codekingpro/portable-devtools

1from abc import ABC2from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Type3 4import numpy as np5from langchain_core.documents import Document6from langchain_core.embeddings import Embeddings7from langchain_core.vectorstores import VectorStore8from pydantic import Field9 10from langchain_community.vectorstores.utils import maximal_marginal_relevance11 12if TYPE_CHECKING:13    from docarray import BaseDoc14    from docarray.index.abstract import BaseDocIndex15 16 17def _check_docarray_import() -> None:18    try:19        import docarray20 21        da_version = docarray.__version__.split(".")22        if int(da_version[0]) == 0 and int(da_version[1]) <= 31:23            raise ImportError(24                f"To use the DocArrayHnswSearch VectorStore the docarray "25                f"version >=0.32.0 is expected, received: {docarray.__version__}."26                f"To upgrade, please run: `pip install -U docarray`."27            )28    except ImportError:29        raise ImportError(30            "Could not import docarray python package. "31            "Please install it with `pip install docarray`."32        )33 34 35class DocArrayIndex(VectorStore, ABC):36    """Base class for `DocArray` based vector stores."""37 38    def __init__(39        self,40        doc_index: "BaseDocIndex",41        embedding: Embeddings,42    ):43        """Initialize a vector store from DocArray's DocIndex."""44        self.doc_index = doc_index45        self.embedding = embedding46 47    @staticmethod48    def _get_doc_cls(**embeddings_params: Any) -> Type["BaseDoc"]:49        """Get docarray Document class describing the schema of DocIndex."""50        from docarray import BaseDoc51        from docarray.typing import NdArray52 53        class DocArrayDoc(BaseDoc):54            text: Optional[str] = Field(default=None)55            embedding: Optional[NdArray] = Field(**embeddings_params)56            metadata: Optional[dict] = Field(default=None)57 58        return DocArrayDoc59 60    @property61    def doc_cls(self) -> Type["BaseDoc"]:62        if self.doc_index._schema is None:63            raise ValueError("doc_index expected to have non-null _schema attribute.")64        return self.doc_index._schema65 66    def add_texts(67        self,68        texts: Iterable[str],69        metadatas: Optional[List[dict]] = None,70        **kwargs: Any,71    ) -> List[str]:72        """Embed texts and add to the vector store.73 74        Args:75            texts: Iterable of strings to add to the vectorstore.76            metadatas: Optional list of metadatas associated with the texts.77 78        Returns:79            List of ids from adding the texts into the vectorstore.80        """81        ids: List[str] = []82        embeddings = self.embedding.embed_documents(list(texts))83        for i, (t, e) in enumerate(zip(texts, embeddings)):84            m = metadatas[i] if metadatas else {}85            doc = self.doc_cls(text=t, embedding=e, metadata=m)86            self.doc_index.index([doc])87            ids.append(str(doc.id))88 89        return ids90 91    def similarity_search_with_score(92        self, query: str, k: int = 4, **kwargs: Any93    ) -> List[Tuple[Document, float]]:94        """Return docs most similar to query.95 96        Args:97            query: Text to look up documents similar to.98            k: Number of Documents to return. Defaults to 4.99 100        Returns:101            List of documents most similar to the query text and102            cosine distance in float for each.103            Lower score represents more similarity.104        """105        query_embedding = self.embedding.embed_query(query)106        query_doc = self.doc_cls(embedding=query_embedding)107        docs, scores = self.doc_index.find(query_doc, search_field="embedding", limit=k)108 109        result = [110            (Document(page_content=doc.text, metadata=doc.metadata), score)111            for doc, score in zip(docs, scores)112        ]113        return result114 115    def similarity_search(116        self, query: str, k: int = 4, **kwargs: Any117    ) -> List[Document]:118        """Return docs most similar to query.119 120        Args:121            query: Text to look up documents similar to.122            k: Number of Documents to return. Defaults to 4.123 124        Returns:125            List of Documents most similar to the query.126        """127        results = self.similarity_search_with_score(query, k=k, **kwargs)128        return [doc for doc, _ in results]129 130    def _similarity_search_with_relevance_scores(131        self,132        query: str,133        k: int = 4,134        **kwargs: Any,135    ) -> List[Tuple[Document, float]]:136        """Return docs and relevance scores, normalized on a scale from 0 to 1.137 138        0 is dissimilar, 1 is most similar.139        """140        raise NotImplementedError()141 142    def similarity_search_by_vector(143        self, embedding: List[float], k: int = 4, **kwargs: Any144    ) -> List[Document]:145        """Return docs most similar to embedding vector.146 147        Args:148            embedding: Embedding to look up documents similar to.149            k: Number of Documents to return. Defaults to 4.150 151        Returns:152            List of Documents most similar to the query vector.153        """154 155        query_doc = self.doc_cls(embedding=embedding)156        docs = self.doc_index.find(157            query_doc, search_field="embedding", limit=k158        ).documents159 160        result = [161            Document(page_content=doc.text, metadata=doc.metadata) for doc in docs162        ]163        return result164 165    def max_marginal_relevance_search(166        self,167        query: str,168        k: int = 4,169        fetch_k: int = 20,170        lambda_mult: float = 0.5,171        **kwargs: Any,172    ) -> List[Document]:173        """Return docs selected using the maximal marginal relevance.174 175        Maximal marginal relevance optimizes for similarity to query AND diversity176        among selected documents.177 178        Args:179            query: Text to look up documents similar to.180            k: Number of Documents to return. Defaults to 4.181            fetch_k: Number of Documents to fetch to pass to MMR algorithm.182            lambda_mult: Number between 0 and 1 that determines the degree183                of diversity among the results with 0 corresponding184                to maximum diversity and 1 to minimum diversity.185                Defaults to 0.5.186        Returns:187            List of Documents selected by maximal marginal relevance.188        """189        query_embedding = self.embedding.embed_query(query)190        query_doc = self.doc_cls(embedding=query_embedding)191 192        docs = self.doc_index.find(193            query_doc, search_field="embedding", limit=fetch_k194        ).documents195 196        mmr_selected = maximal_marginal_relevance(197            np.array(query_embedding), docs.embedding, k=k198        )199        results = [200            Document(page_content=docs[idx].text, metadata=docs[idx].metadata)201            for idx in mmr_selected202        ]203        return results204