futurehouse-client 0.4.2.dev11__py3-none-any.whl → 0.4.3__py3-none-any.whl

futurehouse_client/clients/job_client.py

@@ -19,6 +19,7 @@ from futurehouse_client.models.rest import (
     FinalEnvironmentRequest,
     StoreAgentStatePostRequest,
     StoreEnvironmentFrameRequest,
+    TrajectoryPatchRequest,
 )
 from futurehouse_client.utils.monitoring import (
     external_trace,
@@ -318,3 +319,52 @@ class JobClient:
                 f"Unexpected error storing environment frame for state {state_identifier}",
             )
             raise
+
+    async def patch_trajectory(
+        self,
+        public: bool | None = None,
+        shared_with: list[int] | None = None,
+        notification_enabled: bool | None = None,
+        notification_type: str | None = None,
+        min_estimated_time: float | None = None,
+        max_estimated_time: float | None = None,
+    ) -> None:
+        data = TrajectoryPatchRequest(
+            public=public,
+            shared_with=shared_with,
+            notification_enabled=notification_enabled,
+            notification_type=notification_type,
+            min_estimated_time=min_estimated_time,
+            max_estimated_time=max_estimated_time,
+        )
+        try:
+            async with httpx.AsyncClient(timeout=self.REQUEST_TIMEOUT) as client:
+                url = f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}"
+                headers = {
+                    "Authorization": f"Bearer {self.oauth_jwt}",
+                    "x-trajectory-id": self.trajectory_id,
+                }
+                response = await client.patch(
+                    url=url,
+                    json=data.model_dump(mode="json", exclude_none=True),
+                    headers=headers,
+                )
+                response.raise_for_status()
+                logger.debug("Trajectory updated successfully")
+        except httpx.HTTPStatusError as e:
+            logger.exception(
+                "HTTP error while patching trajectory. "
+                f"Status code: {e.response.status_code}, "
+                f"Response: {e.response.text}",
+            )
+        except httpx.TimeoutException:
+            logger.exception(
+                f"Timeout while patching trajectory after {self.REQUEST_TIMEOUT}s",
+            )
+            raise
+        except httpx.NetworkError:
+            logger.exception("Network error while patching trajectory")
+            raise
+        except Exception:
+            logger.exception("Unexpected error while patching trajectory")
+            raise

futurehouse_client/clients/rest_client.py

@@ -1,3 +1,4 @@
+# ruff: noqa: PLR0915
 import ast
 import asyncio
 import base64
@@ -26,21 +27,14 @@ from httpx import (
     AsyncClient,
     Client,
     CloseError,
-    ConnectError,
-    ConnectTimeout,
     HTTPStatusError,
-    NetworkError,
-    ReadError,
-    ReadTimeout,
     RemoteProtocolError,
     codes,
 )
 from ldp.agent import AgentConfig
-from requests.exceptions import RequestException, Timeout
 from tenacity import (
     before_sleep_log,
     retry,
-    retry_if_exception_type,
     stop_after_attempt,
     wait_exponential,
 )
@@ -48,6 +42,7 @@ from tqdm import tqdm as sync_tqdm
 from tqdm.asyncio import tqdm
 
 from futurehouse_client.clients import JobNames
