futurehouse-client 0.4.5.dev10__py3-none-any.whl → 0.4.5.dev119__py3-none-any.whl

futurehouse_client/clients/data_storage_methods.py

@@ -35,10 +35,12 @@ from futurehouse_client.models.data_storage_methods import (
     DataStorageResponse,
     DataStorageType,
     DirectoryManifest,
+    GetDatasetAndEntriesResponse,
     ManifestEntry,
 )
 from futurehouse_client.models.rest import (
     DataStorageSearchPayload,
+    FilterLogic,
     SearchCriterion,
 )
 from futurehouse_client.utils.general import retry_if_connection_error
@@ -779,6 +781,7 @@ class DataStorageMethods:
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
         project_id: UUID | None = None,
+        tags: list[str] | None = None,
     ) -> DataStorageResponse:
         """Upload a directory as a single zip file collection.
 
@@ -790,6 +793,7 @@ class DataStorageMethods:
             ignore_patterns: List of patterns to ignore when zipping
             ignore_filename: Name of ignore file to read from directory
             project_id: ID of the project this data storage entry belongs to
+            tags: List of tags to associate with the data storage entry
 
         Returns:
             DataStorageResponse for the uploaded zip file
@@ -810,6 +814,7 @@ class DataStorageMethods:
                 path=zip_gcs_path,
                 is_collection=True,
                 project_id=project_id,
+                tags=tags,
             )
 
             logger.debug(
@@ -859,6 +864,7 @@ class DataStorageMethods:
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
         project_id: UUID | None = None,
+        tags: list[str] | None = None,
     ) -> DataStorageResponse:
         """Asynchronously upload a directory as a single zip file.
 
@@ -870,6 +876,7 @@ class DataStorageMethods:
             ignore_patterns: List of patterns to ignore when zipping
             ignore_filename: Name of ignore file to read from directory
             project_id: ID of the project this data storage entry belongs to
+            tags: List of tags to associate with the data storage entry
 
         Returns:
             DataStorageResponse for the uploaded zip file
@@ -890,6 +897,7 @@ class DataStorageMethods:
                 path=zip_gcs_path,
                 is_collection=True,
                 project_id=project_id,
+                tags=tags,
             )
 
             data_storage_response = await self._acreate_data_storage_entry(payload)
@@ -1524,7 +1532,33 @@ class DataStorageMethods:
             project_id: ID of the project this data storage entry belongs to
 
         Returns:
-            DataStorageResponse containing the created data storage entry and storage locations
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
 
         Raises:
             DataStorageCreationError: If there's an error creating the data storage entry
@@ -1571,7 +1605,33 @@ class DataStorageMethods:
             project_id: ID of the project this data storage entry belongs to
 
         Returns:
-            DataStorageResponse containing the created data storage entry and storage locations
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
 
         Raises:
             DataStorageCreationError: If there's an error creating the data storage entry
@@ -1734,6 +1794,7 @@ class DataStorageMethods:
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
         project_id: UUID | None = None,
+        dataset_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Store file or directory content in the data storage system.
 
@@ -1749,13 +1810,45 @@ class DataStorageMethods:
             description: Optional description of the data storage entry
             path: Optional path for the data storage entry
             as_collection: If true, upload directories as a single zip file collection.
-            manifest_filename: Name of manifest file
+            manifest_filename: Name of manifest file (JSON or YAML) containing:
+                - entries - Map of file/directory names to their manifest entries
+                - Each ManifestEntry contains:
+                    - description - Description of the file or directory
+                    - metadata - Additional metadata for the entry
+                - Each DirectoryManifest contains nested entries following the same structure
             ignore_patterns: List of patterns to ignore when zipping directories
             ignore_filename: Name of ignore file to read from directory (default: .gitignore)
             project_id: ID of the project this data storage entry belongs to
+            dataset_id: ID of the dataset this data storage entry belongs to
 
         Returns:
-            DataStorageResponse containing the final data storage entry
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
 
         Raises:
             DataStorageCreationError: If there's an error in the process
@@ -1782,6 +1875,7 @@ class DataStorageMethods:
                     ignore_patterns=ignore_patterns,
                     ignore_filename=ignore_filename,
                     project_id=project_id,
+                    dataset_id=dataset_id,
                 )
                 if not responses:
                     raise DataStorageCreationError(
@@ -1827,15 +1921,47 @@ class DataStorageMethods:
             path: Optional GCS path for the entry.
             as_collection: If uploading a directory, `True` zips it into a single collection,
                            `False` uploads it as a hierarchical structure of individual objects.
-            manifest_filename: Optional manifest file for hierarchical uploads.
+            manifest_filename: Optional manifest file (JSON or YAML) for hierarchical uploads containing:
+                - entries - Map of file/directory names to their manifest entries
+                - Each ManifestEntry contains:
+                    - description - Description of the file or directory
+                    - metadata - Additional metadata for the entry
+                - Each DirectoryManifest contains nested entries following the same structure
             ignore_patterns: List of patterns to ignore when zipping.
             ignore_filename: Name of ignore file to read (default: .gitignore).
             dataset_id: Optional dataset ID to add entry to, or None to create new dataset.
             project_id: ID of the project this data storage entry belongs to
 
         Returns:
-            The `DataStorageResponse` for the created entry. For hierarchical uploads,
-            this is the response for the root directory entry.
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
+
+            For hierarchical uploads, this is the response for the root directory entry.
         """
         file_path = self._validate_file_path(file_path)
 
@@ -1896,7 +2022,12 @@ class DataStorageMethods:
 
         Args:
             name: Name of the data storage entry
-            existing_location: Describes the existing data source location to register
+            existing_location: a pydantic model describing the existing data source location to register, containing:
+                - storage_type - Type of storage (BIGQUERY, GCS, PG_TABLE, RAW_CONTENT, ELASTIC_SEARCH)
+                - content_type - Type of content (BQ_DATASET, BQ_TABLE, TEXT, TEXT_W_EMBEDDINGS, DIRECTORY, FILE, INDEX, INDEX_W_EMBEDDINGS)
+                - content_schema - Content schema for the data
+                - metadata - Additional metadata for the location
+                - location - Location path or identifier
             description: Optional description of the data storage entry
             as_collection: If uploading a directory, `True` creates a single storage entry for
                 the whole directory and multiple storage locations for each file, `False` assumes
@@ -1905,7 +2036,33 @@ class DataStorageMethods:
             project_id: ID of the project this data storage entry belongs to
 
         Returns:
-            DataStorageResponse containing the created data storage entry and storage locations
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
 
         Raises:
             DataStorageCreationError: If there's an error creating the data storage entry
@@ -1951,7 +2108,12 @@ class DataStorageMethods:
 
         Args:
             name: Name of the data storage entry
-            existing_location: Describes the existing data source location to register
+            existing_location: a pydantic model describing the existing data source location to register, containing:
+                - storage_type - Type of storage (BIGQUERY, GCS, PG_TABLE, RAW_CONTENT, ELASTIC_SEARCH)
+                - content_type - Type of content (BQ_DATASET, BQ_TABLE, TEXT, TEXT_W_EMBEDDINGS, DIRECTORY, FILE, INDEX, INDEX_W_EMBEDDINGS)
+                - content_schema - Content schema for the data
+                - metadata - Additional metadata for the location
+                - location - Location path or identifier
             description: Optional description of the data storage entry
             as_collection: If uploading a directory, `True` creates a single storage entry for
                 the whole directory and multiple storage locations for each file, `False` assumes
@@ -1960,7 +2122,33 @@ class DataStorageMethods:
             project_id: ID of the project this data storage entry belongs to
 
         Returns:
-            DataStorageResponse containing the created data storage entry and storage locations
+            DataStorageResponse: A Pydantic model containing:
+                - data_storage: DataStorageEntry with fields:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+                - storage_locations with each location containing:
+                    - id - Unique identifier for the storage location
+                    - data_storage_id - ID of the associated data storage entry
+                    - storage_config pydantic model with fields:
+                        - storage_type - Type of storage (e.g., 'gcs', 'pg_table')
+                        - content_type - Type of content stored
+                        - content_schema - Content schema
+                        - metadata - Location metadata
+                        - location - Location path or identifier
+                        - signed_url - Signed URL for uploading/downloading
 
         Raises:
             DataStorageCreationError: If there's an error creating the data storage entry
@@ -1997,12 +2185,17 @@ class DataStorageMethods:
         self,
         criteria: list[SearchCriterion] | None = None,
         size: int = 10,
+        filter_logic: FilterLogic = FilterLogic.OR,
     ) -> list[dict]:
         """Search data storage objects using structured criteria.
 
         Args:
-            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            criteria: List of SearchCriterion pydantic models with fields:
+                - field - Field name to search on
+                - operator - Search operator (EQUALS, CONTAINS, STARTS_WITH, ENDS_WITH, GREATER_THAN, LESS_THAN, BETWEEN, IN)
+                - value - Value to search for
             size: Number of results to return (1-100)
+            filter_logic: Either "AND" (all criteria must match) or "OR" (at least one must match)
 
         Returns:
             List of search results with scores and data storage information
@@ -2023,6 +2216,7 @@ class DataStorageMethods:
             payload = DataStorageSearchPayload(
                 criteria=criteria or [],
                 size=max(1, min(100, size)),  # Clamp between 1-100
+                filter_logic=filter_logic,
             )
 
             response = self.client.post(
@@ -2053,12 +2247,17 @@ class DataStorageMethods:
         self,
         criteria: list[SearchCriterion] | None = None,
         size: int = 10,
+        filter_logic: FilterLogic = FilterLogic.OR,
     ) -> list[dict]:
         """Asynchronously search data storage objects using structured criteria.
 
         Args:
-            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            criteria: List of SearchCriterion pydantic models with fields:
+                - field - Field name to search on
+                - operator - Search operator (EQUALS, CONTAINS, STARTS_WITH, ENDS_WITH, GREATER_THAN, LESS_THAN, BETWEEN, IN)
+                - value - Value to search for
             size: Number of results to return (1-100)
+            filter_logic: Either "AND" (all criteria must match) or "OR" (at least one must match)
 
         Returns:
             List of search results with scores and data storage information
@@ -2079,6 +2278,7 @@ class DataStorageMethods:
             payload = DataStorageSearchPayload(
                 criteria=criteria or [],
                 size=max(1, min(100, size)),  # Clamp between 1-100
+                filter_logic=filter_logic,
             )
 
             response = await self.async_client.post(
@@ -2118,11 +2318,11 @@ class DataStorageMethods:
         """Search data storage objects using vector similarity.
 
         Args:
-            embedding: Embedding vector for similarity search
+            embedding: List of float values representing the embedding vector for similarity search
             size: Number of results to return (1-100)
             min_score: Minimum similarity score (0.0-1.0)
             dataset_id: Optional dataset ID filter
-            tags: Optional list of tags to filter by
+            tags: Optional list of string tags to filter by
             user_id: Optional user ID filter (admin only)
             project_id: Optional project ID filter
 
@@ -2196,11 +2396,11 @@ class DataStorageMethods:
         """Asynchronously search data storage objects using vector similarity.
 
         Args:
-            embedding: Embedding vector for similarity search
+            embedding: List of float values representing the embedding vector for similarity search
             size: Number of results to return (1-100)
             min_score: Minimum similarity score (0.0-1.0)
             dataset_id: Optional dataset ID filter
-            tags: Optional list of tags to filter by
+            tags: Optional list of string tags to filter by
             user_id: Optional user ID filter (admin only)
             project_id: Optional project ID filter
 
@@ -2268,12 +2468,12 @@ class DataStorageMethods:
         """Fetch data from the storage system (sync version).
 
         Args:
-            data_storage_id: ID of the data storage entry to fetch
+            data_storage_id: UUID of the data storage entry to fetch
 
         Returns:
             For PG_TABLE storage: string content
             For GCS storage: Path to downloaded file (may be unzipped if it was a zip)
-            For multi-location entries: dict of location IDs to dicts with signed URL and file name
+            For multi-location entries: list of downloaded files
             None if not found or error occurred
         """
         if not data_storage_id:
@@ -2342,12 +2542,12 @@ class DataStorageMethods:
         """Fetch data from the storage system.
 
         Args:
-            data_storage_id: ID of the data storage entry to fetch
+            data_storage_id: UUID of the data storage entry to fetch
 
         Returns:
             For PG_TABLE storage: string content
             For GCS storage: Path to downloaded file (may be unzipped if it was a zip)
-            For multi-location entries: dict of location IDs to dicts with signed URL and file name
+            For multi-location entries: list of downloaded files
             None if not found or error occurred
         """
         if not data_storage_id:
@@ -2417,7 +2617,23 @@ class DataStorageMethods:
         name: str,
         description: str | None = None,
         dataset_id: UUID | None = None,
-    ):
+    ) -> CreateDatasetPayload:
+        """Asynchronously create a new dataset.
+
+        Args:
+            name: Name of the dataset to create
+            description: Optional description of the dataset
+            dataset_id: Optional UUID to assign to the dataset, or None to auto-generate
+
+        Returns:
+            CreateDatasetPayload: A Pydantic model containing:
+                - id - ID of the created dataset (None if auto-generated)
+                - name - Name of the dataset
+                - description - Description of the dataset
+
+        Raises:
+            DataStorageCreationError: If there's an error creating the dataset
+        """
         try:
             payload = CreateDatasetPayload(
                 name=name,
@@ -2448,7 +2664,23 @@ class DataStorageMethods:
         name: str,
         description: str | None = None,
         dataset_id: UUID | None = None,
-    ):
+    ) -> CreateDatasetPayload:
+        """Create a new dataset.
+
+        Args:
+            name: Name of the dataset to create
+            description: Optional description of the dataset
+            dataset_id: Optional UUID to assign to the dataset, or None to auto-generate
+
+        Returns:
+            CreateDatasetPayload: A Pydantic model containing:
+                - id - ID of the created dataset (None if auto-generated)
+                - name - Name of the dataset
+                - description - Description of the dataset
+
+        Raises:
+            DataStorageCreationError: If there's an error creating the dataset
+        """
         try:
             payload = CreateDatasetPayload(
                 name=name,
@@ -2522,14 +2754,48 @@ class DataStorageMethods:
         retry=retry_if_connection_error,
         before_sleep=before_sleep_log(logger, logging.WARNING),
     )
-    async def aget_dataset(self, dataset_id: UUID):
+    async def aget_dataset(self, dataset_id: UUID) -> GetDatasetAndEntriesResponse:
+        """Asynchronously retrieve a dataset by ID.
+
+        Args:
+            dataset_id: UUID of the dataset to retrieve
+
+        Returns:
+            GetDatasetAndEntriesResponse: A dict containing:
+                - dataset: DatasetStorage with fields:
+                    - id - Unique identifier for the dataset
+                    - name - Name of the dataset
+                    - user_id - ID of the user who created the dataset
+                    - description - Description of the dataset
+                    - created_at - Timestamp when the dataset was created
+                    - modified_at - Timestamp when the dataset was last modified
+                - data_storage_entries - List of data storage entries in the dataset, each containing:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+
+        Raises:
+            DataStorageError: If there's an error retrieving the dataset
+        """
         try:
             response = await self.async_client.get(
                 f"/v0.1/data-storage/datasets/{dataset_id}"
             )
             response.raise_for_status()
 
-            return response.json()
+            return GetDatasetAndEntriesResponse.model_validate(response.json())
         except HTTPStatusError as e:
             self._handle_http_errors(e, "retrieving")
         except Exception as e:
@@ -2541,12 +2807,46 @@ class DataStorageMethods:
         retry=retry_if_connection_error,
         before_sleep=before_sleep_log(logger, logging.WARNING),
     )
-    def get_dataset(self, dataset_id: UUID):
+    def get_dataset(self, dataset_id: UUID) -> GetDatasetAndEntriesResponse:
+        """Retrieve a dataset by ID.
+
+        Args:
+            dataset_id: UUID of the dataset to retrieve
+
+        Returns:
+            GetDatasetAndEntriesResponse: A dict containing:
+                - dataset: DatasetStorage with fields:
+                    - id - Unique identifier for the dataset
+                    - name - Name of the dataset
+                    - user_id - ID of the user who created the dataset
+                    - description - Description of the dataset
+                    - created_at - Timestamp when the dataset was created
+                    - modified_at - Timestamp when the dataset was last modified
+                - data_storage_entries - List of data storage entries in the dataset, each containing:
+                    - id - Unique identifier for the data storage entry
+                    - name - Name of the data storage entry
+                    - description - Description of the data storage entry
+                    - content - Content of the data storage entry
+                    - embedding - Embedding vector for the content
+                    - is_collection - Whether this entry is a collection
+                    - tags - List of tags associated with the entry
+                    - parent_id - ID of the parent entry for hierarchical storage
+                    - project_id - ID of the project this entry belongs to
+                    - dataset_id - ID of the dataset this entry belongs to
+                    - path - Path in the storage system where this entry is located
+                    - bigquery_schema - Target BigQuery schema for the entry
+                    - user_id - ID of the user who created this entry
+                    - created_at - Timestamp when the entry was created
+                    - modified_at - Timestamp when the entry was last updated
+
+        Raises:
+            DataStorageError: If there's an error retrieving the dataset
+        """
         try:
             response = self.client.get(f"/v0.1/data-storage/datasets/{dataset_id}")
             response.raise_for_status()
 
-            return response.json()
+            return GetDatasetAndEntriesResponse.model_validate(response.json())
         except HTTPStatusError as e:
             self._handle_http_errors(e, "retrieving")
         except Exception as e:
@@ -2622,7 +2922,15 @@ class DataStorageMethods:
         retry=retry_if_connection_error,
         before_sleep=before_sleep_log(logger, logging.WARNING),
     )
-    async def adelete_data_storage_entry(self, data_storage_entry_id: UUID):
+    async def adelete_data_storage_entry(self, data_storage_entry_id: UUID) -> None:
+        """Asynchronously delete a data storage entry.
+
+        Args:
+            data_storage_entry_id: UUID of the data storage entry to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the data storage entry
+        """
         try:
             await self.async_client.delete(
                 f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"
@@ -2638,7 +2946,15 @@ class DataStorageMethods:
         retry=retry_if_connection_error,
         before_sleep=before_sleep_log(logger, logging.WARNING),
     )
-    def delete_data_storage_entry(self, data_storage_entry_id: UUID):
+    def delete_data_storage_entry(self, data_storage_entry_id: UUID) -> None:
+        """Delete a data storage entry.
+
+        Args:
+            data_storage_entry_id: UUID of the data storage entry to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the data storage entry
+        """
         try:
             self.client.delete(
                 f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"

futurehouse_client/models/client.py

@@ -2,6 +2,7 @@ from typing import Any, Generic, TypeAlias, TypeVar
 
 from aviary.message import Message
 from aviary.tools.base import Tool
+from ldp.agent import Agent
 from ldp.data_structures import Transition
 from ldp.graph.ops import OpResult
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
@@ -34,6 +35,7 @@ class ASVState(BaseState, Generic[T]):
     def serialize_action(self, action: OpResult[T]) -> dict:
         return action.to_dict()
 
+
 class EnvResetState(BaseState):
     observations: list[Message] = Field()
     tools: list[Tool] = Field()
@@ -57,6 +59,64 @@ class TransitionState(BaseState):
         }
 
 
+class GlobalState(BaseState):
+    agent: Agent | None = None
+    env: Any | None = None
+    agent_state: Any | None = None
+    next_agent_state: Any | None = None
+    observations: list = []
+    action: Any | None = None
+    value: float = 0.0
+    last_step_state: Transition | None = None
+
+    def update_observations(self, obs: list[Message]) -> list[Message]:
+        previous_observations = self.observations or []
+        self.observations = obs
+        return previous_observations
+
+    def store_step_state(self, step_state: Transition) -> None:
+        self.last_step_state = step_state
+
+    def update_trajectory_data(self, **kwargs) -> None:
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def _get_safe_previous_observations(
+        self, current_obs: list[Message] | None = None
+    ) -> list[Message]:
+        if self.last_step_state:
+            last_step_state = self.last_step_state
+            if last_step_state.next_observation:
+                return last_step_state.next_observation
+        if self.observations:
+            return self.observations
+        return current_obs or []
+
+    def create_step_state(self, callback_type: str, **kwargs) -> Transition:
+        defaults = {
+            "timestep": getattr(self.agent, "_timestep", 0) if self.agent else 0,
+            "agent_state": self.agent_state,
+            "next_agent_state": self.next_agent_state or self.agent_state,
+            "observation": self._get_safe_previous_observations(),
+            "next_observation": self.observations or [],
+            "action": self.action,
+            "reward": 0.0,
+            "truncated": False,
+            "done": False,
+            "value": self.value or 0.0,
+            "metadata": {"callback_type": callback_type},
+        }
+
+        for key, value in kwargs.items():
+            if key == "metadata" and isinstance(value, dict):
+                if isinstance(defaults["metadata"], dict):
+                    defaults["metadata"].update(value)
+            else:
+                defaults[key] = value
+
+        return Transition(**defaults)
+
+
 StateType: TypeAlias = (
     BeforeTransitionState
     | InitialState
@@ -64,4 +124,5 @@ StateType: TypeAlias = (
     | EnvResetState
     | EnvStepState
     | TransitionState
+    | GlobalState
 )

futurehouse_client/models/data_storage_methods.py

@@ -3,10 +3,32 @@ from datetime import datetime
 from enum import StrEnum, auto
 from os import PathLike
 from pathlib import Path
-from typing import Any
+from typing import Annotated, Any
 from uuid import UUID
 
-from pydantic import BaseModel, Field, JsonValue
+from pydantic import (
+    BaseModel,
+    Field,
+    JsonValue,
+    PlainSerializer,
+    PlainValidator,
+    WithJsonSchema,
+)
+from sqlalchemy_utils import Ltree
+
+LtreeField = Annotated[
+    Ltree,
+    PlainValidator(Ltree),
+    PlainSerializer(lambda v: v.path),
+    WithJsonSchema({"type": "string", "examples": ["some.path"]}),
+]
+
+
+class DataStorageEntryStatus(StrEnum):
+    PENDING = auto()
+    ACTIVE = auto()
+    FAILED = auto()
+    DISABLED = auto()
 
 
 class DataStorageEntry(BaseModel):
@@ -20,6 +42,9 @@ class DataStorageEntry(BaseModel):
     content: str | None = Field(
         default=None, description="Content of the data storage entry"
     )
+    status: DataStorageEntryStatus = Field(
+        description="Status of the data storage entry"
+    )
     embedding: list[float] | None = Field(
         default=None, description="Embedding vector for the content"
     )
@@ -151,6 +176,26 @@ class DataStorageRequestPayload(BaseModel):
     existing_location: DataStorageLocationPayload | None = Field(
         default=None, description="Target storage metadata"
     )
+    tags: list[str] | None = Field(
+        default=None,
+        description="List of tags associated with the data storage entry",
+    )
+
+
+class DatasetStorage(BaseModel):
+    """Pydantic model representing a DatasetStorage record."""
+
+    id: UUID
+    name: str
+    user_id: str
+    description: str | None = None
+    created_at: datetime
+    modified_at: datetime
+
+
+class GetDatasetAndEntriesResponse(BaseModel):
+    dataset: DatasetStorage
+    data_storage_entries: list[DataStorageEntry]
 
 
 class CreateDatasetPayload(BaseModel):

futurehouse_client/models/rest.py

@@ -67,7 +67,8 @@ class SearchOperator(StrEnum):
     """Operators for structured search criteria."""
 
     EQUALS = "equals"
-    CONTAINS = "contains"
+    CONTAINS = "contains"  # Exact phrase/substring matching
+    FULLTEXT = "fulltext"  # Tokenized full-text search (match query)
     STARTS_WITH = "starts_with"
     ENDS_WITH = "ends_with"
     GREATER_THAN = "greater_than"
@@ -84,6 +85,11 @@ class SearchCriterion(BaseModel):
     value: str | list[str] | bool
 
 
+class FilterLogic(StrEnum):
+    AND = "AND"
+    OR = "OR"
+
+
 class WorldModelSearchPayload(BaseModel):
     """Payload for structured world model search."""
 
@@ -173,3 +179,4 @@ class DataStorageSearchPayload(BaseModel):
 
     criteria: list[SearchCriterion]
     size: int = 10
+    filter_logic: FilterLogic = FilterLogic.OR

futurehouse_client/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.5.dev10'
-__version_tuple__ = version_tuple = (0, 4, 5, 'dev10')
+__version__ = version = '0.4.5.dev119'
+__version_tuple__ = version_tuple = (0, 4, 5, 'dev119')
 
 __commit_id__ = commit_id = None

{futurehouse_client-0.4.5.dev10.dist-info → futurehouse_client-0.4.5.dev119.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.4.5.dev10
+Version: 0.4.5.dev119
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 License:                                  Apache License
@@ -224,6 +224,7 @@ Requires-Dist: openai<1.100.0,>=1
 Requires-Dist: pydantic
 Requires-Dist: python-dotenv
 Requires-Dist: requests
+Requires-Dist: sqlalchemy-utils>=0.41.2
 Requires-Dist: tenacity
 Requires-Dist: tqdm>=4.62
 Provides-Extra: dev

{futurehouse_client-0.4.5.dev10.dist-info → futurehouse_client-0.4.5.dev119.dist-info}/RECORD

@@ -1,23 +1,23 @@
 futurehouse_client/__init__.py,sha256=PvFTkocA-hobsWoDEBEdrUgLIbuVbDs_0nvMdImJmHk,707
 futurehouse_client/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-futurehouse_client/version.py,sha256=0gN9IqRid_hywJfwdbe8e1dFFTTj5VGTzTuG5nFBWxw,719
+futurehouse_client/version.py,sha256=ygTkTx_4WMa3DIXWf_ZxjBUe9cI-wOluDVgxQX8thfA,721
 futurehouse_client/clients/__init__.py,sha256=-HXNj-XJ3LRO5XM6MZ709iPs29YpApss0Q2YYg1qMZw,280
-futurehouse_client/clients/data_storage_methods.py,sha256=VESdX0J_frITd0QAjQ5UMPiqYvpToo0ooDCrS-U4uH8,99535
+futurehouse_client/clients/data_storage_methods.py,sha256=f8ZsVicEtO50pRXoPzEB2GpiyqosNofyoW8vJeYvFnM,119266
 futurehouse_client/clients/job_client.py,sha256=b5gpzulZpxpv9R337r3UKItnMdtd6CGlI1sV3_VQJso,13985
 futurehouse_client/clients/rest_client.py,sha256=RdyFEipvADDCHyY5XFy565IoL9-N1myJjF0G8x2wlK8,103183
 futurehouse_client/models/__init__.py,sha256=0YlzKGymbY1g4cXxnUc0BUnthTkVBf12bCZlGUcMQqk,701
 futurehouse_client/models/app.py,sha256=UUg17I3zk6cH_7mrdojHGYvQfm_SeDkuUxsPlRyIYz0,31895
-futurehouse_client/models/client.py,sha256=Dg7bYgbdmP5GkDWp5e71vxs5YdkzyqqlwBN-A84Jt8w,1709
-futurehouse_client/models/data_storage_methods.py,sha256=GS1FbuMsUJSh7Evjt86vOri-95hfiLyASBS1xG7erNk,12793
-futurehouse_client/models/rest.py,sha256=Fqw0_ypULzd7IV93PKooSG9W5_g7fGFsdW9jNVVImHA,4514
+futurehouse_client/models/client.py,sha256=3WLS0xdB7CYqHShi_gqyRa6PGj-QvP--0HzD1R93yvY,3868
+futurehouse_client/models/data_storage_methods.py,sha256=cpF2g4y_REECaz--WhaJeLqXA_3m3keRP5XOXiL8GOI,13811
+futurehouse_client/models/rest.py,sha256=SbeXZSPUCM0lQ_gVUPa64vKzMxuUVgqmJ5YThfDWs8g,4726
 futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 futurehouse_client/utils/auth.py,sha256=tgWELjKfg8eWme_qdcRmc8TjQN9DVZuHHaVXZNHLchk,2960
 futurehouse_client/utils/general.py,sha256=PIkGLCSA3kUvc6mwR-prEB7YnMdKILOIm6cPowSZzzs,2532
 futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
 futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
 futurehouse_client/utils/world_model_tools.py,sha256=v2krZGrco0ur2a_pcRMtnQL05SxlIoBXuJ5R1JkQNws,2921
-futurehouse_client-0.4.5.dev10.dist-info/licenses/LICENSE,sha256=oQ9ZHjUi-_6GfP3gs14FlPb0OlGwE1QCCKFGnJ4LD2I,11341
-futurehouse_client-0.4.5.dev10.dist-info/METADATA,sha256=LDAeM4-wOHzn9c7svqmT7rIvXGfq0Gi9lrc5Nk_cUvc,27060
-futurehouse_client-0.4.5.dev10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-futurehouse_client-0.4.5.dev10.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
-futurehouse_client-0.4.5.dev10.dist-info/RECORD,,
+futurehouse_client-0.4.5.dev119.dist-info/licenses/LICENSE,sha256=oQ9ZHjUi-_6GfP3gs14FlPb0OlGwE1QCCKFGnJ4LD2I,11341
+futurehouse_client-0.4.5.dev119.dist-info/METADATA,sha256=_GrNdEBxKiRCI4lXWI8WcIBjVRTEMF_FqvvrQcT3l_E,27101
+futurehouse_client-0.4.5.dev119.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+futurehouse_client-0.4.5.dev119.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
+futurehouse_client-0.4.5.dev119.dist-info/RECORD,,