futurehouse-client 0.4.2.dev11__py3-none-any.whl → 0.4.2.dev274__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

@@ -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,
@@ -160,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
@@ -239,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.
 
@@ -251,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)
@@ -282,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,
                 )
             )

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

@@ -0,0 +1,341 @@
+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",
+    )
+    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",
+    )
+    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"
+    )
+    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",
+    )
+    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/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"

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

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.2.dev11'
-__version_tuple__ = version_tuple = (0, 4, 2, 'dev11')
+__version__ = version = '0.4.2.dev274'
+__version_tuple__ = version_tuple = (0, 4, 2, 'dev274')
 
 __commit_id__ = commit_id = None

{futurehouse_client-0.4.2.dev11.dist-info → futurehouse_client-0.4.2.dev274.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.4.2.dev11
+Version: 0.4.2.dev274
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 License:                                  Apache License
@@ -213,18 +213,22 @@ Classifier: Programming Language :: Python
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: aiofiles
 Requires-Dist: cloudpickle
 Requires-Dist: fhaviary
+Requires-Dist: google-resumable-media[aiohttp]
 Requires-Dist: httpx
 Requires-Dist: ldp>=0.22.0
 Requires-Dist: litellm
 Requires-Dist: pydantic
 Requires-Dist: python-dotenv
+Requires-Dist: requests
 Requires-Dist: tenacity
 Requires-Dist: tqdm>=4.62
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: futurehouse-client[monitoring,typing]; extra == "dev"
+Requires-Dist: ipykernel; extra == "dev"
 Requires-Dist: jupyter; extra == "dev"
 Requires-Dist: jupyterlab; extra == "dev"
 Requires-Dist: mypy; extra == "dev"
@@ -244,6 +248,7 @@ Requires-Dist: setuptools_scm; extra == "dev"
 Provides-Extra: monitoring
 Requires-Dist: newrelic>=8.8.0; extra == "monitoring"
 Provides-Extra: typing
+Requires-Dist: types-PyYAML; extra == "typing"
 Requires-Dist: types-requests; extra == "typing"
 Requires-Dist: types-tqdm; extra == "typing"
 Dynamic: license-file