clean up storage api

blink1073 · blink1073 · commit 919ec27bdc6f · 2025-04-30T09:31:57.000-05:00
diff --git a/pymongo_voyageai/client.py b/pymongo_voyageai/client.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import logging
 from collections.abc import Mapping, Sequence
 from time import monotonic, sleep
@@ -206,7 +207,19 @@ def image_to_storage(self, document: ImageDocument | Image.Image) -> StoredDocum
         """
         if isinstance(document, Image.Image):
             document = ImageDocument(image=document)
-        return self._storage.save_image(document)
+        object_name = f"{ObjectId()}.png"
+        fd = io.BytesIO()
+        document.image.save(fd, "png")
+        fd.seek(0)
+        self._storage.save_data(fd, object_name)
+        return StoredDocument(
+            root_location=self._storage.root_location,
+            object_name=object_name,
+            page_number=document.page_number,
+            source_url=document.source_url,
+            name=document.name,
+            metadata=document.metadata,
+        )
 
     async def aimage_to_storage(self, document: ImageDocument | Image.Image) -> StoredDocument:
         """Convert an image to a stored document.
@@ -232,7 +245,15 @@ def storage_to_image(self, document: StoredDocument | str) -> ImageDocument:
             document = StoredDocument(
                 root_location=self._storage.root_location, object_name=document
             )
