Toggle Menu
Prev Up Next

# Source code for langchain_community.vectorstores.awadb

    from __future__ import annotations

    import logging
    import uuid
    from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Type

    import numpy as np
    from langchain_core.documents import Document
    from langchain_core.embeddings import Embeddings
    from langchain_core.vectorstores import VectorStore

    from langchain_community.vectorstores.utils import maximal_marginal_relevance

    if TYPE_CHECKING:
        import awadb

    logger = logging.getLogger()
    DEFAULT_TOPN = 4

    class AwaDB(VectorStore):
        """`AwaDB` vector store."""

        _DEFAULT_TABLE_NAME = "langchain_awadb"

        def __init__(
            self,
            table_name: str = _DEFAULT_TABLE_NAME,
            embedding: Optional[Embeddings] = None,
            log_and_data_dir: Optional[str] = None,
            client: Optional[awadb.Client] = None,
            **kwargs: Any,
        ) -> None:
            """Initialize with AwaDB client.
               If table_name is not specified,
               a random table name of `_DEFAULT_TABLE_NAME + last segment of uuid`
               would be created automatically.

            Args:
                table_name: Name of the table created, default _DEFAULT_TABLE_NAME.
                embedding: Optional Embeddings initially set.
                log_and_data_dir: Optional the root directory of log and data.
                client: Optional AwaDB client.
                kwargs: Any possible extend parameters in the future.

            Returns:
                None.
            """
            try:
                import awadb
            except ImportError:
                raise ImportError(
                    "Could not import awadb python package. "
                    "Please install it with `pip install awadb`."
                )

            if client is not None:
                self.awadb_client = client
            else:
                if log_and_data_dir is not None:
                    self.awadb_client = awadb.Client(log_and_data_dir)
                else:
                    self.awadb_client = awadb.Client()

            if table_name == self._DEFAULT_TABLE_NAME:
                table_name += "_"
                table_name += str(uuid.uuid4()).split("-")[-1]

            self.awadb_client.Create(table_name)
            self.table2embeddings: dict[str, Embeddings] = {}
            if embedding is not None:
                self.table2embeddings[table_name] = embedding
            self.using_table_name = table_name

        @property
        def embeddings(self) -> Optional[Embeddings]:
            if self.using_table_name in self.table2embeddings:
                return self.table2embeddings[self.using_table_name]
            return None

        def add_texts(
            self,
            texts: Iterable[str],
            metadatas: Optional[List[dict]] = None,
            is_duplicate_texts: Optional[bool] = None,
            **kwargs: Any,
        ) -> List[str]:
            """Run more texts through the embeddings and add to the vectorstore.
            Args:
                texts: Iterable of strings to add to the vectorstore.
                metadatas: Optional list of metadatas associated with the texts.
                is_duplicate_texts: Optional whether to duplicate texts. Defaults to True.
                kwargs: any possible extend parameters in the future.

            Returns:
                List of ids from adding the texts into the vectorstore.
            """
            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            embeddings = None
            if self.using_table_name in self.table2embeddings:
                embeddings = self.table2embeddings[self.using_table_name].embed_documents(
                    list(texts)
                )

            return self.awadb_client.AddTexts(
                "embedding_text",
                "text_embedding",
                texts,
                embeddings,
                metadatas,
                is_duplicate_texts,
            )

        def load_local(
            self,
            table_name: str,
            **kwargs: Any,
        ) -> bool:
            """Load the local specified table.

            Args:
                table_name: Table name
                kwargs: Any possible extend parameters in the future.

            Returns:
                Success or failure of loading the local specified table
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            return self.awadb_client.Load(table_name)

        def similarity_search(
            self,
            query: str,
            k: int = DEFAULT_TOPN,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Return docs most similar to query.

            Args:
                query: Text query.
                k: The maximum number of documents to return.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter (Optional[dict]): Filter by metadata. Defaults to None.
                E.g. `{"color" : "red", "price": 4.20}`. Optional.
                E.g. `{"max_price" : 15.66, "min_price": 4.20}`
                `price` is the metadata field, means range filter(4.20<'price'<15.66).
                E.g. `{"maxe_price" : 15.66, "mine_price": 4.20}`
                `price` is the metadata field, means range filter(4.20<='price'<=15.66).
                kwargs: Any possible extend parameters in the future.

            Returns:
                Returns the k most similar documents to the specified text query.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            embedding = None
            if self.using_table_name in self.table2embeddings:
                embedding = self.table2embeddings[self.using_table_name].embed_query(query)
            else:
                from awadb import AwaEmbedding

                embedding = AwaEmbedding().Embedding(query)

            not_include_fields: Set[str] = {"text_embedding", "_id", "score"}
            return self.similarity_search_by_vector(
                embedding,
                k,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
                not_include_fields_in_metadata=not_include_fields,
            )

        def similarity_search_with_score(
            self,
            query: str,
            k: int = DEFAULT_TOPN,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            **kwargs: Any,
        ) -> List[Tuple[Document, float]]:
            """The most k similar documents and scores of the specified query.

            Args:
                query: Text query.
                k: The k most similar documents to the text query.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter: Filter by metadata. Defaults to None.
                kwargs: Any possible extend parameters in the future.

            Returns:
                The k most similar documents to the specified text query.
                0 is dissimilar, 1 is the most similar.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            embedding = None
            if self.using_table_name in self.table2embeddings:
                embedding = self.table2embeddings[self.using_table_name].embed_query(query)
            else:
                from awadb import AwaEmbedding

                embedding = AwaEmbedding().Embedding(query)

            results: List[Tuple[Document, float]] = []

            not_include_fields: Set[str] = {"text_embedding", "_id"}
            retrieval_docs = self.similarity_search_by_vector(
                embedding,
                k,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
                not_include_fields_in_metadata=not_include_fields,
            )

            for doc in retrieval_docs:
                score = doc.metadata["score"]
                del doc.metadata["score"]
                doc_tuple = (doc, score)
                results.append(doc_tuple)

            return results

        def _similarity_search_with_relevance_scores(
            self,
            query: str,
            k: int = 4,
            **kwargs: Any,
        ) -> List[Tuple[Document, float]]:
            return self.similarity_search_with_score(query, k, **kwargs)

        def similarity_search_by_vector(
            self,
            embedding: Optional[List[float]] = None,
            k: int = DEFAULT_TOPN,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            not_include_fields_in_metadata: Optional[Set[str]] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Return docs most similar to embedding vector.

            Args:
                embedding: Embedding to look up documents similar to.
                k: Number of Documents to return. Defaults to 4.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter: Filter by metadata. Defaults to None.
                not_incude_fields_in_metadata: Not include meta fields of each document.

            Returns:
                List of Documents which are the most similar to the query vector.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            results: List[Document] = []

            if embedding is None:
                return results

            show_results = self.awadb_client.Search(
                embedding,
                k,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
                not_include_fields=not_include_fields_in_metadata,
            )

            if show_results.__len__() == 0:
                return results

            for item_detail in show_results[0]["ResultItems"]:
                content = ""
                meta_data = {}
                for item_key in item_detail:
                    if item_key == "embedding_text":
                        content = item_detail[item_key]
                        continue
                    elif not_include_fields_in_metadata is not None:
                        if item_key in not_include_fields_in_metadata:
                            continue
                    meta_data[item_key] = item_detail[item_key]
                results.append(Document(page_content=content, metadata=meta_data))
            return results

        def max_marginal_relevance_search(
            self,
            query: str,
            k: int = 4,
            fetch_k: int = 20,
            lambda_mult: float = 0.5,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Return docs selected using the maximal marginal relevance.

            Maximal marginal relevance optimizes for similarity to query AND diversity
            among selected documents.

            Args:
                query: Text to look up documents similar to.
                k: Number of Documents to return. Defaults to 4.
                fetch_k: Number of Documents to fetch to pass to MMR algorithm.
                lambda_mult: Number between 0 and 1 that determines the degree
                            of diversity among the results with 0 corresponding
                            to maximum diversity and 1 to minimum diversity.
                            Defaults to 0.5.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter (Optional[dict]): Filter by metadata. Defaults to None.

            Returns:
                List of Documents selected by maximal marginal relevance.
            """
            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            embedding: List[float] = []
            if self.using_table_name in self.table2embeddings:
                embedding = self.table2embeddings[self.using_table_name].embed_query(query)
            else:
                from awadb import AwaEmbedding

                embedding = AwaEmbedding().Embedding(query)

            if embedding.__len__() == 0:
                return []

            results = self.max_marginal_relevance_search_by_vector(
                embedding,
                k,
                fetch_k,
                lambda_mult=lambda_mult,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
            )
            return results

        def max_marginal_relevance_search_by_vector(
            self,
            embedding: List[float],
            k: int = 4,
            fetch_k: int = 20,
            lambda_mult: float = 0.5,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Return docs selected using the maximal marginal relevance.

            Maximal marginal relevance optimizes for similarity to query AND diversity
            among selected documents.

            Args:
                embedding: Embedding to look up documents similar to.
                k: Number of Documents to return. Defaults to 4.
                fetch_k: Number of Documents to fetch to pass to MMR algorithm.
                lambda_mult: Number between 0 and 1 that determines the degree
                            of diversity among the results with 0 corresponding
                            to maximum diversity and 1 to minimum diversity.
                            Defaults to 0.5.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter (Optional[dict]): Filter by metadata. Defaults to None.

            Returns:
                List of Documents selected by maximal marginal relevance.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            results: List[Document] = []

            if embedding is None:
                return results

            not_include_fields: set = {"_id", "score"}
            retrieved_docs = self.similarity_search_by_vector(
                embedding,
                fetch_k,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
                not_include_fields_in_metadata=not_include_fields,
            )

            top_embeddings = []

            for doc in retrieved_docs:
                top_embeddings.append(doc.metadata["text_embedding"])

            selected_docs = maximal_marginal_relevance(
                np.array(embedding, dtype=np.float32), embedding_list=top_embeddings
            )

            for s_id in selected_docs:
                if "text_embedding" in retrieved_docs[s_id].metadata:
                    del retrieved_docs[s_id].metadata["text_embedding"]
                    results.append(retrieved_docs[s_id])
            return results

        def get(
            self,
            ids: Optional[List[str]] = None,
            text_in_page_content: Optional[str] = None,
            meta_filter: Optional[dict] = None,
            not_include_fields: Optional[Set[str]] = None,
            limit: Optional[int] = None,
            **kwargs: Any,
        ) -> Dict[str, Document]:
            """Return docs according ids.

            Args:
                ids: The ids of the embedding vectors.
                text_in_page_content: Filter by the text in page_content of Document.
                meta_filter: Filter by any metadata of the document.
                not_include_fields: Not pack the specified fields of each document.
                limit: The number of documents to return. Defaults to 5. Optional.

            Returns:
                Documents which satisfy the input conditions.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            docs_detail = self.awadb_client.Get(
                ids=ids,
                text_in_page_content=text_in_page_content,
                meta_filter=meta_filter,
                not_include_fields=not_include_fields,
                limit=limit,
            )

            results: Dict[str, Document] = {}
            for doc_detail in docs_detail:
                content = ""
                meta_info = {}
                for field in doc_detail:
                    if field == "embedding_text":
                        content = doc_detail[field]
                        continue
                    elif field == "text_embedding" or field == "_id":
                        continue

                    meta_info[field] = doc_detail[field]

                doc = Document(page_content=content, metadata=meta_info)
                results[doc_detail["_id"]] = doc
            return results

        def delete(
            self,
            ids: Optional[List[str]] = None,
            **kwargs: Any,
        ) -> Optional[bool]:
            """Delete the documents which have the specified ids.

            Args:
                ids: The ids of the embedding vectors.
                **kwargs: Other keyword arguments that subclasses might use.

            Returns:
                Optional[bool]: True if deletion is successful.
                False otherwise, None if not implemented.
            """
            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")
            ret: Optional[bool] = None
            if ids is None or ids.__len__() == 0:
                return ret
            ret = self.awadb_client.Delete(ids)
            return ret

        def update(
            self,
            ids: List[str],
            texts: Iterable[str],
            metadatas: Optional[List[dict]] = None,
            **kwargs: Any,
        ) -> List[str]:
            """Update the documents which have the specified ids.

            Args:
                ids: The id list of the updating embedding vector.
                texts: The texts of the updating documents.
                metadatas: The metadatas of the updating documents.
            Returns:
                the ids of the updated documents.
            """

            if self.awadb_client is None:
                raise ValueError("AwaDB client is None!!!")

            return self.awadb_client.UpdateTexts(
                ids=ids, text_field_name="embedding_text", texts=texts, metadatas=metadatas
            )

        def create_table(
            self,
            table_name: str,
            **kwargs: Any,
        ) -> bool:
            """Create a new table."""

            if self.awadb_client is None:
                return False

            ret = self.awadb_client.Create(table_name)

            if ret:
                self.using_table_name = table_name
            return ret

        def use(
            self,
            table_name: str,
            **kwargs: Any,
        ) -> bool:
            """Use the specified table. Don't know the tables, please invoke list_tables."""

            if self.awadb_client is None:
                return False

            ret = self.awadb_client.Use(table_name)
            if ret:
                self.using_table_name = table_name

            return ret

        def list_tables(
            self,
            **kwargs: Any,
        ) -> List[str]:
            """List all the tables created by the client."""

            if self.awadb_client is None:
                return []

            return self.awadb_client.ListAllTables()

        def get_current_table(
            self,
            **kwargs: Any,
        ) -> str:
            """Get the current table."""

            return self.using_table_name

        @classmethod
        def from_texts(
            cls: Type[AwaDB],
            texts: List[str],
            embedding: Optional[Embeddings] = None,
            metadatas: Optional[List[dict]] = None,
            table_name: str = _DEFAULT_TABLE_NAME,
            log_and_data_dir: Optional[str] = None,
            client: Optional[awadb.Client] = None,
            **kwargs: Any,
        ) -> AwaDB:
            """Create an AwaDB vectorstore from a raw documents.

            Args:
                texts (List[str]): List of texts to add to the table.
                embedding (Optional[Embeddings]): Embedding function. Defaults to None.
                metadatas (Optional[List[dict]]): List of metadatas. Defaults to None.
                table_name (str): Name of the table to create.
                log_and_data_dir (Optional[str]): Directory of logging and persistence.
                client (Optional[awadb.Client]): AwaDB client

            Returns:
                AwaDB: AwaDB vectorstore.
            """
            awadb_client = cls(
                table_name=table_name,
                embedding=embedding,
                log_and_data_dir=log_and_data_dir,
                client=client,
            )
            awadb_client.add_texts(texts=texts, metadatas=metadatas)
            return awadb_client

        @classmethod
        def from_documents(
            cls: Type[AwaDB],
            documents: List[Document],
            embedding: Optional[Embeddings] = None,
            table_name: str = _DEFAULT_TABLE_NAME,
            log_and_data_dir: Optional[str] = None,
            client: Optional[awadb.Client] = None,
            **kwargs: Any,
        ) -> AwaDB:
            """Create an AwaDB vectorstore from a list of documents.

            If a log_and_data_dir specified, the table will be persisted there.

            Args:
                documents (List[Document]): List of documents to add to the vectorstore.
                embedding (Optional[Embeddings]): Embedding function. Defaults to None.
                table_name (str): Name of the table to create.
                log_and_data_dir (Optional[str]): Directory to persist the table.
                client (Optional[awadb.Client]): AwaDB client.
                Any: Any possible parameters in the future

            Returns:
                AwaDB: AwaDB vectorstore.
            """
            texts = [doc.page_content for doc in documents]
            metadatas = [doc.metadata for doc in documents]
            return cls.from_texts(
                texts=texts,
                embedding=embedding,
                metadatas=metadatas,
                table_name=table_name,
                log_and_data_dir=log_and_data_dir,
                client=client,
            )