Add Chroma Cloud config

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",

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")

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:

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
 # ============================================================================