-        return self._storage.load_image(document=document)
+        buffer = self._storage.read_data(document.object_name)
+        image = Image.open(buffer)
+        return ImageDocument(
+            image=image,
+            source_url=document.source_url,
+            page_number=document.page_number,
+            metadata=document.metadata,
+            name=document.name,
+        )
 
     async def astorage_to_image(self, document: StoredDocument | str) -> ImageDocument:
         """Convert a stored document to an image document.
@@ -470,7 +491,7 @@ def delete_many(
                 self._expand_doc(obj, False)
                 for inp in obj["inputs"]:
                     if isinstance(inp, StoredDocument):
-                        self._storage.delete_image(inp)
+                        self._storage.delete_data(inp.object_name)
         return self._coll.delete_many(filter=filter, **kwargs).acknowledged
 
     async def adelete_many(
diff --git a/pymongo_voyageai/storage.py b/pymongo_voyageai/storage.py
@@ -2,40 +2,36 @@
 
 import boto3  # type:ignore[import-untyped]
 import botocore  # type:ignore[import-untyped]
-from bson import ObjectId
-from PIL import Image
-
-from .document import ImageDocument, StoredDocument
 
 
 class ObjectStorage:
-    """A class used store image documents."""
+    """A class used to store binary data."""
 
     root_location: str
-    """The root location to use in the object store."""
+    """The default root location to use in the object store."""
 
     url_prefixes: list[str] | None
     """The url prefixes used by the object store, for reading data from a url."""
 
-    def save_image(self, image: ImageDocument) -> StoredDocument:
-        """Save an image document to the object store."""
+    def save_data(self, data: io.BytesIO, object_name: str) -> None:
+        """Save data to the object store."""
         raise NotImplementedError
 
-    def load_image(self, document: StoredDocument) -> ImageDocument:
-        """Load an image document from the object store."""
+    def read_data(self, object_name: str) -> io.BytesIO:
+        """Read data from the object store."""
         raise NotImplementedError
 
-    def read_from_url(self, url: str) -> io.BytesIO:
-        """Read data from a url into a BytesIO object."""
+    def load_url(self, url: str) -> io.BytesIO:
+        """Load data from a url."""
         raise NotImplementedError
 
-    def delete_image(self, document: StoredDocument) -> None:
-        """Remove an image document from the object store."""
+    def delete_data(self, object_name: str) -> None:
+        """Delete data from the object store."""
         raise NotImplementedError
 
-    def close(self) -> None:
+    def close(self):
         """Close the object store."""
-        raise NotImplementedError
+        pass
 
 
 class S3Storage(ObjectStorage):
@@ -59,41 +55,26 @@ def __init__(
         self.client = client or boto3.client("s3", region_name=region_name)
         self.root_location = bucket_name
 
-    def save_image(self, image: ImageDocument) -> StoredDocument:
-        object_name = f"{ObjectId()}.png"
-        fd = io.BytesIO()
-        image.image.save(fd, "png")
-        fd.seek(0)
-        self.client.upload_fileobj(fd, self.root_location, object_name)
-        return StoredDocument(
-            root_location=self.root_location,
-            object_name=object_name,
-            page_number=image.page_number,
-            source_url=image.source_url,
-            name=image.name,
-            metadata=image.metadata,
-        )
-
-    def load_image(self, document: StoredDocument) -> ImageDocument:
+    def save_data(self, data: io.BytesIO, object_name: str) -> None:
+        """Save data to the object store."""
+        self.client.upload_fileobj(data, self.root_location, object_name)
+
+    def read_data(self, object_name: str) -> io.BytesIO:
+        """Read data using the object store."""
         buffer = io.BytesIO()
-        self.client.download_fileobj(document.root_location, document.object_name, buffer)
-        image = Image.open(buffer)
-        return ImageDocument(
-            image=image,
-            source_url=document.source_url,
-            page_number=document.page_number,
-            metadata=document.metadata,
-            name=document.name,
-        )
-
-    def read_from_url(self, url: str) -> io.BytesIO:
-        bucket, key = url.replace("s3://", "").split("/")
+        self.client.download_fileobj(self.root_location, object_name, buffer)
+        return buffer
+
+    def load_url(self, url: str) -> io.BytesIO:
+        """Load data from a url."""
+        bucket, _, object_name = url.replace("s3://", "").partition("/")
         buffer = io.BytesIO()
-        self.client.download_fileobj(bucket, key, buffer)
+        self.client.download_fileobj(bucket, object_name, buffer)
         return buffer
 
-    def delete_image(self, document: StoredDocument) -> None:
-        self.client.delete_object(Bucket=document.root_location, Key=document.object_name)
+    def delete_data(self, object_name: str) -> None:
+        """Delete data from the object store."""
+        self.client.delete_object(Bucket=self.root_location, Key=object_name)
 
     def close(self) -> None:
         self.client.close()
@@ -106,29 +87,21 @@ class MemoryStorage(ObjectStorage):
 
     def __init__(self) -> None:
         self.root_location = "foo"
-        self.storage: dict[str, ImageDocument] = dict()
-
-    def save_image(self, image: ImageDocument) -> StoredDocument:
-        object_name = str(ObjectId())
-        self.storage[object_name] = image
-        return StoredDocument(
-            root_location=self.root_location,
-            name=image.name,
-            object_name=object_name,
-            source_url=image.source_url,
-            page_number=image.page_number,
-        )
-
-    def load_image(self, document: StoredDocument) -> ImageDocument:
-        return self.storage[document.object_name]
-
-    def read_from_url(self, url: str) -> io.BytesIO:
-        with open(url.replace("file://", ""), "rb") as fid:
-            data = fid.read()
-        return io.BytesIO(data)
+        self.storage: dict[str, io.BytesIO] = dict()
 
-    def delete_image(self, document: StoredDocument) -> None:
-        self.storage.pop(document.object_name, None)
+    def save_data(self, data: io.BytesIO, object_name: str) -> None:
+        """Save data to the object store."""
+        self.storage[object_name] = data
 
-    def close(self):
-        pass
+    def read_data(self, object_name: str) -> io.BytesIO:
+        """Read data using the object store."""
+        return self.storage[object_name]
+
+    def load_url(self, url: str) -> io.BytesIO:
+        """Load data from a url."""
+        with open(url.replace("file://", ""), "rb") as fid:
+            return io.BytesIO(fid.read())
+
+    def delete_data(self, object_name: str) -> None:
+        """Delete data from the object store."""
+        self.storage.pop(object_name, None)
diff --git a/pymongo_voyageai/utils.py b/pymongo_voyageai/utils.py
@@ -92,15 +92,15 @@ def url_to_images(
     if storage and storage.url_prefixes:
         for pattern in storage.url_prefixes:
             if url.startswith(pattern):
-                source = storage.read_from_url(url)
+                source = storage.load_url(url)
                 break
     # For parquet files that are not loaded by the storage object, let pandas handle the download.
     if source is None and url.endswith(".parquet"):
         source = url
     # For s3 files that are not loaded by the storage object, create a temp S3Storage object.
     if source is None and url.startswith("s3://"):
         storage = S3Storage("")
-        source = storage.read_from_url(url)
+        source = storage.load_url(url)
         storage.close()
     # For all other files, use the native download.
     if source is None:
