codekingpro/portable-devtools

1from __future__ import annotations2 3import uuid4import warnings5from itertools import repeat6from typing import (7    TYPE_CHECKING,8    Any,9    Dict,10    Iterable,11    List,12    Optional,13    Tuple,14    Type,15    Union,16)17 18import numpy as np19from langchain_core.documents import Document20from langchain_core.embeddings import Embeddings21from langchain_core.vectorstores import VectorStore22 23from langchain_community.vectorstores.utils import maximal_marginal_relevance24 25if TYPE_CHECKING:26    import supabase27 28 29class SupabaseVectorStore(VectorStore):30    """`Supabase Postgres` vector store.31 32    It assumes you have the `pgvector`33    extension installed and a `match_documents` (or similar) function. For more details:34    https://integrations.langchain.com/vectorstores?integration_name=SupabaseVectorStore35 36    You can implement your own `match_documents` function in order to limit the search37    space to a subset of documents based on your own authorization or business logic.38 39    Note that the Supabase Python client does not yet support async operations.40 41    If you'd like to use `max_marginal_relevance_search`, please review the instructions42    below on modifying the `match_documents` function to return matched embeddings.43 44 45    Examples:46 47    .. code-block:: python48 49        from langchain_community.embeddings.openai import OpenAIEmbeddings50        from langchain_core.documents import Document51        from langchain_community.vectorstores import SupabaseVectorStore52        from supabase.client import create_client53 54        docs = [55            Document(page_content="foo", metadata={"id": 1}),56        ]57        embeddings = OpenAIEmbeddings()58        supabase_client = create_client("my_supabase_url", "my_supabase_key")59        vector_store = SupabaseVectorStore.from_documents(60            docs,61            embeddings,62            client=supabase_client,63            table_name="documents",64            query_name="match_documents",65            chunk_size=500,66        )67 68    To load from an existing table:69 70    .. code-block:: python71 72        from langchain_community.embeddings.openai import OpenAIEmbeddings73        from langchain_community.vectorstores import SupabaseVectorStore74        from supabase.client import create_client75 76 77        embeddings = OpenAIEmbeddings()78        supabase_client = create_client("my_supabase_url", "my_supabase_key")79        vector_store = SupabaseVectorStore(80            client=supabase_client,81            embedding=embeddings,82            table_name="documents",83            query_name="match_documents",84        )85 86    """87 88    def __init__(89        self,90        client: supabase.client.Client,91        embedding: Embeddings,92        table_name: str,93        chunk_size: int = 500,94        query_name: Union[str, None] = None,95    ) -> None:96        """Initialize with supabase client."""97        try:98            import supabase  # noqa: F40199        except ImportError:100            raise ImportError(101                "Could not import supabase python package. "102                "Please install it with `pip install supabase`."103            )104 105        self._client = client106        self._embedding: Embeddings = embedding107        self.table_name = table_name or "documents"108        self.query_name = query_name or "match_documents"109        self.chunk_size = chunk_size or 500110        # According to the SupabaseVectorStore JS implementation, the best chunk size111        # is 500. Though for large datasets it can be too large so it is configurable.112 113    @property114    def embeddings(self) -> Embeddings:115        return self._embedding116 117    def add_texts(118        self,119        texts: Iterable[str],120        metadatas: Optional[List[Dict[Any, Any]]] = None,121        ids: Optional[List[str]] = None,122        **kwargs: Any,123    ) -> List[str]:124        ids = ids or [str(uuid.uuid4()) for _ in texts]125        docs = self._texts_to_documents(texts, metadatas)126 127        vectors = self._embedding.embed_documents(list(texts))128        return self.add_vectors(vectors, docs, ids)129 130    @classmethod131    def from_texts(132        cls: Type["SupabaseVectorStore"],133        texts: List[str],134        embedding: Embeddings,135        metadatas: Optional[List[dict]] = None,136        client: Optional[supabase.client.Client] = None,137        table_name: Optional[str] = "documents",138        query_name: Union[str, None] = "match_documents",139        chunk_size: int = 500,140        ids: Optional[List[str]] = None,141        **kwargs: Any,142    ) -> "SupabaseVectorStore":143        """Return VectorStore initialized from texts and embeddings."""144 145        if not client:146            raise ValueError("Supabase client is required.")147 148        if not table_name:149            raise ValueError("Supabase document table_name is required.")150 151        embeddings = embedding.embed_documents(texts)152        ids = [str(uuid.uuid4()) for _ in texts]153        docs = cls._texts_to_documents(texts, metadatas)154        cls._add_vectors(155            client, table_name, embeddings, docs, ids, chunk_size, **kwargs156        )157 158        return cls(159            client=client,160            embedding=embedding,161            table_name=table_name,162            query_name=query_name,163            chunk_size=chunk_size,164        )165 166    def add_vectors(167        self,168        vectors: List[List[float]],169        documents: List[Document],170        ids: List[str],171    ) -> List[str]:172        return self._add_vectors(173            self._client, self.table_name, vectors, documents, ids, self.chunk_size174        )175 176    def similarity_search(177        self,178        query: str,179        k: int = 4,180        filter: Optional[Dict[str, Any]] = None,181        **kwargs: Any,182    ) -> List[Document]:183        vector = self._embedding.embed_query(query)184        return self.similarity_search_by_vector(vector, k=k, filter=filter, **kwargs)185 186    def similarity_search_by_vector(187        self,188        embedding: List[float],189        k: int = 4,190        filter: Optional[Dict[str, Any]] = None,191        **kwargs: Any,192    ) -> List[Document]:193        result = self.similarity_search_by_vector_with_relevance_scores(194            embedding, k=k, filter=filter, **kwargs195        )196 197        documents = [doc for doc, _ in result]198 199        return documents200 201    def similarity_search_with_relevance_scores(202        self,203        query: str,204        k: int = 4,205        filter: Optional[Dict[str, Any]] = None,206        **kwargs: Any,207    ) -> List[Tuple[Document, float]]:208        vector = self._embedding.embed_query(query)209        return self.similarity_search_by_vector_with_relevance_scores(210            vector, k=k, filter=filter, **kwargs211        )212 213    def match_args(214        self, query: List[float], filter: Optional[Dict[str, Any]]215    ) -> Dict[str, Any]:216        ret: Dict[str, Any] = dict(query_embedding=query)217        if filter:218            ret["filter"] = filter219        return ret220 221    def similarity_search_by_vector_with_relevance_scores(222        self,223        query: List[float],224        k: int,225        filter: Optional[Dict[str, Any]] = None,226        postgrest_filter: Optional[str] = None,227        score_threshold: Optional[float] = None,228    ) -> List[Tuple[Document, float]]:229        # Convert MongoDB-style filter to PostgreSQL syntax if needed230        if filter:231            for key, value in filter.items():232                if isinstance(value, dict) and "$in" in value:233                    # Extract the list of values for the $in operator234                    in_values = value["$in"]235                    # Create a PostgreSQL IN clause236                    values_str = ",".join(f"'{str(v)}'" for v in in_values)237                    new_filter = f"metadata->>{key} IN ({values_str})"238 239                    # Combine with existing postgrest_filter if present240                    if postgrest_filter:241                        postgrest_filter = f"({postgrest_filter}) and ({new_filter})"242                    else:243                        postgrest_filter = new_filter244 245        match_documents_params = self.match_args(query, filter)246        query_builder = self._client.rpc(self.query_name, match_documents_params)247 248        if postgrest_filter:249            query_builder.params = query_builder.params.set(250                "and", f"({postgrest_filter})"251            )252 253        query_builder.params = query_builder.params.set("limit", k)254 255        res = query_builder.execute()256 257        match_result = [258            (259                Document(260                    metadata=search.get("metadata", {}),261                    page_content=search.get("content", ""),262                ),263                search.get("similarity", 0.0),264            )265            for search in res.data266            if search.get("content")267        ]268 269        if score_threshold is not None:270            match_result = [271                (doc, similarity)272                for doc, similarity in match_result273                if similarity >= score_threshold274            ]275            if len(match_result) == 0:276                warnings.warn(277                    "No relevant docs were retrieved using the relevance score"278                    f" threshold {score_threshold}"279                )280 281        return match_result282 283    def similarity_search_by_vector_returning_embeddings(284        self,285        query: List[float],286        k: int,287        filter: Optional[Dict[str, Any]] = None,288        postgrest_filter: Optional[str] = None,289    ) -> List[Tuple[Document, float, np.ndarray]]:290        match_documents_params = self.match_args(query, filter)291        query_builder = self._client.rpc(self.query_name, match_documents_params)292 293        if postgrest_filter:294            query_builder.params = query_builder.params.set(295                "and", f"({postgrest_filter})"296            )297 298        query_builder.params = query_builder.params.set("limit", k)299 300        res = query_builder.execute()301 302        match_result = [303            (304                Document(305                    metadata=search.get("metadata", {}),306                    page_content=search.get("content", ""),307                ),308                search.get("similarity", 0.0),309                # Supabase returns a vector type as its string represation (!).310                # This is a hack to convert the string to numpy array.311                np.fromstring(312                    search.get("embedding", "").strip("[]"), np.float32, sep=","313                ),314            )315            for search in res.data316            if search.get("content")317        ]318 319        return match_result320 321    @staticmethod322    def _texts_to_documents(323        texts: Iterable[str],324        metadatas: Optional[Iterable[Dict[Any, Any]]] = None,325    ) -> List[Document]:326        """Return list of Documents from list of texts and metadatas."""327        if metadatas is None:328            metadatas = repeat({})329 330        docs = [331            Document(page_content=text, metadata=metadata)332            for text, metadata in zip(texts, metadatas)333        ]334 335        return docs336 337    @staticmethod338    def _add_vectors(339        client: supabase.client.Client,340        table_name: str,341        vectors: List[List[float]],342        documents: List[Document],343        ids: List[str],344        chunk_size: int,345        **kwargs: Any,346    ) -> List[str]:347        """Add vectors to Supabase table."""348 349        rows: List[Dict[str, Any]] = [350            {351                "id": ids[idx],352                "content": documents[idx].page_content,353                "embedding": embedding,354                "metadata": documents[idx].metadata,355                **kwargs,356            }357            for idx, embedding in enumerate(vectors)358        ]359        id_list: List[str] = []360        for i in range(0, len(rows), chunk_size):361            chunk = rows[i : i + chunk_size]362 363            result = client.from_(table_name).upsert(chunk).execute()364 365            if len(result.data) == 0:366                raise Exception("Error inserting: No rows added")367 368            # VectorStore.add_vectors returns ids as strings369            ids = [str(i.get("id")) for i in result.data if i.get("id")]370 371            id_list.extend(ids)372 373        return id_list374 375    def max_marginal_relevance_search_by_vector(376        self,377        embedding: List[float],378        k: int = 4,379        fetch_k: int = 20,380        lambda_mult: float = 0.5,381        **kwargs: Any,382    ) -> List[Document]:383        """Return docs selected using the maximal marginal relevance.384 385        Maximal marginal relevance optimizes for similarity to query AND diversity386        among selected documents.387 388        Args:389            embedding: Embedding to look up documents similar to.390            k: Number of Documents to return. Defaults to 4.391            fetch_k: Number of Documents to fetch to pass to MMR algorithm.392            lambda_mult: Number between 0 and 1 that determines the degree393                        of diversity among the results with 0 corresponding394                        to maximum diversity and 1 to minimum diversity.395                        Defaults to 0.5.396        Returns:397            List of Documents selected by maximal marginal relevance.398        """399        result = self.similarity_search_by_vector_returning_embeddings(400            embedding, fetch_k401        )402 403        matched_documents = [doc_tuple[0] for doc_tuple in result]404        matched_embeddings = [doc_tuple[2] for doc_tuple in result]405 406        mmr_selected = maximal_marginal_relevance(407            np.array([embedding], dtype=np.float32),408            matched_embeddings,409            k=k,410            lambda_mult=lambda_mult,411        )412 413        filtered_documents = [matched_documents[i] for i in mmr_selected]414 415        return filtered_documents416 417    def max_marginal_relevance_search(418        self,419        query: str,420        k: int = 4,421        fetch_k: int = 20,422        lambda_mult: float = 0.5,423        **kwargs: Any,424    ) -> List[Document]:425        """Return docs selected using the maximal marginal relevance.426 427        Maximal marginal relevance optimizes for similarity to query AND diversity428        among selected documents.429 430        Args:431            query: Text to look up documents similar to.432            k: Number of Documents to return. Defaults to 4.433            fetch_k: Number of Documents to fetch to pass to MMR algorithm.434            lambda_mult: Number between 0 and 1 that determines the degree435                        of diversity among the results with 0 corresponding436                        to maximum diversity and 1 to minimum diversity.437                        Defaults to 0.5.438        Returns:439            List of Documents selected by maximal marginal relevance.440 441        `max_marginal_relevance_search` requires that `query_name` returns matched442        embeddings alongside the match documents. The following function443        demonstrates how to do this:444 445        ```sql446        CREATE FUNCTION match_documents_embeddings(query_embedding vector(1536),447                                                   match_count int)448            RETURNS TABLE(449                id uuid,450                content text,451                metadata jsonb,452                embedding vector(1536),453                similarity float)454            LANGUAGE plpgsql455            AS $$456            # variable_conflict use_column457        BEGIN458            RETURN query459            SELECT460                id,461                content,462                metadata,463                embedding,464                1 -(docstore.embedding <=> query_embedding) AS similarity465            FROM466                docstore467            ORDER BY468                docstore.embedding <=> query_embedding469            LIMIT match_count;470        END;471        $$;472        ```473        """474        embedding = self._embedding.embed_query(query)475        docs = self.max_marginal_relevance_search_by_vector(476            embedding, k, fetch_k, lambda_mult=lambda_mult477        )478        return docs479 480    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> None:481        """Delete by vector IDs.482 483        Args:484            ids: List of ids to delete.485        """486 487        if ids is None:488            raise ValueError("No ids provided to delete.")489 490        rows: List[Dict[str, Any]] = [491            {492                "id": id,493            }494            for id in ids495        ]496 497        # TODO: Check if this can be done in bulk498        for row in rows:499            self._client.from_(self.table_name).delete().eq("id", row["id"]).execute()500