+from futurehouse_client.clients.data_storage_methods import DataStorageMethods
 from futurehouse_client.models.app import (
     AuthType,
     JobDeploymentConfig,
@@ -60,15 +55,20 @@ from futurehouse_client.models.app import (
 from futurehouse_client.models.rest import (
     DiscoveryResponse,
     ExecutionStatus,
+    SearchCriterion,
     UserAgentRequest,
     UserAgentRequestPostPayload,
     UserAgentRequestStatus,
     UserAgentResponsePayload,
     WorldModel,
     WorldModelResponse,
+    WorldModelSearchPayload,
 )
 from futurehouse_client.utils.auth import RefreshingJWT
-from futurehouse_client.utils.general import gather_with_concurrency
+from futurehouse_client.utils.general import (
+    create_retry_if_connection_error,
+    gather_with_concurrency,
+)
 from futurehouse_client.utils.module_utils import (
     OrganizationSelector,
     fetch_environment_function_docstring,
@@ -160,28 +160,14 @@ class FileUploadError(RestClientError):
     """Raised when there's an error uploading a file."""
 
 
-retry_if_connection_error = retry_if_exception_type((
-    # From requests
-    Timeout,
-    ConnectionError,
-    RequestException,
-    # From httpx
-    ConnectError,
-    ConnectTimeout,
-    ReadTimeout,
-    ReadError,
-    NetworkError,
-    RemoteProtocolError,
-    CloseError,
-    FileUploadError,
-))
+retry_if_connection_error = create_retry_if_connection_error(FileUploadError)
 
 DEFAULT_AGENT_TIMEOUT: int = 2400  # seconds
 
 
-# pylint: disable=too-many-public-methods
-class RestClient:
-    REQUEST_TIMEOUT: ClassVar[float] = 30.0  # sec
+class RestClient(DataStorageMethods):
+    REQUEST_TIMEOUT: ClassVar[float] = 30.0  # sec - for general API calls
+    FILE_UPLOAD_TIMEOUT: ClassVar[float] = 600.0  # 10 minutes - for file uploads
     MAX_RETRY_ATTEMPTS: ClassVar[int] = 3
     RETRY_MULTIPLIER: ClassVar[int] = 1
     MAX_RETRY_WAIT: ClassVar[int] = 10
@@ -239,11 +225,35 @@ class RestClient:
         """Authenticated HTTP client for multipart uploads."""
         return cast(Client, self.get_client(None, authenticated=True))
 
+    @property
+    def file_upload_client(self) -> Client:
+        """Authenticated HTTP client with extended timeout for file uploads."""
+        return cast(
+            Client,
+            self.get_client(
+                "application/json", authenticated=True, timeout=self.FILE_UPLOAD_TIMEOUT
+            ),
+        )
+
+    @property
+    def async_file_upload_client(self) -> AsyncClient:
+        """Authenticated async HTTP client with extended timeout for file uploads."""
+        return cast(
+            AsyncClient,
+            self.get_client(
+                "application/json",
+                authenticated=True,
+                async_client=True,
+                timeout=self.FILE_UPLOAD_TIMEOUT,
+            ),
+        )
+
     def get_client(
         self,
         content_type: str | None = "application/json",
         authenticated: bool = True,
         async_client: bool = False,
+        timeout: float | None = None,
     ) -> Client | AsyncClient:
         """Return a cached HTTP client or create one if needed.
 
@@ -251,12 +261,13 @@ class RestClient:
             content_type: The desired content type header. Use None for multipart uploads.
             authenticated: Whether the client should include authentication.
             async_client: Whether to use an async client.
+            timeout: Custom timeout in seconds. Uses REQUEST_TIMEOUT if not provided.
 
         Returns:
             An HTTP client configured with the appropriate headers.
         """
-        # Create a composite key based on content type and auth flag
-        key = f"{content_type or 'multipart'}_{authenticated}_{async_client}"
+        client_timeout = timeout or self.REQUEST_TIMEOUT
+        key = f"{content_type or 'multipart'}_{authenticated}_{async_client}_{client_timeout}"
 
         if key not in self._clients:
             headers = copy.deepcopy(self.headers)
@@ -282,14 +293,14 @@ class RestClient:
                 AsyncClient(
                     base_url=self.base_url,
                     headers=headers,
-                    timeout=self.REQUEST_TIMEOUT,
+                    timeout=client_timeout,
                     auth=auth,
                 )
                 if async_client
                 else Client(
                     base_url=self.base_url,
                     headers=headers,
-                    timeout=self.REQUEST_TIMEOUT,
+                    timeout=client_timeout,
                     auth=auth,
                 )
             )
@@ -323,16 +334,23 @@ class RestClient:
             raise ValueError(f"Organization '{organization}' not found.")
         return filtered_orgs
 
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
     def _check_job(self, name: str, organization: str) -> dict[str, Any]:
-        try:
-            response = self.client.get(
-                f"/v0.1/crows/{name}/organizations/{organization}"
-            )
-            response.raise_for_status()
-            return response.json()
-        except Exception as e:
-            raise JobFetchError(f"Error checking job: {e!r}.") from e
+        response = self.client.get(f"/v0.1/crows/{name}/organizations/{organization}")
+        response.raise_for_status()
+        return response.json()
 
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
     def _fetch_my_orgs(self) -> list[str]:
         response = self.client.get(f"/v0.1/organizations?filter={True}")
         response.raise_for_status()
@@ -690,10 +708,12 @@ class RestClient:
 
     async def arun_tasks_until_done(
         self,
-        task_data: TaskRequest
-        | dict[str, Any]
-        | Collection[TaskRequest]
-        | Collection[dict[str, Any]],
+        task_data: (
+            TaskRequest
+            | dict[str, Any]
+            | Collection[TaskRequest]
+            | Collection[dict[str, Any]]
+        ),
         verbose: bool = False,
         progress_bar: bool = False,
         concurrency: int = 10,
@@ -761,10 +781,12 @@ class RestClient:
 
     def run_tasks_until_done(
         self,
-        task_data: TaskRequest
-        | dict[str, Any]
-        | Collection[TaskRequest]
-        | Collection[dict[str, Any]],
+        task_data: (
+            TaskRequest
+            | dict[str, Any]
+            | Collection[TaskRequest]
+            | Collection[dict[str, Any]]
+        ),
         verbose: bool = False,
         progress_bar: bool = False,
         timeout: int = DEFAULT_AGENT_TIMEOUT,
@@ -837,12 +859,9 @@ class RestClient:
     )
     def get_build_status(self, build_id: UUID | None = None) -> dict[str, Any]:
         """Get the status of a build."""
-        try:
-            build_id = build_id or self.build_id
-            response = self.client.get(f"/v0.1/builds/{build_id}")
-            response.raise_for_status()
-        except Exception as e:
-            raise JobFetchError(f"Error getting build status: {e!r}.") from e
+        build_id = build_id or self.build_id
+        response = self.client.get(f"/v0.1/builds/{build_id}")
+        response.raise_for_status()
         return response.json()
 
     # TODO: Refactor later so we don't have to ignore PLR0915
@@ -852,7 +871,7 @@ class RestClient:
         retry=retry_if_connection_error,
         before_sleep=before_sleep_log(logger, logging.WARNING),
     )
-    def create_job(self, config: JobDeploymentConfig) -> dict[str, Any]:  # noqa: PLR0915
+    def create_job(self, config: JobDeploymentConfig) -> dict[str, Any]:
         """Creates a futurehouse job deployment from the environment and environment files.
 
         Args:
@@ -1597,6 +1616,56 @@ class RestClient:
         except Exception as e:
             raise WorldModelFetchError(f"An unexpected error occurred: {e!r}.") from e
 
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+    )
+    def list_world_models(
+        self,
+        name: str | None = None,
+        project_id: UUID | str | None = None,
+        limit: int = 150,
+        offset: int = 0,
+        sort_order: str = "asc",
+    ) -> list[WorldModelResponse]:
+        """List world models with different behavior based on filters.
+
+        When filtering by name: returns only the latest version for that name.
+        When filtering by project_id (without name): returns all versions for that project.
+        When no filters: returns latest version of each world model.
+
+        Args:
+            name: Filter by world model name.
+            project_id: Filter by project ID.
+            limit: The maximum number of models to return.
+            offset: Number of results to skip for pagination.
+            sort_order: Sort order 'asc' or 'desc'.
+
+        Returns:
+            A list of world model dictionaries.
+        """
+        try:
+            params: dict[str, str | int] = {
+                "limit": limit,
+                "offset": offset,
+                "sort_order": sort_order,
+            }
+            if name:
+                params["name"] = name
+            if project_id:
+                params["project_id"] = str(project_id)
+
+            response = self.client.get("/v0.1/world-models", params=params)
+            response.raise_for_status()
+            return response.json()
+        except HTTPStatusError as e:
+            raise WorldModelFetchError(
+                f"Error listing world models: {e.response.status_code} - {e.response.text}"
+            ) from e
+        except Exception as e:
+            raise WorldModelFetchError(f"An unexpected error occurred: {e!r}.") from e
+
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -1604,34 +1673,43 @@ class RestClient:
     )
     def search_world_models(
         self,
-        query: str,
+        criteria: list[SearchCriterion] | None = None,
         size: int = 10,
-        total_search_size: int = 50,
+        project_id: UUID | str | None = None,
         search_all_versions: bool = False,
-    ) -> list[str]:
-        """Search for world models.
+    ) -> list[WorldModelResponse]:
+        """Search world models using structured criteria.
 
         Args:
-            query: The search query.
+            criteria: List of SearchCriterion objects with field, operator, and value.
             size: The number of results to return.
-            total_search_size: The number of results to search for.
-            search_all_versions: Whether to search all versions of the world model or just the latest one.
+            project_id: Optional filter by project ID.
+            search_all_versions: Whether to search all versions or just latest.
 
         Returns:
-            A list of world model names.
+            A list of world model responses.
+
+        Example:
+            from futurehouse_client.models.rest import SearchCriterion, SearchOperator
+            criteria = [
+                SearchCriterion(field="name", operator=SearchOperator.CONTAINS, value="chemistry"),
+                SearchCriterion(field="email", operator=SearchOperator.CONTAINS, value="tyler"),
+            ]
+            results = client.search_world_models(criteria=criteria, size=20)
         """
         try:
-            # Use the consolidated endpoint with search parameters
-            response = self.client.get(
-                "/v0.1/world-models",
-                params={
-                    "q": query,
-                    "size": size,
-                    "search_all_versions": search_all_versions,
-                },
+            payload = WorldModelSearchPayload(
+                criteria=criteria or [],
+                size=size,
+                project_id=project_id,
+                search_all_versions=search_all_versions,
+            )
+
+            response = self.client.post(
+                "/v0.1/world-models/search",
+                json=payload.model_dump(mode="json"),
             )
             response.raise_for_status()
-            # The new endpoint returns a list of models directly
             return response.json()
         except HTTPStatusError as e:
             raise WorldModelFetchError(

futurehouse_client/models/__init__.py

@@ -13,7 +13,7 @@ from .app import (
     TaskResponse,
     TaskResponseVerbose,
 )
-from .rest import WorldModel, WorldModelResponse
+from .rest import TrajectoryPatchRequest, WorldModel, WorldModelResponse
 
 __all__ = [
     "AuthType",
@@ -29,6 +29,7 @@ __all__ = [
     "TaskRequest",
     "TaskResponse",
     "TaskResponseVerbose",
+    "TrajectoryPatchRequest",
     "WorldModel",
     "WorldModelResponse",
 ]

futurehouse_client/models/client.py

@@ -27,13 +27,17 @@ class InitialState(BaseState):
 
 class ASVState(BaseState, Generic[T]):
     action: OpResult[T] = Field()
-    next_agent_state: Any = Field()
+    next_state: Any = Field()
     value: float = Field()
 
     @field_serializer("action")
     def serialize_action(self, action: OpResult[T]) -> dict:
         return action.to_dict()
 
+    @field_serializer("next_state")
+    def serialize_next_state(self, state: Any) -> str:
+        return str(state)
+
 
 class EnvResetState(BaseState):
     observations: list[Message] = Field()