futurehouse-client 0.3.20.dev266__py3-none-any.whl → 0.3.20.dev411__py3-none-any.whl

futurehouse_client/clients/rest_client.py

@@ -26,21 +26,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 +41,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,
@@ -58,6 +52,7 @@ from futurehouse_client.models.app import (
     TrajectoryQueryParams,
 )
 from futurehouse_client.models.rest import (
+    DiscoveryResponse,
     ExecutionStatus,
     UserAgentRequest,
     UserAgentRequestPostPayload,
@@ -67,7 +62,10 @@ from futurehouse_client.models.rest import (
     WorldModelResponse,
 )
 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,
@@ -135,10 +133,22 @@ class WorldModelCreationError(RestClientError):
     """Raised when there's an error creating a world model."""
 
 
+class WorldModelDeletionError(RestClientError):
+    """Raised when there's an error deleting a world model."""
+
+
 class ProjectError(RestClientError):
     """Raised when there's an error with trajectory group operations."""
 
 
+class DiscoveryCreationError(RestClientError):
+    """Raised when there's an error creating a discovery."""
+
+
+class DiscoveryFetchError(RestClientError):
+    """Raised when there's an error fetching a discovery."""
+
+
 class InvalidTaskDescriptionError(Exception):
     """Raised when the task description is invalid or empty."""
 
@@ -147,28 +157,15 @@ 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
@@ -226,11 +223,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.
 
@@ -238,12 +259,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)
@@ -269,14 +291,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,
                 )
             )
@@ -1584,6 +1606,48 @@ 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 search_world_models(
+        self,
+        query: str,
+        size: int = 10,
+        total_search_size: int = 50,
+        search_all_versions: bool = False,
+    ) -> list[str]:
+        """Search for world models.
+
+        Args:
+            query: The search query.
+            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.
+
+        Returns:
+            A list of world model names.
+        """
+        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,
+                },
+            )
+            response.raise_for_status()
+            return response.json()
+        except HTTPStatusError as e:
+            raise WorldModelFetchError(
+                f"Error searching 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),
@@ -1659,6 +1723,32 @@ class RestClient:
                 f"An unexpected error occurred during world model creation: {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,
+    )
+    async def delete_world_model(self, world_model_id: UUID) -> None:
+        """Delete a world model snapshot by its ID.
+
+        Args:
+            world_model_id: The unique ID of the world model snapshot to delete.
+
+        Raises:
+            WorldModelDeletionError: If the API call fails.
+        """
+        try:
+            response = await self.async_client.delete(
+                f"/v0.1/world-models/{world_model_id}"
+            )
+            response.raise_for_status()
+        except HTTPStatusError as e:
+            raise WorldModelDeletionError(
+                f"Error deleting world model: {e.response.status_code} - {e.response.text}"
+            ) from e
+        except Exception as e:
+            raise WorldModelDeletionError(f"An unexpected error occurred: {e}") from e
+
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -1737,11 +1827,18 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
-    def create_project(self, name: str) -> str:
+    def create_project(
+        self,
+        name: str,
+        shared_with_organizations: list[int] | None = None,
+        shared_with_users: list[str] | None = None,
+    ) -> str:
         """Create a new project.
 
         Args:
             name: The name for the project
+            shared_with_organizations: List of organization IDs to share the project with
+            shared_with_users: List of user emails to share the project with
 
         Returns:
             UUID of the created project
@@ -1750,7 +1847,11 @@ class RestClient:
             ProjectError: If there's an error creating the project
         """
         try:
-            data = {"name": name}
+            data = {
+                "name": name,
+                "shared_with_organizations": shared_with_organizations,
+                "shared_with_users": shared_with_users,
+            }
             response = self.client.post("/v0.1/projects", json=data)
             response.raise_for_status()
             return response.json()
@@ -1802,11 +1903,18 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
-    async def acreate_project(self, name: str) -> str:
+    async def acreate_project(
+        self,
+        name: str,
+        shared_with_organizations: list[int] | None = None,
+        shared_with_users: list[str] | None = None,
+    ) -> str:
         """Asynchronously create a new project.
 
         Args:
             name: The name for the project
+            shared_with_organizations: List of organization IDs to share the project with
+            shared_with_users: List of user emails to share the project with
 
         Returns:
             UUID of the created project
@@ -1815,7 +1923,11 @@ class RestClient:
             ProjectError: If there's an error creating the project
         """
         try:
