futurehouse-client 0.4.2.dev274__py3-none-any.whl → 0.4.3.dev3__py3-none-any.whl

futurehouse_client/clients/data_storage_methods.py

@@ -7,7 +7,7 @@ import tempfile
 import zipfile
 from os import PathLike
 from pathlib import Path
-from typing import NoReturn
+from typing import Any, NoReturn
 from uuid import UUID
 
 import aiofiles
@@ -15,6 +15,8 @@ import aiohttp
 import requests as requests_lib
 from google.resumable_media import requests as resumable_requests
 from httpx import AsyncClient, Client, HTTPStatusError, codes
+from lmi.utils import gather_with_concurrency
+from pydantic import HttpUrl
 from requests.adapters import HTTPAdapter
 from tenacity import (
     before_sleep_log,
@@ -26,12 +28,19 @@ from tqdm import tqdm
 from urllib3.util.retry import Retry
 
 from futurehouse_client.models.data_storage_methods import (
+    CreateDatasetPayload,
+    DataContentType,
     DataStorageLocationPayload,
     DataStorageRequestPayload,
     DataStorageResponse,
+    DataStorageType,
     DirectoryManifest,
     ManifestEntry,
 )
+from futurehouse_client.models.rest import (
+    DataStorageSearchPayload,
+    SearchCriterion,
+)
 from futurehouse_client.utils.general import retry_if_connection_error
 
 # this is only required if they're using a yaml manifest
@@ -54,6 +63,7 @@ INITIATE_HEADERS = {
     "x-goog-resumable": "start",
     "Content-Length": "0",
 }
+DOWNLOAD_CONCURRENCY = 3
 
 
 def _should_ignore_file(
@@ -438,6 +448,10 @@ class DataStorageCreationError(DataStorageError):
     """Raised when there's an error creating a data storage entry."""
 
 
+class DataStorageRetrievalError(DataStorageError):
+    """Raised when there's an error retrieving a data storage entry."""
+
+
 class ProgressWrapper:
     """Common progress wrapper for file uploads."""
 
@@ -462,7 +476,7 @@ class ProgressWrapper:
         return self.file_obj.tell()
 
 
-class DataStorageMethods:
+class DataStorageMethods:  # pylint: disable=too-many-public-methods
     """Data storage methods for RestClient.
 
     This class contains methods for interacting with the data storage API endpoints.
@@ -472,14 +486,12 @@ class DataStorageMethods:
     def _handle_http_errors(self, e: HTTPStatusError, operation: str) -> NoReturn:
         """Handle common HTTP errors for data storage operations."""
         if e.response.status_code == codes.FORBIDDEN:
-            raise DataStorageCreationError(
+            raise DataStorageError(
                 f"Error {operation} data storage entry, not authorized"
             ) from e
         if e.response.status_code == codes.UNPROCESSABLE_ENTITY:
-            raise DataStorageCreationError(
-                f"Invalid request payload: {e.response.text}"
-            ) from e
-        raise DataStorageCreationError(
+            raise DataStorageError(f"Invalid request payload: {e.response.text}") from e
+        raise DataStorageError(
             f"Error {operation} data storage entry: {e.response.status_code} - {e.response.text}"
         ) from e
 
@@ -487,7 +499,7 @@ class DataStorageMethods:
         """Validate file path exists and return Path object."""
         file_path = Path(file_path)
         if not file_path.exists():
-            raise DataStorageCreationError(f"File or directory not found: {file_path}")
+            raise DataStorageError(f"File or directory not found: {file_path}")
         return file_path
 
     def _build_zip_path(self, name: str, path: str | None) -> str:
@@ -529,19 +541,24 @@ class DataStorageMethods:
             return extracted_items[0]
         return extract_dir
 
-    async def _adownload_from_gcs(self, signed_url: str) -> Path:
+    async def _adownload_from_gcs(
+        self, signed_url: str, file_name: str | None = None
+    ) -> Path:
         """Download file from GCS using signed URL and handle unzipping if needed.
 
         Args:
             signed_url: The signed URL to download from
+            file_name: The name of the file to download
 
         Returns:
             Path to the downloaded file (or unzipped directory if it was a zip)
         """
+        file_name = file_name or "downloaded_file"
+
         try:
             with tempfile.TemporaryDirectory() as temp_dir_str:
                 temp_dir = Path(temp_dir_str)
-                temp_file = temp_dir / "downloaded_file"
+                temp_file = temp_dir / file_name
 
                 async with self.async_client.stream("GET", signed_url) as response:
                     response.raise_for_status()
@@ -549,11 +566,11 @@ class DataStorageMethods:
                     content_disposition = response.headers.get(
                         "content-disposition", ""
                     )
-                    filename = "downloaded_file"
+                    filename = file_name
                     if "filename=" in content_disposition:
                         filename = content_disposition.split("filename=")[-1].strip('"')
 
-                    if filename != "downloaded_file":
+                    if filename != file_name:
                         temp_file = temp_dir / filename
 
                     async with aiofiles.open(temp_file, "wb") as f:
@@ -583,21 +600,23 @@ class DataStorageMethods:
                     return final_file
 
         except Exception as e:
-            raise DataStorageCreationError(f"Failed to download from GCS: {e}") from e
+            raise DataStorageError(f"Failed to download from GCS: {e}") from e
 
-    def _download_from_gcs(self, signed_url: str) -> Path:
+    def _download_from_gcs(self, signed_url: str, file_name: str | None = None) -> Path:
         """Download file from GCS using signed URL and handle unzipping if needed (sync version).
 
         Args:
             signed_url: The signed URL to download from
-
+            file_name: The name of the file to download
         Returns:
             Path to the downloaded file (or unzipped directory if it was a zip)
         """
+        file_name = file_name or "downloaded_file"
+
         try:
             with tempfile.TemporaryDirectory() as temp_dir_str:
                 temp_dir = Path(temp_dir_str)
-                temp_file = temp_dir / "downloaded_file"
+                temp_file = temp_dir / file_name
 
                 with requests_lib.get(signed_url, stream=True, timeout=30) as response:
                     response.raise_for_status()
@@ -605,11 +624,11 @@ class DataStorageMethods:
                     content_disposition = response.headers.get(
                         "content-disposition", ""
                     )
-                    filename = "downloaded_file"
+                    filename = file_name
                     if "filename=" in content_disposition:
                         filename = content_disposition.split("filename=")[-1].strip('"')
 
-                    if filename != "downloaded_file":
+                    if filename != file_name:
                         temp_file = temp_dir / filename
 
                     with open(temp_file, "wb") as f:
@@ -639,9 +658,7 @@ class DataStorageMethods:
                     return final_file
 
         except Exception as e:
-            raise DataStorageCreationError(f"Failed to download from GCS: {e}") from e
-
-    # =====================================
+            raise DataStorageError(f"Failed to download from GCS: {e}") from e
 
     def _prepare_single_file_upload(
         self, name: str, file_path: Path, description: str | None, path: str | None
@@ -676,7 +693,7 @@ class DataStorageMethods:
     ) -> DataStorageResponse:
         """Create data storage entry via API (sync version)."""
         response = self.client.post(
-            "/v0.1/data-storage",
+            "/v0.1/data-storage/data-entries",
             json=payload.model_dump(mode="json", exclude_none=True),
         )
         response.raise_for_status()
@@ -687,7 +704,7 @@ class DataStorageMethods:
     ) -> DataStorageResponse:
         """Create data storage entry via API (async version)."""
         response = await self.async_client.post(
-            "/v0.1/data-storage",
+            "/v0.1/data-storage/data-entries",
             json=payload.model_dump(mode="json", exclude_none=True),
         )
         response.raise_for_status()
@@ -800,24 +817,30 @@ class DataStorageMethods:
             )
             data_storage_response = self._create_data_storage_entry(payload)
 
-            if not data_storage_response.signed_url:
-                raise DataStorageCreationError("No signed URL returned for zip upload")
+            for storage_location in data_storage_response.storage_locations:
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageCreationError(
+                        "No signed URL returned for zip upload"
+                    )
 
-            with tqdm(
-                total=zip_size,
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                desc=f"Uploading {dir_path.name} (zipped)",
-                miniters=1,
-                mininterval=0.1,
-            ) as pbar:
-                _upload_file_with_progress(
-                    data_storage_response.signed_url, temp_zip_path, pbar, zip_size
-                )
+                with tqdm(
+                    total=zip_size,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    desc=f"Uploading {dir_path.name} (zipped)",
+                    miniters=1,
+                    mininterval=0.1,
+                ) as pbar:
+                    _upload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        temp_zip_path,
+                        pbar,
+                        zip_size,
+                    )
 
             status_response = self.client.patch(
-                f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+                f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
                 json={"status": "active"},
             )
             status_response.raise_for_status()
@@ -871,24 +894,30 @@ class DataStorageMethods:
 
             data_storage_response = await self._acreate_data_storage_entry(payload)
 
-            if not data_storage_response.signed_url:
-                raise DataStorageCreationError("No signed URL returned for zip upload")
+            for storage_location in data_storage_response.storage_locations:
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageCreationError(
+                        "No signed URL returned for zip upload"
+                    )
 
-            with tqdm(
-                total=zip_size,
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                desc=f"Uploading {dir_path.name} (zipped)",
-                miniters=1,
-                mininterval=0.1,
-            ) as pbar:
-                await _aupload_file_with_progress(
-                    data_storage_response.signed_url, temp_zip_path, pbar, zip_size
-                )
+                with tqdm(
+                    total=zip_size,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    desc=f"Uploading {dir_path.name} (zipped)",
+                    miniters=1,
+                    mininterval=0.1,
+                ) as pbar:
+                    await _aupload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        temp_zip_path,
+                        pbar,
+                        zip_size,
+                    )
 
             status_response = await self.async_client.patch(
-                f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+                f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
                 json={"status": "active"},
             )
             status_response.raise_for_status()
@@ -951,30 +980,34 @@ class DataStorageMethods:
 
         data_storage_response = self._create_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for storage_location in data_storage_response.storage_locations:
+            if not storage_location.storage_config.signed_url:
+                raise DataStorageCreationError("No signed URL returned from server")
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-        ) as pbar:
-            try:
-                _upload_file_with_progress(
-                    data_storage_response.signed_url, file_path, pbar, file_size
-                )
-                logger.debug("File upload to signed URL completed successfully")
-            except Exception as e:
-                logger.error(f"Failed to upload file to signed URL: {e}")
-                raise
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+            ) as pbar:
+                try:
+                    _upload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        file_path,
+                        pbar,
+                        file_size,
+                    )
+                    logger.debug("File upload to signed URL completed successfully")
+                except Exception as e:
+                    logger.error(f"Failed to upload file to signed URL: {e}")
+                    raise
 
         logger.debug("Updating data storage status to active")
         status_response = self.client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1015,24 +1048,28 @@ class DataStorageMethods:
 
         data_storage_response = await self._acreate_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for location in data_storage_response.storage_locations:
+            if not location.storage_config.signed_url:
+                raise DataStorageCreationError(
+                    f"No signed URL returned from server for location: {location.id}"
+                )
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-        ) as pbar:
-            await _aupload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
-            )
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+                leave=False,
+            ) as pbar:
+                await _aupload_file_with_progress(
+                    location.storage_config.signed_url, file_path, pbar, file_size
+                )
 
         status_response = await self.async_client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1075,25 +1112,26 @@ class DataStorageMethods:
         )
         data_storage_response = self._create_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for location in data_storage_response.storage_locations:
+            if not location.storage_config.signed_url:
+                raise DataStorageCreationError("No signed URL returned from server")
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-            leave=False,
-        ) as pbar:
-            _upload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
-            )
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+                leave=False,
+            ) as pbar:
+                _upload_file_with_progress(
+                    location.storage_config.signed_url, file_path, pbar, file_size
+                )
 
         status_response = self.client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1295,7 +1333,9 @@ class DataStorageMethods:
         )
         data_storage_response = await self._acreate_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
+        storage_location = data_storage_response.storage_locations[0]
+
+        if not storage_location.storage_config.signed_url:
             raise DataStorageCreationError("No signed URL returned from server")
 
         with tqdm(
@@ -1308,11 +1348,11 @@ class DataStorageMethods:
             mininterval=0.1,
         ) as pbar:
             await _aupload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
+                storage_location.storage_config.signed_url, file_path, pbar, file_size
             )
 
         status_response = await self.async_client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1553,6 +1593,130 @@ class DataStorageMethods:
                 f"An unexpected error occurred: {e!r}"
             ) from e
 
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def astore_link(
+        self,
+        name: str,
+        url: HttpUrl,
+        description: str,
+        instructions: str,
+        api_key: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
+    ) -> DataStorageResponse:
+        """Asynchronously store a link/URL in the data storage system.
+
+        Args:
+            name: Name of the link entry
+            url: The URL/link to store
+            description: Searchable details of the link
+            instructions: Instructions for how to consume the link or api
+            api_key: Any authentication key to access the api. If this is included, you should also include
+                details of how the key should be consumed in the instructions.
+            metadata: Any additional metadata about the link
+            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:
+            DataStorageResponse containing the created link storage entry
+
+        Raises:
+            DataStorageCreationError: If there's an error creating the link storage entry
+        """
+        try:
+            link_metadata = metadata.copy() if metadata else {}
+            link_metadata["instructions"] = instructions
+            if api_key:
+                link_metadata["api_key"] = api_key
+
+            existing_location = DataStorageLocationPayload(
+                storage_type=DataStorageType.LINK,
+                content_type=DataContentType.TEXT,
+                location=url,
+                metadata=link_metadata or None,
+            )
+
+            payload = DataStorageRequestPayload(
+                name=name,
+                content=url,
+                description=description,
+                dataset_id=dataset_id,
+                project_id=project_id,
+                existing_location=existing_location,
+            )
+            return await self._acreate_data_storage_entry(payload)
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
+    def store_link(
+        self,
+        name: str,
+        url: HttpUrl,
+        description: str,
+        instructions: str,
+        api_key: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
+    ) -> DataStorageResponse:
+        """Store a link/URL in the data storage system.
+
+        Args:
+            name: Name of the link entry
+            url: The URL/link to store
+            description: Searchable details of the link
+            instructions: Instructions for how to consume the link or api
+            api_key: Any authentication key to access the api. If this is included, you should also include
+                details of how the key should be consumed in the instructions.
+            metadata: Any additional metadata about the link
+            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:
+            DataStorageResponse containing the created link storage entry
+
+        Raises:
+            DataStorageCreationError: If there's an error creating the link storage entry
+        """
+        try:
+            link_metadata = metadata.copy() if metadata else {}
+            link_metadata["instructions"] = instructions
+            if api_key:
+                link_metadata["api_key"] = api_key
+
+            existing_location = DataStorageLocationPayload(
+                storage_type=DataStorageType.LINK,
+                content_type=DataContentType.TEXT,
+                location=url,
+                metadata=link_metadata or None,
+            )
+
+            payload = DataStorageRequestPayload(
+                name=name,
+                content=url,
+                description=description,
+                dataset_id=dataset_id,
+                project_id=project_id,
+                existing_location=existing_location,
+            )
+            return self._create_data_storage_entry(payload)
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
     @retry(
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, max=10),
@@ -1724,6 +1888,7 @@ class DataStorageMethods:
         name: str,
         existing_location: DataStorageLocationPayload,
         description: str | None = None,
+        as_collection: bool = False,
         path: str | None = None,
         project_id: UUID | None = None,
     ) -> DataStorageResponse:
@@ -1733,6 +1898,9 @@ class DataStorageMethods:
             name: Name of the data storage entry
             existing_location: Describes the existing data source location to register
             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
+                you are uploading a single file.
             path: Optional path for the data storage entry
             project_id: ID of the project this data storage entry belongs to
 
@@ -1749,9 +1917,11 @@ class DataStorageMethods:
                 path=path,
                 existing_location=existing_location,
                 project_id=project_id,
+                is_collection=as_collection,
             )
             response = self.client.post(
-                "/v0.1/data-storage", json=payload.model_dump(exclude_none=True)
+                "/v0.1/data-storage/data-entries",
+                json=payload.model_dump(exclude_none=True),
             )
             response.raise_for_status()
             return DataStorageResponse.model_validate(response.json())
@@ -1772,6 +1942,7 @@ class DataStorageMethods:
         self,
         name: str,
         existing_location: DataStorageLocationPayload,
+        as_collection: bool = False,
         description: str | None = None,
         path: str | None = None,
         project_id: UUID | None = None,
@@ -1782,6 +1953,9 @@ class DataStorageMethods:
             name: Name of the data storage entry
             existing_location: Describes the existing data source location to register
             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
+                you are uploading a single file.
             path: Optional path for the data storage entry
             project_id: ID of the project this data storage entry belongs to
 
@@ -1798,9 +1972,11 @@ class DataStorageMethods:
                 path=path,
                 existing_location=existing_location,
                 project_id=project_id,
+                is_collection=as_collection,
             )
             response = await self.async_client.post(
-                "/v0.1/data-storage", json=payload.model_dump(exclude_none=True)
+                "/v0.1/data-storage/data-entries",
+                json=payload.model_dump(exclude_none=True),
             )
             response.raise_for_status()
             return DataStorageResponse.model_validate(response.json())
@@ -1811,8 +1987,274 @@ class DataStorageMethods:
                 f"An unexpected error occurred: {e!r}"
             ) from e
 
-    # TODO: EVERYTHING BELOW THIS LINE SHOULD BE MOVED TO FH_TOOLS REPO
-    # =================================================
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def search_data_storage(
+        self,
+        criteria: list[SearchCriterion] | None = None,
+        size: int = 10,
+    ) -> list[dict]:
+        """Search data storage objects using structured criteria.
+
+        Args:
+            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            size: Number of results to return (1-100)
+
+        Returns:
+            List of search results with scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error searching data storage entries
+
+        Example:
+            from futurehouse_client.models.rest import SearchCriterion, SearchOperator
+            criteria = [
+                SearchCriterion(field="name", operator=SearchOperator.CONTAINS, value="document"),
+                SearchCriterion(field="project_id", operator=SearchOperator.EQUALS, value="my-project-id"),
+                SearchCriterion(field="status", operator=SearchOperator.EQUALS, value="active"),
+            ]
+            results = client.search_data_storage(criteria=criteria, size=20)
+        """
+        try:
+            payload = DataStorageSearchPayload(
+                criteria=criteria or [],
+                size=max(1, min(100, size)),  # Clamp between 1-100
+            )
+
+            response = self.client.post(
+                "/v0.1/data-storage/search",
+                json=payload.model_dump(mode="json"),
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "searching")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def asearch_data_storage(
+        self,
+        criteria: list[SearchCriterion] | None = None,
+        size: int = 10,
+    ) -> list[dict]:
+        """Asynchronously search data storage objects using structured criteria.
+
+        Args:
+            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            size: Number of results to return (1-100)
+
+        Returns:
+            List of search results with scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error searching data storage entries
+
+        Example:
+            from futurehouse_client.models.rest import SearchCriterion, SearchOperator
+            criteria = [
+                SearchCriterion(field="name", operator=SearchOperator.CONTAINS, value="document"),
+                SearchCriterion(field="project_id", operator=SearchOperator.EQUALS, value="my-project-id"),
+                SearchCriterion(field="status", operator=SearchOperator.EQUALS, value="active"),
+            ]
+            results = await client.asearch_data_storage(criteria=criteria, size=20)
+        """
+        try:
+            payload = DataStorageSearchPayload(
+                criteria=criteria or [],
+                size=max(1, min(100, size)),  # Clamp between 1-100
+            )
+
+            response = await self.async_client.post(
+                "/v0.1/data-storage/search",
+                json=payload.model_dump(mode="json"),
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "searching")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during async search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def similarity_search_data_storage(
+        self,
+        embedding: list[float],
+        size: int = 10,
+        min_score: float = 0.7,
+        dataset_id: UUID | None = None,
+        tags: list[str] | None = None,
+        user_id: str | None = None,
+        project_id: str | None = None,
+    ) -> list[dict]:
+        """Search data storage objects using vector similarity.
+
+        Args:
+            embedding: 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
+            user_id: Optional user ID filter (admin only)
+            project_id: Optional project ID filter
+
+        Returns:
+            List of search results with similarity scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error performing similarity search
+        """
+        try:
+            # Validate inputs
+            if not embedding:
+                raise DataStorageCreationError("Embedding vector is required")
+
+            if not all(isinstance(x, int | float) for x in embedding):
+                raise DataStorageCreationError("Embedding must be a list of numbers")
+
+            size = max(1, min(100, size))  # Clamp between 1-100
+            min_score = max(0.0, min(1.0, min_score))  # Clamp between 0.0-1.0
+
+            # Build request payload
+            payload = {
+                "embedding": embedding,
+                "size": size,
+                "min_score": min_score,
+            }
+
+            # Add optional filters
+            if dataset_id is not None:
+                payload["dataset_id"] = str(dataset_id)
+            if tags is not None:
+                payload["tags"] = tags
+            if user_id is not None:
+                payload["user_id"] = user_id
+            if project_id is not None:
+                payload["project_id"] = project_id
+
+            response = self.client.post(
+                "/v0.1/data-storage/similarity-search", json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Similarity search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "performing similarity search")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during similarity search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def asimilarity_search_data_storage(
+        self,
+        embedding: list[float],
+        size: int = 10,
+        min_score: float = 0.7,
+        dataset_id: UUID | None = None,
+        tags: list[str] | None = None,
+        user_id: str | None = None,
+        project_id: str | None = None,
+    ) -> list[dict]:
+        """Asynchronously search data storage objects using vector similarity.
+
+        Args:
+            embedding: 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
+            user_id: Optional user ID filter (admin only)
+            project_id: Optional project ID filter
+
+        Returns:
+            List of search results with similarity scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error performing similarity search
+        """
+        try:
+            # Validate inputs
+            if not embedding:
+                raise DataStorageCreationError("Embedding vector is required")
+
+            if not all(isinstance(x, int | float) for x in embedding):
+                raise DataStorageCreationError("Embedding must be a list of numbers")
+
+            size = max(1, min(100, size))  # Clamp between 1-100
+            min_score = max(0.0, min(1.0, min_score))  # Clamp between 0.0-1.0
+
+            # Build request payload
+            payload = {
+                "embedding": embedding,
+                "size": size,
+                "min_score": min_score,
+            }
+
+            # Add optional filters
+            if dataset_id is not None:
+                payload["dataset_id"] = str(dataset_id)
+            if tags is not None:
+                payload["tags"] = tags
+            if user_id is not None:
+                payload["user_id"] = user_id
+            if project_id is not None:
+                payload["project_id"] = project_id
+
+            response = await self.async_client.post(
+                "/v0.1/data-storage/similarity-search", json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Similarity search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "performing similarity search")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during async similarity search: {e!r}"
+            ) from e
+
     @retry(
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, max=10),
@@ -1822,7 +2264,7 @@ class DataStorageMethods:
     def fetch_data_from_storage(
         self,
         data_storage_id: UUID | None = None,
-    ) -> str | Path | None:
+    ) -> str | Path | list[Path] | None:
         """Fetch data from the storage system (sync version).
 
         Args:
@@ -1831,27 +2273,43 @@ class DataStorageMethods:
         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
             None if not found or error occurred
         """
         if not data_storage_id:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 "data_storage_id must be provided at this time"
             )
 
         try:
-            response = self.client.get(f"/v0.1/data-storage/{data_storage_id}")
+            response = self.client.get(
+                f"/v0.1/data-storage/data-entries/{data_storage_id}", timeout=100
+            )
             response.raise_for_status()
             result = DataStorageResponse.model_validate(response.json())
 
-            storage_type = result.storage_location.storage_config.storage_type
+            if len(result.storage_locations) > 1:
+                return [
+                    self._download_from_gcs(
+                        location.storage_config.signed_url or "",
+                        (location.storage_config.location or "").split("/")[-1],
+                    )
+                    for location in result.storage_locations
+                ]
+
+            # Most scenarios will only have one location
+            storage_location = result.storage_locations[0]
+            storage_type = storage_location.storage_config.storage_type
 
             if storage_type == "gcs":
-                if not result.signed_url:
-                    raise DataStorageCreationError(
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageRetrievalError(
                         "No signed URL available for GCS download"
                     )
 
-                return self._download_from_gcs(result.signed_url)
+                return self._download_from_gcs(
+                    storage_location.storage_config.signed_url
+                )
 
             if storage_type in {"raw_content", "pg_table"}:
                 content = result.data_storage.content
@@ -1862,12 +2320,12 @@ class DataStorageMethods:
                     return None
                 return content
 
-            raise DataStorageCreationError(f"Unsupported storage type: {storage_type}")
+            raise DataStorageRetrievalError(f"Unsupported storage type: {storage_type}")
 
         except HTTPStatusError as e:
             self._handle_http_errors(e, "retrieving")
         except Exception as e:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 f"An unexpected error occurred: {e!r}"
             ) from e
 
@@ -1880,7 +2338,7 @@ class DataStorageMethods:
     async def afetch_data_from_storage(
         self,
         data_storage_id: UUID | None = None,
-    ) -> str | Path | None:
+    ) -> str | Path | list[Path] | None:
         """Fetch data from the storage system.
 
         Args:
@@ -1889,29 +2347,46 @@ class DataStorageMethods:
         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
             None if not found or error occurred
         """
         if not data_storage_id:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 "data_storage_id must be provided at this time"
             )
 
         try:
             response = await self.async_client.get(
-                f"/v0.1/data-storage/{data_storage_id}"
+                f"/v0.1/data-storage/data-entries/{data_storage_id}", timeout=100
             )
             response.raise_for_status()
             result = DataStorageResponse.model_validate(response.json())
 
-            storage_type = result.storage_location.storage_config.storage_type
+            if len(result.storage_locations) > 1:
+                return await gather_with_concurrency(
+                    DOWNLOAD_CONCURRENCY,
+                    [
+                        self._adownload_from_gcs(
+                            location.storage_config.signed_url or "",
+                            (location.storage_config.location or "").split("/")[-1],
+                        )
+                        for location in result.storage_locations
+                    ],
+                )
+
+            # Most scenarios will only have one location
+            storage_location = result.storage_locations[0]
+            storage_type = storage_location.storage_config.storage_type
 
             if storage_type == "gcs":
-                if not result.signed_url:
-                    raise DataStorageCreationError(
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageRetrievalError(
                         "No signed URL available for GCS download"
                     )
 
-                return await self._adownload_from_gcs(result.signed_url)
+                return await self._adownload_from_gcs(
+                    storage_location.storage_config.signed_url
+                )
 
             if storage_type in {"raw_content", "pg_table"}:
                 content = result.data_storage.content
@@ -1922,11 +2397,189 @@ class DataStorageMethods:
                     return None
                 return content
 
-            raise DataStorageCreationError(f"Unsupported storage type: {storage_type}")
+            raise DataStorageRetrievalError(f"Unsupported storage type: {storage_type}")
 
         except HTTPStatusError as e:
             self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageRetrievalError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def acreate_dataset(
+        self,
+        name: str,
+        description: str | None = None,
+        dataset_id: UUID | None = None,
+    ):
+        try:
+            payload = CreateDatasetPayload(
+                name=name,
+                description=description,
+                id=dataset_id,
+            )
+            response = await self.async_client.post(
+                "/v0.1/data-storage/datasets",
+                json=payload.model_dump(exclude_none=True),
+            )
+            response.raise_for_status()
+            return CreateDatasetPayload.model_validate(response.json())
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
             ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def create_dataset(
+        self,
+        name: str,
+        description: str | None = None,
+        dataset_id: UUID | None = None,
+    ):
+        try:
+            payload = CreateDatasetPayload(
+                name=name,
+                description=description,
+                id=dataset_id,
+            )
+            response = self.client.post(
+                "/v0.1/data-storage/datasets",
+                json=payload.model_dump(exclude_none=True),
+            )
+            response.raise_for_status()
+            return CreateDatasetPayload.model_validate(response.json())
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def adelete_dataset(self, dataset_id: UUID):
+        """Delete a dataset.
+
+        Note: This will delete all data storage entries associated with the dataset.
+
+        Args:
+            dataset_id: ID of the dataset to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the dataset
+        """
+        try:
+            await self.async_client.delete(f"/v0.1/data-storage/datasets/{dataset_id}")
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def delete_dataset(self, dataset_id: UUID):
+        """Delete a dataset.
+
+        Note: This will delete all data storage entries associated with the dataset.
+
+        Args:
+            dataset_id: ID of the dataset to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the dataset
+        """
+        try:
+            self.client.delete(f"/v0.1/data-storage/datasets/{dataset_id}")
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def aget_dataset(self, dataset_id: UUID):
+        try:
+            response = await self.async_client.get(
+                f"/v0.1/data-storage/datasets/{dataset_id}"
+            )
+            response.raise_for_status()
+
+            return response.json()
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def get_dataset(self, dataset_id: UUID):
+        try:
+            response = self.client.get(f"/v0.1/data-storage/datasets/{dataset_id}")
+            response.raise_for_status()
+
+            return response.json()
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        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):
+        try:
+            await self.async_client.delete(
+                f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"
+            )
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def delete_data_storage_entry(self, data_storage_entry_id: UUID):
+        try:
+            self.client.delete(
+                f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"
+            )
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e