From f5e0fb9e4b3f64f0b1c3d85ae0b986d9db1e1a60 Mon Sep 17 00:00:00 2001 From: Parshva Daftari <89991302+parshvadaftari@users.noreply.github.com> Date: Tue, 9 Sep 2025 22:22:09 +0530 Subject: [PATCH] Added support for chromadb cloud (#3436) --- docs/components/vectordbs/dbs/chroma.mdx | 11 ++++++++-- mem0/configs/vector_stores/chroma.py | 28 ++++++++++++++++++++---- mem0/vector_stores/chroma.py | 13 +++++++++++ 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/docs/components/vectordbs/dbs/chroma.mdx b/docs/components/vectordbs/dbs/chroma.mdx index ed798ebbe..2e546b883 100644 --- a/docs/components/vectordbs/dbs/chroma.mdx +++ b/docs/components/vectordbs/dbs/chroma.mdx @@ -1,7 +1,9 @@ -[Chroma](https://www.trychroma.com/) is an AI-native open-source vector database that simplifies building LLM apps by providing tools for storing, embedding, and searching embeddings with a focus on simplicity and speed. +[Chroma](https://www.trychroma.com/) is an AI-native open-source vector database that simplifies building LLM apps by providing tools for storing, embedding, and searching embeddings with a focus on simplicity and speed. It supports both local deployment and cloud hosting through ChromaDB Cloud. ### Usage +#### Local Installation + ```python import os from mem0 import Memory @@ -14,6 +16,9 @@ config = { "config": { "collection_name": "test", "path": "db", + # Optional: ChromaDB Cloud configuration + # "api_key": "your-chroma-cloud-api-key", + # "tenant": "your-chroma-cloud-tenant-id", } } } @@ -38,4 +43,6 @@ Here are the parameters available for configuring Chroma: | `client` | Custom client for Chroma | `None` | | `path` | Path for the Chroma database | `db` | | `host` | The host where the Chroma server is running | `None` | -| `port` | The port where the Chroma server is running | `None` | \ No newline at end of file +| `port` | The port where the Chroma server is running | `None` | +| `api_key` | ChromaDB Cloud API key (for cloud usage) | `None` | +| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` | \ No newline at end of file diff --git a/mem0/configs/vector_stores/chroma.py b/mem0/configs/vector_stores/chroma.py index cf9104237..764e6a381 100644 --- a/mem0/configs/vector_stores/chroma.py +++ b/mem0/configs/vector_stores/chroma.py @@ -10,17 +10,37 @@ class ChromaDbConfig(BaseModel): raise ImportError("The 'chromadb' library is required. Please install it using 'pip install chromadb'.") Client: ClassVar[type] = Client - collection_name: str = Field("mem0", description="Default name for the collection") + collection_name: str = Field("mem0", description="Default name for the collection/database") client: Optional[Client] = Field(None, description="Existing ChromaDB client instance") path: Optional[str] = Field(None, description="Path to the database directory") host: Optional[str] = Field(None, description="Database connection remote host") port: Optional[int] = Field(None, description="Database connection remote port") + # ChromaDB Cloud configuration + api_key: Optional[str] = Field(None, description="ChromaDB Cloud API key") + tenant: Optional[str] = Field(None, description="ChromaDB Cloud tenant ID") @model_validator(mode="before") - def check_host_port_or_path(cls, values): + def check_connection_config(cls, values): host, port, path = values.get("host"), values.get("port"), values.get("path") - if not path and not (host and port): - raise ValueError("Either 'host' and 'port' or 'path' must be provided.") + api_key, tenant = values.get("api_key"), values.get("tenant") + + # Check if cloud configuration is provided + cloud_config = bool(api_key and tenant) + + # If cloud configuration is provided, remove any default path that might have been added + if cloud_config and path == "/tmp/chroma": + values.pop("path", None) + return values + + # Check if local/server configuration is provided (excluding default tmp path for cloud config) + local_config = bool(path and path != "/tmp/chroma") or bool(host and port) + + if not cloud_config and not local_config: + raise ValueError("Either ChromaDB Cloud configuration (api_key, tenant) or local configuration (path or host/port) must be provided.") + + if cloud_config and local_config: + raise ValueError("Cannot specify both cloud configuration and local configuration. Choose one.") + return values @model_validator(mode="before") diff --git a/mem0/vector_stores/chroma.py b/mem0/vector_stores/chroma.py index 681d4626c..8d231719c 100644 --- a/mem0/vector_stores/chroma.py +++ b/mem0/vector_stores/chroma.py @@ -28,6 +28,8 @@ class ChromaDB(VectorStoreBase): host: Optional[str] = None, port: Optional[int] = None, path: Optional[str] = None, + api_key: Optional[str] = None, + tenant: Optional[str] = None, ): """ Initialize the Chromadb vector store. @@ -38,10 +40,21 @@ class ChromaDB(VectorStoreBase): host (str, optional): Host address for chromadb server. Defaults to None. port (int, optional): Port for chromadb server. Defaults to None. path (str, optional): Path for local chromadb database. Defaults to None. + api_key (str, optional): ChromaDB Cloud API key. Defaults to None. + tenant (str, optional): ChromaDB Cloud tenant ID. Defaults to None. """ if client: self.client = client + elif api_key and tenant: + # Initialize ChromaDB Cloud client + logger.info("Initializing ChromaDB Cloud client") + self.client = chromadb.CloudClient( + api_key=api_key, + tenant=tenant, + database="mem0" # Use fixed database name for cloud + ) else: + # Initialize local or server client self.settings = Settings(anonymized_telemetry=False) if host and port: