futurehouse-client 0.4.5.dev10__tar.gz → 0.4.5.dev119__tar.gz

{futurehouse_client-0.4.5.dev10/src/futurehouse_client.egg-info → futurehouse_client-0.4.5.dev119}/PKG-INFO

@@ -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 → futurehouse_client-0.4.5.dev119}/pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "pydantic",
     "python-dotenv",
     "requests",
+    "sqlalchemy-utils>=0.41.2",
     "tenacity",
     "tqdm>=4.62",  # For tqdm.asyncio stdlib consistency
 ]

{futurehouse_client-0.4.5.dev10 → futurehouse_client-0.4.5.dev119}/src/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}"