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
|
### Usage
|
||||||
|
|
||||||
|
#### Local Installation
|
||||||
|
|
||||||
```python
|
```python
|
||||||
import os
|
import os
|
||||||
from mem0 import Memory
|
from mem0 import Memory
|
||||||
@@ -14,6 +16,9 @@ config = {
|
|||||||
"config": {
|
"config": {
|
||||||
"collection_name": "test",
|
"collection_name": "test",
|
||||||
"path": "db",
|
"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` |
|
| `path` | Path for the Chroma database | `db` |
|
||||||
| `host` | The host where the Chroma server is running | `None` |
|
| `host` | The host where the Chroma server is running | `None` |
|
||||||
| `port` | The port 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'.")
|
raise ImportError("The 'chromadb' library is required. Please install it using 'pip install chromadb'.")
|
||||||
Client: ClassVar[type] = Client
|
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")
|
client: Optional[Client] = Field(None, description="Existing ChromaDB client instance")
|
||||||
path: Optional[str] = Field(None, description="Path to the database directory")
|
path: Optional[str] = Field(None, description="Path to the database directory")
|
||||||
host: Optional[str] = Field(None, description="Database connection remote host")
|
host: Optional[str] = Field(None, description="Database connection remote host")
|
||||||
port: Optional[int] = Field(None, description="Database connection remote port")
|
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")
|
@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")
|
host, port, path = values.get("host"), values.get("port"), values.get("path")
|
||||||
if not path and not (host and port):
|
api_key, tenant = values.get("api_key"), values.get("tenant")
|
||||||
raise ValueError("Either 'host' and 'port' or 'path' must be provided.")
|
|
||||||
|
# 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
|
return values
|
||||||
|
|
||||||
@model_validator(mode="before")
|
@model_validator(mode="before")
|
||||||
|
|||||||
@@ -28,6 +28,8 @@ class ChromaDB(VectorStoreBase):
|
|||||||
host: Optional[str] = None,
|
host: Optional[str] = None,
|
||||||
port: Optional[int] = None,
|
port: Optional[int] = None,
|
||||||
path: Optional[str] = None,
|
path: Optional[str] = None,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
tenant: Optional[str] = None,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initialize the Chromadb vector store.
|
Initialize the Chromadb vector store.
|
||||||
@@ -38,10 +40,21 @@ class ChromaDB(VectorStoreBase):
|
|||||||
host (str, optional): Host address for chromadb server. Defaults to None.
|
host (str, optional): Host address for chromadb server. Defaults to None.
|
||||||
port (int, optional): Port 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.
|
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:
|
if client:
|
||||||
self.client = 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:
|
else:
|
||||||
|
# Initialize local or server client
|
||||||
self.settings = Settings(anonymized_telemetry=False)
|
self.settings = Settings(anonymized_telemetry=False)
|
||||||
|
|
||||||
if host and port:
|
if host and port:
|
||||||
|
|||||||
Reference in New Issue
Block a user