langroid 0.33.6__py3-none-any.whl → 0.33.7__py3-none-any.whl

langroid/vector_store/base.py

@@ -0,0 +1,359 @@
+import copy
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Sequence, Tuple, Type
+
+import numpy as np
+import pandas as pd
+
+from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
+from langroid.embedding_models.models import OpenAIEmbeddingsConfig
+from langroid.mytypes import DocMetaData, Document
+from langroid.pydantic_v1 import BaseSettings
+from langroid.utils.algorithms.graph import components, topological_sort
+from langroid.utils.configuration import settings
+from langroid.utils.object_registry import ObjectRegistry
+from langroid.utils.output.printing import print_long_text
+from langroid.utils.pandas_utils import stringify
+from langroid.utils.pydantic_utils import flatten_dict
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStoreConfig(BaseSettings):
+    type: str = ""  # deprecated, keeping it for backward compatibility
+    collection_name: str | None = "temp"
+    replace_collection: bool = False  # replace collection if it already exists
+    storage_path: str = ".qdrant/data"
+    cloud: bool = False
+    batch_size: int = 200
+    embedding: EmbeddingModelsConfig = OpenAIEmbeddingsConfig(
+        model_type="openai",
+    )
+    embedding_model: Optional[EmbeddingModel] = None
+    timeout: int = 60
+    host: str = "127.0.0.1"
+    port: int = 6333
+    # used when parsing search results back as Document objects
+    document_class: Type[Document] = Document
+    metadata_class: Type[DocMetaData] = DocMetaData
+    # compose_file: str = "langroid/vector_store/docker-compose-qdrant.yml"
+
+
+class VectorStore(ABC):
+    """
+    Abstract base class for a vector store.
+    """
+
+    def __init__(self, config: VectorStoreConfig):
+        self.config = config
+        if config.embedding_model is None:
+            self.embedding_model = EmbeddingModel.create(config.embedding)
+        else:
+            self.embedding_model = config.embedding_model
+
+    @staticmethod
+    def create(config: VectorStoreConfig) -> Optional["VectorStore"]:
+        from langroid.vector_store.chromadb import ChromaDB, ChromaDBConfig
+        from langroid.vector_store.lancedb import LanceDB, LanceDBConfig
+        from langroid.vector_store.meilisearch import MeiliSearch, MeiliSearchConfig
+        from langroid.vector_store.momento import MomentoVI, MomentoVIConfig
+        from langroid.vector_store.qdrantdb import QdrantDB, QdrantDBConfig
+
+        if isinstance(config, QdrantDBConfig):
+            return QdrantDB(config)
+        elif isinstance(config, ChromaDBConfig):
+            return ChromaDB(config)
+        elif isinstance(config, MomentoVIConfig):
+            return MomentoVI(config)
+        elif isinstance(config, LanceDBConfig):
+            return LanceDB(config)
+        elif isinstance(config, MeiliSearchConfig):
+            return MeiliSearch(config)
+
+        else:
+            logger.warning(
+                f"""
+                Unknown vector store config: {config.__repr_name__()},
+                so skipping vector store creation!
+                If you intended to use a vector-store, please set a specific 
+                vector-store in your script, typically in the `vecdb` field of a 
+                `ChatAgentConfig`, otherwise set it to None.
+                """
+            )
+            return None
+
+    @abstractmethod
+    def clear_empty_collections(self) -> int:
+        """Clear all empty collections in the vector store.
+        Returns the number of collections deleted.
+        """
+        pass
+
+    @abstractmethod
+    def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
+        """
+        Clear all collections in the vector store.
+
+        Args:
+            really (bool, optional): Whether to really clear all collections.
+                Defaults to False.
+            prefix (str, optional): Prefix of collections to clear.
+        Returns:
+            int: Number of collections deleted.
+        """
+        pass
+
+    @abstractmethod
+    def list_collections(self, empty: bool = False) -> List[str]:
+        """List all collections in the vector store
+        (only non empty collections if empty=False).
+        """
+        pass
+
+    def set_collection(self, collection_name: str, replace: bool = False) -> None:
+        """
+        Set the current collection to the given collection name.
+        Args:
+            collection_name (str): Name of the collection.
+            replace (bool, optional): Whether to replace the collection if it
+                already exists. Defaults to False.
+        """
+
+        self.config.collection_name = collection_name
+        self.config.replace_collection = replace
+        if replace:
+            self.create_collection(collection_name, replace=True)
+
+    @abstractmethod
+    def create_collection(self, collection_name: str, replace: bool = False) -> None:
+        """Create a collection with the given name.
+        Args:
+            collection_name (str): Name of the collection.
+            replace (bool, optional): Whether to replace the
+                collection if it already exists. Defaults to False.
+        """
+        pass
+
+    @abstractmethod
+    def add_documents(self, documents: Sequence[Document]) -> None:
+        pass
+
+    def compute_from_docs(self, docs: List[Document], calc: str) -> str:
+        """Compute a result on a set of documents,
+        using a dataframe calc string like `df.groupby('state')['income'].mean()`.
+        """
+        # convert each doc to a dict, using dotted paths for nested fields
+        dicts = [flatten_dict(doc.dict(by_alias=True)) for doc in docs]
+        df = pd.DataFrame(dicts)
+
+        try:
+            result = pd.eval(  # safer than eval but limited to single expression
+                calc,
+                engine="python",
+                parser="pandas",
+                local_dict={"df": df},
+            )
+        except Exception as e:
+            # return error message so LLM can fix the calc string if needed
+            err = f"""
+            Error encountered in pandas eval: {str(e)}
+            """
+            if isinstance(e, KeyError) and "not in index" in str(e):
+                # Pd.eval sometimes fails on a perfectly valid exprn like
+                # df.loc[..., 'column'] with a KeyError.
+                err += """
+                Maybe try a different way, e.g. 
+                instead of df.loc[..., 'column'], try df.loc[...]['column']
+                """
+            return err
+        return stringify(result)
+
+    def maybe_add_ids(self, documents: Sequence[Document]) -> None:
+        """Add ids to metadata if absent, since some
+        vecdbs don't like having blank ids."""
+        for d in documents:
+            if d.metadata.id in [None, ""]:
+                d.metadata.id = ObjectRegistry.new_id()
+
+    @abstractmethod
+    def similar_texts_with_scores(
+        self,
+        text: str,
+        k: int = 1,
+        where: Optional[str] = None,
+    ) -> List[Tuple[Document, float]]:
+        """
+        Find k most similar texts to the given text, in terms of vector distance metric
+        (e.g., cosine similarity).
+
+        Args:
+            text (str): The text to find similar texts for.
+            k (int, optional): Number of similar texts to retrieve. Defaults to 1.
+            where (Optional[str], optional): Where clause to filter the search.
+
+        Returns:
+            List[Tuple[Document,float]]: List of (Document, score) tuples.
+
+        """
+        pass
+
+    def add_context_window(
+        self, docs_scores: List[Tuple[Document, float]], neighbors: int = 0
+    ) -> List[Tuple[Document, float]]:
+        """
+        In each doc's metadata, there may be a window_ids field indicating
+        the ids of the chunks around the current chunk.
+        These window_ids may overlap, so we
+        - coalesce each overlapping groups into a single window (maintaining ordering),
+        - create a new document for each part, preserving metadata,
+
+        We may have stored a longer set of window_ids than we need during chunking.
+        Now, we just want `neighbors` on each side of the center of the window_ids list.
+
+        Args:
+            docs_scores (List[Tuple[Document, float]]): List of pairs of documents
+                to add context windows to together with their match scores.
+            neighbors (int, optional): Number of neighbors on "each side" of match to
+                retrieve. Defaults to 0.
+                "Each side" here means before and after the match,
+                in the original text.
+
+        Returns:
+            List[Tuple[Document, float]]: List of (Document, score) tuples.
+        """
+        # We return a larger context around each match, i.e.
+        # a window of `neighbors` on each side of the match.
+        docs = [d for d, s in docs_scores]
+        scores = [s for d, s in docs_scores]
+        if neighbors == 0:
+            return docs_scores
+        doc_chunks = [d for d in docs if d.metadata.is_chunk]
+        if len(doc_chunks) == 0:
+            return docs_scores
+        window_ids_list = []
+        id2metadata = {}
+        # id -> highest score of a doc it appears in
+        id2max_score: Dict[int | str, float] = {}
+        for i, d in enumerate(docs):
+            window_ids = d.metadata.window_ids
+            if len(window_ids) == 0:
+                window_ids = [d.id()]
+            id2metadata.update({id: d.metadata for id in window_ids})
+
+            id2max_score.update(
+                {id: max(id2max_score.get(id, 0), scores[i]) for id in window_ids}
+            )
+            n = len(window_ids)
+            chunk_idx = window_ids.index(d.id())
+            neighbor_ids = window_ids[
+                max(0, chunk_idx - neighbors) : min(n, chunk_idx + neighbors + 1)
+            ]
+            window_ids_list += [neighbor_ids]
+
+        # window_ids could be from different docs,
+        # and they may overlap, so we coalesce overlapping groups into
+        # separate windows.
+        window_ids_list = self.remove_overlaps(window_ids_list)
+        final_docs = []
+        final_scores = []
+        for w in window_ids_list:
+            metadata = copy.deepcopy(id2metadata[w[0]])
+            metadata.window_ids = w
+            document = Document(
+                content=" ".join([d.content for d in self.get_documents_by_ids(w)]),
+                metadata=metadata,
+            )
+            # make a fresh id since content is in general different
+            document.metadata.id = ObjectRegistry.new_id()
+            final_docs += [document]
+            final_scores += [max(id2max_score[id] for id in w)]
+        return list(zip(final_docs, final_scores))
+
+    @staticmethod
+    def remove_overlaps(windows: List[List[str]]) -> List[List[str]]:
+        """
+        Given a collection of windows, where each window is a sequence of ids,
+        identify groups of overlapping windows, and for each overlapping group,
+        order the chunk-ids using topological sort so they appear in the original
+        order in the text.
+
+        Args:
+            windows (List[int|str]): List of windows, where each window is a
+                sequence of ids.
+
+        Returns:
+            List[int|str]: List of windows, where each window is a sequence of ids,
+                and no two windows overlap.
+        """
+        ids = set(id for w in windows for id in w)
+        # id -> {win -> # pos}
+        id2win2pos: Dict[str, Dict[int, int]] = {id: {} for id in ids}
+
+        for i, w in enumerate(windows):
+            for j, id in enumerate(w):
+                id2win2pos[id][i] = j
+
+        n = len(windows)
+        # relation between windows:
+        order = np.zeros((n, n), dtype=np.int8)
+        for i, w in enumerate(windows):
+            for j, x in enumerate(windows):
+                if i == j:
+                    continue
+                if len(set(w).intersection(x)) == 0:
+                    continue
+                id = list(set(w).intersection(x))[0]  # any common id
+                if id2win2pos[id][i] > id2win2pos[id][j]:
+                    order[i, j] = -1  # win i is before win j
+                else:
+                    order[i, j] = 1  # win i is after win j
+
+        # find groups of windows that overlap, like connected components in a graph
+        groups = components(np.abs(order))
+
+        # order the chunk-ids in each group using topological sort
+        new_windows = []
+        for g in groups:
+            # find total ordering among windows in group based on order matrix
+            # (this is a topological sort)
+            _g = np.array(g)
+            order_matrix = order[_g][:, _g]
+            ordered_window_indices = topological_sort(order_matrix)
+            ordered_window_ids = [windows[i] for i in _g[ordered_window_indices]]
+            flattened = [id for w in ordered_window_ids for id in w]
+            flattened_deduped = list(dict.fromkeys(flattened))
+            # Note we are not going to split these, and instead we'll return
+            # larger windows from concatenating the connected groups.
+            # This ensures context is retained for LLM q/a
+            new_windows += [flattened_deduped]
+
+        return new_windows
+
+    @abstractmethod
+    def get_all_documents(self, where: str = "") -> List[Document]:
+        """
+        Get all documents in the current collection, possibly filtered by `where`.
+        """
+        pass
+
+    @abstractmethod
+    def get_documents_by_ids(self, ids: List[str]) -> List[Document]:
+        """
+        Get documents by their ids.
+        Args:
+            ids (List[str]): List of document ids.
+
+        Returns:
+            List[Document]: List of documents
+        """
+        pass
+
+    @abstractmethod
+    def delete_collection(self, collection_name: str) -> None:
+        pass
+
+    def show_if_debug(self, doc_score_pairs: List[Tuple[Document, float]]) -> None:
+        if settings.debug:
+            for i, (d, s) in enumerate(doc_score_pairs):
+                print_long_text("red", "italic red", f"\nMATCH-{i}\n", d.content)

