codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from enum import Enum5from typing import (6    TYPE_CHECKING,7    Any,8    Dict,9    Generator,10    Iterable,11    List,12    Optional,13    TypeVar,14    Union,15)16 17from langchain_core.documents import Document18from langchain_core.vectorstores import VectorStore19 20if TYPE_CHECKING:21    from langchain_core.embeddings import Embeddings22    from pymongo.collection import Collection23 24 25# Before Python 3.11 native StrEnum is not available26class DocumentDBSimilarityType(str, Enum):27    """DocumentDB Similarity Type as enumerator."""28 29    COS = "cosine"30    """Cosine similarity"""31    DOT = "dotProduct"32    """Dot product"""33    EUC = "euclidean"34    """Euclidean distance"""35 36 37DocumentDBDocumentType = TypeVar("DocumentDBDocumentType", bound=Dict[str, Any])38 39logger = logging.getLogger(__name__)40 41DEFAULT_INSERT_BATCH_SIZE = 12842 43 44class DocumentDBVectorSearch(VectorStore):45    """`Amazon DocumentDB (with MongoDB compatibility)` vector store.46    Please refer to the official Vector Search documentation for more details:47    https://docs.aws.amazon.com/documentdb/latest/developerguide/vector-search.html48 49    To use, you should have both:50    - the ``pymongo`` python package installed51    - a connection string and credentials associated with a DocumentDB cluster52 53    Example:54        . code-block:: python55 56            from langchain_community.vectorstores import DocumentDBVectorSearch57            from langchain_community.embeddings.openai import OpenAIEmbeddings58            from pymongo import MongoClient59 60            mongo_client = MongoClient("<YOUR-CONNECTION-STRING>")61            collection = mongo_client["<db_name>"]["<collection_name>"]62            embeddings = OpenAIEmbeddings()63            vectorstore = DocumentDBVectorSearch(collection, embeddings)64    """65 66    def __init__(67        self,68        collection: Collection[DocumentDBDocumentType],69        embedding: Embeddings,70        *,71        index_name: str = "vectorSearchIndex",72        text_key: str = "textContent",73        embedding_key: str = "vectorContent",74    ):75        """Constructor for DocumentDBVectorSearch76 77        Args:78            collection: MongoDB collection to add the texts to.79            embedding: Text embedding model to use.80            index_name: Name of the Vector Search index.81            text_key: MongoDB field that will contain the text82                for each document.83            embedding_key: MongoDB field that will contain the embedding84                for each document.85        """86        self._collection = collection87        self._embedding = embedding88        self._index_name = index_name89        self._text_key = text_key90        self._embedding_key = embedding_key91        self._similarity_type = DocumentDBSimilarityType.COS92 93    @property94    def embeddings(self) -> Embeddings:95        return self._embedding96 97    def get_index_name(self) -> str:98        """Returns the index name99 100        Returns:101            Returns the index name102 103        """104        return self._index_name105 106    @classmethod107    def from_connection_string(108        cls,109        connection_string: str,110        namespace: str,111        embedding: Embeddings,112        **kwargs: Any,113    ) -> DocumentDBVectorSearch:114        """Creates an Instance of DocumentDBVectorSearch from a Connection String115 116        Args:117            connection_string: The DocumentDB cluster endpoint connection string118            namespace: The namespace (database.collection)119            embedding: The embedding utility120            **kwargs: Dynamic keyword arguments121 122        Returns:123            an instance of the vector store124 125        """126        try:127            from pymongo import MongoClient128        except ImportError:129            raise ImportError(130                "Could not import pymongo, please install it with "131                "`pip install pymongo`."132            )133        client: MongoClient = MongoClient(connection_string)134        db_name, collection_name = namespace.split(".")135        collection = client[db_name][collection_name]136        return cls(collection, embedding, **kwargs)137 138    def index_exists(self) -> bool:139        """Verifies if the specified index name during instance140            construction exists on the collection141 142        Returns:143          Returns True on success and False if no such index exists144            on the collection145        """146        cursor = self._collection.list_indexes()147        index_name = self._index_name148 149        for res in cursor:150            current_index_name = res.pop("name")151            if current_index_name == index_name:152                return True153 154        return False155 156    def delete_index(self) -> None:157        """Deletes the index specified during instance construction if it exists"""158        if self.index_exists():159            self._collection.drop_index(self._index_name)160            # Raises OperationFailure on an error (e.g. trying to drop161            # an index that does not exist)162 163    def create_index(164        self,165        dimensions: int = 1536,166        similarity: DocumentDBSimilarityType = DocumentDBSimilarityType.COS,167        m: int = 16,168        ef_construction: int = 64,169    ) -> dict[str, Any]:170        """Creates an index using the index name specified at171            instance construction172 173        Args:174            dimensions: Number of dimensions for vector similarity.175                The maximum number of supported dimensions is 2000176 177            similarity: Similarity algorithm to use with the HNSW index.178                 Possible options are:179                    - DocumentDBSimilarityType.COS (cosine distance),180                    - DocumentDBSimilarityType.EUC (Euclidean distance), and181                    - DocumentDBSimilarityType.DOT (dot product).182 183            m: Specifies the max number of connections for an HNSW index.184                Large impact on memory consumption.185 186            ef_construction: Specifies the size of the dynamic candidate list187                for constructing the graph for HNSW index. Higher values lead188                to more accurate results but slower indexing speed.189 190 191        Returns:192            An object describing the created index193 194        """195        self._similarity_type = similarity196 197        # prepare the command198        create_index_commands = {199            "createIndexes": self._collection.name,200            "indexes": [201                {202                    "name": self._index_name,203                    "key": {self._embedding_key: "vector"},204                    "vectorOptions": {205                        "type": "hnsw",206                        "similarity": similarity,207                        "dimensions": dimensions,208                        "m": m,209                        "efConstruction": ef_construction,210                    },211                }212            ],213        }214 215        # retrieve the database object216        current_database = self._collection.database217 218        # invoke the command from the database object219        create_index_responses: dict[str, Any] = current_database.command(220            create_index_commands221        )222 223        return create_index_responses224 225    def add_texts(226        self,227        texts: Iterable[str],228        metadatas: Optional[List[Dict[str, Any]]] = None,229        **kwargs: Any,230    ) -> List:231        batch_size = kwargs.get("batch_size", DEFAULT_INSERT_BATCH_SIZE)232        _metadatas: Union[List, Generator] = metadatas or ({} for _ in texts)233        texts_batch = []234        metadatas_batch = []235        result_ids = []236        for i, (text, metadata) in enumerate(zip(texts, _metadatas)):237            texts_batch.append(text)238            metadatas_batch.append(metadata)239            if (i + 1) % batch_size == 0:240                result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))241                texts_batch = []242                metadatas_batch = []243        if texts_batch:244            result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))245        return result_ids246 247    def _insert_texts(self, texts: List[str], metadatas: List[Dict[str, Any]]) -> List:248        """Used to Load Documents into the collection249 250        Args:251            texts: The list of documents strings to load252            metadatas: The list of metadata objects associated with each document253 254        Returns:255 256        """257        # If the text is empty, then exit early258        if not texts:259            return []260 261        # Embed and create the documents262        embeddings = self._embedding.embed_documents(texts)263        to_insert = [264            {self._text_key: t, self._embedding_key: embedding, **m}265            for t, m, embedding in zip(texts, metadatas, embeddings)266        ]267        # insert the documents in DocumentDB268        insert_result = self._collection.insert_many(to_insert)269        return insert_result.inserted_ids270 271    @classmethod272    def from_texts(273        cls,274        texts: List[str],275        embedding: Embeddings,276        metadatas: Optional[List[dict]] = None,277        collection: Optional[Collection[DocumentDBDocumentType]] = None,278        **kwargs: Any,279    ) -> DocumentDBVectorSearch:280        if collection is None:281            raise ValueError("Must provide 'collection' named parameter.")282        vectorstore = cls(collection, embedding, **kwargs)283        vectorstore.add_texts(texts, metadatas=metadatas)284        return vectorstore285 286    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:287        if ids is None:288            raise ValueError("No document ids provided to delete.")289 290        for document_id in ids:291            self.delete_document_by_id(document_id)292        return True293 294    def delete_document_by_id(self, document_id: Optional[str] = None) -> None:295        """Removes a Specific Document by Id296 297        Args:298            document_id: The document identifier299        """300        try:301            from bson.objectid import ObjectId302        except ImportError as e:303            raise ImportError(304                "Unable to import bson, please install with `pip install bson`."305            ) from e306        if document_id is None:307            raise ValueError("No document id provided to delete.")308 309        self._collection.delete_one({"_id": ObjectId(document_id)})310 311    def _similarity_search_without_score(312        self,313        embeddings: List[float],314        k: int = 4,315        ef_search: int = 40,316        filter: Optional[Dict[str, Any]] = None,317    ) -> List[Document]:318        """Returns a list of documents.319 320        Args:321            embeddings: The query vector322            k: the number of documents to return323            ef_search: Specifies the size of the dynamic candidate list324                that HNSW index uses during search. A higher value of325                efSearch provides better recall at cost of speed.326            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.327        Returns:328            A list of documents closest to the query vector329        """330        # $match can't be null, so initializes to {} when None to avoid331        # "the match filter must be an expression in an object"332        if not filter:333            filter = {}334        pipeline: List[dict[str, Any]] = [335            {"$match": filter},336            {337                "$search": {338                    "vectorSearch": {339                        "vector": embeddings,340                        "path": self._embedding_key,341                        "similarity": self._similarity_type,342                        "k": k,343                        "efSearch": ef_search,344                    }345                },346            },347        ]348 349        cursor = self._collection.aggregate(pipeline)350 351        docs = []352 353        for res in cursor:354            text = res.pop(self._text_key)355            docs.append(Document(page_content=text, metadata=res))356 357        return docs358 359    def similarity_search(360        self,361        query: str,362        k: int = 4,363        ef_search: int = 40,364        *,365        filter: Optional[Dict[str, Any]] = None,366        **kwargs: Any,367    ) -> List[Document]:368        embeddings = self._embedding.embed_query(query)369        docs = self._similarity_search_without_score(370            embeddings=embeddings, k=k, ef_search=ef_search, filter=filter371        )372        return [doc for doc in docs]373