-            data = {"name": name}
+            data = {
+                "name": name,
+                "shared_with_organizations": shared_with_organizations,
+                "shared_with_users": shared_with_users,
+            }
             response = await self.async_client.post("/v0.1/projects", json=data)
             response.raise_for_status()
             return response.json()
@@ -2326,6 +2438,108 @@ class RestClient:
                 f"An unexpected error occurred while responding to the request: {e!r}."
             ) from e
 
+    def create_discovery(
+        self,
+        project_id: UUID,
+        world_model_id: UUID,
+        dataset_id: UUID,
+        description: str,
+        associated_trajectories: list[UUID],
+        validation_level: int,
+    ) -> UUID:
+        """Create a new discovery.
+
+        Args:
+            project_id: The ID of the project the discovery is associated with.
+            world_model_id: The ID of the world model the discovery is associated with.
+            dataset_id: The ID of the dataset the discovery is associated with.
+            description: The description of the discovery.
+            associated_trajectories: The IDs of the trajectories to associate with the discovery.
+            validation_level: The validation level of the discovery.
+
+        Returns:
+            The ID of the created discovery.
+
+        Raises:
+            DiscoveryCreationError: If there's an error creating the discovery.
+        """
+        try:
+            data = {
+                "project_id": str(project_id),
+                "world_model_id": str(world_model_id),
+                "dataset_id": str(dataset_id),
+                "description": description,
+                "associated_trajectories": [
+                    str(trajectory) for trajectory in associated_trajectories
+                ],
+                "validation_level": validation_level,
+            }
+            response = self.client.post("/v0.1/discoveries", json=data)
+            response.raise_for_status()
+            return UUID(response.json())
+        except HTTPStatusError as e:
+            raise DiscoveryCreationError(
+                f"Error creating discovery: {e.response.status_code} - {e.response.text}"
+            ) from e
+        except Exception as e:
+            raise DiscoveryCreationError(f"Error creating discovery: {e!r}") from e
+
+    def get_discovery(self, discovery_id: UUID) -> DiscoveryResponse:
+        """Get a discovery by its ID.
+
+        Args:
+            discovery_id: The ID of the discovery to get.
+
+        Returns:
+            The discovery.
+
+        Raises:
+            DiscoveryFetchError: If there's an error fetching the discovery.
+        """
+        try:
+            response = self.client.get(f"/v0.1/discoveries/{discovery_id}")
+            response.raise_for_status()
+            return DiscoveryResponse.model_validate(response.json())
+        except HTTPStatusError as e:
+            raise DiscoveryFetchError(
+                f"Error fetching discovery: {e.response.status_code} - {e.response.text}"
+            ) from e
+        except Exception as e:
+            raise DiscoveryFetchError(f"Error fetching discovery: {e!r}") from e
+
+    def list_discoveries_for_project(
+        self, project_id: UUID, limit: int = 50, offset: int = 0
+    ) -> list[DiscoveryResponse]:
+        """List discoveries for a specific project.
+
+        Args:
+            project_id: The ID of the project to get discoveries for.
+            limit: The maximum number of discoveries to return.
+            offset: The number of discoveries to skip.
+
+        Returns:
+            A list of discoveries for the specified project.
+
+        Raises:
+            DiscoveryFetchError: If there's an error fetching the discoveries.
+        """
+        try:
+            response = self.client.get(
+                f"/v0.1/projects/{project_id}/discoveries",
+                params={"limit": limit, "offset": offset},
+            )
+            response.raise_for_status()
+            data = response.json()
+            return [DiscoveryResponse.model_validate(d) for d in data]
+        except HTTPStatusError as e:
+            raise DiscoveryFetchError(
+                f"Error fetching discoveries for project: {e.response.status_code} - {e.response.text}"
+            ) from e
+        except Exception as e:
+            raise DiscoveryFetchError(
+                f"Error fetching discoveries for project: {e!r}"
+            ) from e
+
 
 def get_installed_packages() -> dict[str, str]:
     """Returns a dictionary of installed packages and their versions."""