codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4import uuid5import warnings6from typing import (7    TYPE_CHECKING,8    Any,9    Callable,10    Iterable,11    List,12    Optional,13    Tuple,14    TypeVar,15    Union,16)17 18import numpy as np19from langchain_core.documents import Document20from langchain_core.embeddings import Embeddings21from langchain_core.vectorstores import VectorStore22 23from langchain_community.vectorstores.utils import (24    DistanceStrategy,25    maximal_marginal_relevance,26)27 28if TYPE_CHECKING:29    from aerospike_vector_search import Client30    from aerospike_vector_search.types import Neighbor, VectorDistanceMetric31 32logger = logging.getLogger(__name__)33 34 35def _import_aerospike() -> Any:36    try:37        from aerospike_vector_search import Client38    except ImportError as e:39        raise ImportError(40            "Could not import aerospike_vector_search python package. "41            "Please install it with `pip install aerospike_vector`."42        ) from e43    return Client44 45 46AVST = TypeVar("AVST", bound="Aerospike")47 48 49class Aerospike(VectorStore):50    """`Aerospike` vector store.51 52    To use, you should have the ``aerospike_vector_search`` python package installed.53    """54 55    def __init__(56        self,57        client: Client,58        embedding: Union[Embeddings, Callable],59        namespace: str,60        index_name: Optional[str] = None,61        vector_key: str = "_vector",62        text_key: str = "_text",63        id_key: str = "_id",64        set_name: Optional[str] = None,65        distance_strategy: Optional[66            Union[DistanceStrategy, VectorDistanceMetric]67        ] = DistanceStrategy.EUCLIDEAN_DISTANCE,68    ):69        """Initialize with Aerospike client.70 71        Args:72            client: Aerospike client.73            embedding: Embeddings object or Callable (deprecated) to embed text.74            namespace: Namespace to use for storing vectors. This should match75            index_name: Name of the index previously created in Aerospike. This76            vector_key: Key to use for vector in metadata. This should match the77                key used during index creation.78            text_key: Key to use for text in metadata.79            id_key: Key to use for id in metadata.80            set_name: Default set name to use for storing vectors.81            distance_strategy: Distance strategy to use for similarity search82                This should match the distance strategy used during index creation.83        """84 85        aerospike = _import_aerospike()86 87        if not isinstance(embedding, Embeddings):88            warnings.warn(89                "Passing in `embedding` as a Callable is deprecated. Please pass in an"90                " Embeddings object instead."91            )92 93        if not isinstance(client, aerospike):94            raise ValueError(95                f"client should be an instance of aerospike_vector_search.Client, "96                f"got {type(client)}"97            )98 99        self._client = client100        self._embedding = embedding101        self._text_key = text_key102        self._vector_key = vector_key103        self._id_key = id_key104        self._index_name = index_name105        self._namespace = namespace106        self._set_name = set_name107        self._distance_strategy = self.convert_distance_strategy(distance_strategy)108 109    @property110    def embeddings(self) -> Optional[Embeddings]:111        """Access the query embedding object if available."""112        if isinstance(self._embedding, Embeddings):113            return self._embedding114        return None115 116    def _embed_documents(self, texts: Iterable[str]) -> List[List[float]]:117        """Embed search docs."""118        if isinstance(self._embedding, Embeddings):119            return self._embedding.embed_documents(list(texts))120        return [self._embedding(t) for t in texts]121 122    def _embed_query(self, text: str) -> List[float]:123        """Embed query text."""124        if isinstance(self._embedding, Embeddings):125            return self._embedding.embed_query(text)126        return self._embedding(text)127 128    @staticmethod129    def convert_distance_strategy(130        distance_strategy: Union[VectorDistanceMetric, DistanceStrategy],131    ) -> DistanceStrategy:132        """133        Convert Aerospikes distance strategy to langchains DistanceStrategy134        enum. This is a convenience method to allow users to pass in the same135        distance metric used to create the index.136        """137        from aerospike_vector_search.types import VectorDistanceMetric138 139        if isinstance(distance_strategy, DistanceStrategy):140            return distance_strategy141 142        if distance_strategy == VectorDistanceMetric.COSINE:143            return DistanceStrategy.COSINE144 145        if distance_strategy == VectorDistanceMetric.DOT_PRODUCT:146            return DistanceStrategy.DOT_PRODUCT147 148        if distance_strategy == VectorDistanceMetric.SQUARED_EUCLIDEAN:149            return DistanceStrategy.EUCLIDEAN_DISTANCE150 151        raise ValueError(152            "Unknown distance strategy, must be cosine, dot_product, or euclidean"153        )154 155    def add_texts(156        self,157        texts: Iterable[str],158        metadatas: Optional[List[dict]] = None,159        ids: Optional[List[str]] = None,160        set_name: Optional[str] = None,161        embedding_chunk_size: int = 1000,162        index_name: Optional[str] = None,163        wait_for_index: bool = True,164        **kwargs: Any,165    ) -> List[str]:166        """Run more texts through the embeddings and add to the vectorstore.167 168 169        Args:170            texts: Iterable of strings to add to the vectorstore.171            metadatas: Optional list of metadata associated with the texts.172            ids: Optional list of ids to associate with the texts.173            set_name: Optional aerospike set name to add the texts to.174            batch_size: Batch size to use when adding the texts to the vectorstore.175            embedding_chunk_size: Chunk size to use when embedding the texts.176            index_name: Optional aerospike index name used for waiting for index177                completion. If not provided, the default index_name will be used.178            wait_for_index: If True, wait for the all the texts to be indexed179                before returning. Requires index_name to be provided. Defaults180                to True.181            kwargs: Additional keyword arguments to pass to the client upsert call.182 183        Returns:184            List of ids from adding the texts into the vectorstore.185 186        """187        if set_name is None:188            set_name = self._set_name189 190        if index_name is None:191            index_name = self._index_name192 193        if wait_for_index and index_name is None:194            raise ValueError("if wait_for_index is True, index_name must be provided")195 196        texts = list(texts)197        ids = ids or [str(uuid.uuid4()) for _ in texts]198 199        # We need to shallow copy so that we can add the vector and text keys200        if metadatas:201            metadatas = [m.copy() for m in metadatas]202        else:203            metadatas = metadatas or [{} for _ in texts]204 205        for i in range(0, len(texts), embedding_chunk_size):206            chunk_texts = texts[i : i + embedding_chunk_size]207            chunk_ids = ids[i : i + embedding_chunk_size]208            chunk_metadatas = metadatas[i : i + embedding_chunk_size]209            embeddings = self._embed_documents(chunk_texts)210 211            for metadata, embedding, text in zip(212                chunk_metadatas, embeddings, chunk_texts213            ):214                metadata[self._vector_key] = embedding215                metadata[self._text_key] = text216 217            for id, metadata in zip(chunk_ids, chunk_metadatas):218                metadata[self._id_key] = id219                self._client.upsert(220                    namespace=self._namespace,221                    key=id,222                    set_name=set_name,223                    record_data=metadata,224                    **kwargs,225                )226 227        if wait_for_index:228            self._client.wait_for_index_completion(229                namespace=self._namespace,230                name=index_name,231            )232 233        return ids234 235    def delete(236        self,237        ids: Optional[List[str]] = None,238        set_name: Optional[str] = None,239        **kwargs: Any,240    ) -> Optional[bool]:241        """Delete by vector ID or other criteria.242 243        Args:244            ids: List of ids to delete.245            **kwargs: Other keyword arguments to pass to client delete call.246 247        Returns:248            Optional[bool]: True if deletion is successful,249            False otherwise, None if not implemented.250        """251        from aerospike_vector_search import AVSServerError252 253        if ids:254            for id in ids:255                try:256                    self._client.delete(257                        namespace=self._namespace,258                        key=id,259                        set_name=set_name,260                        **kwargs,261                    )262                except AVSServerError:263                    return False264 265        return True266 267    def similarity_search_with_score(268        self,269        query: str,270        k: int = 4,271        metadata_keys: Optional[List[str]] = None,272        index_name: Optional[str] = None,273        **kwargs: Any,274    ) -> List[Tuple[Document, float]]:275        """Return aerospike documents most similar to query, along with scores.276 277        Args:278            query: Text to look up documents similar to.279            k: Number of Documents to return. Defaults to 4.280            metadata_keys: List of metadata keys to return with the documents.281                If None, all metadata keys will be returned. Defaults to None.282            index_name: Name of the index to search. Overrides the default283                index_name.284            kwargs: Additional keyword arguments to pass to the search method.285 286        Returns:287            List of Documents most similar to the query and associated scores.288        """289 290        return self.similarity_search_by_vector_with_score(291            self._embed_query(query),292            k=k,293            metadata_keys=metadata_keys,294            index_name=index_name,295            **kwargs,296        )297 298    def similarity_search_by_vector_with_score(299        self,300        embedding: List[float],301        k: int = 4,302        metadata_keys: Optional[List[str]] = None,303        index_name: Optional[str] = None,304        **kwargs: Any,305    ) -> List[Tuple[Document, float]]:306        """Return aerospike documents most similar to embedding, along with scores.307 308        Args:309            embedding: Embedding to look up documents similar to.310            k: Number of Documents to return. Defaults to 4.311            metadata_keys: List of metadata keys to return with the documents.312                If None, all metadata keys will be returned. Defaults to None.313            index_name: Name of the index to search. Overrides the default314                index_name.315            kwargs: Additional keyword arguments to pass to the client316                vector_search method.317 318        Returns:319            List of Documents most similar to the query and associated scores.320 321        """322 323        docs = []324 325        if metadata_keys and self._text_key not in metadata_keys:326            metadata_keys = [self._text_key] + metadata_keys327 328        if index_name is None:329            index_name = self._index_name330 331        if index_name is None:332            raise ValueError("index_name must be provided")333 334        results: list[Neighbor] = self._client.vector_search(335            index_name=index_name,336            namespace=self._namespace,337            query=embedding,338            limit=k,339            field_names=metadata_keys,340            **kwargs,341        )342 343        for result in results:344            metadata = result.fields345 346            if self._text_key in metadata:347                text = metadata.pop(self._text_key)348                score = result.distance349                docs.append((Document(page_content=text, metadata=metadata), score))350            else:351                logger.warning(352                    f"Found document with no `{self._text_key}` key. Skipping."353                )354                continue355 356        return docs357 358    def similarity_search_by_vector(359        self,360        embedding: List[float],361        k: int = 4,362        metadata_keys: Optional[List[str]] = None,363        index_name: Optional[str] = None,364        **kwargs: Any,365    ) -> List[Document]:366        """Return docs most similar to embedding vector.367 368        Args:369            embedding: Embedding to look up documents similar to.370            k: Number of Documents to return. Defaults to 4.371            metadata_keys: List of metadata keys to return with the documents.372                If None, all metadata keys will be returned. Defaults to None.373            index_name: Name of the index to search. Overrides the default374                index_name.375            kwargs: Additional keyword arguments to pass to the search method.376 377 378        Returns:379            List of Documents most similar to the query vector.380        """381        return [382            doc383            for doc, _ in self.similarity_search_by_vector_with_score(384                embedding,385                k=k,386                metadata_keys=metadata_keys,387                index_name=index_name,388                **kwargs,389            )390        ]391 392    def similarity_search(393        self,394        query: str,395        k: int = 4,396        metadata_keys: Optional[List[str]] = None,397        index_name: Optional[str] = None,398        **kwargs: Any,399    ) -> List[Document]:400        """Return aerospike documents most similar to query.401 402        Args:403            query: Text to look up documents similar to.404            k: Number of Documents to return. Defaults to 4.405            metadata_keys: List of metadata keys to return with the documents.406                If None, all metadata keys will be returned. Defaults to None.407            index_name: Optional name of the index to search. Overrides the408                default index_name.409 410        Returns:411            List of Documents most similar to the query and score for each412        """413        docs_and_scores = self.similarity_search_with_score(414            query, k=k, metadata_keys=metadata_keys, index_name=index_name, **kwargs415        )416        return [doc for doc, _ in docs_and_scores]417 418    def _select_relevance_score_fn(self) -> Callable[[float], float]:419        """420        The 'correct' relevance function421        may differ depending on a few things, including:422        - the distance / similarity metric used by the VectorStore423        - the scale of your embeddings (OpenAI's are unit normed. Many others are not!)424        - embedding dimensionality425        - etc.426 427        0 is dissimilar, 1 is similar.428 429        Aerospike's relevance_fn assume euclidean and dot product embeddings are430        normalized to unit norm.431        """432        if self._distance_strategy == DistanceStrategy.COSINE:433            return self._cosine_relevance_score_fn434        elif self._distance_strategy == DistanceStrategy.DOT_PRODUCT:435            return self._max_inner_product_relevance_score_fn436        elif self._distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:437            return self._euclidean_relevance_score_fn438        else:439            raise ValueError(440                "Unknown distance strategy, must be cosine, dot_product, or euclidean"441            )442 443    @staticmethod444    def _cosine_relevance_score_fn(score: float) -> float:445        """Aerospike returns cosine distance scores between [0,2]446 447        0 is dissimilar, 1 is similar.448        """449        return 1 - (score / 2)450 451    def max_marginal_relevance_search_by_vector(452        self,453        embedding: List[float],454        k: int = 4,455        fetch_k: int = 20,456        lambda_mult: float = 0.5,457        metadata_keys: Optional[List[str]] = None,458        index_name: Optional[str] = None,459        **kwargs: Any,460    ) -> List[Document]:461        """Return docs selected using the maximal marginal relevance.462 463        Maximal marginal relevance optimizes for similarity to query AND diversity464        among selected documents.465 466        Args:467            embedding: Embedding to look up documents similar to.468            k: Number of Documents to return. Defaults to 4.469            fetch_k: Number of Documents to fetch to pass to MMR algorithm.470            lambda_mult: Number between 0 and 1 that determines the degree of471                diversity among the results with 0 corresponding to maximum472                diversity and 1 to minimum diversity. Defaults to 0.5.473            metadata_keys: List of metadata keys to return with the documents.474                If None, all metadata keys will be returned. Defaults to None.475            index_name: Optional name of the index to search. Overrides the476                default index_name.477        Returns:478            List of Documents selected by maximal marginal relevance.479        """480 481        if metadata_keys and self._vector_key not in metadata_keys:482            metadata_keys = [self._vector_key] + metadata_keys483 484        docs = self.similarity_search_by_vector(485            embedding,486            k=fetch_k,487            metadata_keys=metadata_keys,488            index_name=index_name,489            **kwargs,490        )491        mmr_selected = maximal_marginal_relevance(492            np.array([embedding], dtype=np.float32),493            [doc.metadata[self._vector_key] for doc in docs],494            k=k,495            lambda_mult=lambda_mult,496        )497 498        if metadata_keys and self._vector_key in metadata_keys:499            for i in mmr_selected:500                docs[i].metadata.pop(self._vector_key)501 502        return [docs[i] for i in mmr_selected]503 504    def max_marginal_relevance_search(505        self,506        query: str,507        k: int = 4,508        fetch_k: int = 20,509        lambda_mult: float = 0.5,510        metadata_keys: Optional[List[str]] = None,511        index_name: Optional[str] = None,512        **kwargs: Any,513    ) -> List[Document]:514        """Return docs selected using the maximal marginal relevance.515 516        Maximal marginal relevance optimizes for similarity to query AND diversity517        among selected documents.518 519        Args:520            query: Text to look up documents similar to.521            k: Number of Documents to return. Defaults to 4.522            fetch_k: Number of Documents to fetch to pass to MMR algorithm.523            lambda_mult: Number between 0 and 1 that determines the degree524                        of diversity among the results with 0 corresponding525                        to maximum diversity and 1 to minimum diversity.526                        Defaults to 0.5.527            index_name: Name of the index to search.528        Returns:529            List of Documents selected by maximal marginal relevance.530        """531        embedding = self._embed_query(query)532        return self.max_marginal_relevance_search_by_vector(533            embedding,534            k,535            fetch_k,536            lambda_mult,537            metadata_keys=metadata_keys,538            index_name=index_name,539            **kwargs,540        )541 542    @classmethod543    def from_texts(544        cls,545        texts: List[str],546        embedding: Embeddings,547        metadatas: Optional[List[dict]] = None,548        client: Client = None,549        namespace: str = "test",550        index_name: Optional[str] = None,551        ids: Optional[List[str]] = None,552        embeddings_chunk_size: int = 1000,553        client_kwargs: Optional[dict] = None,554        **kwargs: Any,555    ) -> Aerospike:556        """557        This is a user friendly interface that:558            1. Embeds text.559            2. Converts the texts into documents.560            3. Adds the documents to a provided Aerospike index561 562        This is intended to be a quick way to get started.563 564        Example:565            .. code-block:: python566 567                from langchain_community.vectorstores import Aerospike568                from langchain_openai import OpenAIEmbeddings569                from aerospike_vector_search import Client, HostPort570 571                client = Client(seeds=HostPort(host="localhost", port=5000))572                aerospike = Aerospike.from_texts(573                    ["foo", "bar", "baz"],574                    embedder,575                    client,576                    "namespace",577                    index_name="index",578                    vector_key="vector",579                    distance_strategy=MODEL_DISTANCE_CALC,580                )581        """582        aerospike = cls(583            client,584            embedding,585            namespace,586            **kwargs,587        )588 589        aerospike.add_texts(590            texts,591            metadatas=metadatas,592            ids=ids,593            index_name=index_name,594            embedding_chunk_size=embeddings_chunk_size,595            **(client_kwargs or {}),596        )597        return aerospike598