langroid 0.31.1__py3-none-any.whl → 0.33.3__py3-none-any.whl

langroid/vector_store/base.py

@@ -1,357 +0,0 @@
-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
-
-    @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

@@ -1,214 +0,0 @@
-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