langroid/vector_store/chromadb.py

@@ -0,0 +1,214 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+from langroid.embedding_models.base import (
+    EmbeddingModelsConfig,
+)
+from langroid.embedding_models.models import OpenAIEmbeddingsConfig
+from langroid.exceptions import LangroidImportError
+from langroid.mytypes import Document
+from langroid.utils.configuration import settings
+from langroid.utils.output.printing import print_long_text
+from langroid.vector_store.base import VectorStore, VectorStoreConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaDBConfig(VectorStoreConfig):
+    collection_name: str = "temp"
+    storage_path: str = ".chroma/data"
+    embedding: EmbeddingModelsConfig = OpenAIEmbeddingsConfig()
+    host: str = "127.0.0.1"
+    port: int = 6333
+
+
+class ChromaDB(VectorStore):
+    def __init__(self, config: ChromaDBConfig = ChromaDBConfig()):
+        super().__init__(config)
+        try:
+            import chromadb
+        except ImportError:
+            raise LangroidImportError("chromadb", "chromadb")
+        self.config = config
+        self.embedding_fn = self.embedding_model.embedding_fn()
+        self.client = chromadb.Client(
+            chromadb.config.Settings(
+                # chroma_db_impl="duckdb+parquet",
+                # is_persistent=bool(config.storage_path),
+                persist_directory=config.storage_path,
+            )
+        )
+        if self.config.collection_name is not None:
+            self.create_collection(
+                self.config.collection_name,
+                replace=self.config.replace_collection,
+            )
+
+    def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
+        """Clear all collections in the vector store with the given prefix."""
+
+        if not really:
+            logger.warning("Not deleting all collections, set really=True to confirm")
+            return 0
+        coll = [c for c in self.client.list_collections() if c.name.startswith(prefix)]
+        if len(coll) == 0:
+            logger.warning(f"No collections found with prefix {prefix}")
+            return 0
+        n_empty_deletes = 0
+        n_non_empty_deletes = 0
+        for c in coll:
+            n_empty_deletes += c.count() == 0
+            n_non_empty_deletes += c.count() > 0
+            self.client.delete_collection(name=c.name)
+        logger.warning(
+            f"""
+            Deleted {n_empty_deletes} empty collections and
+            {n_non_empty_deletes} non-empty collections.
+            """
+        )
+        return n_empty_deletes + n_non_empty_deletes
+
+    def clear_empty_collections(self) -> int:
+        colls = self.client.list_collections()
+        n_deletes = 0
+        for coll in colls:
+            if coll.count() == 0:
+                n_deletes += 1
+                self.client.delete_collection(name=coll.name)
+        return n_deletes
+
+    def list_collections(self, empty: bool = False) -> List[str]:
+        """
+        List non-empty collections in the vector store.
+        Args:
+            empty (bool, optional): Whether to list empty collections.
+        Returns:
+            List[str]: List of non-empty collection names.
+        """
+        colls = self.client.list_collections()
+        if empty:
+            return [coll.name for coll in colls]
+        return [coll.name for coll in colls if coll.count() > 0]
+
+    def create_collection(self, collection_name: str, replace: bool = False) -> None:
+        """
+        Create a collection in the vector store, optionally replacing an existing
+            collection if `replace` is True.
+        Args:
+            collection_name (str): Name of the collection to create or replace.
+            replace (bool, optional): Whether to replace an existing collection.
+                Defaults to False.
+
+        """
+        self.config.collection_name = collection_name
+        if collection_name in self.list_collections(empty=True) and replace:
+            logger.warning(f"Replacing existing collection {collection_name}")
+            self.client.delete_collection(collection_name)
+        self.collection = self.client.create_collection(
+            name=self.config.collection_name,
+            embedding_function=self.embedding_fn,
+            get_or_create=not replace,
+        )
+
+    def add_documents(self, documents: Sequence[Document]) -> None:
+        super().maybe_add_ids(documents)
+        if documents is None:
+            return
+        contents: List[str] = [document.content for document in documents]
+        # convert metadatas to dicts so chroma can handle them
+        metadata_dicts: List[dict[str, Any]] = [
+            d.metadata.dict_bool_int() for d in documents
+        ]
+        for m in metadata_dicts:
+            # chroma does not handle non-atomic types in metadata
+            m["window_ids"] = ",".join(m["window_ids"])
+
+        ids = [str(d.id()) for d in documents]
+
+        colls = self.list_collections(empty=True)
+        if self.config.collection_name is None:
+            raise ValueError("No collection name set, cannot ingest docs")
+        if self.config.collection_name not in colls:
+            self.create_collection(self.config.collection_name, replace=True)
+
+        self.collection.add(
+            # embedding_models=embedding_models,
+            documents=contents,
+            metadatas=metadata_dicts,
+            ids=ids,
+        )
+
+    def get_all_documents(self, where: str = "") -> List[Document]:
+        filter = json.loads(where) if where else None
+        results = self.collection.get(
+            include=["documents", "metadatas"],
+            where=filter,
+        )
+        results["documents"] = [results["documents"]]
+        results["metadatas"] = [results["metadatas"]]
+        return self._docs_from_results(results)
+
+    def get_documents_by_ids(self, ids: List[str]) -> List[Document]:
+        # get them one by one since chroma mangles the order of the results
+        # when fetched from a list of ids.
+        results = [
+            self.collection.get(ids=[id], include=["documents", "metadatas"])
+            for id in ids
+        ]
+        final_results = {}
+        final_results["documents"] = [[r["documents"][0] for r in results]]
+        final_results["metadatas"] = [[r["metadatas"][0] for r in results]]
+        return self._docs_from_results(final_results)
+
+    def delete_collection(self, collection_name: str) -> None:
+        try:
+            self.client.delete_collection(name=collection_name)
+        except Exception:
+            pass
+
+    def similar_texts_with_scores(
+        self, text: str, k: int = 1, where: Optional[str] = None
+    ) -> List[Tuple[Document, float]]:
+        n = self.collection.count()
+        filter = json.loads(where) if where else None
+        results = self.collection.query(
+            query_texts=[text],
+            n_results=min(n, k),
+            where=filter,
+            include=["documents", "distances", "metadatas"],
+        )
+        docs = self._docs_from_results(results)
+        # chroma distances are 1 - cosine.
+        scores = [1 - s for s in results["distances"][0]]
+        return list(zip(docs, scores))
+
+    def _docs_from_results(self, results: Dict[str, Any]) -> List[Document]:
+        """
+        Helper function to convert results from ChromaDB to a list of Documents
+        Args:
+            results (dict): results from ChromaDB
+
+        Returns:
+            List[Document]: list of Documents
+        """
+        if len(results["documents"][0]) == 0:
+            return []
+        contents = results["documents"][0]
+        if settings.debug:
+            for i, c in enumerate(contents):
+                print_long_text("red", "italic red", f"MATCH-{i}", c)
+        metadatas = results["metadatas"][0]
+        for m in metadatas:
+            # restore the stringified list of window_ids into the original List[str]
+            if m["window_ids"].strip() == "":
+                m["window_ids"] = []
+            else:
+                m["window_ids"] = m["window_ids"].split(",")
+        docs = [
+            self.config.document_class(
+                content=d, metadata=self.config.metadata_class(**m)
+            )
+            for d, m in zip(contents, metadatas)
+        ]
+        return docs