codekingpro/portable-devtools

1# System imports2from __future__ import annotations3 4import logging5import time6import uuid7from typing import Any, Dict, List, Optional, Sequence, Tuple, Type8 9# Third-party imports10import numpy as np11 12# Local imports13from langchain_core.documents import Document14from langchain_core.embeddings import Embeddings15from langchain_core.indexing.base import UpsertResponse16from langchain_core.vectorstores import VectorStore17from typing_extensions import override18 19# Configure some defaults20ENGINE = "HNSW"21METRIC = "CS"22DESCRIPTOR_SET = "langchain"23BATCHSIZE = 100024PROPERTY_PREFIX = "lc_"  # Prefix for properties that are in the client metadata25TEXT_PROPERTY = "text"  # Property name for the text26UNIQUEID_PROPERTY = "uniqueid"  # Property name for the unique id27 28 29class ApertureDB(VectorStore):30    @override31    def __init__(32        self,33        embeddings: Embeddings,34        descriptor_set: str = DESCRIPTOR_SET,35        dimensions: Optional[int] = None,36        engine: Optional[str] = None,37        metric: Optional[str] = None,38        log_level: int = logging.WARN,39        properties: Optional[Dict] = None,40        **kwargs: Any,41    ) -> None:42        """Create a vectorstore backed by ApertureDB43 44        A single ApertureDB instance can support many vectorstores,45        distinguished by 'descriptor_set' name.  The descriptor set is created46        if it does not exist.  Different descriptor sets can use different47        engines and metrics, be supplied by different embedding models, and have48        different dimensions.49 50        See ApertureDB documentation on `AddDescriptorSet`51        https://docs.aperturedata.io/query_language/Reference/descriptor_commands/desc_set_commands/AddDescriptorSet52        for more information on the engine and metric options.53 54        Args:55            embeddings (Embeddings): Embeddings object56            descriptor_set (str, optional): Descriptor set name. Defaults to57                "langchain".58            dimensions (Optional[int], optional): Number of dimensions of the59                embeddings. Defaults to None.60            engine (str, optional): Engine to use. Defaults to "HNSW" for new61                descriptorsets.62            metric (str, optional): Metric to use. Defaults to "CS" for new63                descriptorsets.64            log_level (int, optional): Logging level. Defaults to logging.WARN.65        """66        # ApertureDB imports67        try:68            from aperturedb.Utils import Utils, create_connector69        except ImportError:70            raise ImportError(71                "ApertureDB is not installed. Please install it using "72                "'pip install aperturedb'"73            )74 75        super().__init__(**kwargs)76        self.logger = logging.getLogger(__name__)77        self.logger.setLevel(log_level)78        self.descriptor_set = descriptor_set79 80        self.embedding_function = embeddings81        self.dimensions = dimensions82        self.engine = engine83        self.metric = metric84        self.properties = properties85        if embeddings is None:86            self.logger.fatal("No embedding function provided.")87            raise ValueError("No embedding function provided.")88 89        try:90            from aperturedb.Utils import Utils, create_connector91        except ImportError:92            self.logger.exception(93                "ApertureDB is not installed. Please install it using "94                "'pip install aperturedb'"95            )96            raise97 98        self.connection = create_connector()99        self.utils = Utils(self.connection)100        try:101            self.utils.status()102        except Exception:103            self.logger.exception("Failed to connect to ApertureDB")104            raise105 106        self._find_or_add_descriptor_set()107 108    def _find_or_add_descriptor_set(self) -> None:109        descriptor_set = self.descriptor_set110        """Checks if the descriptor set exists, if not, creates it"""111        find_ds_query = [112            {113                "FindDescriptorSet": {114                    "with_name": descriptor_set,115                    "engines": True,116                    "metrics": True,117                    "dimensions": True,118                    "results": {"all_properties": True},119                }120            }121        ]122        r, b = self.connection.query(find_ds_query)123        assert self.connection.last_query_ok(), r124        n_entities = (125            len(r[0]["FindDescriptorSet"]["entities"])126            if "entities" in r[0]["FindDescriptorSet"]127            else 0128        )129        assert n_entities <= 1, "Multiple descriptor sets with the same name"130 131        if n_entities == 1:  # Descriptor set exists already132            e = r[0]["FindDescriptorSet"]["entities"][0]133            self.logger.info(f"Descriptor set {descriptor_set} already exists")134 135            engines = e["_engines"]136            assert len(engines) == 1, "Only one engine is supported"137 138            if self.engine is None:139                self.engine = engines[0]140            elif self.engine != engines[0]:141                self.logger.error(f"Engine mismatch: {self.engine} != {engines[0]}")142 143            metrics = e["_metrics"]144            assert len(metrics) == 1, "Only one metric is supported"145            if self.metric is None:146                self.metric = metrics[0]147            elif self.metric != metrics[0]:148                self.logger.error(f"Metric mismatch: {self.metric} != {metrics[0]}")149 150            dimensions = e["_dimensions"]151            if self.dimensions is None:152                self.dimensions = dimensions153            elif self.dimensions != dimensions:154                self.logger.error(155                    f"Dimensions mismatch: {self.dimensions} != {dimensions}"156                )157 158            self.properties = {159                k[len(PROPERTY_PREFIX) :]: v160                for k, v in e.items()161                if k.startswith(PROPERTY_PREFIX)162            }163 164        else:165            self.logger.info(166                f"Descriptor set {descriptor_set} does not exist. Creating it"167            )168            if self.engine is None:169                self.engine = ENGINE170            if self.metric is None:171                self.metric = METRIC172            if self.dimensions is None:173                self.dimensions = len(self.embedding_function.embed_query("test"))174 175            properties = (176                {PROPERTY_PREFIX + k: v for k, v in self.properties.items()}177                if self.properties is not None178                else None179            )180 181            self.utils.add_descriptorset(182                name=descriptor_set,183                dim=self.dimensions,184                engine=self.engine,185                metric=self.metric,186                properties=properties,187            )188 189            # Create indexes190            self.utils.create_entity_index("_Descriptor", "_create_txn")191            self.utils.create_entity_index("_DescriptorSet", "_name")192            self.utils.create_entity_index("_Descriptor", UNIQUEID_PROPERTY)193 194    @override195    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:196        """Delete documents from the vectorstore by id.197 198        Args:199            ids: List of ids to delete from the vectorstore.200 201        Returns:202            True if the deletion was successful, False otherwise203        """204        assert ids is not None, "ids must be provided"205        query = [206            {207                "DeleteDescriptor": {208                    "set": self.descriptor_set,209                    "constraints": {UNIQUEID_PROPERTY: ["in", ids]},210                }211            }212        ]213 214        result, _ = self.utils.execute(query)215        return result216 217    @override218    def get_by_ids(self, ids: Sequence[str], /) -> List[Document]:219        """Find documents in the vectorstore by id.220 221        Args:222            ids: List of ids to find in the vectorstore.223 224        Returns:225            documents: List of Document objects found in the vectorstore.226        """227        query = [228            {229                "FindDescriptor": {230                    "set": self.descriptor_set,231                    "constraints": {UNIQUEID_PROPERTY: ["in", ids]},232                    "results": {"all_properties": True},233                }234            }235        ]236 237        results, _ = self.utils.execute(query)238        docs = [239            self._descriptor_to_document(d)240            for d in results[0]["FindDescriptor"].get("entities", [])241        ]242        return docs243 244    @override245    def similarity_search(246        self, query: str, k: int = 4, *args: Any, **kwargs: Any247    ) -> List[Document]:248        """Search for documents similar to the query using the vectorstore249 250        Args:251            query: Query string to search for.252            k: Number of results to return.253 254        Returns:255            List of Document objects ordered by decreasing similarity to the query.256        """257        assert self.embedding_function is not None, "Embedding function is not set"258        embedding = self.embedding_function.embed_query(query)259        return self.similarity_search_by_vector(embedding, k, *args, **kwargs)260 261    @override262    def similarity_search_with_score(263        self, query: str, *args: Any, **kwargs: Any264    ) -> List[Tuple[Document, float]]:265        embedding = self.embedding_function.embed_query(query)266        return self._similarity_search_with_score_by_vector(embedding, *args, **kwargs)267 268    def _descriptor_to_document(self, d: dict) -> Document:269        metadata = {}270        for k, v in d.items():271            if k.startswith(PROPERTY_PREFIX):272                metadata[k[len(PROPERTY_PREFIX) :]] = v273        text = d[TEXT_PROPERTY]274        uniqueid = d[UNIQUEID_PROPERTY]275        doc = Document(page_content=text, metadata=metadata, id=uniqueid)276        return doc277 278    def _similarity_search_with_score_by_vector(279        self, embedding: List[float], k: int = 4, vectors: bool = False280    ) -> List[Tuple[Document, float]]:281        from aperturedb.Descriptors import Descriptors282 283        descriptors = Descriptors(self.connection)284        start_time = time.time()285        descriptors.find_similar(286            set=self.descriptor_set, vector=embedding, k_neighbors=k, distances=True287        )288        self.logger.info(289            f"ApertureDB similarity search took {time.time() - start_time} seconds"290        )291        return [(self._descriptor_to_document(d), d["_distance"]) for d in descriptors]292 293    @override294    def similarity_search_by_vector(295        self, embedding: List[float], k: int = 4, **kwargs: Any296    ) -> List[Document]:297        """Returns the k most similar documents to the given embedding vector298 299        Args:300            embedding: The embedding vector to search for301            k: The number of similar documents to return302 303        Returns:304            List of Document objects ordered by decreasing similarity to the query.305        """306        from aperturedb.Descriptors import Descriptors307 308        descriptors = Descriptors(self.connection)309        start_time = time.time()310        descriptors.find_similar(311            set=self.descriptor_set, vector=embedding, k_neighbors=k312        )313        self.logger.info(314            f"ApertureDB similarity search took {time.time() - start_time} seconds"315        )316        return [self._descriptor_to_document(d) for d in descriptors]317 318    @override319    def max_marginal_relevance_search(320        self,321        query: str,322        k: int = 4,323        fetch_k: int = 20,324        lambda_mult: float = 0.5,325        **kwargs: Any,326    ) -> List[Document]:327        """Returns similar documents to the query that also have diversity328 329        This algorithm balances relevance and diversity in the search results.330 331        Args:332            query: Query string to search for.333            k: Number of results to return.334            fetch_k: Number of results to fetch.335            lambda_mult: Lambda multiplier for MMR.336 337        Returns:338            List of Document objects ordered by decreasing similarity/diversty.339        """340        self.logger.info(f"Max Marginal Relevance search for query: {query}")341        embedding = self.embedding_function.embed_query(query)342        return self.max_marginal_relevance_search_by_vector(343            embedding, k, fetch_k, lambda_mult, **kwargs344        )345 346    @override347    def max_marginal_relevance_search_by_vector(348        self,349        embedding: List[float],350        k: int = 4,351        fetch_k: int = 20,352        lambda_mult: float = 0.5,353        **kwargs: Any,354    ) -> List[Document]:355        """Returns similar documents to the vector that also have diversity356 357        This algorithm balances relevance and diversity in the search results.358 359        Args:360            embedding: Embedding vector to search for.361            k: Number of results to return.362            fetch_k: Number of results to fetch.363            lambda_mult: Lambda multiplier for MMR.364 365        Returns:366            List of Document objects ordered by decreasing similarity/diversty.367        """368        from aperturedb.Descriptors import Descriptors369 370        descriptors = Descriptors(self.connection)371        start_time = time.time()372        descriptors.find_similar_mmr(373            set=self.descriptor_set,374            vector=embedding,375            k_neighbors=k,376            fetch_k=fetch_k,377            lambda_mult=lambda_mult,378        )379        self.logger.info(380            f"ApertureDB similarity search mmr took {time.time() - start_time} seconds"381        )382        return [self._descriptor_to_document(d) for d in descriptors]383 384    @classmethod385    @override386    def from_texts(387        cls: Type[ApertureDB],388        texts: List[str],389        embedding: Embeddings,390        metadatas: Optional[List[dict]] = None,391        **kwargs: Any,392    ) -> ApertureDB:393        """Creates a new vectorstore from a list of texts394 395        Args:396            texts: List of text strings397            embedding: Embeddings object as for constructing the vectorstore398            metadatas: Optional list of metadatas associated with the texts.399            kwargs: Additional arguments to pass to the constructor400        """401        store = cls(embeddings=embedding, **kwargs)402        store.add_texts(texts, metadatas)403        return store404 405    @classmethod406    @override407    def from_documents(408        cls: Type[ApertureDB],409        documents: List[Document],410        embedding: Embeddings,411        **kwargs: Any,412    ) -> ApertureDB:413        """Creates a new vectorstore from a list of documents414 415        Args:416            documents: List of Document objects417            embedding: Embeddings object as for constructing the vectorstore418            metadatas: Optional list of metadatas associated with the texts.419            kwargs: Additional arguments to pass to the constructor420        """421        store = cls(embeddings=embedding, **kwargs)422        store.add_documents(documents)423        return store424 425    @classmethod426    def delete_vectorstore(class_, descriptor_set: str) -> None:427        """Deletes a vectorstore and all its data from the database428 429        Args:430            descriptor_set: The name of the descriptor set to delete431        """432        from aperturedb.Utils import Utils, create_connector433 434        db = create_connector()435        utils = Utils(db)436        utils.remove_descriptorset(descriptor_set)437 438    @classmethod439    def list_vectorstores(class_) -> None:440        """Returns a list of all vectorstores in the database441 442        Returns:443            List of descriptor sets with properties444        """445        from aperturedb.Utils import create_connector446 447        db = create_connector()448        query = [449            {450                "FindDescriptorSet": {451                    # Return all properties452                    "results": {"all_properties": True},453                    "engines": True,454                    "metrics": True,455                    "dimensions": True,456                }457            }458        ]459        response, _ = db.query(query)460        assert db.last_query_ok(), response461        return response[0]["FindDescriptorSet"]["entities"]462 463    def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]:464        """Add or update documents in the vectorstore.465 466        Args:467            documents: Documents to add to the vectorstore.468            kwargs: Additional keyword arguments.469                if kwargs contains ids and documents contain ids,470                the ids in the kwargs will receive precedence.471 472        Returns:473            List of IDs of the added texts.474 475        Raises:476            ValueError: If the number of ids does not match the number of documents.477        """478 479        if "ids" in kwargs:480            ids = kwargs.pop("ids")481            if ids and len(ids) != len(documents):482                raise ValueError(483                    "The number of ids must match the number of documents. "484                    "Got {len(ids)} ids and {len(documents)} documents."485                )486 487            documents_ = []488 489            for id_, document in zip(ids, documents):490                doc_with_id = Document(491                    page_content=document.page_content,492                    metadata=document.metadata,493                    id=id_,494                )495                documents_.append(doc_with_id)496        else:497            documents_ = documents498 499        # If upsert has been implemented, we can use it to add documents500        return self.upsert(documents_, **kwargs)["succeeded"]501 502    def upsert(self, items: Sequence[Document], /, **kwargs: Any) -> UpsertResponse:503        """Insert or update items504 505        Updating documents is dependent on the documents' `id` attribute.506 507        Args:508            items: List of Document objects to upsert509 510        Returns:511            UpsertResponse object with succeeded and failed512        """513        # For now, simply delete and add514        # We could do something more efficient to update metadata,515        # but we don't support changing the embedding of a descriptor.516 517        from aperturedb.ParallelLoader import ParallelLoader518 519        ids_to_delete: List[str] = [520            item.id for item in items if hasattr(item, "id") and item.id is not None521        ]522        if ids_to_delete:523            self.delete(ids_to_delete)524 525        texts = [doc.page_content for doc in items]526        metadatas = [527            doc.metadata if getattr(doc, "metadata", None) is not None else {}528            for doc in items529        ]530        embeddings = self.embedding_function.embed_documents(texts)531        ids: List[str] = [532            doc.id if hasattr(doc, "id") and doc.id is not None else str(uuid.uuid4())533            for doc in items534        ]535 536        data = []537        for text, embedding, metadata, unique_id in zip(538            texts, embeddings, metadatas, ids539        ):540            properties = {PROPERTY_PREFIX + k: v for k, v in metadata.items()}541            properties[TEXT_PROPERTY] = text542            properties[UNIQUEID_PROPERTY] = unique_id543            command = {544                "AddDescriptor": {545                    "set": self.descriptor_set,546                    "properties": properties,547                }548            }549            query = [command]550            blobs = [np.array(embedding, dtype=np.float32).tobytes()]551            data.append((query, blobs))552        loader = ParallelLoader(self.connection)553        loader.ingest(data, batchsize=BATCHSIZE)554        return UpsertResponse(succeeded=ids, failed=[])555