codekingpro/portable-devtools

1from __future__ import annotations2 3import uuid4from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union5 6from langchain_core.documents import Document7from langchain_core.embeddings import Embeddings8from langchain_core.vectorstores import VectorStore9 10 11class PGVecto_rs(VectorStore):12    """VectorStore backed by pgvecto_rs."""13 14    _store = None15    _embedding: Embeddings16 17    def __init__(18        self,19        embedding: Embeddings,20        dimension: int,21        db_url: str,22        collection_name: str,23        new_table: bool = False,24    ) -> None:25        """Initialize a PGVecto_rs vectorstore.26 27        Args:28            embedding: Embeddings to use.29            dimension: Dimension of the embeddings.30            db_url: Database URL.31            collection_name: Name of the collection.32            new_table: Whether to create a new table or connect to an existing one.33            If true, the table will be dropped if exists, then recreated.34            Defaults to False.35        """36        try:37            from pgvecto_rs.sdk import PGVectoRs38        except ImportError as e:39            raise ImportError(40                "Unable to import pgvector_rs.sdk , please install with "41                '`pip install "pgvecto_rs[sdk]"`.'42            ) from e43        self._store = PGVectoRs(44            db_url=db_url,45            collection_name=collection_name,46            dimension=dimension,47            recreate=new_table,48        )49        self._embedding = embedding50 51    # ================ Create interface =================52    @classmethod53    def from_texts(54        cls,55        texts: List[str],56        embedding: Embeddings,57        metadatas: Optional[List[dict]] = None,58        db_url: str = "",59        collection_name: str = str(uuid.uuid4().hex),60        **kwargs: Any,61    ) -> PGVecto_rs:62        """Return VectorStore initialized from texts and optional metadatas."""63        sample_embedding = embedding.embed_query("Hello pgvecto_rs!")64        dimension = len(sample_embedding)65        if db_url is None:66            raise ValueError("db_url must be provided")67        _self: PGVecto_rs = cls(68            embedding=embedding,69            dimension=dimension,70            db_url=db_url,71            collection_name=collection_name,72        )73        _self.add_texts(texts, metadatas, **kwargs)74        return _self75 76    @classmethod77    def from_documents(78        cls,79        documents: List[Document],80        embedding: Embeddings,81        db_url: str = "",82        collection_name: str = str(uuid.uuid4().hex),83        **kwargs: Any,84    ) -> PGVecto_rs:85        """Return VectorStore initialized from documents."""86        texts = [document.page_content for document in documents]87        metadatas = [document.metadata for document in documents]88        return cls.from_texts(89            texts, embedding, metadatas, db_url, collection_name, **kwargs90        )91 92    @classmethod93    def from_collection_name(94        cls,95        embedding: Embeddings,96        db_url: str,97        collection_name: str,98    ) -> PGVecto_rs:99        """Create new empty vectorstore with collection_name.100        Or connect to an existing vectorstore in database if exists.101        Arguments should be the same as when the vectorstore was created."""102        sample_embedding = embedding.embed_query("Hello pgvecto_rs!")103        return cls(104            embedding=embedding,105            dimension=len(sample_embedding),106            db_url=db_url,107            collection_name=collection_name,108        )109 110    # ================ Insert interface =================111 112    def add_texts(113        self,114        texts: Iterable[str],115        metadatas: Optional[List[dict]] = None,116        **kwargs: Any,117    ) -> List[str]:118        """Run more texts through the embeddings and add to the vectorstore.119 120        Args:121            texts: Iterable of strings to add to the vectorstore.122            metadatas: Optional list of metadatas associated with the texts.123            kwargs: vectorstore specific parameters124 125        Returns:126            List of ids of the added texts.127 128        """129        from pgvecto_rs.sdk import Record130 131        embeddings = self._embedding.embed_documents(list(texts))132        records = [133            Record.from_text(text, embedding, meta)134            for text, embedding, meta in zip(texts, embeddings, metadatas or [])135        ]136        self._store.insert(records)  # type: ignore[union-attr]137        return [str(record.id) for record in records]138 139    def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]:140        """Run more documents through the embeddings and add to the vectorstore.141 142        Args:143            documents (List[Document]): List of documents to add to the vectorstore.144 145        Returns:146            List of ids of the added documents.147        """148        return self.add_texts(149            [document.page_content for document in documents],150            [document.metadata for document in documents],151            **kwargs,152        )153 154    # ================ Query interface =================155    def similarity_search_with_score_by_vector(156        self,157        query_vector: List[float],158        k: int = 4,159        distance_func: Literal[160            "sqrt_euclid", "neg_dot_prod", "ned_cos"161        ] = "sqrt_euclid",162        filter: Union[None, Dict[str, Any], Any] = None,163        **kwargs: Any,164    ) -> List[Tuple[Document, float]]:165        """Return docs most similar to query vector, with its score."""166 167        from pgvecto_rs.sdk.filters import meta_contains168 169        distance_func_map = {170            "sqrt_euclid": "<->",171            "neg_dot_prod": "<#>",172            "ned_cos": "<=>",173        }174        if filter is None:175            real_filter = None176        elif isinstance(filter, dict):177            real_filter = meta_contains(filter)178        else:179            real_filter = filter180        results = self._store.search(  # type: ignore[union-attr]181            query_vector,182            distance_func_map[distance_func],183            k,184            filter=real_filter,185        )186 187        return [188            (189                Document(190                    page_content=res[0].text,191                    metadata=res[0].meta,192                ),193                res[1],194            )195            for res in results196        ]197 198    def similarity_search_by_vector(199        self,200        embedding: List[float],201        k: int = 4,202        distance_func: Literal[203            "sqrt_euclid", "neg_dot_prod", "ned_cos"204        ] = "sqrt_euclid",205        filter: Optional[Any] = None,206        **kwargs: Any,207    ) -> List[Document]:208        return [209            doc210            for doc, _score in self.similarity_search_with_score_by_vector(211                embedding, k, distance_func, **kwargs212            )213        ]214 215    def similarity_search_with_score(216        self,217        query: str,218        k: int = 4,219        distance_func: Literal[220            "sqrt_euclid", "neg_dot_prod", "ned_cos"221        ] = "sqrt_euclid",222        **kwargs: Any,223    ) -> List[Tuple[Document, float]]:224        query_vector = self._embedding.embed_query(query)225        return self.similarity_search_with_score_by_vector(226            query_vector, k, distance_func, **kwargs227        )228 229    def similarity_search(230        self,231        query: str,232        k: int = 4,233        distance_func: Literal[234            "sqrt_euclid", "neg_dot_prod", "ned_cos"235        ] = "sqrt_euclid",236        **kwargs: Any,237    ) -> List[Document]:238        """Return docs most similar to query."""239        query_vector = self._embedding.embed_query(query)240        return [241            doc242            for doc, _score in self.similarity_search_with_score_by_vector(243                query_vector, k, distance_func, **kwargs244            )245        ]246