Skip to content

Commit

Permalink
fix: Update RAG ChromaDB UTs v0.6.1 (#803)
Browse files Browse the repository at this point in the history
**Reason for Change**:
ChromaDB has been updated 0.5.18->0.6.1. Fix addresses slight library
changes.
  • Loading branch information
ishaansehgal99 authored Jan 4, 2025
1 parent 1567afa commit 47143e5
Showing 1 changed file with 5 additions and 6 deletions.
11 changes: 5 additions & 6 deletions presets/ragengine/vector_store/chromadb_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]),
}
Expand All @@ -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.")

Expand Down

0 comments on commit 47143e5

Please sign in to comment.