codekingpro/portable-devtools

1"""Wrapper around in-memory storage."""2 3from __future__ import annotations4 5from typing import Any, Dict, List, Literal, Optional6 7from langchain_core.embeddings import Embeddings8 9from langchain_community.vectorstores.docarray.base import (10    DocArrayIndex,11    _check_docarray_import,12)13 14 15class DocArrayInMemorySearch(DocArrayIndex):16    """In-memory `DocArray` storage for exact search.17 18    To use it, you should have the ``docarray`` package with version >=0.32.0 installed.19    You can install it with `pip install docarray`.20    """21 22    @classmethod23    def from_params(24        cls,25        embedding: Embeddings,26        metric: Literal[27            "cosine_sim", "euclidian_dist", "sgeuclidean_dist"28        ] = "cosine_sim",29        **kwargs: Any,30    ) -> DocArrayInMemorySearch:31        """Initialize DocArrayInMemorySearch store.32 33        Args:34            embedding (Embeddings): Embedding function.35            metric (str): metric for exact nearest-neighbor search.36                Can be one of: "cosine_sim", "euclidean_dist" and "sqeuclidean_dist".37                Defaults to "cosine_sim".38            **kwargs: Other keyword arguments to be passed to the get_doc_cls method.39        """40        _check_docarray_import()41        from docarray.index import InMemoryExactNNIndex42 43        doc_cls = cls._get_doc_cls(space=metric, **kwargs)44        doc_index = InMemoryExactNNIndex[doc_cls]()45        return cls(doc_index, embedding)46 47    @classmethod48    def from_texts(49        cls,50        texts: List[str],51        embedding: Embeddings,52        metadatas: Optional[List[Dict[Any, Any]]] = None,53        **kwargs: Any,54    ) -> DocArrayInMemorySearch:55        """Create an DocArrayInMemorySearch store and insert data.56 57        Args:58            texts (List[str]): Text data.59            embedding (Embeddings): Embedding function.60            metadatas (Optional[List[Dict[Any, Any]]]): Metadata for each text61                if it exists. Defaults to None.62            **kwargs: Other keyword arguments to be passed to the from_params method.63 64        Returns:65            DocArrayInMemorySearch Vector Store66        """67        store = cls.from_params(embedding, **kwargs)68        store.add_texts(texts=texts, metadatas=metadatas)69        return store70