diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py index 1d6ad04a0285..877f37afccea 100644 --- a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/__init__.py @@ -1,4 +1,5 @@ from ._chroma_configs import ( + ChromaCloudVectorMemoryConfig, ChromaDBVectorMemoryConfig, CustomEmbeddingFunctionConfig, DefaultEmbeddingFunctionConfig, @@ -12,6 +13,7 @@ __all__ = [ "ChromaDBVectorMemory", "ChromaDBVectorMemoryConfig", + "ChromaCloudVectorMemoryConfig", "PersistentChromaDBVectorMemoryConfig", "HttpChromaDBVectorMemoryConfig", "DefaultEmbeddingFunctionConfig", diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py index 3f30caacdbde..e48e0212254a 100644 --- a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chroma_configs.py @@ -107,7 +107,7 @@ class ChromaDBVectorMemoryConfig(BaseModel): Added support for custom embedding functions via embedding_function_config. """ - client_type: Literal["persistent", "http"] + client_type: Literal["persistent", "http", "cloud"] collection_name: str = Field(default="memory_store", description="Name of the ChromaDB collection") distance_metric: str = Field(default="cosine", description="Distance metric for similarity search") k: int = Field(default=3, description="Number of results to return in queries") @@ -123,15 +123,22 @@ class ChromaDBVectorMemoryConfig(BaseModel): class PersistentChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): """Configuration for persistent ChromaDB memory.""" - client_type: Literal["persistent", "http"] = "persistent" + client_type: Literal["persistent", "http", "cloud"] = "persistent" persistence_path: str = Field(default="./chroma_db", description="Path for persistent storage") class HttpChromaDBVectorMemoryConfig(ChromaDBVectorMemoryConfig): """Configuration for HTTP ChromaDB memory.""" - client_type: Literal["persistent", "http"] = "http" + client_type: Literal["persistent", "http", "cloud"] = "http" host: str = Field(default="localhost", description="Host of the remote server") port: int = Field(default=8000, description="Port of the remote server") ssl: bool = Field(default=False, description="Whether to use HTTPS") headers: Dict[str, str] | None = Field(default=None, description="Headers to send to the server") + + +class ChromaCloudVectorMemoryConfig(ChromaDBVectorMemoryConfig): + """Configuration for Chroma Cloud memory.""" + + client_type: Literal["persistent", "http", "cloud"] = "cloud" + api_key: str = Field(description="API key for Chroma Cloud authentication") diff --git a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py index 6664b404407c..209481d9f140 100644 --- a/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py +++ b/python/packages/autogen-ext/src/autogen_ext/memory/chromadb/_chromadb.py @@ -6,12 +6,13 @@ from autogen_core.memory import Memory, MemoryContent, MemoryMimeType, MemoryQueryResult, UpdateContextResult from autogen_core.model_context import ChatCompletionContext from autogen_core.models import SystemMessage -from chromadb import HttpClient, PersistentClient +from chromadb import CloudClient, HttpClient, PersistentClient from chromadb.api.models.Collection import Collection from chromadb.api.types import Document, Metadata from typing_extensions import Self from ._chroma_configs import ( + ChromaCloudVectorMemoryConfig, ChromaDBVectorMemoryConfig, CustomEmbeddingFunctionConfig, DefaultEmbeddingFunctionConfig, @@ -54,9 +55,10 @@ class ChromaDBVectorMemory(Memory, Component[ChromaDBVectorMemoryConfig]): Args: config (ChromaDBVectorMemoryConfig | None): Configuration for the ChromaDB memory. If None, defaults to a PersistentChromaDBVectorMemoryConfig with default values. - Two config types are supported: + Three config types are supported: * PersistentChromaDBVectorMemoryConfig: For local storage * HttpChromaDBVectorMemoryConfig: For connecting to a remote ChromaDB server + * ChromaCloudVectorMemoryConfig: For connecting to Chroma Cloud Example: @@ -70,6 +72,7 @@ class ChromaDBVectorMemory(Memory, Component[ChromaDBVectorMemoryConfig]): from autogen_core.memory import MemoryContent, MemoryMimeType from autogen_ext.memory.chromadb import ( ChromaDBVectorMemory, + ChromaCloudVectorMemoryConfig, PersistentChromaDBVectorMemoryConfig, SentenceTransformerEmbeddingFunctionConfig, OpenAIEmbeddingFunctionConfig, @@ -118,6 +121,19 @@ async def main() -> None: ) ) + # Using Chroma Cloud + cloud_memory = ChromaDBVectorMemory( + config=ChromaCloudVectorMemoryConfig( + collection_name="cloud_memory", + tenant="your-tenant", + database="your-database", + api_key=os.environ["CHROMA_API_KEY"], + embedding_function_config=OpenAIEmbeddingFunctionConfig( + api_key=os.environ["OPENAI_API_KEY"], model_name="text-embedding-3-small" + ), + ) + ) + # Add user preferences to memory await openai_memory.add( MemoryContent( @@ -148,6 +164,7 @@ async def main() -> None: await default_memory.close() await custom_memory.close() await openai_memory.close() + await cloud_memory.close() asyncio.run(main()) @@ -260,6 +277,12 @@ def _ensure_initialized(self) -> None: tenant=self._config.tenant, database=self._config.database, ) + elif isinstance(self._config, ChromaCloudVectorMemoryConfig): + self._client = CloudClient( + tenant=self._config.tenant, + database=self._config.database, + api_key=self._config.api_key, + ) else: raise ValueError(f"Unsupported config type: {type(self._config)}") except Exception as e: diff --git a/python/packages/autogen-ext/tests/memory/test_chroma_memory.py b/python/packages/autogen-ext/tests/memory/test_chroma_memory.py index f62c91c7dde1..53fa8ac120b0 100644 --- a/python/packages/autogen-ext/tests/memory/test_chroma_memory.py +++ b/python/packages/autogen-ext/tests/memory/test_chroma_memory.py @@ -5,6 +5,7 @@ from autogen_core.model_context import BufferedChatCompletionContext from autogen_core.models import UserMessage from autogen_ext.memory.chromadb import ( + ChromaCloudVectorMemoryConfig, ChromaDBVectorMemory, CustomEmbeddingFunctionConfig, DefaultEmbeddingFunctionConfig, @@ -274,6 +275,22 @@ def test_http_config(tmp_path: Path) -> None: assert config.headers == {"Authorization": "Bearer test-token"} +@pytest.mark.asyncio +def test_cloud_config(tmp_path: Path) -> None: + """Test ChromaCloud configuration.""" + config = ChromaCloudVectorMemoryConfig( + collection_name="test_cloud", + api_key="test-token", + tenant="test-tenant", + database="test-database", + ) + + assert config.client_type == "cloud" + assert config.api_key == "test-token" + assert config.tenant == "test-tenant" + assert config.database == "test-database" + + # ============================================================================ # Embedding Function Configuration Tests # ============================================================================