codekingpro/portable-devtools

1import uuid2from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple3 4from langchain_core.documents import Document5from langchain_core.embeddings import Embeddings6from langchain_core.vectorstores import VectorStore7 8DEFAULT_DISTANCE_STRATEGY = "cosine"  # or "l2"9DEFAULT_TiDB_VECTOR_TABLE_NAME = "langchain_vector"10 11 12class TiDBVectorStore(VectorStore):13    """TiDB Vector Store."""14 15    def __init__(16        self,17        connection_string: str,18        embedding_function: Embeddings,19        table_name: str = DEFAULT_TiDB_VECTOR_TABLE_NAME,20        distance_strategy: str = DEFAULT_DISTANCE_STRATEGY,21        *,22        engine_args: Optional[Dict[str, Any]] = None,23        drop_existing_table: bool = False,24        **kwargs: Any,25    ) -> None:26        """27        Initialize a TiDB Vector Store in Langchain with a flexible28        and standardized table structure for storing vector data29        which remains fixed regardless of the dynamic table name setting.30 31        The vector table schema includes:32        - 'id': a UUID for each entry.33        - 'embedding': stores vector data in a VectorType column.34        - 'document': a Text column for the original data or additional information.35        - 'meta': a JSON column for flexible metadata storage.36        - 'create_time' and 'update_time': timestamp columns for tracking data changes.37 38        This table structure caters to general use cases and39        complex scenarios where the table serves as a semantic layer for advanced40        data integration and analysis, leveraging SQL for join queries.41 42        Args:43            connection_string (str): The connection string for the TiDB database,44                format: "mysql+pymysql://root@34.212.137.91:4000/test".45            embedding_function: The embedding function used to generate embeddings.46            table_name (str, optional): The name of the table that will be used to47                store vector data. If you do not provide a table name,48                a default table named `langchain_vector` will be created automatically.49            distance_strategy: The strategy used for similarity search,50                defaults to "cosine", valid values: "l2", "cosine".51            engine_args (Optional[Dict]): Additional arguments for the database engine,52                defaults to None.53            drop_existing_table: Drop the existing TiDB table before initializing,54                defaults to False.55            **kwargs (Any): Additional keyword arguments.56 57        Examples:58            .. code-block:: python59 60            from langchain_community.vectorstores import TiDBVectorStore61            from langchain_openai import OpenAIEmbeddings62 63            embeddingFunc = OpenAIEmbeddings()64            CONNECTION_STRING = "mysql+pymysql://root@34.212.137.91:4000/test"65 66            vs = TiDBVector.from_texts(67                embedding=embeddingFunc,68                texts = [..., ...],69                connection_string=CONNECTION_STRING,70                distance_strategy="l2",71                table_name="tidb_vector_langchain",72            )73 74            query = "What did the president say about Ketanji Brown Jackson"75            docs = db.similarity_search_with_score(query)76 77        """78 79        super().__init__(**kwargs)80        self._connection_string = connection_string81        self._embedding_function = embedding_function82        self._distance_strategy = distance_strategy83        self._vector_dimension = self._get_dimension()84 85        try:86            from tidb_vector.integrations import TiDBVectorClient87        except ImportError:88            raise ImportError(89                "Could not import tidbvec python package. "90                "Please install it with `pip install tidb-vector`."91            )92 93        self._tidb = TiDBVectorClient(94            connection_string=connection_string,95            table_name=table_name,96            distance_strategy=distance_strategy,97            vector_dimension=self._vector_dimension,98            engine_args=engine_args,99            drop_existing_table=drop_existing_table,100            **kwargs,101        )102 103    @property104    def embeddings(self) -> Embeddings:105        """Return the function used to generate embeddings."""106        return self._embedding_function107 108    @property109    def tidb_vector_client(self) -> Any:110        """Return the TiDB Vector Client."""111        return self._tidb112 113    @property114    def distance_strategy(self) -> Any:115        """116        Returns the current distance strategy.117        """118        return self._distance_strategy119 120    def _get_dimension(self) -> int:121        """122        Get the dimension of the vector using embedding functions.123        """124        return len(self._embedding_function.embed_query("test embedding length"))125 126    @classmethod127    def from_texts(128        cls,129        texts: List[str],130        embedding: Embeddings,131        metadatas: Optional[List[dict]] = None,132        **kwargs: Any,133    ) -> "TiDBVectorStore":134        """135        Create a VectorStore from a list of texts.136 137        Args:138            texts (List[str]): The list of texts to be added to the TiDB Vector.139            embedding (Embeddings): The function to use for generating embeddings.140            metadatas: The list of metadata dictionaries corresponding to each text,141                defaults to None.142            **kwargs (Any): Additional keyword arguments.143                connection_string (str): The connection string for the TiDB database,144                    format: "mysql+pymysql://root@34.212.137.91:4000/test".145                table_name (str, optional): The name of table used to store vector data,146                    defaults to "langchain_vector".147                distance_strategy: The distance strategy used for similarity search,148                    defaults to "cosine", allowed: "l2", "cosine".149                ids (Optional[List[str]]): The list of IDs corresponding to each text,150                    defaults to None.151                engine_args: Additional arguments for the underlying database engine,152                    defaults to None.153                drop_existing_table: Drop the existing TiDB table before initializing,154                    defaults to False.155 156        Returns:157            VectorStore: The created TiDB Vector Store.158 159        """160 161        # Extract arguments from kwargs with default values162        connection_string = kwargs.pop("connection_string", None)163        if connection_string is None:164            raise ValueError("please provide your tidb connection_url")165        table_name = kwargs.pop("table_name", "langchain_vector")166        distance_strategy = kwargs.pop("distance_strategy", "cosine")167        ids = kwargs.pop("ids", None)168        engine_args = kwargs.pop("engine_args", None)169        drop_existing_table = kwargs.pop("drop_existing_table", False)170 171        embeddings = embedding.embed_documents(list(texts))172 173        vs = cls(174            connection_string=connection_string,175            table_name=table_name,176            embedding_function=embedding,177            distance_strategy=distance_strategy,178            engine_args=engine_args,179            drop_existing_table=drop_existing_table,180            **kwargs,181        )182 183        vs._tidb.insert(184            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs185        )186 187        return vs188 189    @classmethod190    def from_existing_vector_table(191        cls,192        embedding: Embeddings,193        connection_string: str,194        table_name: str,195        distance_strategy: str = DEFAULT_DISTANCE_STRATEGY,196        *,197        engine_args: Optional[Dict[str, Any]] = None,198        **kwargs: Any,199    ) -> VectorStore:200        """201        Create a VectorStore instance from an existing TiDB Vector Store in TiDB.202 203        Args:204            embedding (Embeddings): The function to use for generating embeddings.205            connection_string (str): The connection string for the TiDB database,206                format: "mysql+pymysql://root@34.212.137.91:4000/test".207            table_name (str, optional): The name of table used to store vector data,208                defaults to "langchain_vector".209            distance_strategy: The distance strategy used for similarity search,210                defaults to "cosine", allowed: "l2", "cosine".211            engine_args: Additional arguments for the underlying database engine,212                defaults to None.213            **kwargs (Any): Additional keyword arguments.214        Returns:215            VectorStore: The VectorStore instance.216 217        Raises:218            NoSuchTableError: If the specified table does not exist in the TiDB.219        """220 221        try:222            from tidb_vector.integrations import check_table_existence223        except ImportError:224            raise ImportError(225                "Could not import tidbvec python package. "226                "Please install it with `pip install tidb-vector`."227            )228 229        if check_table_existence(connection_string, table_name):230            return cls(231                connection_string=connection_string,232                table_name=table_name,233                embedding_function=embedding,234                distance_strategy=distance_strategy,235                engine_args=engine_args,236                **kwargs,237            )238        else:239            raise ValueError(f"Table {table_name} does not exist in the TiDB database.")240 241    def drop_vectorstore(self) -> None:242        """243        Drop the Vector Store from the TiDB database.244        """245        self._tidb.drop_table()246 247    def add_texts(248        self,249        texts: Iterable[str],250        metadatas: Optional[List[dict]] = None,251        ids: Optional[List[str]] = None,252        **kwargs: Any,253    ) -> List[str]:254        """255        Add texts to TiDB Vector Store.256 257        Args:258            texts (Iterable[str]): The texts to be added.259            metadatas (Optional[List[dict]]): The metadata associated with each text,260                Defaults to None.261            ids (Optional[List[str]]): The IDs to be assigned to each text,262                Defaults to None, will be generated if not provided.263 264        Returns:265            List[str]: The IDs assigned to the added texts.266        """267 268        embeddings = self._embedding_function.embed_documents(list(texts))269        if ids is None:270            ids = [str(uuid.uuid4()) for _ in texts]271        if not metadatas:272            metadatas = [{} for _ in texts]273 274        return self._tidb.insert(275            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs276        )277 278    def delete(279        self,280        ids: Optional[List[str]] = None,281        **kwargs: Any,282    ) -> None:283        """284        Delete vector data from the TiDB Vector Store.285 286        Args:287            ids (Optional[List[str]]): A list of vector IDs to delete.288            kwargs: Additional keyword arguments.289        """290 291        self._tidb.delete(ids=ids, **kwargs)292 293    def similarity_search(294        self,295        query: str,296        k: int = 4,297        filter: Optional[dict] = None,298        **kwargs: Any,299    ) -> List[Document]:300        """301        Perform a similarity search using the given query.302 303        Args:304            query: The query string.305            k: The number of results to retrieve. Defaults to 4.306            filter: A filter to apply to the search results.307            kwargs: Additional keyword arguments.308 309        Returns:310            A list of `Document` objects representing the search results.311        """312        result = self.similarity_search_with_score(query, k, filter, **kwargs)313        return [doc for doc, _ in result]314 315    def similarity_search_with_score(316        self,317        query: str,318        k: int = 5,319        filter: Optional[dict] = None,320        **kwargs: Any,321    ) -> List[Tuple[Document, float]]:322        """323        Perform a similarity search with score based on the given query.324 325        Args:326            query (str): The query string.327            k (int, optional): The number of results to return. Defaults to 5.328            filter (dict, optional): A filter to apply to the search results.329                Defaults to None.330            kwargs: Additional keyword arguments.331 332        Returns:333            A list of tuples containing relevant documents and their similarity scores.334        """335        query_vector = self._embedding_function.embed_query(query)336        relevant_docs = self._tidb.query(337            query_vector=query_vector, k=k, filter=filter, **kwargs338        )339        return [340            (341                Document(342                    page_content=doc.document,343                    metadata=doc.metadata,344                ),345                doc.distance,346            )347            for doc in relevant_docs348        ]349 350    def _select_relevance_score_fn(self) -> Callable[[float], float]:351        """352        Select the relevance score function based on the distance strategy.353        """354        if self._distance_strategy == "cosine":355            return self._cosine_relevance_score_fn356        elif self._distance_strategy == "l2":357            return self._euclidean_relevance_score_fn358        else:359            raise ValueError(360                "No supported normalization function"361                f" for distance_strategy of {self._distance_strategy}."362                "Consider providing relevance_score_fn to PGVector constructor."363            )364