futurehouse-client 0.4.0__py3-none-any.whl → 0.4.1.dev95__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,
@@ -68,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,
@@ -136,6 +133,10 @@ 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."""
 
@@ -156,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
@@ -235,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.
 
@@ -247,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)
@@ -278,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,
                 )
             )
@@ -1593,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),
@@ -1668,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),

futurehouse_client/models/data_storage_methods.py

@@ -0,0 +1,333 @@
+import contextlib
+from datetime import datetime
+from enum import StrEnum, auto
+from os import PathLike
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+from pydantic import BaseModel, Field, JsonValue
+
+
+class DataStorageEntry(BaseModel):
+    """Model representing a data storage entry."""
+
+    id: UUID = Field(description="Unique identifier for the data storage entry")
+    name: str = Field(description="Name of the data storage entry")
+    description: str | None = Field(
+        default=None, description="Description of the data storage entry"
+    )
+    content: str | None = Field(
+        default=None, description="Content of the data storage entry"
+    )
+    embedding: list[float] | None = Field(
+        default=None, description="Embedding vector for the content"
+    )
+    is_collection: bool = Field(
+        default=False, description="Whether this entry is a collection"
+    )
+    tags: list[str] | None = Field(
+        default=None,
+        description="List of tags associated with the data storage entry",
+    )
+    parent_id: UUID | None = Field(
+        default=None,
+        description="ID of the parent entry if this is a sub-entry for hierarchical storage",
+    )
+    dataset_id: UUID | None = Field(
+        default=None,
+        description="ID of the dataset this entry belongs to",
+    )
+    path: str | None = Field(
+        default=None,
+        description="Path in the storage system where this entry is located, if a file.",
+    )
+    bigquery_schema: Any | None = Field(
+        default=None, description="Target BigQuery schema for the data storage entry"
+    )
+    user_id: str = Field(description="ID of the user who created this entry")
+    created_at: datetime = Field(description="Timestamp when the entry was created")
+    modified_at: datetime = Field(
+        description="Timestamp when the entry was last updated"
+    )
+
+
+class DataStorageType(StrEnum):
+    BIGQUERY = auto()
+    GCS = auto()
+    PG_TABLE = auto()
+    RAW_CONTENT = auto()
+    ELASTIC_SEARCH = auto()
+
+
+class DataContentType(StrEnum):
+    BQ_DATASET = auto()
+    BQ_TABLE = auto()
+    TEXT = auto()
+    TEXT_W_EMBEDDINGS = auto()
+    DIRECTORY = auto()
+    FILE = auto()
+    INDEX = auto()
+    INDEX_W_EMBEDDINGS = auto()
+
+
+class DataStorageLocationPayload(BaseModel):
+    storage_type: DataStorageType
+    content_type: DataContentType
+    content_schema: JsonValue | None = None
+    metadata: JsonValue | None = None
+    location: str | None = None
+
+
+class DataStorageLocationDetails(BaseModel):
+    """Model representing the location details 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")
+    content_schema: JsonValue | None = Field(default=None, description="Content schema")
+    metadata: JsonValue | None = Field(default=None, description="Location metadata")
+    location: str | None = Field(
+        default=None, description="Location path or identifier"
+    )
+
+
+class DataStorageLocations(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"
+    )
+    created_at: datetime = Field(description="Timestamp when the location was created")
+
+
+class DataStorageResponse(BaseModel):
+    """Response model for data storage operations."""
+
+    data_storage: DataStorageEntry = Field(description="The created data storage entry")
+    storage_location: DataStorageLocations = 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):
+    """Payload for creating a data storage entry."""
+
+    name: str = Field(description="Name of the data storage entry")
+    description: str | None = Field(
+        default=None, description="Description of the data storage entry"
+    )
+    content: str | None = Field(
+        default=None, description="Content of the data storage entry"
+    )
+    is_collection: bool = Field(
+        default=False, description="Whether this entry is a collection"
+    )
+    parent_id: UUID | None = Field(
+        default=None, description="ID of the parent entry for hierarchical storage"
+    )
+    dataset_id: UUID | None = Field(
+        default=None,
+        description="ID of existing dataset to add entry to, or None to create new dataset",
+    )
+    path: PathLike | str | None = Field(
+        default=None,
+        description="Path to store in the GCS bucket.",
+    )
+    existing_location: DataStorageLocationPayload | None = Field(
+        default=None, description="Target storage metadata"
+    )
+
+
+class ManifestEntry(BaseModel):
+    """Model representing a single entry in a manifest file."""
+
+    description: str | None = Field(
+        default=None, description="Description of the file or directory"
+    )
+    metadata: dict[str, Any] | None = Field(
+        default=None, description="Additional metadata for the entry"
+    )
+
+
+class DirectoryManifest(BaseModel):
+    """Model representing the structure of a manifest file."""
+
+    entries: dict[str, "ManifestEntry | DirectoryManifest"] = Field(
+        default_factory=dict,
+        description="Map of file/directory names to their manifest entries",
+    )
+
+    def get_entry_description(self, name: str) -> str | None:
+        """Get description for a specific entry."""
+        entry = self.entries.get(name)
+        if isinstance(entry, ManifestEntry):
+            return entry.description
+        if isinstance(entry, DirectoryManifest):
+            # For nested directories, could derive description from contents
+            return None
+        return None
+
+    def get_entry_metadata(self, name: str) -> dict[str, Any] | None:
+        """Get metadata for a specific entry."""
+        entry = self.entries.get(name)
+        if isinstance(entry, ManifestEntry):
+            return entry.metadata
+        return None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "DirectoryManifest":
+        """Create DirectoryManifest from a dictionary (loaded from JSON/YAML)."""
+        entries: dict[str, ManifestEntry | DirectoryManifest] = {}
+        for name, value in data.items():
+            if isinstance(value, dict):
+                if "description" in value or "metadata" in value:
+                    # This looks like a ManifestEntry
+                    entries[name] = ManifestEntry(**value)
+                else:
+                    # This looks like a nested directory
+                    entries[name] = cls.from_dict(value)
+            else:
+                # Simple string description
+                entries[name] = ManifestEntry(description=str(value))
+
+        return cls(entries=entries)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert back to dictionary format."""
+        result = {}
+        for name, entry in self.entries.items():
+            if isinstance(entry, ManifestEntry):
+                if entry.description is not None or entry.metadata is not None:
+                    entry_dict = {}
+                    if entry.description is not None:
+                        entry_dict["description"] = entry.description
+                    if entry.metadata is not None:
+                        entry_dict.update(entry.metadata)
+                    result[name] = entry_dict
+            elif isinstance(entry, DirectoryManifest):
+                result[name] = entry.to_dict()
+        return result
+
+
+class FileMetadata(BaseModel):
+    """Model representing metadata for a file being processed."""
+
+    path: Path = Field(description="Path to the file")
+    name: str = Field(description="Name of the file")
+    size: int | None = Field(default=None, description="Size of the file in bytes")
+    description: str | None = Field(
+        default=None, description="Description from manifest or generated"
+    )
+    is_directory: bool = Field(default=False, description="Whether this is a directory")
+    parent_id: UUID | None = Field(
+        default=None, description="Parent directory ID in the storage system"
+    )
+    dataset_id: UUID | None = Field(
+        default=None, description="Dataset ID this file belongs to"
+    )
+
+    @classmethod
+    def from_path(
+        cls,
+        path: Path,
+        description: str | None = None,
+        parent_id: UUID | None = None,
+        dataset_id: UUID | None = None,
+    ) -> "FileMetadata":
+        """Create FileMetadata from a Path object."""
+        size = None
+        is_directory = path.is_dir()
+
+        if not is_directory:
+            with contextlib.suppress(OSError):
+                size = path.stat().st_size
+
+        return cls(
+            path=path,
+            name=path.name,
+            size=size,
+            description=description,
+            is_directory=is_directory,
+            parent_id=parent_id,
+            dataset_id=dataset_id,
+        )
+
+
+class UploadProgress(BaseModel):
+    """Model for tracking upload progress."""
+
+    total_files: int = Field(description="Total number of files to upload")
+    uploaded_files: int = Field(default=0, description="Number of files uploaded")
+    total_bytes: int | None = Field(default=None, description="Total bytes to upload")
+    uploaded_bytes: int = Field(default=0, description="Number of bytes uploaded")
+    current_file: str | None = Field(
+        default=None, description="Currently uploading file"
+    )
+    errors: list[str] = Field(
+        default_factory=list, description="List of error messages"
+    )
+
+    @property
+    def progress_percentage(self) -> float:
+        """Calculate progress percentage based on files."""
+        if self.total_files == 0:
+            return 0.0
+        return (self.uploaded_files / self.total_files) * 100.0
+
+    @property
+    def bytes_percentage(self) -> float | None:
+        """Calculate progress percentage based on bytes."""
+        if not self.total_bytes or self.total_bytes == 0:
+            return None
+        return (self.uploaded_bytes / self.total_bytes) * 100.0
+
+    def add_error(self, error: str) -> None:
+        """Add an error message."""
+        self.errors.append(error)
+
+    def increment_files(self, bytes_uploaded: int = 0) -> None:
+        """Increment the uploaded files counter."""
+        self.uploaded_files += 1
+        self.uploaded_bytes += bytes_uploaded
+
+
+class DirectoryUploadConfig(BaseModel):
+    """Configuration for directory uploads."""
+
+    name: str = Field(description="Name for the directory upload")
+    description: str | None = Field(
+        default=None, description="Description for the directory"
+    )
+    as_collection: bool = Field(
+        default=False, description="Upload as single collection or hierarchically"
+    )
+    manifest_filename: str | None = Field(
+        default=None, description="Name of manifest file to use"
+    )
+    ignore_patterns: list[str] = Field(
+        default_factory=list, description="Patterns to ignore"
+    )
+    ignore_filename: str = Field(
+        default=".gitignore", description="Name of ignore file to read"
+    )
+    base_path: str | None = Field(default=None, description="Base path for storage")
+    parent_id: UUID | None = Field(default=None, description="Parent directory ID")
+    dataset_id: UUID | None = Field(default=None, description="Dataset ID to use")
+
+    def with_parent(
+        self, parent_id: UUID, dataset_id: UUID | None = None
+    ) -> "DirectoryUploadConfig":
+        """Create a new config with parent and dataset IDs set."""
+        return self.model_copy(
+            update={"parent_id": parent_id, "dataset_id": dataset_id or self.dataset_id}
+        )
+
+
+# Forward reference resolution for DirectoryManifest
+DirectoryManifest.model_rebuild()

futurehouse_client/utils/general.py

@@ -2,11 +2,45 @@ import asyncio
 from collections.abc import Awaitable, Iterable
 from typing import TypeVar
 
+from httpx import (
+    CloseError,
+    ConnectError,
+    ConnectTimeout,
+    NetworkError,
+    ReadError,
+    ReadTimeout,
+    RemoteProtocolError,
+)
+from requests.exceptions import RequestException, Timeout
+from tenacity import retry_if_exception_type
 from tqdm.asyncio import tqdm
 
 T = TypeVar("T")
 
 
+_BASE_CONNECTION_ERRORS = (
+    # From requests
+    Timeout,
+    ConnectionError,
+    RequestException,
+    # From httpx
+    ConnectError,
+    ConnectTimeout,
+    ReadTimeout,
+    ReadError,
+    NetworkError,
+    RemoteProtocolError,
+    CloseError,
+)
+
+retry_if_connection_error = retry_if_exception_type(_BASE_CONNECTION_ERRORS)
+
+
+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)
+
+
 async def gather_with_concurrency(
     n: int | asyncio.Semaphore, coros: Iterable[Awaitable[T]], progress: bool = False
 ) -> list[T]:

futurehouse_client/utils/world_model_tools.py

@@ -0,0 +1,69 @@
+import os
+from uuid import UUID
+
+from aviary.core import Tool
+
+from futurehouse_client.clients.rest_client import RestClient
+from futurehouse_client.models.app import Stage
+from futurehouse_client.models.rest import WorldModel
+
+
+class WorldModelTools:
+    _client: RestClient | None = None
+
+    @classmethod
+    def _get_client(cls) -> RestClient:
+        """Lazy initialization of the RestClient to avoid validation errors during import."""
+        if cls._client is None:
+            api_key = os.getenv("FH_PLATFORM_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "FH_PLATFORM_API_KEY environment variable is required for WorldModelTools"
+                )
+            cls._client = RestClient(
+                stage=Stage.from_string(os.getenv("CROW_ENV", "dev")),
+                api_key=api_key,
+            )
+        return cls._client
+
+    @staticmethod
+    def create_world_model(name: str, description: str, content: str) -> UUID:
+        """Create a new world model.
+
+        Args:
+            name: The name of the world model.
+            description: A description of the world model.
+            content: The content/data of the world model.
+
+        Returns:
+            UUID: The ID of the newly created world model.
+        """
+        world_model = WorldModel(
+            name=name,
+            description=description,
+            content=content,
+        )
+        return WorldModelTools._get_client().create_world_model(world_model)
+
+    @staticmethod
+    def search_world_models(query: str) -> list[str]:
+        """Search for world models using a text query.
+
+        Args:
+            query: The search query string to match against world model content.
+
+        Returns:
+            list[str]: A list of world model IDs that match the search query.
+        """
+        return WorldModelTools._get_client().search_world_models(query, size=1)
+
+
+create_world_model_tool = Tool.from_function(WorldModelTools.create_world_model)
+search_world_model_tool = Tool.from_function(WorldModelTools.search_world_models)
+
+
+def make_world_model_tools() -> list[Tool]:
+    return [
+        search_world_model_tool,
+        create_world_model_tool,
+    ]