Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
base.py1112 linesDownload Raw Back to vectorstores
1"""A vector store stores embedded data and performs vector search.2 3One of the most common ways to store and search over unstructured data is to4embed it and store the resulting embedding vectors, and then query the store5and retrieve the data that are 'most similar' to the embedded query.6"""7 8from __future__ import annotations9 10import logging11import math12import warnings13from abc import ABC, abstractmethod14from itertools import cycle15from typing import (16    TYPE_CHECKING,17    Any,18    ClassVar,19    TypeVar,20)21 22from pydantic import ConfigDict, Field, model_validator23from typing_extensions import Self, override24 25from langchain_core.documents import Document26from langchain_core.embeddings import Embeddings27from langchain_core.retrievers import BaseRetriever, LangSmithRetrieverParams28from langchain_core.runnables.config import run_in_executor29 30if TYPE_CHECKING:31    from collections.abc import Callable, Collection, Iterable, Iterator, Sequence32 33    from langchain_core.callbacks.manager import (34        AsyncCallbackManagerForRetrieverRun,35        CallbackManagerForRetrieverRun,36    )37 38logger = logging.getLogger(__name__)39 40VST = TypeVar("VST", bound="VectorStore")41 42 43class VectorStore(ABC):44    """Interface for vector store."""45 46    def add_texts(47        self,48        texts: Iterable[str],49        metadatas: list[dict] | None = None,50        *,51        ids: list[str] | None = None,52        **kwargs: Any,53    ) -> list[str]:54        """Run more texts through the embeddings and add to the `VectorStore`.55 56        Args:57            texts: Iterable of strings to add to the `VectorStore`.58            metadatas: Optional list of metadatas associated with the texts.59            ids: Optional list of IDs associated with the texts.60            **kwargs: `VectorStore` specific parameters.61 62                One of the kwargs should be `ids` which is a list of ids63                associated with the texts.64 65        Returns:66            List of IDs from adding the texts into the `VectorStore`.67 68        Raises:69            ValueError: If the number of metadatas does not match the number of texts.70            ValueError: If the number of IDs does not match the number of texts.71        """72        if type(self).add_documents != VectorStore.add_documents:73            # This condition is triggered if the subclass has provided74            # an implementation of the upsert method.75            # The existing add_texts76            texts_: Sequence[str] = (77                texts if isinstance(texts, (list, tuple)) else list(texts)78            )79            if metadatas and len(metadatas) != len(texts_):80                msg = (81                    "The number of metadatas must match the number of texts."82                    f"Got {len(metadatas)} metadatas and {len(texts_)} texts."83                )84                raise ValueError(msg)85            metadatas_ = iter(metadatas) if metadatas else cycle([{}])86            ids_: Iterator[str | None] = iter(ids) if ids else cycle([None])87            docs = [88                Document(id=id_, page_content=text, metadata=metadata_)89                for text, metadata_, id_ in zip(texts, metadatas_, ids_, strict=False)90            ]91            if ids is not None:92                # For backward compatibility93                kwargs["ids"] = ids94 95            return self.add_documents(docs, **kwargs)96        msg = f"`add_texts` has not been implemented for {self.__class__.__name__} "97        raise NotImplementedError(msg)98 99    @property100    def embeddings(self) -> Embeddings | None:101        """Access the query embedding object if available."""102        logger.debug(103            "The embeddings property has not been implemented for %s",104            self.__class__.__name__,105        )106        return None107 108    def delete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:109        """Delete by vector ID or other criteria.110 111        Args:112            ids: List of IDs to delete. If `None`, delete all.113            **kwargs: Other keyword arguments that subclasses might use.114 115        Returns:116            `True` if deletion is successful, `False` otherwise, `None` if not117                implemented.118        """119        msg = "delete method must be implemented by subclass."120        raise NotImplementedError(msg)121 122    def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:123        """Get documents by their IDs.124 125        The returned documents are expected to have the ID field set to the ID of the126        document in the vector store.127 128        Fewer documents may be returned than requested if some IDs are not found or129        if there are duplicated IDs.130 131        Users should not assume that the order of the returned documents matches132        the order of the input IDs. Instead, users should rely on the ID field of the133        returned documents.134 135        This method should **NOT** raise exceptions if no documents are found for136        some IDs.137 138        Args:139            ids: List of IDs to retrieve.140 141        Returns:142            List of `Document` objects.143        """144        msg = f"{self.__class__.__name__} does not yet support get_by_ids."145        raise NotImplementedError(msg)146 147    # Implementations should override this method to provide an async native version.148    async def aget_by_ids(self, ids: Sequence[str], /) -> list[Document]:149        """Async get documents by their IDs.150 151        The returned documents are expected to have the ID field set to the ID of the152        document in the vector store.153 154        Fewer documents may be returned than requested if some IDs are not found or155        if there are duplicated IDs.156 157        Users should not assume that the order of the returned documents matches158        the order of the input IDs. Instead, users should rely on the ID field of the159        returned documents.160 161        This method should **NOT** raise exceptions if no documents are found for162        some IDs.163 164        Args:165            ids: List of IDs to retrieve.166 167        Returns:168            List of `Document` objects.169        """170        return await run_in_executor(None, self.get_by_ids, ids)171 172    async def adelete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:173        """Async delete by vector ID or other criteria.174 175        Args:176            ids: List of IDs to delete. If `None`, delete all.177            **kwargs: Other keyword arguments that subclasses might use.178 179        Returns:180            `True` if deletion is successful, `False` otherwise, `None` if not181                implemented.182        """183        return await run_in_executor(None, self.delete, ids, **kwargs)184 185    async def aadd_texts(186        self,187        texts: Iterable[str],188        metadatas: list[dict] | None = None,189        *,190        ids: list[str] | None = None,191        **kwargs: Any,192    ) -> list[str]:193        """Async run more texts through the embeddings and add to the `VectorStore`.194 195        Args:196            texts: Iterable of strings to add to the `VectorStore`.197            metadatas: Optional list of metadatas associated with the texts.198            ids: Optional list199            **kwargs: `VectorStore` specific parameters.200 201        Returns:202            List of IDs from adding the texts into the `VectorStore`.203 204        Raises:205            ValueError: If the number of metadatas does not match the number of texts.206            ValueError: If the number of IDs does not match the number of texts.207        """208        if ids is not None:209            # For backward compatibility210            kwargs["ids"] = ids211        if type(self).aadd_documents != VectorStore.aadd_documents:212            # This condition is triggered if the subclass has provided213            # an implementation of the upsert method.214            # The existing add_texts215            texts_: Sequence[str] = (216                texts if isinstance(texts, (list, tuple)) else list(texts)217            )218            if metadatas and len(metadatas) != len(texts_):219                msg = (220                    "The number of metadatas must match the number of texts."221                    f"Got {len(metadatas)} metadatas and {len(texts_)} texts."222                )223                raise ValueError(msg)224            metadatas_ = iter(metadatas) if metadatas else cycle([{}])225            ids_: Iterator[str | None] = iter(ids) if ids else cycle([None])226 227            docs = [228                Document(id=id_, page_content=text, metadata=metadata_)229                for text, metadata_, id_ in zip(texts, metadatas_, ids_, strict=False)230            ]231            return await self.aadd_documents(docs, **kwargs)232        return await run_in_executor(None, self.add_texts, texts, metadatas, **kwargs)233 234    def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:235        """Add or update documents in the `VectorStore`.236 237        Args:238            documents: Documents to add to the `VectorStore`.239            **kwargs: Additional keyword arguments.240 241                If kwargs contains IDs and documents contain ids, the IDs in the kwargs242                will receive precedence.243 244        Returns:245            List of IDs of the added texts.246        """247        if type(self).add_texts != VectorStore.add_texts:248            if "ids" not in kwargs:249                ids = [doc.id for doc in documents]250 251                # If there's at least one valid ID, we'll assume that IDs252                # should be used.253                if any(ids):254                    kwargs["ids"] = ids255 256            texts = [doc.page_content for doc in documents]257            metadatas = [doc.metadata for doc in documents]258            return self.add_texts(texts, metadatas, **kwargs)259        msg = (260            f"`add_documents` and `add_texts` has not been implemented "261            f"for {self.__class__.__name__} "262        )263        raise NotImplementedError(msg)264 265    async def aadd_documents(266        self, documents: list[Document], **kwargs: Any267    ) -> list[str]:268        """Async run more documents through the embeddings and add to the `VectorStore`.269 270        Args:271            documents: Documents to add to the `VectorStore`.272            **kwargs: Additional keyword arguments.273 274        Returns:275            List of IDs of the added texts.276        """277        # If the async method has been overridden, we'll use that.278        if type(self).aadd_texts != VectorStore.aadd_texts:279            if "ids" not in kwargs:280                ids = [doc.id for doc in documents]281 282                # If there's at least one valid ID, we'll assume that IDs283                # should be used.284                if any(ids):285                    kwargs["ids"] = ids286 287            texts = [doc.page_content for doc in documents]288            metadatas = [doc.metadata for doc in documents]289            return await self.aadd_texts(texts, metadatas, **kwargs)290 291        return await run_in_executor(None, self.add_documents, documents, **kwargs)292 293    def search(self, query: str, search_type: str, **kwargs: Any) -> list[Document]:294        """Return docs most similar to query using a specified search type.295 296        Args:297            query: Input text.298            search_type: Type of search to perform.299 300                Can be `'similarity'`, `'mmr'`, or `'similarity_score_threshold'`.301            **kwargs: Arguments to pass to the search method.302 303        Returns:304            List of `Document` objects most similar to the query.305 306        Raises:307            ValueError: If `search_type` is not one of `'similarity'`,308                `'mmr'`, or `'similarity_score_threshold'`.309        """310        if search_type == "similarity":311            return self.similarity_search(query, **kwargs)312        if search_type == "similarity_score_threshold":313            docs_and_similarities = self.similarity_search_with_relevance_scores(314                query, **kwargs315            )316            return [doc for doc, _ in docs_and_similarities]317        if search_type == "mmr":318            return self.max_marginal_relevance_search(query, **kwargs)319        msg = (320            f"search_type of {search_type} not allowed. Expected "321            "search_type to be 'similarity', 'similarity_score_threshold'"322            " or 'mmr'."323        )324        raise ValueError(msg)325 326    async def asearch(327        self, query: str, search_type: str, **kwargs: Any328    ) -> list[Document]:329        """Async return docs most similar to query using a specified search type.330 331        Args:332            query: Input text.333            search_type: Type of search to perform.334 335                Can be `'similarity'`, `'mmr'`, or `'similarity_score_threshold'`.336            **kwargs: Arguments to pass to the search method.337 338        Returns:339            List of `Document` objects most similar to the query.340 341        Raises:342            ValueError: If `search_type` is not one of `'similarity'`,343                `'mmr'`, or `'similarity_score_threshold'`.344        """345        if search_type == "similarity":346            return await self.asimilarity_search(query, **kwargs)347        if search_type == "similarity_score_threshold":348            docs_and_similarities = await self.asimilarity_search_with_relevance_scores(349                query, **kwargs350            )351            return [doc for doc, _ in docs_and_similarities]352        if search_type == "mmr":353            return await self.amax_marginal_relevance_search(query, **kwargs)354        msg = (355            f"search_type of {search_type} not allowed. Expected "356            "search_type to be 'similarity', 'similarity_score_threshold' or 'mmr'."357        )358        raise ValueError(msg)359 360    @abstractmethod361    def similarity_search(362        self, query: str, k: int = 4, **kwargs: Any363    ) -> list[Document]:364        """Return docs most similar to query.365 366        Args:367            query: Input text.368            k: Number of `Document` objects to return.369            **kwargs: Arguments to pass to the search method.370 371        Returns:372            List of `Document` objects most similar to the query.373        """374 375    @staticmethod376    def _euclidean_relevance_score_fn(distance: float) -> float:377        """Return a similarity score on a scale [0, 1]."""378        # The 'correct' relevance function379        # may differ depending on a few things, including:380        # - the distance / similarity metric used by the VectorStore381        # - the scale of your embeddings (OpenAI's are unit normed. Many382        #  others are not!)383        # - embedding dimensionality384        # - etc.385        # This function converts the Euclidean norm of normalized embeddings386        # (0 is most similar, sqrt(2) most dissimilar)387        # to a similarity function (0 to 1)388        return 1.0 - distance / math.sqrt(2)389 390    @staticmethod391    def _cosine_relevance_score_fn(distance: float) -> float:392        """Normalize the distance to a score on a scale [0, 1]."""393        return 1.0 - distance394 395    @staticmethod396    def _max_inner_product_relevance_score_fn(distance: float) -> float:397        """Normalize the distance to a score on a scale [0, 1]."""398        if distance > 0:399            return 1.0 - distance400 401        return -1.0 * distance402 403    def _select_relevance_score_fn(self) -> Callable[[float], float]:404        """The 'correct' relevance function.405 406        May differ depending on a few things, including:407 408        - The distance / similarity metric used by the VectorStore409        - The scale of your embeddings (OpenAI's are unit normed. Many others are not!)410        - Embedding dimensionality411        - etc.412 413        Vectorstores should define their own selection-based method of relevance.414        """415        raise NotImplementedError416 417    def similarity_search_with_score(418        self, *args: Any, **kwargs: Any419    ) -> list[tuple[Document, float]]:420        """Run similarity search with distance.421 422        Args:423            *args: Arguments to pass to the search method.424            **kwargs: Arguments to pass to the search method.425 426        Returns:427            List of tuples of `(doc, similarity_score)`.428        """429        raise NotImplementedError430 431    async def asimilarity_search_with_score(432        self, *args: Any, **kwargs: Any433    ) -> list[tuple[Document, float]]:434        """Async run similarity search with distance.435 436        Args:437            *args: Arguments to pass to the search method.438            **kwargs: Arguments to pass to the search method.439 440        Returns:441            List of tuples of `(doc, similarity_score)`.442        """443        # This is a temporary workaround to make the similarity search444        # asynchronous. The proper solution is to make the similarity search445        # asynchronous in the vector store implementations.446        return await run_in_executor(447            None, self.similarity_search_with_score, *args, **kwargs448        )449 450    def _similarity_search_with_relevance_scores(451        self,452        query: str,453        k: int = 4,454        **kwargs: Any,455    ) -> list[tuple[Document, float]]:456        """Default similarity search with relevance scores.457 458        Modify if necessary in subclass.459        Return docs and relevance scores in the range `[0, 1]`.460 461        `0` is dissimilar, `1` is most similar.462 463        Args:464            query: Input text.465            k: Number of `Document` objects to return.466            **kwargs: Kwargs to be passed to similarity search.467 468                Should include `score_threshold`, an optional floating point value469                between `0` to `1` to filter the resulting set of retrieved docs.470 471        Returns:472            List of tuples of `(doc, similarity_score)`473        """474        relevance_score_fn = self._select_relevance_score_fn()475        docs_and_scores = self.similarity_search_with_score(query, k, **kwargs)476        return [(doc, relevance_score_fn(score)) for doc, score in docs_and_scores]477 478    async def _asimilarity_search_with_relevance_scores(479        self,480        query: str,481        k: int = 4,482        **kwargs: Any,483    ) -> list[tuple[Document, float]]:484        """Default similarity search with relevance scores.485 486        Modify if necessary in subclass.487        Return docs and relevance scores in the range `[0, 1]`.488 489        `0` is dissimilar, `1` is most similar.490 491        Args:492            query: Input text.493            k: Number of `Document` objects to return.494            **kwargs: Kwargs to be passed to similarity search.495 496                Should include `score_threshold`, an optional floating point value497                between `0` to `1` to filter the resulting set of retrieved docs.498 499        Returns:500            List of tuples of `(doc, similarity_score)`501        """502        relevance_score_fn = self._select_relevance_score_fn()503        docs_and_scores = await self.asimilarity_search_with_score(query, k, **kwargs)504        return [(doc, relevance_score_fn(score)) for doc, score in docs_and_scores]505 506    def similarity_search_with_relevance_scores(507        self,508        query: str,509        k: int = 4,510        **kwargs: Any,511    ) -> list[tuple[Document, float]]:512        """Return docs and relevance scores in the range `[0, 1]`.513 514        `0` is dissimilar, `1` is most similar.515 516        Args:517            query: Input text.518            k: Number of `Document` objects to return.519            **kwargs: Kwargs to be passed to similarity search.520 521                Should include `score_threshold`, an optional floating point value522                between `0` to `1` to filter the resulting set of retrieved docs.523 524        Returns:525            List of tuples of `(doc, similarity_score)`.526        """527        score_threshold = kwargs.pop("score_threshold", None)528 529        docs_and_similarities = self._similarity_search_with_relevance_scores(530            query, k=k, **kwargs531        )532        if any(533            similarity < 0.0 or similarity > 1.0534            for _, similarity in docs_and_similarities535        ):536            warnings.warn(537                "Relevance scores must be between"538                f" 0 and 1, got {docs_and_similarities}",539                stacklevel=2,540            )541 542        if score_threshold is not None:543            docs_and_similarities = [544                (doc, similarity)545                for doc, similarity in docs_and_similarities546                if similarity >= score_threshold547            ]548            if len(docs_and_similarities) == 0:549                logger.warning(550                    "No relevant docs were retrieved using the "551                    "relevance score threshold %s",552                    score_threshold,553                )554        return docs_and_similarities555 556    async def asimilarity_search_with_relevance_scores(557        self,558        query: str,559        k: int = 4,560        **kwargs: Any,561    ) -> list[tuple[Document, float]]:562        """Async return docs and relevance scores in the range `[0, 1]`.563 564        `0` is dissimilar, `1` is most similar.565 566        Args:567            query: Input text.568            k: Number of `Document` objects to return.569            **kwargs: Kwargs to be passed to similarity search.570 571                Should include `score_threshold`, an optional floating point value572                between `0` to `1` to filter the resulting set of retrieved docs.573 574        Returns:575            List of tuples of `(doc, similarity_score)`576        """577        score_threshold = kwargs.pop("score_threshold", None)578 579        docs_and_similarities = await self._asimilarity_search_with_relevance_scores(580            query, k=k, **kwargs581        )582        if any(583            similarity < 0.0 or similarity > 1.0584            for _, similarity in docs_and_similarities585        ):586            warnings.warn(587                "Relevance scores must be between"588                f" 0 and 1, got {docs_and_similarities}",589                stacklevel=2,590            )591 592        if score_threshold is not None:593            docs_and_similarities = [594                (doc, similarity)595                for doc, similarity in docs_and_similarities596                if similarity >= score_threshold597            ]598            if len(docs_and_similarities) == 0:599                logger.warning(600                    "No relevant docs were retrieved using the "601                    "relevance score threshold %s",602                    score_threshold,603                )604        return docs_and_similarities605 606    async def asimilarity_search(607        self, query: str, k: int = 4, **kwargs: Any608    ) -> list[Document]:609        """Async return docs most similar to query.610 611        Args:612            query: Input text.613            k: Number of `Document` objects to return.614            **kwargs: Arguments to pass to the search method.615 616        Returns:617            List of `Document` objects most similar to the query.618        """619        # This is a temporary workaround to make the similarity search620        # asynchronous. The proper solution is to make the similarity search621        # asynchronous in the vector store implementations.622        return await run_in_executor(None, self.similarity_search, query, k=k, **kwargs)623 624    def similarity_search_by_vector(625        self, embedding: list[float], k: int = 4, **kwargs: Any626    ) -> list[Document]:627        """Return docs most similar to embedding vector.628 629        Args:630            embedding: Embedding to look up documents similar to.631            k: Number of `Document` objects to return.632            **kwargs: Arguments to pass to the search method.633 634        Returns:635            List of `Document` objects most similar to the query vector.636        """637        raise NotImplementedError638 639    async def asimilarity_search_by_vector(640        self, embedding: list[float], k: int = 4, **kwargs: Any641    ) -> list[Document]:642        """Async return docs most similar to embedding vector.643 644        Args:645            embedding: Embedding to look up documents similar to.646            k: Number of `Document` objects to return.647            **kwargs: Arguments to pass to the search method.648 649        Returns:650            List of `Document` objects most similar to the query vector.651        """652        # This is a temporary workaround to make the similarity search653        # asynchronous. The proper solution is to make the similarity search654        # asynchronous in the vector store implementations.655        return await run_in_executor(656            None, self.similarity_search_by_vector, embedding, k=k, **kwargs657        )658 659    def max_marginal_relevance_search(660        self,661        query: str,662        k: int = 4,663        fetch_k: int = 20,664        lambda_mult: float = 0.5,665        **kwargs: Any,666    ) -> list[Document]:667        """Return docs selected using the maximal marginal relevance.668 669        Maximal marginal relevance optimizes for similarity to query AND diversity670        among selected documents.671 672        Args:673            query: Text to look up documents similar to.674            k: Number of `Document` objects to return.675            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.676            lambda_mult: Number between `0` and `1` that determines the degree of677                diversity among the results with `0` corresponding to maximum diversity678                and `1` to minimum diversity.679            **kwargs: Arguments to pass to the search method.680 681        Returns:682            List of `Document` objects selected by maximal marginal relevance.683        """684        raise NotImplementedError685 686    async def amax_marginal_relevance_search(687        self,688        query: str,689        k: int = 4,690        fetch_k: int = 20,691        lambda_mult: float = 0.5,692        **kwargs: Any,693    ) -> list[Document]:694        """Async return docs selected using the maximal marginal relevance.695 696        Maximal marginal relevance optimizes for similarity to query AND diversity697        among selected documents.698 699        Args:700            query: Text to look up documents similar to.701            k: Number of `Document` objects to return.702            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.703            lambda_mult: Number between `0` and `1` that determines the degree of704                diversity among the results with `0` corresponding to maximum diversity705                and `1` to minimum diversity.706            **kwargs: Arguments to pass to the search method.707 708        Returns:709            List of `Document` objects selected by maximal marginal relevance.710        """711        # This is a temporary workaround to make the similarity search712        # asynchronous. The proper solution is to make the similarity search713        # asynchronous in the vector store implementations.714        return await run_in_executor(715            None,716            self.max_marginal_relevance_search,717            query,718            k=k,719            fetch_k=fetch_k,720            lambda_mult=lambda_mult,721            **kwargs,722        )723 724    def max_marginal_relevance_search_by_vector(725        self,726        embedding: list[float],727        k: int = 4,728        fetch_k: int = 20,729        lambda_mult: float = 0.5,730        **kwargs: Any,731    ) -> list[Document]:732        """Return docs selected using the maximal marginal relevance.733 734        Maximal marginal relevance optimizes for similarity to query AND diversity735        among selected documents.736 737        Args:738            embedding: Embedding to look up documents similar to.739            k: Number of `Document` objects to return.740            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.741            lambda_mult: Number between `0` and `1` that determines the degree of742                diversity among the results with `0` corresponding to maximum diversity743                and `1` to minimum diversity.744            **kwargs: Arguments to pass to the search method.745 746        Returns:747            List of `Document` objects selected by maximal marginal relevance.748        """749        raise NotImplementedError750 751    async def amax_marginal_relevance_search_by_vector(752        self,753        embedding: list[float],754        k: int = 4,755        fetch_k: int = 20,756        lambda_mult: float = 0.5,757        **kwargs: Any,758    ) -> list[Document]:759        """Async return docs selected using the maximal marginal relevance.760 761        Maximal marginal relevance optimizes for similarity to query AND diversity762        among selected documents.763 764        Args:765            embedding: Embedding to look up documents similar to.766            k: Number of `Document` objects to return.767            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.768            lambda_mult: Number between `0` and `1` that determines the degree of769                diversity among the results with `0` corresponding to maximum diversity770                and `1` to minimum diversity.771            **kwargs: Arguments to pass to the search method.772 773        Returns:774            List of `Document` objects selected by maximal marginal relevance.775        """776        return await run_in_executor(777            None,778            self.max_marginal_relevance_search_by_vector,779            embedding,780            k=k,781            fetch_k=fetch_k,782            lambda_mult=lambda_mult,783            **kwargs,784        )785 786    @classmethod787    def from_documents(788        cls,789        documents: list[Document],790        embedding: Embeddings,791        **kwargs: Any,792    ) -> Self:793        """Return `VectorStore` initialized from documents and embeddings.794 795        Args:796            documents: List of `Document` objects to add to the `VectorStore`.797            embedding: Embedding function to use.798            **kwargs: Additional keyword arguments.799 800        Returns:801            `VectorStore` initialized from documents and embeddings.802        """803        texts = [d.page_content for d in documents]804        metadatas = [d.metadata for d in documents]805 806        if "ids" not in kwargs:807            ids = [doc.id for doc in documents]808 809            # If there's at least one valid ID, we'll assume that IDs810            # should be used.811            if any(ids):812                kwargs["ids"] = ids813 814        return cls.from_texts(texts, embedding, metadatas=metadatas, **kwargs)815 816    @classmethod817    async def afrom_documents(818        cls,819        documents: list[Document],820        embedding: Embeddings,821        **kwargs: Any,822    ) -> Self:823        """Async return `VectorStore` initialized from documents and embeddings.824 825        Args:826            documents: List of `Document` objects to add to the `VectorStore`.827            embedding: Embedding function to use.828            **kwargs: Additional keyword arguments.829 830        Returns:831            `VectorStore` initialized from documents and embeddings.832        """833        texts = [d.page_content for d in documents]834        metadatas = [d.metadata for d in documents]835 836        if "ids" not in kwargs:837            ids = [doc.id for doc in documents]838 839            # If there's at least one valid ID, we'll assume that IDs840            # should be used.841            if any(ids):842                kwargs["ids"] = ids843 844        return await cls.afrom_texts(texts, embedding, metadatas=metadatas, **kwargs)845 846    @classmethod847    @abstractmethod848    def from_texts(849        cls: type[VST],850        texts: list[str],851        embedding: Embeddings,852        metadatas: list[dict] | None = None,853        *,854        ids: list[str] | None = None,855        **kwargs: Any,856    ) -> VST:857        """Return `VectorStore` initialized from texts and embeddings.858 859        Args:860            texts: Texts to add to the `VectorStore`.861            embedding: Embedding function to use.862            metadatas: Optional list of metadatas associated with the texts.863            ids: Optional list of IDs associated with the texts.864            **kwargs: Additional keyword arguments.865 866        Returns:867            `VectorStore` initialized from texts and embeddings.868        """869 870    @classmethod871    async def afrom_texts(872        cls,873        texts: list[str],874        embedding: Embeddings,875        metadatas: list[dict] | None = None,876        *,877        ids: list[str] | None = None,878        **kwargs: Any,879    ) -> Self:880        """Async return `VectorStore` initialized from texts and embeddings.881 882        Args:883            texts: Texts to add to the `VectorStore`.884            embedding: Embedding function to use.885            metadatas: Optional list of metadatas associated with the texts.886            ids: Optional list of IDs associated with the texts.887            **kwargs: Additional keyword arguments.888 889        Returns:890            `VectorStore` initialized from texts and embeddings.891        """892        if ids is not None:893            kwargs["ids"] = ids894        return await run_in_executor(895            None, cls.from_texts, texts, embedding, metadatas, **kwargs896        )897 898    def _get_retriever_tags(self) -> list[str]:899        """Get tags for retriever."""900        tags = [self.__class__.__name__]901        if self.embeddings:902            tags.append(self.embeddings.__class__.__name__)903        return tags904 905    def as_retriever(self, **kwargs: Any) -> VectorStoreRetriever:906        """Return `VectorStoreRetriever` initialized from this `VectorStore`.907 908        Args:909            **kwargs: Keyword arguments to pass to the search function.910 911                Can include:912 913                * `search_type`: Defines the type of search that the Retriever should914                    perform. Can be `'similarity'` (default), `'mmr'`, or915                    `'similarity_score_threshold'`.916                * `search_kwargs`: Keyword arguments to pass to the search function.917 918                    Can include things like:919 920                    * `k`: Amount of documents to return (Default: `4`)921                    * `score_threshold`: Minimum relevance threshold922                        for `similarity_score_threshold`923                    * `fetch_k`: Amount of documents to pass to MMR algorithm924                        (Default: `20`)925                    * `lambda_mult`: Diversity of results returned by MMR;926                        `1` for minimum diversity and 0 for maximum. (Default: `0.5`)927                    * `filter`: Filter by document metadata928 929        Returns:930            Retriever class for `VectorStore`.931 932        Examples:933        ```python934        # Retrieve more documents with higher diversity935        # Useful if your dataset has many similar documents936        docsearch.as_retriever(937            search_type="mmr", search_kwargs={"k": 6, "lambda_mult": 0.25}938        )939 940        # Fetch more documents for the MMR algorithm to consider941        # But only return the top 5942        docsearch.as_retriever(search_type="mmr", search_kwargs={"k": 5, "fetch_k": 50})943 944        # Only retrieve documents that have a relevance score945        # Above a certain threshold946        docsearch.as_retriever(947            search_type="similarity_score_threshold",948            search_kwargs={"score_threshold": 0.8},949        )950 951        # Only get the single most similar document from the dataset952        docsearch.as_retriever(search_kwargs={"k": 1})953 954        # Use a filter to only retrieve documents from a specific paper955        docsearch.as_retriever(956            search_kwargs={"filter": {"paper_title": "GPT-4 Technical Report"}}957        )958        ```959        """960        tags = kwargs.pop("tags", None) or [*self._get_retriever_tags()]961        return VectorStoreRetriever(vectorstore=self, tags=tags, **kwargs)962 963 964class VectorStoreRetriever(BaseRetriever):965    """Base Retriever class for VectorStore."""966 967    vectorstore: VectorStore968    """VectorStore to use for retrieval."""969 970    search_type: str = "similarity"971    """Type of search to perform."""972 973    search_kwargs: dict = Field(default_factory=dict)974    """Keyword arguments to pass to the search function."""975 976    allowed_search_types: ClassVar[Collection[str]] = (977        "similarity",978        "similarity_score_threshold",979        "mmr",980    )981 982    model_config = ConfigDict(983        arbitrary_types_allowed=True,984    )985 986    @model_validator(mode="before")987    @classmethod988    def validate_search_type(cls, values: dict) -> Any:989        """Validate search type.990 991        Args:992            values: Values to validate.993 994        Returns:995            Validated values.996 997        Raises:998            ValueError: If `search_type` is not one of the allowed search types.999            ValueError: If `score_threshold` is not specified with a float value(`0~1`)1000        """1001        search_type = values.get("search_type", "similarity")1002        if search_type not in cls.allowed_search_types:1003            msg = (1004                f"search_type of {search_type} not allowed. Valid values are: "1005                f"{cls.allowed_search_types}"1006            )1007            raise ValueError(msg)1008        if search_type == "similarity_score_threshold":1009            score_threshold = values.get("search_kwargs", {}).get("score_threshold")1010            if (score_threshold is None) or (not isinstance(score_threshold, float)):1011                msg = (1012                    "`score_threshold` is not specified with a float value(0~1) "1013                    "in `search_kwargs`."1014                )1015                raise ValueError(msg)1016        return values1017 1018    def _get_ls_params(self, **kwargs: Any) -> LangSmithRetrieverParams:1019        """Get standard params for tracing."""1020        kwargs_ = self.search_kwargs | kwargs1021 1022        ls_params = super()._get_ls_params(**kwargs_)1023 1024        ls_params["ls_vector_store_provider"] = self.vectorstore.__class__.__name__1025 1026        if self.vectorstore.embeddings:1027            ls_params["ls_embedding_provider"] = (1028                self.vectorstore.embeddings.__class__.__name__1029            )1030        elif hasattr(self.vectorstore, "embedding") and isinstance(1031            self.vectorstore.embedding, Embeddings1032        ):1033            ls_params["ls_embedding_provider"] = (1034                self.vectorstore.embedding.__class__.__name__1035            )1036 1037        return ls_params1038 1039    @override1040    def _get_relevant_documents(1041        self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any1042    ) -> list[Document]:1043        kwargs_ = self.search_kwargs | kwargs1044        if self.search_type == "similarity":1045            docs = self.vectorstore.similarity_search(query, **kwargs_)1046        elif self.search_type == "similarity_score_threshold":1047            docs_and_similarities = (1048                self.vectorstore.similarity_search_with_relevance_scores(1049                    query, **kwargs_1050                )1051            )1052            docs = [doc for doc, _ in docs_and_similarities]1053        elif self.search_type == "mmr":1054            docs = self.vectorstore.max_marginal_relevance_search(query, **kwargs_)1055        else:1056            msg = f"search_type of {self.search_type} not allowed."1057            raise ValueError(msg)1058        return docs1059 1060    @override1061    async def _aget_relevant_documents(1062        self,1063        query: str,1064        *,1065        run_manager: AsyncCallbackManagerForRetrieverRun,1066        **kwargs: Any,1067    ) -> list[Document]:1068        kwargs_ = self.search_kwargs | kwargs1069        if self.search_type == "similarity":1070            docs = await self.vectorstore.asimilarity_search(query, **kwargs_)1071        elif self.search_type == "similarity_score_threshold":1072            docs_and_similarities = (1073                await self.vectorstore.asimilarity_search_with_relevance_scores(1074                    query, **kwargs_1075                )1076            )1077            docs = [doc for doc, _ in docs_and_similarities]1078        elif self.search_type == "mmr":1079            docs = await self.vectorstore.amax_marginal_relevance_search(1080                query, **kwargs_1081            )1082        else:1083            msg = f"search_type of {self.search_type} not allowed."1084            raise ValueError(msg)1085        return docs1086 1087    def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:1088        """Add documents to the `VectorStore`.1089 1090        Args:1091            documents: Documents to add to the `VectorStore`.1092            **kwargs: Other keyword arguments that subclasses might use.1093 1094        Returns:1095            List of IDs of the added texts.1096        """1097        return self.vectorstore.add_documents(documents, **kwargs)1098 1099    async def aadd_documents(1100        self, documents: list[Document], **kwargs: Any1101    ) -> list[str]:1102        """Async add documents to the `VectorStore`.1103 1104        Args:1105            documents: Documents to add to the `VectorStore`.1106            **kwargs: Other keyword arguments that subclasses might use.1107 1108        Returns:1109            List of IDs of the added texts.1110        """1111        return await self.vectorstore.aadd_documents(documents, **kwargs)1112 
codekingpro/portable-devtools · Team Ai