codekingpro/portable-devtools

1from __future__ import annotations2 3import json4import logging5import time6import uuid7from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Type8 9from langchain_core._api.deprecation import deprecated10from langchain_core.documents import Document11from langchain_core.embeddings import Embeddings12from langchain_core.vectorstores import VectorStore13 14from langchain_community.utilities.vertexai import get_client_info15 16if TYPE_CHECKING:17    from google.cloud import storage18    from google.cloud.aiplatform import MatchingEngineIndex, MatchingEngineIndexEndpoint19    from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import (20        Namespace,21    )22    from google.oauth2.service_account import Credentials23 24    from langchain_community.embeddings import TensorflowHubEmbeddings25 26logger = logging.getLogger(__name__)27 28 29@deprecated(30    since="0.0.12",31    removal="1.0",32    alternative_import="langchain_google_vertexai.VectorSearchVectorStore",33)34class MatchingEngine(VectorStore):35    """`Google Vertex AI Vector Search` (previously Matching Engine) vector store.36 37    While the embeddings are stored in the Matching Engine, the embedded38    documents will be stored in GCS.39 40    An existing Index and corresponding Endpoint are preconditions for41    using this module.42 43    See usage in docs/integrations/vectorstores/google_vertex_ai_vector_search.ipynb44 45    Note that this implementation is mostly meant for reading if you are46    planning to do a real time implementation. While reading is a real time47    operation, updating the index takes close to one hour."""48 49    def __init__(50        self,51        project_id: str,52        index: MatchingEngineIndex,53        endpoint: MatchingEngineIndexEndpoint,54        embedding: Embeddings,55        gcs_client: storage.Client,56        gcs_bucket_name: str,57        credentials: Optional[Credentials] = None,58        *,59        document_id_key: Optional[str] = None,60    ):61        """Google Vertex AI Vector Search (previously Matching Engine)62         implementation of the vector store.63 64        While the embeddings are stored in the Matching Engine, the embedded65        documents will be stored in GCS.66 67        An existing Index and corresponding Endpoint are preconditions for68        using this module.69 70        See usage in71        docs/integrations/vectorstores/google_vertex_ai_vector_search.ipynb.72 73        Note that this implementation is mostly meant for reading if you are74        planning to do a real time implementation. While reading is a real time75        operation, updating the index takes close to one hour.76 77        Attributes:78            project_id: The GCS project id.79            index: The created index class. See80                ~:func:`MatchingEngine.from_components`.81            endpoint: The created endpoint class. See82                ~:func:`MatchingEngine.from_components`.83            embedding: A :class:`Embeddings` that will be used for84                embedding the text sent. If none is sent, then the85                multilingual Tensorflow Universal Sentence Encoder will be used.86            gcs_client: The GCS client.87            gcs_bucket_name: The GCS bucket name.88            credentials (Optional): Created GCP credentials.89            document_id_key (Optional): Key for storing document ID in document90                metadata. If None, document ID will not be returned in document91                metadata.92        """93        super().__init__()94        self._validate_google_libraries_installation()95 96        self.project_id = project_id97        self.index = index98        self.endpoint = endpoint99        self.embedding = embedding100        self.gcs_client = gcs_client101        self.credentials = credentials102        self.gcs_bucket_name = gcs_bucket_name103        self.document_id_key = document_id_key104 105    @property106    def embeddings(self) -> Embeddings:107        return self.embedding108 109    def _validate_google_libraries_installation(self) -> None:110        """Validates that Google libraries that are needed are installed."""111        try:112            from google.cloud import aiplatform, storage  # noqa: F401113            from google.oauth2 import service_account  # noqa: F401114        except ImportError:115            raise ImportError(116                "You must run `pip install --upgrade "117                "google-cloud-aiplatform google-cloud-storage`"118                "to use the MatchingEngine Vectorstore."119            )120 121    def add_texts(122        self,123        texts: Iterable[str],124        metadatas: Optional[List[dict]] = None,125        **kwargs: Any,126    ) -> List[str]:127        """Run more texts through the embeddings and add to the vectorstore.128 129        Args:130            texts: Iterable of strings to add to the vectorstore.131            metadatas: Optional list of metadatas associated with the texts.132            kwargs: vectorstore specific parameters.133 134        Returns:135            List of ids from adding the texts into the vectorstore.136        """137        texts = list(texts)138        if metadatas is not None and len(texts) != len(metadatas):139            raise ValueError(140                "texts and metadatas do not have the same length. Received "141                f"{len(texts)} texts and {len(metadatas)} metadatas."142            )143        logger.debug("Embedding documents.")144        embeddings = self.embedding.embed_documents(texts)145        jsons = []146        ids = []147        # Could be improved with async.148        for idx, (embedding, text) in enumerate(zip(embeddings, texts)):149            id = str(uuid.uuid4())150            ids.append(id)151            json_: dict = {"id": id, "embedding": embedding}152            if metadatas is not None:153                json_["metadata"] = metadatas[idx]154            jsons.append(json_)155            self._upload_to_gcs(text, f"documents/{id}")156 157        logger.debug(f"Uploaded {len(ids)} documents to GCS.")158 159        # Creating json lines from the embedded documents.160        result_str = "\n".join([json.dumps(x) for x in jsons])161 162        filename_prefix = f"indexes/{uuid.uuid4()}"163        filename = f"{filename_prefix}/{time.time()}.json"164        self._upload_to_gcs(result_str, filename)165        logger.debug(166            f"Uploaded updated json with embeddings to "167            f"{self.gcs_bucket_name}/{filename}."168        )169 170        self.index = self.index.update_embeddings(171            contents_delta_uri=f"gs://{self.gcs_bucket_name}/{filename_prefix}/"172        )173 174        logger.debug("Updated index with new configuration.")175 176        return ids177 178    def _upload_to_gcs(self, data: str, gcs_location: str) -> None:179        """Uploads data to gcs_location.180 181        Args:182            data: The data that will be stored.183            gcs_location: The location where the data will be stored.184        """185        bucket = self.gcs_client.get_bucket(self.gcs_bucket_name)186        blob = bucket.blob(gcs_location)187        blob.upload_from_string(data)188 189    def similarity_search_with_score(190        self,191        query: str,192        k: int = 4,193        filter: Optional[List[Namespace]] = None,194    ) -> List[Tuple[Document, float]]:195        """Return docs most similar to query and their cosine distance from the query.196 197        Args:198            query: String query look up documents similar to.199            k: Number of Documents to return. Defaults to 4.200            filter: Optional. A list of Namespaces for filtering201                the matching results.202                For example:203                [Namespace("color", ["red"], []), Namespace("shape", [], ["squared"])]204                will match datapoints that satisfy "red color" but not include205                datapoints with "squared shape". Please refer to206                https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#json207                for more detail.208 209        Returns:210            List[Tuple[Document, float]]: List of documents most similar to211            the query text and cosine distance in float for each.212            Lower score represents more similarity.213        """214        logger.debug(f"Embedding query {query}.")215        embedding_query = self.embedding.embed_query(query)216        return self.similarity_search_by_vector_with_score(217            embedding_query, k=k, filter=filter218        )219 220    def similarity_search_by_vector_with_score(221        self,222        embedding: List[float],223        k: int = 4,224        filter: Optional[List[Namespace]] = None,225    ) -> List[Tuple[Document, float]]:226        """Return docs most similar to the embedding and their cosine distance.227 228        Args:229            embedding: Embedding to look up documents similar to.230            k: Number of Documents to return. Defaults to 4.231            filter: Optional. A list of Namespaces for filtering232                the matching results.233                For example:234                [Namespace("color", ["red"], []), Namespace("shape", [], ["squared"])]235                will match datapoints that satisfy "red color" but not include236                datapoints with "squared shape". Please refer to237                https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#json238                for more detail.239 240        Returns:241            List[Tuple[Document, float]]: List of documents most similar to242            the query text and cosine distance in float for each.243            Lower score represents more similarity.244 245        """246        filter = filter or []247 248        # If the endpoint is public we use the find_neighbors function.249        if hasattr(self.endpoint, "_public_match_client") and (250            self.endpoint._public_match_client251        ):252            response = self.endpoint.find_neighbors(253                deployed_index_id=self._get_index_id(),254                queries=[embedding],255                num_neighbors=k,256                filter=filter,257            )258        else:259            response = self.endpoint.match(260                deployed_index_id=self._get_index_id(),261                queries=[embedding],262                num_neighbors=k,263                filter=filter,264            )265 266        logger.debug(f"Found {len(response)} matches.")267 268        if len(response) == 0:269            return []270 271        docs: List[Tuple[Document, float]] = []272 273        # I'm only getting the first one because queries receives an array274        # and the similarity_search method only receives one query. This275        # means that the match method will always return an array with only276        # one element.277        for result in response[0]:278            page_content = self._download_from_gcs(f"documents/{result.id}")279            # TODO: return all metadata.280            metadata = {}281            if self.document_id_key is not None:282                metadata[self.document_id_key] = result.id283            document = Document(284                page_content=page_content,285                metadata=metadata,286            )287            docs.append((document, result.distance))288 289        logger.debug("Downloaded documents for query.")290 291        return docs292 293    def similarity_search(294        self,295        query: str,296        k: int = 4,297        filter: Optional[List[Namespace]] = None,298        **kwargs: Any,299    ) -> List[Document]:300        """Return docs most similar to query.301 302        Args:303            query: The string that will be used to search for similar documents.304            k: The amount of neighbors that will be retrieved.305            filter: Optional. A list of Namespaces for filtering the matching results.306                For example:307                [Namespace("color", ["red"], []), Namespace("shape", [], ["squared"])]308                will match datapoints that satisfy "red color" but not include309                datapoints with "squared shape". Please refer to310                https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#json311                 for more detail.312 313        Returns:314            A list of k matching documents.315        """316        docs_and_scores = self.similarity_search_with_score(317            query, k=k, filter=filter, **kwargs318        )319 320        return [doc for doc, _ in docs_and_scores]321 322    def similarity_search_by_vector(323        self,324        embedding: List[float],325        k: int = 4,326        filter: Optional[List[Namespace]] = None,327        **kwargs: Any,328    ) -> List[Document]:329        """Return docs most similar to the embedding.330 331        Args:332            embedding: Embedding to look up documents similar to.333            k: The amount of neighbors that will be retrieved.334            filter: Optional. A list of Namespaces for filtering the matching results.335                For example:336                [Namespace("color", ["red"], []), Namespace("shape", [], ["squared"])]337                will match datapoints that satisfy "red color" but not include338                datapoints with "squared shape". Please refer to339                https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#json340                 for more detail.341 342        Returns:343            A list of k matching documents.344        """345        docs_and_scores = self.similarity_search_by_vector_with_score(346            embedding, k=k, filter=filter, **kwargs347        )348 349        return [doc for doc, _ in docs_and_scores]350 351    def _get_index_id(self) -> str:352        """Gets the correct index id for the endpoint.353 354        Returns:355            The index id if found (which should be found) or throws356            ValueError otherwise.357        """358        for index in self.endpoint.deployed_indexes:359            if index.index == self.index.resource_name:360                return index.id361 362        raise ValueError(363            f"No index with id {self.index.resource_name} "364            f"deployed on endpoint "365            f"{self.endpoint.display_name}."366        )367 368    def _download_from_gcs(self, gcs_location: str) -> str:369        """Downloads from GCS in text format.370 371        Args:372            gcs_location: The location where the file is located.373 374        Returns:375            The string contents of the file.376        """377        bucket = self.gcs_client.get_bucket(self.gcs_bucket_name)378        blob = bucket.blob(gcs_location)379        return blob.download_as_string()380 381    @classmethod382    def from_texts(383        cls: Type["MatchingEngine"],384        texts: List[str],385        embedding: Embeddings,386        metadatas: Optional[List[dict]] = None,387        **kwargs: Any,388    ) -> "MatchingEngine":389        """Use from components instead."""390        raise NotImplementedError(391            "This method is not implemented. Instead, you should initialize the class"392            " with `MatchingEngine.from_components(...)` and then call "393            "`add_texts`"394        )395 396    @classmethod397    def from_components(398        cls: Type["MatchingEngine"],399        project_id: str,400        region: str,401        gcs_bucket_name: str,402        index_id: str,403        endpoint_id: str,404        credentials_path: Optional[str] = None,405        embedding: Optional[Embeddings] = None,406        **kwargs: Any,407    ) -> "MatchingEngine":408        """Takes the object creation out of the constructor.409 410        Args:411            project_id: The GCP project id.412            region: The default location making the API calls. It must have413            the same location as the GCS bucket and must be regional.414            gcs_bucket_name: The location where the vectors will be stored in415            order for the index to be created.416            index_id: The id of the created index.417            endpoint_id: The id of the created endpoint.418            credentials_path: (Optional) The path of the Google credentials on419            the local file system.420            embedding: The :class:`Embeddings` that will be used for421            embedding the texts.422            kwargs: Additional keyword arguments to pass to MatchingEngine.__init__().423 424        Returns:425            A configured MatchingEngine with the texts added to the index.426        """427        gcs_bucket_name = cls._validate_gcs_bucket(gcs_bucket_name)428        credentials = cls._create_credentials_from_file(credentials_path)429        index = cls._create_index_by_id(index_id, project_id, region, credentials)430        endpoint = cls._create_endpoint_by_id(431            endpoint_id,432            project_id,433            region,434            credentials,435        )436 437        gcs_client = cls._get_gcs_client(credentials, project_id)438        cls._init_aiplatform(project_id, region, gcs_bucket_name, credentials)439 440        return cls(441            project_id=project_id,442            index=index,443            endpoint=endpoint,444            embedding=embedding or cls._get_default_embeddings(),445            gcs_client=gcs_client,446            credentials=credentials,447            gcs_bucket_name=gcs_bucket_name,448            **kwargs,449        )450 451    @classmethod452    def _validate_gcs_bucket(cls, gcs_bucket_name: str) -> str:453        """Validates the gcs_bucket_name as a bucket name.454 455        Args:456              gcs_bucket_name: The received bucket uri.457 458        Returns:459              A valid gcs_bucket_name or throws ValueError if full path is460              provided.461        """462        gcs_bucket_name = gcs_bucket_name.replace("gs://", "")463        if "/" in gcs_bucket_name:464            raise ValueError(465                f"The argument gcs_bucket_name should only be "466                f"the bucket name. Received {gcs_bucket_name}"467            )468        return gcs_bucket_name469 470    @classmethod471    def _create_credentials_from_file(472        cls, json_credentials_path: Optional[str]473    ) -> Optional[Credentials]:474        """Creates credentials for GCP.475 476        Args:477             json_credentials_path: The path on the file system where the478             credentials are stored.479 480         Returns:481             An optional of Credentials or None, in which case the default482             will be used.483        """484 485        from google.oauth2 import service_account486 487        credentials = None488        if json_credentials_path is not None:489            credentials = service_account.Credentials.from_service_account_file(490                json_credentials_path491            )492 493        return credentials494 495    @classmethod496    def _create_index_by_id(497        cls, index_id: str, project_id: str, region: str, credentials: "Credentials"498    ) -> MatchingEngineIndex:499        """Creates a MatchingEngineIndex object by id.500 501        Args:502            index_id: The created index id.503            project_id: The project to retrieve index from.504            region: Location to retrieve index from.505            credentials: GCS credentials.506 507        Returns:508            A configured MatchingEngineIndex.509        """510 511        from google.cloud import aiplatform512 513        logger.debug(f"Creating matching engine index with id {index_id}.")514        return aiplatform.MatchingEngineIndex(515            index_name=index_id,516            project=project_id,517            location=region,518            credentials=credentials,519        )520 521    @classmethod522    def _create_endpoint_by_id(523        cls, endpoint_id: str, project_id: str, region: str, credentials: "Credentials"524    ) -> MatchingEngineIndexEndpoint:525        """Creates a MatchingEngineIndexEndpoint object by id.526 527        Args:528            endpoint_id: The created endpoint id.529            project_id: The project to retrieve index from.530            region: Location to retrieve index from.531            credentials: GCS credentials.532 533        Returns:534            A configured MatchingEngineIndexEndpoint.535        """536 537        from google.cloud import aiplatform538 539        logger.debug(f"Creating endpoint with id {endpoint_id}.")540        return aiplatform.MatchingEngineIndexEndpoint(541            index_endpoint_name=endpoint_id,542            project=project_id,543            location=region,544            credentials=credentials,545        )546 547    @classmethod548    def _get_gcs_client(549        cls, credentials: "Credentials", project_id: str550    ) -> "storage.Client":551        """Lazily creates a GCS client.552 553        Returns:554            A configured GCS client.555        """556 557        from google.cloud import storage558 559        return storage.Client(560            credentials=credentials,561            project=project_id,562            client_info=get_client_info(module="vertex-ai-matching-engine"),563        )564 565    @classmethod566    def _init_aiplatform(567        cls,568        project_id: str,569        region: str,570        gcs_bucket_name: str,571        credentials: "Credentials",572    ) -> None:573        """Configures the aiplatform library.574 575        Args:576            project_id: The GCP project id.577            region: The default location making the API calls. It must have578            the same location as the GCS bucket and must be regional.579            gcs_bucket_name: GCS staging location.580            credentials: The GCS Credentials object.581        """582 583        from google.cloud import aiplatform584 585        logger.debug(586            f"Initializing AI Platform for project {project_id} on "587            f"{region} and for {gcs_bucket_name}."588        )589        aiplatform.init(590            project=project_id,591            location=region,592            staging_bucket=gcs_bucket_name,593            credentials=credentials,594        )595 596    @classmethod597    def _get_default_embeddings(cls) -> "TensorflowHubEmbeddings":598        """This function returns the default embedding.599 600        Returns:601            Default TensorflowHubEmbeddings to use.602        """603 604        from langchain_community.embeddings import TensorflowHubEmbeddings605 606        return TensorflowHubEmbeddings()607