Added support for chromadb cloud (#3436)
This commit is contained in:
@@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -39,3 +44,5 @@ Here are the parameters available for configuring Chroma:
|
||||
| `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` |
|
||||
| `api_key` | ChromaDB Cloud API key (for cloud usage) | `None` |
|
||||
| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` |
|
||||
@@ -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")
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user