feat: Adding multi modal support for PGVectorStore

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,51 @@ 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.__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.__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.__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.__vs.asimilarity_search_image(image_uri, k, filter, **kwargs)  # type: ignore
+        )