# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors # # SPDX-License-Identifier: Apache-2.0 # # Portions derived from https://github.com/microsoft/autogen are under the MIT License. # SPDX-License-Identifier: MIT import os from typing import Any, Callable from ....import_utils import optional_import_block, require_optional_import from .base import Document, ItemID, QueryResults, VectorDB from .utils import chroma_results_to_query_results, filter_results_by_distance, get_logger with optional_import_block() as result: import chromadb import chromadb.errors import chromadb.utils.embedding_functions as ef from chromadb.api.models.Collection import Collection if result.is_successful and chromadb.__version__ < "0.4.15": raise ImportError("Please upgrade chromadb to version 0.4.15 or later.") CHROMADB_MAX_BATCH_SIZE = os.environ.get("CHROMADB_MAX_BATCH_SIZE", 40000) logger = get_logger(__name__) @require_optional_import("chromadb", "retrievechat") class ChromaVectorDB(VectorDB): """A vector database that uses ChromaDB as the backend.""" def __init__( self, *, client=None, path: str = "tmp/db", embedding_function: Callable = None, metadata: dict = None, **kwargs ) -> None: """Initialize the vector database. Args: client: chromadb.Client | The client object of the vector database. Default is None. If provided, it will use the client object directly and ignore other arguments. path: str | The path to the vector database. Default is `tmp/db`. The default was `None` for version `<=0.2.24`. embedding_function: Callable | The embedding function used to generate the vector representation of the documents. Default is None, SentenceTransformerEmbeddingFunction("all-MiniLM-L6-v2") will be used. metadata: dict | The metadata of the vector database. Default is None. If None, it will use this setting: `{"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}`. For more details of the metadata, please refer to [distances](https://github.com/nmslib/hnswlib#supported-distances), [hnsw](https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184), and [ALGO_PARAMS](https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md). kwargs: dict | Additional keyword arguments. Returns: None """ self.client = client self.path = path self.embedding_function = ( ef.SentenceTransformerEmbeddingFunction("all-MiniLM-L6-v2") if embedding_function is None else embedding_function ) self.metadata = metadata if metadata else {"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32} if not self.client: if self.path is not None: self.client = chromadb.PersistentClient(path=self.path, **kwargs) else: self.client = chromadb.Client(**kwargs) self.active_collection = None self.type = "chroma" def create_collection( self, collection_name: str, overwrite: bool = False, get_or_create: bool = True ) -> "Collection": """Create a collection in the vector database. Case 1. if the collection does not exist, create the collection. Case 2. the collection exists, if overwrite is True, it will overwrite the collection. Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection, otherwise it raise a ValueError. Args: collection_name: str | The name of the collection. overwrite: bool | Whether to overwrite the collection if it exists. Default is False. get_or_create: bool | Whether to get the collection if it exists. Default is True. Returns: Collection | The collection object. """ try: if self.active_collection and self.active_collection.name == collection_name: collection = self.active_collection else: collection = self.client.get_collection(collection_name, embedding_function=self.embedding_function) except (ValueError, chromadb.errors.ChromaError): collection = None if collection is None: return self.client.create_collection( collection_name, embedding_function=self.embedding_function, get_or_create=get_or_create, metadata=self.metadata, ) elif overwrite: self.client.delete_collection(collection_name) return self.client.create_collection( collection_name, embedding_function=self.embedding_function, get_or_create=get_or_create, metadata=self.metadata, ) elif get_or_create: return collection else: raise ValueError(f"Collection {collection_name} already exists.") def get_collection(self, collection_name: str = None) -> "Collection": """Get the collection from the vector database. Args: collection_name: str | The name of the collection. Default is None. If None, return the current active collection. Returns: Collection | The collection object. """ if collection_name is None: if self.active_collection is None: raise ValueError("No collection is specified.") else: logger.info( f"No collection is specified. Using current active collection {self.active_collection.name}." ) else: if not (self.active_collection and self.active_collection.name == collection_name): self.active_collection = self.client.get_collection( name=collection_name, embedding_function=self.embedding_function ) return self.active_collection def delete_collection(self, collection_name: str) -> None: """Delete the collection from the vector database. Args: collection_name: str | The name of the collection. Returns: None """ self.client.delete_collection(collection_name) if self.active_collection and self.active_collection.name == collection_name: self.active_collection = None def _batch_insert( self, collection: "Collection", embeddings=None, ids=None, metadatas=None, documents=None, upsert=False ) -> None: batch_size = int(CHROMADB_MAX_BATCH_SIZE) for i in range(0, len(documents), min(batch_size, len(documents))): end_idx = i + min(batch_size, len(documents) - i) collection_kwargs = { "documents": documents[i:end_idx], "ids": ids[i:end_idx], "metadatas": metadatas[i:end_idx] if metadatas else None, "embeddings": embeddings[i:end_idx] if embeddings else None, } if upsert: collection.upsert(**collection_kwargs) else: collection.add(**collection_kwargs) def insert_docs(self, docs: list[Document], collection_name: str = None, upsert: bool = False) -> None: """Insert documents into the collection of the vector database. Args: docs: List[Document] | A list of documents. Each document is a TypedDict `Document`. collection_name: str | The name of the collection. Default is None. upsert: bool | Whether to update the document if it exists. Default is False. kwargs: Dict | Additional keyword arguments. Returns: None """ if not docs: return if docs[0].get("content") is None: raise ValueError("The document content is required.") if docs[0].get("id") is None: raise ValueError("The document id is required.") documents = [doc.get("content") for doc in docs] ids = [doc.get("id") for doc in docs] collection = self.get_collection(collection_name) if docs[0].get("embedding") is None: logger.info( "No content embedding is provided. Will use the VectorDB's embedding function to generate the content embedding." ) embeddings = None else: embeddings = [doc.get("embedding") for doc in docs] metadatas = None if docs[0].get("metadata") is None else [doc.get("metadata") for doc in docs] self._batch_insert(collection, embeddings, ids, metadatas, documents, upsert) def update_docs(self, docs: list[Document], collection_name: str = None) -> None: """Update documents in the collection of the vector database. Args: docs: List[Document] | A list of documents. collection_name: str | The name of the collection. Default is None. Returns: None """ self.insert_docs(docs, collection_name, upsert=True) def delete_docs(self, ids: list[ItemID], collection_name: str = None, **kwargs) -> None: """Delete documents from the collection of the vector database. Args: ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`. collection_name: str | The name of the collection. Default is None. kwargs: Dict | Additional keyword arguments. Returns: None """ collection = self.get_collection(collection_name) collection.delete(ids, **kwargs) def retrieve_docs( self, queries: list[str], collection_name: str = None, n_results: int = 10, distance_threshold: float = -1, **kwargs: Any, ) -> QueryResults: """Retrieve documents from the collection of the vector database based on the queries. Args: queries: List[str] | A list of queries. Each query is a string. collection_name: str | The name of the collection. Default is None. n_results: int | The number of relevant documents to return. Default is 10. distance_threshold: float | The threshold for the distance score, only distance smaller than it will be returned. Don't filter with it if `< 0`. Default is -1. kwargs: Dict | Additional keyword arguments. Returns: QueryResults | The query results. Each query result is a list of list of tuples containing the document and the distance. """ collection = self.get_collection(collection_name) if isinstance(queries, str): queries = [queries] results = collection.query( query_texts=queries, n_results=n_results, **kwargs, ) results["contents"] = results.pop("documents") results = chroma_results_to_query_results(results) results = filter_results_by_distance(results, distance_threshold) return results @staticmethod def _chroma_get_results_to_list_documents(data_dict) -> list[Document]: """Converts a dictionary with list values to a list of Document. Args: data_dict: A dictionary where keys map to lists or None. Returns: List[Document] | The list of Document. Example: ```python data_dict = { "key1s": [1, 2, 3], "key2s": ["a", "b", "c"], "key3s": None, "key4s": ["x", "y", "z"], } results = [ {"key1": 1, "key2": "a", "key4": "x"}, {"key1": 2, "key2": "b", "key4": "y"}, {"key1": 3, "key2": "c", "key4": "z"}, ] ``` """ results = [] keys = [key for key in data_dict if data_dict[key] is not None] for i in range(len(data_dict[keys[0]])): sub_dict = {} for key in data_dict: if data_dict[key] is not None and len(data_dict[key]) > i: sub_dict[key[:-1]] = data_dict[key][i] results.append(sub_dict) return results def get_docs_by_ids( self, ids: list[ItemID] = None, collection_name: str = None, include=None, **kwargs ) -> list[Document]: """Retrieve documents from the collection of the vector database based on the ids. Args: ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None. collection_name: str | The name of the collection. Default is None. include: List[str] | The fields to include. Default is None. If None, will include ["metadatas", "documents"], ids will always be included. kwargs: dict | Additional keyword arguments. Returns: List[Document] | The results. """ collection = self.get_collection(collection_name) include = include if include else ["metadatas", "documents"] results = collection.get(ids, include=include, **kwargs) results = self._chroma_get_results_to_list_documents(results) return results