feat(aadd_images): add support for store_uri_only parameter (#428)

* feat(aadd_images): add support for store_uri_only parameter

* chore: balck reformating

* Delete get-pip.py

* fix: tests for store_uri_only

---------

Co-authored-by: Vishwaraj Anand <[email protected]>

Author: zaherweb-2

src/langchain_google_alloydb_pg/async_vectorstore.py

@@ -397,29 +397,42 @@ async def aadd_images(
         uris: list[str],
         metadatas: Optional[list[dict]] = None,
         ids: Optional[list[str]] = None,
+        store_uri_only: bool = False,
         **kwargs: Any,
     ) -> list[str]:
         """Embed images and add to the table.
 
         Args:
-            uris (list[str]): List of local image URIs to add to the table.
+            uris (list[str]): List of image URIs to add to the table.
             metadatas (Optional[list[dict]]): List of metadatas to add to table records.
             ids: (Optional[list[str]]): List of IDs to add to table records.
+            store_uri_only (bool): If True, stores the URI in the content column
+                                   instead of the base64 encoded image. Defaults to False.
+            **kwargs: Any other arguments to pass to the embedding service.
 
         Returns:
             List of record IDs added.
         """
-        encoded_images = []
         if metadatas is None:
+            # Ensure URI is always in metadata if not explicitly provided elsewhere
             metadatas = [{"image_uri": uri} for uri in uris]
+        elif store_uri_only:
+            # If storing URI only and metadatas are provided, ensure image_uri is present
+            for i, m in enumerate(metadatas):
+                if "image_uri" not in m:  # Add if not already provided by user
+                    m["image_uri"] = uris[i]
+
+        texts_for_content_column: list[str]
+        if store_uri_only:
+            texts_for_content_column = uris
+        else:
+            texts_for_content_column = [self._encode_image(uri) for uri in uris]
 
-        for uri in uris:
-            encoded_image = self._encode_image(uri)
-            encoded_images.append(encoded_image)
-
+        # Embeddings are always generated from the actual image content via URIs
         embeddings = self._images_embedding_helper(uris)
+
         ids = await self.aadd_embeddings(
-            encoded_images, embeddings, metadatas=metadatas, ids=ids, **kwargs
+            texts_for_content_column, embeddings, metadatas=metadatas, ids=ids, **kwargs
         )
         return ids
 

src/langchain_google_alloydb_pg/vectorstore.py

@@ -231,11 +231,14 @@ async def aadd_images(
         uris: list[str],
         metadatas: Optional[list[dict]] = None,
         ids: Optional[list[str]] = None,
+        store_uri_only: bool = False,
         **kwargs: Any,
     ) -> list[str]:
         """Embed images and add to the table."""
         return await self._engine._run_as_async(
-            self.__vs.aadd_images(uris, metadatas, ids, **kwargs)
+            self.__vs.aadd_images(
+                uris, metadatas, ids, store_uri_only=store_uri_only, **kwargs
+            )
         )
 
     def add_embeddings(
@@ -287,11 +290,14 @@ def add_images(
         uris: list[str],
         metadatas: Optional[list[dict]] = None,
         ids: Optional[list[str]] = None,
+        store_uri_only: bool = False,
         **kwargs: Any,
     ) -> list[str]:
         """Embed images and add to the table."""
         return self._engine._run_as_sync(
-            self.__vs.aadd_images(uris, metadatas, ids, **kwargs)
+            self.__vs.aadd_images(
+                uris, metadatas, ids, store_uri_only=store_uri_only, **kwargs
+            )
         )
 
     async def adelete(

tests/test_async_vectorstore.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
 import uuid
 from typing import Sequence
@@ -48,7 +49,7 @@
 class FakeImageEmbedding(DeterministicFakeEmbedding):
 
     def embed_image(self, image_paths: list[str]) -> list[list[float]]:
-        return [self.embed_query(path) for path in image_paths]
+        return [self.embed_query(f"Image Path: {path}") for path in image_paths]
 
 
 image_embedding_service = FakeImageEmbedding(size=VECTOR_SIZE)
@@ -154,6 +155,7 @@ async def image_vs(self, engine):
                 Column("image_id", "TEXT"),
                 Column("source", "TEXT"),
             ],
+            metadata_json_column="mymeta",
         )
         vs = await AsyncAlloyDBVectorStore.create(
             engine,
@@ -278,6 +280,32 @@ async def test_aadd_images(self, engine, image_vs, image_uris):
         assert results[0]["source"] == "google.com"
         await aexecute(engine, (f'TRUNCATE TABLE "{IMAGE_TABLE}"'))
 
+    async def test_aadd_images_store_uri_only(self, engine, image_vs, image_uris):
+        ids = [str(uuid.uuid4()) for i in range(len(image_uris))]
+        metadatas = [
+            {"image_id": str(i), "source": "google.com"} for i in range(len(image_uris))
+        ]
+        await image_vs.aadd_images(image_uris, metadatas, ids, store_uri_only=True)
+        results = await afetch(engine, (f'SELECT * FROM "{IMAGE_TABLE}"'))
+        assert len(results) == len(image_uris)
+        # Check that content column stores the URI
+        for i, result_row in enumerate(results):
+            assert result_row[image_vs.content_column] == image_uris[i]
+            # Check that embedding is not an embedding of the URI string itself (basic check)
+            uri_embedding = embeddings_service.embed_query(image_uris[i])
+            image_embedding = image_embedding_service.embed_image([image_uris[i]])[0]
+            actual_embedding = json.loads(result_row[image_vs.embedding_column])
+            assert actual_embedding != pytest.approx(uri_embedding)
+            assert actual_embedding == pytest.approx(image_embedding)
+            assert result_row["image_id"] == str(i)
+            assert result_row["source"] == "google.com"
+            # Check that the original URI is also in the metadata (json column)
+            assert (
+                result_row[image_vs.metadata_json_column]["image_uri"] == image_uris[i]
+            )
+
+        await aexecute(engine, (f'TRUNCATE TABLE "{IMAGE_TABLE}"'))
+
     async def test_adelete(self, engine, vs):
         ids = [str(uuid.uuid4()) for i in range(len(texts))]
         await vs.aadd_texts(texts, ids=ids)

tests/test_vectorstore.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import os
 import uuid
 from threading import Thread
@@ -52,7 +53,7 @@
 class FakeImageEmbedding(DeterministicFakeEmbedding):
 
     def embed_image(self, image_paths: list[str]) -> list[list[float]]:
-        return [self.embed_query(path) for path in image_paths]
+        return [self.embed_query(f"Image Path: {path}") for path in image_paths]
 
 
 image_embedding_service = FakeImageEmbedding(size=VECTOR_SIZE)
@@ -357,7 +358,52 @@ async def test_aadd_images(self, engine_sync, image_uris):
         assert len(results) == len(image_uris)
         assert results[0]["image_id"] == "0"
         assert results[0]["source"] == "google.com"
-        await aexecute(engine_sync, f'TRUNCATE TABLE "{IMAGE_TABLE}"')
+        await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{IMAGE_TABLE}"')
+
+    async def test_aadd_images_store_uri_only(self, engine_sync, image_uris):
+        table_name = IMAGE_TABLE_SYNC + "_store_uri_only"
+        engine_sync.init_vectorstore_table(
+            table_name,
+            VECTOR_SIZE,
+            metadata_columns=[
+                Column("image_id", "TEXT"),
+                Column("source", "TEXT"),
+            ],
+            metadata_json_column="mymeta",
+        )
+        vs = AlloyDBVectorStore.create_sync(
+            engine_sync,
+            embedding_service=image_embedding_service,
+            table_name=table_name,
+            metadata_columns=["image_id", "source"],
+            metadata_json_column="mymeta",
+        )
+        ids = [str(uuid.uuid4()) for i in range(len(image_uris))]
+        metadatas = [
+            {"image_id": str(i), "source": "google.com"} for i in range(len(image_uris))
+        ]
+        # Test the async method on the sync class
+        await vs.aadd_images(image_uris, metadatas, ids, store_uri_only=True)
+        results = await afetch(engine_sync, f'SELECT * FROM "{table_name}"')
+        assert len(results) == len(image_uris)
+        for i, result_row in enumerate(results):
+            assert (
+                result_row[vs._AlloyDBVectorStore__vs.content_column] == image_uris[i]
+            )
+            uri_embedding = embeddings_service.embed_query(image_uris[i])
+            image_embedding = image_embedding_service.embed_image([image_uris[i]])[0]
+            actual_embedding = json.loads(
+                result_row[vs._AlloyDBVectorStore__vs.embedding_column]
+            )
+            assert actual_embedding != pytest.approx(uri_embedding)
+            assert actual_embedding == pytest.approx(image_embedding)
+            assert result_row["image_id"] == str(i)
+            assert result_row["source"] == "google.com"
+            assert (
+                result_row[vs._AlloyDBVectorStore__vs.metadata_json_column]["image_uri"]
+                == image_uris[i]
+            )
+        await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{table_name}"')
 
     async def test_adelete_custom(self, engine, vs_custom):
         ids = [str(uuid.uuid4()) for i in range(len(texts))]
@@ -405,6 +451,49 @@ async def test_add_images(self, engine_sync, image_uris):
         await vs.adelete(ids)
         await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{IMAGE_TABLE_SYNC}"')
 
+    async def test_add_images_store_uri_only(self, engine_sync, image_uris):
+        table_name = IMAGE_TABLE_SYNC + "_store_uri_only"
+        engine_sync.init_vectorstore_table(
+            table_name,
+            VECTOR_SIZE,
+            metadata_columns=[Column("image_id", "TEXT"), Column("source", "TEXT")],
+            metadata_json_column="mymeta",
+        )
+        vs = AlloyDBVectorStore.create_sync(
+            engine_sync,
+            embedding_service=image_embedding_service,
+            table_name=table_name,
+            metadata_columns=["image_id", "source"],
+            metadata_json_column="mymeta",
+        )
+
+        ids = [str(uuid.uuid4()) for i in range(len(image_uris))]
+        metadatas = [
+            {"image_id": str(i), "source": "google.com"} for i in range(len(image_uris))
+        ]
+        vs.add_images(image_uris, metadatas, ids, store_uri_only=True)
+        results = await afetch(engine_sync, (f'SELECT * FROM "{table_name}"'))
+        assert len(results) == len(image_uris)
+        for i, result_row in enumerate(results):
+            assert (
+                result_row[vs._AlloyDBVectorStore__vs.content_column] == image_uris[i]
+            )
+            uri_embedding = embeddings_service.embed_query(image_uris[i])
+            image_embedding = image_embedding_service.embed_image([image_uris[i]])[0]
+            actual_embedding = json.loads(
+                result_row[vs._AlloyDBVectorStore__vs.embedding_column]
+            )
+            assert actual_embedding != pytest.approx(uri_embedding)
+            assert actual_embedding == pytest.approx(image_embedding)
+            assert result_row["image_id"] == str(i)
+            assert result_row["source"] == "google.com"
+            assert (
+                result_row[vs._AlloyDBVectorStore__vs.metadata_json_column]["image_uri"]
+                == image_uris[i]
+            )
+        await vs.adelete(ids)
+        await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{table_name}"')
+
     async def test_cross_env(self, engine_sync, vs_sync):
         ids = [str(uuid.uuid4()) for i in range(len(texts))]
         await vs_sync.aadd_texts(texts, ids=ids)