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
commit 59d564cbe73399ddd82fe8329dca8a9e9c0be683 Author: youjin <[email protected]> AuthorDate: Wed Jan 14 09:29:55 2026 +0800 [hotfix] Improve create_collection_if_not_exists configuration range for ChromaVectorStore --- .../flink_agents/api/vector_stores/vector_store.py | 4 ++-- .../vector_stores/chroma/chroma_vector_store.py | 21 ++++++++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/python/flink_agents/api/vector_stores/vector_store.py b/python/flink_agents/api/vector_stores/vector_store.py index 7f49ef3..edc778c 100644 --- a/python/flink_agents/api/vector_stores/vector_store.py +++ b/python/flink_agents/api/vector_stores/vector_store.py @@ -69,7 +69,7 @@ class VectorStoreQuery(BaseModel): description="Text query to be converted to embedding for semantic search." ) limit: int = Field(default=10, description="Maximum number of results to return.") - collection_name: str = Field( + collection_name: str | None = Field( default=None, description="The collection to apply the query." ) extra_args: Dict[str, Any] = Field( @@ -323,7 +323,7 @@ class CollectionManageableVectorStore(BaseVectorStore, ABC): @abstractmethod def get_or_create_collection( - self, name: str, metadata: Dict[str, Any] + self, name: str, metadata: Dict[str, Any] | None = None ) -> Collection: """Get a collection, or create it if it doesn't exist. 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 b782bcf..4b742a7 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 @@ -302,6 +302,25 @@ class ChromaVectorStore(CollectionManageableVectorStore): collection_name: str | None = None, **kwargs: Any, ) -> List[str]: + collection_name = collection_name or self.collection + collection_metadata = kwargs.get( + "collection_metadata", self.collection_metadata + ) + create_collection_if_not_exists = kwargs.get( + "create_collection_if_not_exists", self.create_collection_if_not_exists + ) + + # Get or create collection based on configuration + if create_collection_if_not_exists: + # ChromaDB doesn't accept empty metadata, pass None instead + metadata = collection_metadata if collection_metadata else None + collection = self.client.get_or_create_collection( + name=collection_name, + metadata=metadata, + ) + else: + collection = self.client.get_collection(name=collection_name) + documents_chunks = chunk_list(documents, MAX_CHUNK_SIZE) all_ids = [] @@ -311,7 +330,7 @@ class ChromaVectorStore(CollectionManageableVectorStore): embeddings = [doc.embedding for doc in chunk] metadatas = [doc.metadata for doc in chunk] - self.client.get_collection(name=collection_name or self.collection).add( + collection.add( ids=ids, documents=docs, embeddings=embeddings, metadatas=metadatas ) all_ids.extend(ids)
