From 47143e520511f7cf44bdaa52075cd6c6e067c27c Mon Sep 17 00:00:00 2001 From: Ishaan Sehgal Date: Fri, 3 Jan 2025 16:05:56 -0800 Subject: [PATCH] fix: Update RAG ChromaDB UTs v0.6.1 (#803) **Reason for Change**: ChromaDB has been updated 0.5.18->0.6.1. Fix addresses slight library changes. --- presets/ragengine/vector_store/chromadb_store.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/presets/ragengine/vector_store/chromadb_store.py b/presets/ragengine/vector_store/chromadb_store.py index a40c4cdfc..493c87c6e 100644 --- a/presets/ragengine/vector_store/chromadb_store.py +++ b/presets/ragengine/vector_store/chromadb_store.py @@ -35,10 +35,10 @@ def document_exists(self, index_name: str, doc: Document, doc_id: str) -> bool: def list_all_indexed_documents(self) -> Dict[str, Dict[str, Dict[str, str]]]: indexed_docs = {} # Accumulate documents across all indexes try: - for collection in self.chroma_client.list_collections(): - collection_info = collection.get() + for collection_name in self.chroma_client.list_collections(): + collection_info = self.chroma_client.get_collection(collection_name).get() for doc in zip(collection_info["ids"], collection_info["documents"], collection_info["metadatas"]): - indexed_docs.setdefault(collection.name, {})[doc[0]] = { + indexed_docs.setdefault(collection_name, {})[doc[0]] = { "text": doc[1], "metadata": json.dumps(doc[2]), } @@ -54,11 +54,10 @@ def _clear_collection_and_indexes(self): """ try: # Get all collections - collections = self.chroma_client.list_collections() + collection_names = self.chroma_client.list_collections() # Delete each collection - for collection in collections: - collection_name = collection.name + for collection_name in collection_names: self.chroma_client.delete_collection(name=collection_name) print(f"Collection '{collection_name}' has been deleted.")