This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 0d3fb8c [Feature] Support Chroma Cloud (#188)
0d3fb8c is described below
commit 0d3fb8c3731a40b2588c27b76952c2eb892dcb25
Author: Alan Z. <shuai....@gmail.com>
AuthorDate: Sat Sep 20 07:18:49 2025 -0700
[Feature] Support Chroma Cloud (#188)
---
.../vector_stores/chroma/chroma_vector_store.py | 20 +++++++++-
.../chroma/tests/test_chroma_vector_store.py | 43 +++++++++++++++++++++-
2 files changed, 60 insertions(+), 3 deletions(-)
diff --git
a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
index 572424b..9c4fc33 100644
---
a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
+++
b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
@@ -19,6 +19,7 @@ from typing import Any, Dict, List
import chromadb
from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient
from chromadb.config import Settings
from pydantic import Field
@@ -40,6 +41,7 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
- In-memory: No configuration needed (default)
- Persistent: Provide persist_directory
- Server: Provide host and port
+ - Cloud: Provide api_key for Chroma Cloud
Attributes:
----------
@@ -49,6 +51,8 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
Host for ChromaDB server connection.
port : Optional[int]
Port for ChromaDB server connection.
+ api_key : Optional[str]
+ API key for Chroma Cloud connection.
client_settings : Optional[Settings]
ChromaDB client settings for advanced configuration.
tenant : str
@@ -69,6 +73,10 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
default=8000,
description="Port for ChromaDB server connection.",
)
+ api_key: str | None = Field(
+ default=None,
+ description="API key for Chroma Cloud connection.",
+ )
client_settings: Settings | None = Field(
default=None,
description="ChromaDB client settings for advanced configuration.",
@@ -89,6 +97,7 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
persist_directory: str | None = None,
host: str | None = None,
port: int | None = 8000,
+ api_key: str | None = None,
client_settings: Settings | None = None,
tenant: str = "default_tenant",
database: str = "default_database",
@@ -99,6 +108,7 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
persist_directory=persist_directory,
host=host,
port=port,
+ api_key=api_key,
client_settings=client_settings,
tenant=tenant,
database=database,
@@ -111,8 +121,14 @@ class
ChromaVectorStoreConnection(BaseVectorStoreConnection):
if self.__client is None:
# Choose client type based on configuration
- # TODO: support cloud client:
https://docs.trychroma.com/docs/run-chroma/cloud-client
- if self.host is not None:
+ if self.api_key is not None:
+ # Cloud mode
+ self.__client = CloudClient(
+ tenant=self.tenant,
+ database=self.database,
+ api_key=self.api_key,
+ )
+ elif self.host is not None:
# Client-Server Mode
self.__client = chromadb.HttpClient(
host=self.host,
diff --git
a/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py
b/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py
index afff634..9f0b468 100644
---
a/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py
+++
b/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py
@@ -15,6 +15,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
+import os
from typing import Any, Dict
import pytest
@@ -35,6 +36,10 @@ from
flink_agents.integrations.vector_stores.chroma.chroma_vector_store import (
ChromaVectorStoreSetup,
)
+api_key = os.environ.get("TEST_API_KEY")
+tenant = os.environ.get("TEST_TENANT")
+database = os.environ.get("TEST_DATABASE")
+
class MockEmbeddingModel(Resource): # noqa: D101
@classmethod
@@ -77,7 +82,7 @@ def _populate_test_data(connection:
ChromaVectorStoreConnection) -> None:
@pytest.mark.skipif(
not chromadb_available, reason="ChromaDB is not available"
)
-def test_chroma_vector_store_setup() -> None:
+def test_local_chroma_vector_store_setup() -> None:
"""Test ChromaDB vector store setup with embedding model integration."""
connection = ChromaVectorStoreConnection(name="chroma_conn")
embedding_model = MockEmbeddingModel(name="mock_embeddings")
@@ -110,3 +115,39 @@ def test_chroma_vector_store_setup() -> None:
assert result is not None
assert len(result.documents) == 1
assert result.documents[0].id == "doc2"
+
+
+@pytest.mark.skipif(api_key is None, reason="TEST_API_KEY is not set")
+def test_cloud_chroma_vector_store_setup() -> None:
+ """Test cloud ChromaDB vector store setup with embedding model
integration."""
+ connection = ChromaVectorStoreConnection(name="cloud_chroma_conn",
api_key=api_key, tenant=tenant, database=database)
+ embedding_model = MockEmbeddingModel(name="mock_embeddings")
+
+ def get_resource(name: str, resource_type: ResourceType) -> Resource:
+ if resource_type == ResourceType.VECTOR_STORE_CONNECTION:
+ return connection
+ elif resource_type == ResourceType.EMBEDDING_MODEL:
+ return embedding_model
+ else:
+ msg = f"Unknown resource type: {resource_type}"
+ raise ValueError(msg)
+
+ setup = ChromaVectorStoreSetup(
+ name="chroma_setup",
+ connection="cloud_chroma_conn",
+ embedding_model="mock_embeddings",
+ collection="test_collection",
+ get_resource=get_resource
+ )
+
+ _populate_test_data(connection)
+
+ query = VectorStoreQuery(
+ query_text="What is Flink Agent?",
+ limit=1
+ )
+
+ result = setup.query(query)
+ assert result is not None
+ assert len(result.documents) == 1
+ assert result.documents[0].id == "doc2"
