futurehouse-client 0.4.1.dev95__py3-none-any.whl → 0.4.2__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
@@ -54,12 +55,14 @@ 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 (
@@ -332,16 +335,23 @@ class RestClient(DataStorageMethods):
             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()
@@ -699,10 +709,12 @@ class RestClient(DataStorageMethods):
 
     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,
@@ -770,10 +782,12 @@ class RestClient(DataStorageMethods):
 
     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,
@@ -846,12 +860,9 @@ class RestClient(DataStorageMethods):
     )
     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
@@ -861,7 +872,7 @@ class RestClient(DataStorageMethods):
         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:
@@ -1606,6 +1617,56 @@ class RestClient(DataStorageMethods):
         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),
@@ -1613,31 +1674,41 @@ class RestClient(DataStorageMethods):
     )
     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:
-            response = self.client.get(
-                "/v0.1/world-models/search/",
-                params={
-                    "query": query,
-                    "size": size,
-                    "total_search_size": total_search_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()
             return response.json()
@@ -1754,22 +1825,19 @@ class RestClient(DataStorageMethods):
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
-    def get_project_by_name(self, name: str) -> UUID:
+    def get_project_by_name(self, name: str, limit: int = 2) -> UUID | list[UUID]:
         """Get a project UUID by name.
 
         Args:
             name: The name of the project to find
+            limit: Maximum number of projects to return
 
         Returns:
-            UUID of the project as a string
-
-        Raises:
-            ProjectError: If no project is found, multiple projects are found, or there's an error
+            UUID of the project as a string or a list of UUIDs if multiple projects are found
         """
         try:
-            # Get projects filtered by name (backend now filters by name and owner)
             response = self.client.get(
-                "/v0.1/projects", params={"limit": 2, "name": name}
+                "/v0.1/projects", params={"limit": limit, "name": name}
             )
             response.raise_for_status()
             projects = response.json()
@@ -1782,32 +1850,33 @@ class RestClient(DataStorageMethods):
         if len(projects) == 0:
             raise ProjectError(f"No project found with name '{name}'")
         if len(projects) > 1:
-            raise ProjectError(
+            logger.warning(
                 f"Multiple projects found with name '{name}'. Found {len(projects)} projects."
             )
 
-        return UUID(projects[0]["id"])
+        ids = [UUID(project["id"]) for project in projects]
+        return ids[0] if len(ids) == 1 else ids
 
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
-    async def aget_project_by_name(self, name: str) -> UUID:
+    async def aget_project_by_name(
+        self, name: str, limit: int = 2
+    ) -> UUID | list[UUID]:
         """Asynchronously get a project UUID by name.
 
         Args:
             name: The name of the project to find
+            limit: Maximum number of projects to return
 
         Returns:
-            UUID of the project as a string
-
-        Raises:
-            ProjectError: If no project is found, multiple projects are found, or there's an error
+            UUID of the project as a string or a list of UUIDs if multiple projects are found
         """
         try:
             response = await self.async_client.get(
-                "/v0.1/projects", params={"limit": 2, "name": name}
+                "/v0.1/projects", params={"limit": limit, "name": name}
             )
             response.raise_for_status()
             projects = response.json()
@@ -1816,11 +1885,12 @@ class RestClient(DataStorageMethods):
         if len(projects) == 0:
             raise ProjectError(f"No project found with name '{name}'")
         if len(projects) > 1:
-            raise ProjectError(
+            logger.warning(
                 f"Multiple projects found with name '{name}'. Found {len(projects)} projects."
             )
 
-        return UUID(projects[0]["id"])
+        ids = [UUID(project["id"]) for project in projects]
+        return ids[0] if len(ids) == 1 else ids
 
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),

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/app.py

@@ -27,10 +27,7 @@ if TYPE_CHECKING:
 MAX_CROW_JOB_RUN_TIMEOUT = 60 * 60 * 24  # 24 hours in sec
 MIN_CROW_JOB_RUN_TIMEOUT = 0  # sec
 
-
-class PythonVersion(StrEnum):
-    V3_11 = "3.11"
-    V3_12 = "3.12"
+DEFAULT_PYTHON_VERSION_USED_FOR_JOB_BUILDS = "3.13"
 
 
 class AuthType(StrEnum):
@@ -420,9 +417,9 @@ class JobDeploymentConfig(BaseModel):
         description="The configuration for the cloud run container.",
     )
 
-    python_version: PythonVersion = Field(
-        default=PythonVersion.V3_12,
-        description="The python version your docker image should build with.",
+    python_version: str = Field(
+        default=DEFAULT_PYTHON_VERSION_USED_FOR_JOB_BUILDS,
+        description="The python version your docker image should build with (e.g., '3.11', '3.12', '3.13').",
     )
 
     agent: Agent | AgentConfig | str = Field(

futurehouse_client/models/data_storage_methods.py

@@ -34,6 +34,10 @@ class DataStorageEntry(BaseModel):
         default=None,
         description="ID of the parent entry if this is a sub-entry for hierarchical storage",
     )
+    project_id: UUID | None = Field(
+        default=None,
+        description="ID of the project this data storage entry belongs to",
+    )
     dataset_id: UUID | None = Field(
         default=None,
         description="ID of the dataset this entry belongs to",
@@ -79,8 +83,8 @@ class DataStorageLocationPayload(BaseModel):
     location: str | None = None
 
 
-class DataStorageLocationDetails(BaseModel):
-    """Model representing the location details within a DataStorageLocations object."""
+class DataStorageLocationConfig(BaseModel):
+    """Model representing the location configuration within a DataStorageLocations object."""
 
     storage_type: str = Field(description="Type of storage (e.g., 'gcs', 'pg_table')")
     content_type: str = Field(description="Type of content stored")
@@ -89,15 +93,19 @@ class DataStorageLocationDetails(BaseModel):
     location: str | None = Field(
         default=None, description="Location path or identifier"
     )
+    signed_url: str | None = Field(
+        default=None,
+        description="Signed URL for uploading/downloading the file to/from GCS",
+    )
 
 
-class DataStorageLocations(BaseModel):
+class DataStorageLocation(BaseModel):
     """Model representing storage locations for a data storage entry."""
 
     id: UUID = Field(description="Unique identifier for the storage locations")
     data_storage_id: UUID = Field(description="ID of the associated data storage entry")
-    storage_config: DataStorageLocationDetails = Field(
-        description="Storage configuration details"
+    storage_config: DataStorageLocationConfig = Field(
+        description="Storage location configuration"
     )
     created_at: datetime = Field(description="Timestamp when the location was created")
 
@@ -106,13 +114,9 @@ class DataStorageResponse(BaseModel):
     """Response model for data storage operations."""
 
     data_storage: DataStorageEntry = Field(description="The created data storage entry")
-    storage_location: DataStorageLocations = Field(
+    storage_locations: list[DataStorageLocation] = Field(
         description="Storage location for this data entry"
     )
-    signed_url: str | None = Field(
-        default=None,
-        description="Signed URL for uploading/downloading the file to/from GCS",
-    )
 
 
 class DataStorageRequestPayload(BaseModel):
@@ -131,6 +135,10 @@ class DataStorageRequestPayload(BaseModel):
     parent_id: UUID | None = Field(
         default=None, description="ID of the parent entry for hierarchical storage"
     )
+    project_id: UUID | None = Field(
+        default=None,
+        description="ID of the project this data storage entry belongs to",
+    )
     dataset_id: UUID | None = Field(
         default=None,
         description="ID of existing dataset to add entry to, or None to create new dataset",
@@ -144,6 +152,19 @@ class DataStorageRequestPayload(BaseModel):
     )
 
 
+class CreateDatasetPayload(BaseModel):
+    """Payload for creating a dataset."""
+
+    id: UUID | None = Field(
+        default=None,
+        description="ID of the dataset to create, or None to create a new dataset",
+    )
+    name: str = Field(description="Name of the dataset")
+    description: str | None = Field(
+        default=None, description="Description of the dataset"
+    )
+
+
 class ManifestEntry(BaseModel):
     """Model representing a single entry in a manifest file."""
 

futurehouse_client/models/rest.py

@@ -23,6 +23,15 @@ class StoreEnvironmentFrameRequest(BaseModel):
     trajectory_timestep: int
 
 
+class TrajectoryPatchRequest(BaseModel):
+    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
+
+
 class ExecutionStatus(StrEnum):
     QUEUED = auto()
     IN_PROGRESS = "in progress"
@@ -54,7 +63,37 @@ class WorldModel(BaseModel):
     project_id: UUID | str | None = None
 
 
-class WorldModelResponse(BaseModel):
+class SearchOperator(StrEnum):
+    """Operators for structured search criteria."""
+
+    EQUALS = "equals"
+    CONTAINS = "contains"
+    STARTS_WITH = "starts_with"
+    ENDS_WITH = "ends_with"
+    GREATER_THAN = "greater_than"
+    LESS_THAN = "less_than"
+    BETWEEN = "between"
+    IN = "in"
+
+
+class SearchCriterion(BaseModel):
+    """A single search criterion with field, operator, and value."""
+
+    field: str
+    operator: SearchOperator
+    value: str | list[str] | bool
+
+
+class WorldModelSearchPayload(BaseModel):
+    """Payload for structured world model search."""
+
+    criteria: list[SearchCriterion]
+    size: int = 10
+    project_id: UUID | str | None = None
+    search_all_versions: bool = False
+
+
+class WorldModelResponse(WorldModel):
     """
     Response model for a world model snapshot.
 
@@ -62,13 +101,8 @@ class WorldModelResponse(BaseModel):
     """
 
     id: UUID | str
-    prior: UUID | str | None
-    name: str
-    description: str | None
-    content: str
-    trajectory_id: UUID | str | None
+    name: str  # type: ignore[mutable-override]  # The API always returns a non-optional name, overriding the base model's optional field.
     email: str | None
-    model_metadata: JsonValue | None
     enabled: bool
     created_at: datetime
 
@@ -132,3 +166,10 @@ class DiscoveryResponse(BaseModel):
     associated_trajectories: list[UUID | str]
     validation_level: int
     created_at: datetime
+
+
+class DataStorageSearchPayload(BaseModel):
+    """Payload for structured data storage search."""
+
+    criteria: list[SearchCriterion]
+    size: int = 10

futurehouse_client/utils/general.py

@@ -1,22 +1,31 @@
 import asyncio
-from collections.abc import Awaitable, Iterable
+from collections.abc import Awaitable, Callable, Iterable
 from typing import TypeVar
 
 from httpx import (
     CloseError,
     ConnectError,
     ConnectTimeout,
+    HTTPStatusError,
     NetworkError,
     ReadError,
     ReadTimeout,
     RemoteProtocolError,
+    codes,
 )
 from requests.exceptions import RequestException, Timeout
-from tenacity import retry_if_exception_type
+from tenacity import RetryCallState
 from tqdm.asyncio import tqdm
 
 T = TypeVar("T")
 
+RETRYABLE_HTTP_STATUS_CODES = {
+    codes.TOO_MANY_REQUESTS,
+    codes.INTERNAL_SERVER_ERROR,
+    codes.BAD_GATEWAY,
+    codes.SERVICE_UNAVAILABLE,
+    codes.GATEWAY_TIMEOUT,
+}
 
 _BASE_CONNECTION_ERRORS = (
     # From requests
@@ -33,12 +42,32 @@ _BASE_CONNECTION_ERRORS = (
     CloseError,
 )
 
-retry_if_connection_error = retry_if_exception_type(_BASE_CONNECTION_ERRORS)
 
+def create_retry_if_connection_error(
+    *additional_exceptions,
+) -> Callable[[RetryCallState], bool]:
+    """Create a retry condition with base connection errors, HTTP status errors, plus additional exceptions."""
 
-def create_retry_if_connection_error(*additional_exceptions):
-    """Create a retry condition with base connection errors plus additional exceptions."""
-    return retry_if_exception_type(_BASE_CONNECTION_ERRORS + additional_exceptions)
+    def status_retries_with_exceptions(retry_state: RetryCallState) -> bool:
+        if retry_state.outcome is not None and hasattr(
+            retry_state.outcome, "exception"
+        ):
+            exception = retry_state.outcome.exception()
+            # connection errors
+            if isinstance(exception, _BASE_CONNECTION_ERRORS):
+                return True
+            # custom exceptions provided
+            if additional_exceptions and isinstance(exception, additional_exceptions):
+                return True
+            # any http exceptions
+            if isinstance(exception, HTTPStatusError):
+                return exception.response.status_code in RETRYABLE_HTTP_STATUS_CODES
+        return False
+
+    return status_retries_with_exceptions
+
+
+retry_if_connection_error = create_retry_if_connection_error()
 
 
 async def gather_with_concurrency(