codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4import uuid5from typing import Any, Dict, Iterable, List, Optional, Tuple, Type6 7from langchain_core.documents import Document8from langchain_core.embeddings import Embeddings9from langchain_core.utils import get_from_dict_or_env10from langchain_core.vectorstores import VectorStore11 12ADA_TOKEN_COUNT = 153613_LANGCHAIN_DEFAULT_TABLE_NAME = "langchain_pg_embedding"14 15 16class Hologres(VectorStore):17    """`Hologres API` vector store.18 19    - `connection_string` is a hologres connection string.20    - `embedding_function` any embedding function implementing21        `langchain.embeddings.base.Embeddings` interface.22    - `ndims` is the number of dimensions of the embedding output.23    - `table_name` is the name of the table to store embeddings and data.24        (default: langchain_pg_embedding)25        - NOTE: The table will be created when initializing the store (if not exists)26            So, make sure the user has the right permissions to create tables.27    - `pre_delete_table` if True, will delete the table if it exists.28        (default: False)29        - Useful for testing.30    """31 32    def __init__(33        self,34        connection_string: str,35        embedding_function: Embeddings,36        ndims: int = ADA_TOKEN_COUNT,37        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,38        pre_delete_table: bool = False,39        logger: Optional[logging.Logger] = None,40    ) -> None:41        self.connection_string = connection_string42        self.ndims = ndims43        self.table_name = table_name44        self.embedding_function = embedding_function45        self.pre_delete_table = pre_delete_table46        self.logger = logger or logging.getLogger(__name__)47        self.__post_init__()48 49    def __post_init__(50        self,51    ) -> None:52        """53        Initialize the store.54        """55        from hologres_vector import HologresVector56 57        self.storage = HologresVector(58            self.connection_string,59            ndims=self.ndims,60            table_name=self.table_name,61            table_schema={"document": "text"},62            pre_delete_table=self.pre_delete_table,63        )64 65    @property66    def embeddings(self) -> Embeddings:67        return self.embedding_function68 69    @classmethod70    def __from(71        cls,72        texts: List[str],73        embeddings: List[List[float]],74        embedding_function: Embeddings,75        metadatas: Optional[List[dict]] = None,76        ids: Optional[List[str]] = None,77        ndims: int = ADA_TOKEN_COUNT,78        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,79        pre_delete_table: bool = False,80        **kwargs: Any,81    ) -> Hologres:82        if ids is None:83            ids = [str(uuid.uuid4()) for _ in texts]84 85        if not metadatas:86            metadatas = [{} for _ in texts]87 88        connection_string = cls.get_connection_string(kwargs)89 90        store = cls(91            connection_string=connection_string,92            embedding_function=embedding_function,93            ndims=ndims,94            table_name=table_name,95            pre_delete_table=pre_delete_table,96        )97 98        store.add_embeddings(99            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs100        )101 102        return store103 104    def add_embeddings(105        self,106        texts: Iterable[str],107        embeddings: List[List[float]],108        metadatas: List[dict],109        ids: List[str],110        **kwargs: Any,111    ) -> None:112        """Add embeddings to the vectorstore.113 114        Args:115            texts: Iterable of strings to add to the vectorstore.116            embeddings: List of list of embedding vectors.117            metadatas: List of metadatas associated with the texts.118            kwargs: vectorstore specific parameters119        """120        try:121            schema_datas = [{"document": t} for t in texts]122            self.storage.upsert_vectors(embeddings, ids, metadatas, schema_datas)123        except Exception as e:124            self.logger.exception(e)125 126    def add_texts(127        self,128        texts: Iterable[str],129        metadatas: Optional[List[dict]] = None,130        ids: Optional[List[str]] = None,131        **kwargs: Any,132    ) -> List[str]:133        """Run more texts through the embeddings and add to the vectorstore.134 135        Args:136            texts: Iterable of strings to add to the vectorstore.137            metadatas: Optional list of metadatas associated with the texts.138            kwargs: vectorstore specific parameters139 140        Returns:141            List of ids from adding the texts into the vectorstore.142        """143        if ids is None:144            ids = [str(uuid.uuid4()) for _ in texts]145 146        embeddings = self.embedding_function.embed_documents(list(texts))147 148        if not metadatas:149            metadatas = [{} for _ in texts]150 151        self.add_embeddings(texts, embeddings, metadatas, ids, **kwargs)152 153        return ids154 155    def similarity_search(156        self,157        query: str,158        k: int = 4,159        filter: Optional[dict] = None,160        **kwargs: Any,161    ) -> List[Document]:162        """Run similarity search with Hologres with distance.163 164        Args:165            query (str): Query text to search for.166            k (int): Number of results to return. Defaults to 4.167            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.168 169        Returns:170            List of Documents most similar to the query.171        """172        embedding = self.embedding_function.embed_query(text=query)173        return self.similarity_search_by_vector(174            embedding=embedding,175            k=k,176            filter=filter,177        )178 179    def similarity_search_by_vector(180        self,181        embedding: List[float],182        k: int = 4,183        filter: Optional[dict] = None,184        **kwargs: Any,185    ) -> List[Document]:186        """Return docs most similar to embedding vector.187 188        Args:189            embedding: Embedding to look up documents similar to.190            k: Number of Documents to return. Defaults to 4.191            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.192 193        Returns:194            List of Documents most similar to the query vector.195        """196        docs_and_scores = self.similarity_search_with_score_by_vector(197            embedding=embedding, k=k, filter=filter198        )199        return [doc for doc, _ in docs_and_scores]200 201    def similarity_search_with_score(202        self,203        query: str,204        k: int = 4,205        filter: Optional[dict] = None,206    ) -> List[Tuple[Document, float]]:207        """Return docs most similar to query.208 209        Args:210            query: Text to look up documents similar to.211            k: Number of Documents to return. Defaults to 4.212            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.213 214        Returns:215            List of Documents most similar to the query and score for each216        """217        embedding = self.embedding_function.embed_query(query)218        docs = self.similarity_search_with_score_by_vector(219            embedding=embedding, k=k, filter=filter220        )221        return docs222 223    def similarity_search_with_score_by_vector(224        self,225        embedding: List[float],226        k: int = 4,227        filter: Optional[dict] = None,228    ) -> List[Tuple[Document, float]]:229        results: List[dict[str, Any]] = self.storage.search(230            embedding, k=k, select_columns=["document"], metadata_filters=filter231        )232 233        docs = [234            (235                Document(236                    page_content=result["document"],237                    metadata=result["metadata"],238                ),239                result["distance"],240            )241            for result in results242        ]243        return docs244 245    @classmethod246    def from_texts(247        cls: Type[Hologres],248        texts: List[str],249        embedding: Embeddings,250        metadatas: Optional[List[dict]] = None,251        ndims: int = ADA_TOKEN_COUNT,252        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,253        ids: Optional[List[str]] = None,254        pre_delete_table: bool = False,255        **kwargs: Any,256    ) -> Hologres:257        """258        Return VectorStore initialized from texts and embeddings.259        Hologres connection string is required260        "Either pass it as a parameter261        or set the HOLOGRES_CONNECTION_STRING environment variable.262        Create the connection string by calling263        HologresVector.connection_string_from_db_params264        """265        embeddings = embedding.embed_documents(list(texts))266 267        return cls.__from(268            texts,269            embeddings,270            embedding,271            metadatas=metadatas,272            ids=ids,273            ndims=ndims,274            table_name=table_name,275            pre_delete_table=pre_delete_table,276            **kwargs,277        )278 279    @classmethod280    def from_embeddings(281        cls,282        text_embeddings: List[Tuple[str, List[float]]],283        embedding: Embeddings,284        metadatas: Optional[List[dict]] = None,285        ndims: int = ADA_TOKEN_COUNT,286        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,287        ids: Optional[List[str]] = None,288        pre_delete_table: bool = False,289        **kwargs: Any,290    ) -> Hologres:291        """Construct Hologres wrapper from raw documents and pre-292        generated embeddings.293 294        Return VectorStore initialized from documents and embeddings.295        Hologres connection string is required296        "Either pass it as a parameter297        or set the HOLOGRES_CONNECTION_STRING environment variable.298        Create the connection string by calling299        HologresVector.connection_string_from_db_params300 301        Example:302            .. code-block:: python303 304                from langchain_community.vectorstores import Hologres305                from langchain_community.embeddings import OpenAIEmbeddings306                embeddings = OpenAIEmbeddings()307                text_embeddings = embeddings.embed_documents(texts)308                text_embedding_pairs = list(zip(texts, text_embeddings))309                faiss = Hologres.from_embeddings(text_embedding_pairs, embeddings)310        """311        texts = [t[0] for t in text_embeddings]312        embeddings = [t[1] for t in text_embeddings]313 314        return cls.__from(315            texts,316            embeddings,317            embedding,318            metadatas=metadatas,319            ids=ids,320            ndims=ndims,321            table_name=table_name,322            pre_delete_table=pre_delete_table,323            **kwargs,324        )325 326    @classmethod327    def from_existing_index(328        cls: Type[Hologres],329        embedding: Embeddings,330        ndims: int = ADA_TOKEN_COUNT,331        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,332        pre_delete_table: bool = False,333        **kwargs: Any,334    ) -> Hologres:335        """336        Get instance of an existing Hologres store.This method will337        return the instance of the store without inserting any new338        embeddings339        """340 341        connection_string = cls.get_connection_string(kwargs)342 343        store = cls(344            connection_string=connection_string,345            ndims=ndims,346            table_name=table_name,347            embedding_function=embedding,348            pre_delete_table=pre_delete_table,349        )350 351        return store352 353    @classmethod354    def get_connection_string(cls, kwargs: Dict[str, Any]) -> str:355        connection_string: str = get_from_dict_or_env(356            data=kwargs,357            key="connection_string",358            env_key="HOLOGRES_CONNECTION_STRING",359        )360 361        if not connection_string:362            raise ValueError(363                "Hologres connection string is required"364                "Either pass it as a parameter"365                "or set the HOLOGRES_CONNECTION_STRING environment variable."366                "Create the connection string by calling"367                "HologresVector.connection_string_from_db_params"368            )369 370        return connection_string371 372    @classmethod373    def from_documents(374        cls: Type[Hologres],375        documents: List[Document],376        embedding: Embeddings,377        ndims: int = ADA_TOKEN_COUNT,378        table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,379        ids: Optional[List[str]] = None,380        pre_delete_collection: bool = False,381        **kwargs: Any,382    ) -> Hologres:383        """384        Return VectorStore initialized from documents and embeddings.385        Hologres connection string is required386        "Either pass it as a parameter387        or set the HOLOGRES_CONNECTION_STRING environment variable.388        Create the connection string by calling389        HologresVector.connection_string_from_db_params390        """391 392        texts = [d.page_content for d in documents]393        metadatas = [d.metadata for d in documents]394        connection_string = cls.get_connection_string(kwargs)395 396        kwargs["connection_string"] = connection_string397 398        return cls.from_texts(399            texts=texts,400            pre_delete_collection=pre_delete_collection,401            embedding=embedding,402            metadatas=metadatas,403            ids=ids,404            ndims=ndims,405            table_name=table_name,406            **kwargs,407        )408 409    @classmethod410    def connection_string_from_db_params(411        cls,412        host: str,413        port: int,414        database: str,415        user: str,416        password: str,417    ) -> str:418        """Return connection string from database parameters."""419        return (420            f"dbname={database} user={user} password={password} host={host} port={port}"421        )422