update example

Author: blink1073

pymongo_voyageai/utils.py

@@ -5,7 +5,7 @@
 from PIL import Image
 
 from .document import ImageDocument
-from .storage import ObjectStorage
+from .storage import ObjectStorage, S3Storage
 
 try:
     import fitz  # type:ignore[import-untyped]
@@ -97,6 +97,11 @@ def url_to_images(
     # 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)
+        storage.close()
     # For all other files, use the native download.
     if source is None:
         with urllib.request.urlopen(url) as response:

tests/test_client_integration.py

@@ -6,7 +6,7 @@
 import pytest
 from bson import ObjectId
 
-from pymongo_voyageai import PyMongoVoyageAI
+from pymongo_voyageai import MemoryStorage, PyMongoVoyageAI
 
 if "VOYAGEAI_API_KEY" not in os.environ:
     pytest.skip("Requires VoyageAI API Key.", allow_module_level=True)
@@ -95,6 +95,25 @@ def test_pdf_pages_storage(client: PyMongoVoyageAI):
     storage.client.delete_object(Bucket=storage.root_location, Key=object_name)
 
 
+def test_pdf_pages_custom_storage(client: PyMongoVoyageAI):
+    query = "The consequences of a dictator's peace"
+    url = "https://www.fdrlibrary.org/documents/356632/390886/readingcopy.pdf"
+    storage = client._storage
+    object_name = f"{ObjectId()}.pdf"
+    with urllib.request.urlopen(url) as response:
+        storage.client.upload_fileobj(response, storage.root_location, object_name)
+    url = f"s3://{storage.root_location}/{object_name}"
+    client._storage = MemoryStorage()
+    images = client.url_to_images(url)
+    resp = client.add_documents(images)
+    client.wait_for_indexing()
+    data = client.similarity_search(query, extract_images=True)
+    assert len(data[0]["inputs"][0].image.tobytes()) > 0
+    assert len(client.get_by_ids([d["_id"] for d in resp])) == len(resp)
+    client.delete_by_ids([d["_id"] for d in resp])
+    storage.client.delete_object(Bucket=storage.root_location, Key=object_name)
+
+
 @pytest.mark.asyncio
 async def test_image_set_async(client: PyMongoVoyageAI):
     url = "hf://datasets/princeton-nlp/CharXiv/val.parquet"