diff --git a/main.py b/main.py deleted file mode 100644 index 27c02a6..0000000 --- a/main.py +++ /dev/null @@ -1,6 +0,0 @@ -def main(): - print("Hello from roe-python!") - - -if __name__ == "__main__": - main() diff --git a/src/roe/api/agents.py b/src/roe/api/agents.py index 3c1cfcd..41e6e1e 100644 --- a/src/roe/api/agents.py +++ b/src/roe/api/agents.py @@ -109,6 +109,14 @@ from roe.utils.generated_request import request_json, request_raw +_MAX_BATCH_SIZE = 1000 + + +def _iter_chunks(items, chunk_size: int): + for i in range(0, len(items), chunk_size): + yield items[i : i + chunk_size] + + def _build_run_headers( idempotency_key: str | None = None, skip_cache: bool = False, @@ -273,8 +281,6 @@ def delete(self, agent_id: str, version_id: str) -> None: class AgentJobsAPI: """Nested API for agent job operations.""" - _MAX_BATCH_SIZE = 1000 - def __init__(self, agents_api: "AgentsAPI"): self._agents_api = agents_api @@ -286,11 +292,6 @@ def _raw(self) -> AuthenticatedClient: def _org_id(self) -> UUID: return UUID(str(self._agents_api.config.organization_id)) - @staticmethod - def _iter_chunks(items, chunk_size: int): - for i in range(0, len(items), chunk_size): - yield items[i : i + chunk_size] - def list( self, agent_id: str, @@ -349,9 +350,7 @@ def retrieve_result(self, job_id: str) -> AgentJobResultResponse: def retrieve_status_many(self, job_ids: list[str]) -> list[AgentJobStatus]: results: list[AgentJobStatus] = [] is_first_chunk = True - for chunk in self._iter_chunks(job_ids, self._MAX_BATCH_SIZE): - if not chunk: - continue + for chunk in _iter_chunks(job_ids, _MAX_BATCH_SIZE): if not is_first_chunk: time.sleep(self._agents_api.config.batch_chunk_delay) is_first_chunk = False @@ -374,9 +373,7 @@ def retrieve_result_many( ) -> list[Any]: # AgentJobResultItem results: list[Any] = [] is_first_chunk = True - for chunk in self._iter_chunks(job_ids, self._MAX_BATCH_SIZE): - if not chunk: - continue + for chunk in _iter_chunks(job_ids, _MAX_BATCH_SIZE): if not is_first_chunk: time.sleep(self._agents_api.config.batch_chunk_delay) is_first_chunk = False @@ -475,8 +472,6 @@ def delete_data(self, job_id: str) -> AgentJobDeleteDataResponse: class AgentsAPI: """API for managing and running agents.""" - _MAX_BATCH_SIZE = 1000 - def __init__(self, config: RoeConfig, raw_client: AuthenticatedClient): self.config = config self._raw = raw_client @@ -495,11 +490,6 @@ def versions(self) -> AgentVersionsAPI: def jobs(self) -> AgentJobsAPI: return self._jobs - @staticmethod - def _iter_chunks(items, chunk_size: int): - for i in range(0, len(items), chunk_size): - yield items[i : i + chunk_size] - def list( self, page: int | None = None, @@ -666,9 +656,7 @@ def run_many( """ all_job_ids: list[str] = [] is_first_chunk = True - for chunk in self._iter_chunks(batch_inputs, self._MAX_BATCH_SIZE): - if not chunk: - continue + for chunk in _iter_chunks(batch_inputs, _MAX_BATCH_SIZE): if not is_first_chunk: time.sleep(self.config.batch_chunk_delay) is_first_chunk = False diff --git a/src/roe/api/policies.py b/src/roe/api/policies.py index f02c177..024abe4 100644 --- a/src/roe/api/policies.py +++ b/src/roe/api/policies.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any from uuid import UUID from roe._generated.api.policies import ( @@ -57,10 +57,6 @@ def _parse_policy_version(data: dict[str, Any]) -> PolicyVersion: return PolicyVersion.from_dict(_normalize_policy_version_wire(data)) -if TYPE_CHECKING: - pass - - class PolicyVersionsAPI: """Nested API for policy version operations.""" diff --git a/src/roe/models/job.py b/src/roe/models/job.py index 3d3fd80..a7ba97a 100644 --- a/src/roe/models/job.py +++ b/src/roe/models/job.py @@ -16,6 +16,7 @@ from roe._generated.models.agent_job_result_response import AgentJobResultResponse from roe._generated.models.agent_job_single_status import AgentJobSingleStatus from roe._generated.models.agent_job_status import AgentJobStatus +from roe._generated.types import UNSET from roe.exceptions import NotFoundError, RoeAPIException if TYPE_CHECKING: @@ -277,7 +278,7 @@ def retrieve_status(self) -> dict[str, AgentJobStatus]: timestamp=cached.get("timestamp", 0), error_message=cached["error_message"] if cached["error_message"] is not None - else _UNSET_SENTINEL(), + else UNSET, ) else: jobs_to_query.append(job_id) @@ -298,7 +299,7 @@ def retrieve_status(self) -> dict[str, AgentJobStatus]: created_at=None, last_updated_at=None, timestamp=self._extract_timestamp(status_item), - error_message=err if err is not None else _UNSET_SENTINEL(), + error_message=err if err is not None else UNSET, ) self._job_statuses[job_id] = { "status": stat_code, @@ -348,10 +349,3 @@ def _extract_timestamp(status_item: Any) -> int: if isinstance(value, Unset) or value is None: return 0 return int(value) - - -def _UNSET_SENTINEL(): # noqa: N802 — match generated naming style - """Return the ``Unset`` sentinel from the generated types module.""" - from roe._generated.types import UNSET - - return UNSET