feat: Adding multi modal support for PGVectorStore

Author: dishaprakash

langchain_postgres/v2/async_vectorstore.py

@@ -1,12 +1,16 @@
 # TODO: Remove below import when minimum supported Python version is 3.10
 from __future__ import annotations
 
+import base64
 import copy
 import json
+import re
 import uuid
 from typing import Any, Callable, Iterable, Optional, Sequence
 
 import numpy as np
+import requests
+from google.cloud import storage  # type: ignore
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore, utils
@@ -365,6 +369,92 @@ async def aadd_documents(
         ids = await self.aadd_texts(texts, metadatas=metadatas, ids=ids, **kwargs)
         return ids
 
+    def _encode_image(self, uri: str) -> str:
+        """Get base64 string from a image URI."""
+        gcs_uri = re.match("gs://(.*?)/(.*)", uri)
+        if gcs_uri:
+            bucket_name, object_name = gcs_uri.groups()
+            storage_client = storage.Client()
+            bucket = storage_client.bucket(bucket_name)
+            blob = bucket.blob(object_name)
+            return base64.b64encode(blob.download_as_bytes()).decode("utf-8")
+
+        web_uri = re.match(r"^(https?://).*", uri)
+        if web_uri:
+            response = requests.get(uri, stream=True)
+            response.raise_for_status()
+            return base64.b64encode(response.content).decode("utf-8")
+
+        with open(uri, "rb") as image_file:
+            return base64.b64encode(image_file.read()).decode("utf-8")
+
+    async def aadd_images(
+        self,
+        uris: list[str],
+        metadatas: Optional[list[dict]] = None,
+        ids: Optional[list[str]] = None,
+        **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.
+            metadatas (Optional[list[dict]]): List of metadatas to add to table records.
+            ids: (Optional[list[str]]): List of IDs to add to table records.
+
+        Returns:
+            List of record IDs added.
+        """
+        encoded_images = []
+        if metadatas is None:
+            metadatas = [{"image_uri": uri} for uri in uris]
+
+        for uri in uris:
+            encoded_image = self._encode_image(uri)
+            encoded_images.append(encoded_image)
+
+        embeddings = self._images_embedding_helper(uris)
+        ids = await self.aadd_embeddings(
+            encoded_images, embeddings, metadatas=metadatas, ids=ids, **kwargs
+        )
+        return ids
+
+    def _images_embedding_helper(self, image_uris: list[str]) -> list[list[float]]:
+        # check if either `embed_images()` or `embed_image()` API is supported by the embedding service used
+        if hasattr(self.embedding_service, "embed_images"):
+            try:
+                embeddings = self.embedding_service.embed_images(image_uris)
+            except Exception as e:
+                raise Exception(
+                    f"Make sure your selected embedding model supports list of image URIs as input. {str(e)}"
+                )
+        elif hasattr(self.embedding_service, "embed_image"):
+            try:
+                embeddings = self.embedding_service.embed_image(image_uris)
+            except Exception as e:
+                raise Exception(
+                    f"Make sure your selected embedding model supports list of image URIs as input. {str(e)}"
+                )
+        else:
+            raise ValueError(
+                "Please use an embedding model that supports image embedding."
+            )
+        return embeddings
+
+    async def asimilarity_search_image(
+        self,
+        image_uri: str,
+        k: Optional[int] = None,
+        filter: Optional[dict] = None,
+        **kwargs: Any,
+    ) -> list[Document]:
+        """Return docs selected by similarity search on query."""
+        embedding = self._images_embedding_helper([image_uri])[0]
+
+        return await self.asimilarity_search_by_vector(
+            embedding=embedding, k=k, filter=filter, **kwargs
+        )
+
     async def adelete(
         self,
         ids: Optional[list] = None,
@@ -1268,3 +1358,25 @@ def max_marginal_relevance_search_with_score_by_vector(
         raise NotImplementedError(
             "Sync methods are not implemented for AsyncPGVectorStore. Use PGVectorStore interface instead."
         )
+
+    def add_images(
+        self,
+        uris: list[str],
+        metadatas: Optional[list[dict]] = None,
+        ids: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> list[str]:
+        raise NotImplementedError(
+            "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead."
+        )
+
+    def similarity_search_image(
+        self,
+        image_uri: str,
+        k: Optional[int] = None,
+        filter: Optional[dict] = None,
+        **kwargs: Any,
+    ) -> list[Document]:
+        raise NotImplementedError(
+            "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead."
+        )

langchain_postgres/v2/vectorstores.py

@@ -840,3 +840,55 @@ def get_by_ids(self, ids: Sequence[str]) -> list[Document]:
 
     def get_table_name(self) -> str:
         return self.__vs.table_name
+
+    async def aadd_images(
+        self,
+        uris: list[str],
+        metadatas: Optional[list[dict]] = None,
+        ids: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Embed images and add to the table."""
+        return await self._engine._run_as_async(
+            self._PGVectorStore__vs.aadd_images(uris, metadatas, ids, **kwargs)  # type: ignore
+        )
+
+    def add_images(
+        self,
+        uris: list[str],
+        metadatas: Optional[list[dict]] = None,
+        ids: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Embed images and add to the table."""
+        return self._engine._run_as_sync(
+            self._PGVectorStore__vs.aadd_images(uris, metadatas, ids, **kwargs)  # type: ignore
+        )
+
+    def similarity_search_image(
+        self,
+        image_uri: str,
+        k: Optional[int] = None,
+        filter: Optional[dict] = None,
+        **kwargs: Any,
+    ) -> list[Document]:
+        """Return docs selected by similarity search on image."""
+        return self._engine._run_as_sync(
+            self._PGVectorStore__vs.asimilarity_search_image(
+                image_uri, k, filter, **kwargs
+            )  # type: ignore
+        )
+
+    async def asimilarity_search_image(
+        self,
+        image_uri: str,
+        k: Optional[int] = None,
+        filter: Optional[dict] = None,
+        **kwargs: Any,
+    ) -> list[Document]:
+        """Return docs selected by similarity search on image_uri."""
+        return await self._engine._run_as_async(
+            self._PGVectorStore__vs.asimilarity_search_image(
+                image_uri, k, filter, **kwargs
+            )  # type: ignore
+        )

pyproject.toml

@@ -19,6 +19,7 @@ sqlalchemy = "^2"
 pgvector = ">=0.2.5,<0.4"
 numpy = "^1.21"
 asyncpg = "^0.30.0"
+google-cloud-storage = ">=2.18.2, <4.0.0"
 
 [tool.poetry.group.docs.dependencies]
 
@@ -37,6 +38,7 @@ pytest-socket = "^0.7.0"
 pytest-cov = "^5.0.0"
 pytest-timeout = "^2.3.1"
 langchain-tests = "0.3.7"
+pillow = "11.1.0"
 
 [tool.poetry.group.codespell]
 optional = true

tests/unit_tests/v2/test_async_pg_vectorstore.py

@@ -1,10 +1,12 @@
+import os
 import uuid
 from typing import AsyncIterator, Sequence
 
 import pytest
 import pytest_asyncio
 from langchain_core.documents import Document
 from langchain_core.embeddings import DeterministicFakeEmbedding
+from PIL import Image
 from sqlalchemy import text
 from sqlalchemy.engine.row import RowMapping
 
@@ -15,6 +17,7 @@
 DEFAULT_TABLE = "default" + str(uuid.uuid4())
 DEFAULT_TABLE_SYNC = "default_sync" + str(uuid.uuid4())
 CUSTOM_TABLE = "custom" + str(uuid.uuid4())
+IMAGE_TABLE = "image_table" + str(uuid.uuid4())
 VECTOR_SIZE = 768
 
 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE)
@@ -28,6 +31,14 @@
 embeddings = [embeddings_service.embed_query(texts[i]) for i in range(len(texts))]
 
 
+class FakeImageEmbedding(DeterministicFakeEmbedding):
+    def embed_image(self, image_paths: list[str]) -> list[list[float]]:
+        return [self.embed_query(path) for path in image_paths]
+
+
+image_embedding_service = FakeImageEmbedding(size=VECTOR_SIZE)
+
+
 async def aexecute(engine: PGEngine, query: str) -> None:
     async with engine._pool.connect() as conn:
         await conn.execute(text(query))
@@ -52,6 +63,7 @@ async def engine(self) -> AsyncIterator[PGEngine]:
         yield engine
         await engine.adrop_table(DEFAULT_TABLE)
         await engine.adrop_table(CUSTOM_TABLE)
+        await engine.adrop_table(IMAGE_TABLE)
         await engine.close()
 
     @pytest_asyncio.fixture(scope="class")
@@ -87,6 +99,45 @@ async def vs_custom(self, engine: PGEngine) -> AsyncIterator[AsyncPGVectorStore]
         )
         yield vs
 
+    @pytest_asyncio.fixture(scope="class")
+    async def image_vs(self, engine: PGEngine) -> AsyncIterator[AsyncPGVectorStore]:
+        await engine._ainit_vectorstore_table(
+            IMAGE_TABLE,
+            VECTOR_SIZE,
+            metadata_columns=[
+                Column("image_id", "TEXT"),
+                Column("source", "TEXT"),
+            ],
+        )
+        vs = await AsyncPGVectorStore.create(
+            engine,
+            embedding_service=image_embedding_service,
+            table_name=IMAGE_TABLE,
+            metadata_columns=["image_id", "source"],
+            metadata_json_column="mymeta",
+        )
+        yield vs
+
+    @pytest_asyncio.fixture(scope="class")
+    async def image_uris(self) -> AsyncIterator[list[str]]:
+        red_uri = str(uuid.uuid4()).replace("-", "_") + "test_image_red.jpg"
+        green_uri = str(uuid.uuid4()).replace("-", "_") + "test_image_green.jpg"
+        blue_uri = str(uuid.uuid4()).replace("-", "_") + "test_image_blue.jpg"
+        gcs_uri = "gs://github-repo/img/vision/google-cloud-next.jpeg"
+        image = Image.new("RGB", (100, 100), color="red")
+        image.save(red_uri)
+        image = Image.new("RGB", (100, 100), color="green")
+        image.save(green_uri)
+        image = Image.new("RGB", (100, 100), color="blue")
+        image.save(blue_uri)
+        image_uris = [red_uri, green_uri, blue_uri, gcs_uri]
+        yield image_uris
+        for uri in image_uris:
+            try:
+                os.remove(uri)
+            except FileNotFoundError:
+                pass
+
     async def test_init_with_constructor(self, engine: PGEngine) -> None:
         with pytest.raises(Exception):
             AsyncPGVectorStore(
@@ -165,6 +216,20 @@ async def test_adelete(self, engine: PGEngine, vs: AsyncPGVectorStore) -> None:
         assert result == False
         await aexecute(engine, f'TRUNCATE TABLE "{DEFAULT_TABLE}"')
 
+    async def test_aadd_images(
+        self, engine: PGEngine, image_vs: AsyncPGVectorStore, image_uris: list[str]
+    ) -> None:
+        ids = [str(uuid.uuid4()) for i in range(len(image_uris))]
+        metadatas = [
+            {"image_id": str(i), "source": "postgres"} for i in range(len(image_uris))
+        ]
+        await image_vs.aadd_images(image_uris, metadatas, ids)
+        results = await afetch(engine, (f'SELECT * FROM "{IMAGE_TABLE}"'))
+        assert len(results) == len(image_uris)
+        assert results[0]["image_id"] == "0"
+        assert results[0]["source"] == "postgres"
+        await aexecute(engine, (f'TRUNCATE TABLE "{IMAGE_TABLE}"'))
+
     ##### Custom Vector Store  #####
     async def test_aadd_embeddings(
         self, engine: PGEngine, vs_custom: AsyncPGVectorStore