From 62857ba27e7494cd1bcb4fada83ef895065b1730 Mon Sep 17 00:00:00 2001 From: aamirahmed2004 Date: Sat, 21 Sep 2024 16:31:01 -0700 Subject: [PATCH 1/2] create() now gets existing collection, improved logging for consistency with pgvector2 class, implemented clear() --- phi/vectordb/chroma/chromadb.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/phi/vectordb/chroma/chromadb.py b/phi/vectordb/chroma/chromadb.py index 8fdac53252..c5e915e25b 100644 --- a/phi/vectordb/chroma/chromadb.py +++ b/phi/vectordb/chroma/chromadb.py @@ -74,6 +74,10 @@ def create(self) -> None: self._collection = self.client.create_collection( name=self.collection, metadata={"hnsw:space": self.distance.value} ) + + else: + logger.debug(f"Collection already exists: {self.collection}") + self._collection = self.client.get_collection(name=self.collection) def doc_exists(self, document: Document) -> bool: """Check if a document exists in the collection. @@ -125,10 +129,11 @@ def insert(self, documents: List[Document]) -> None: docs_embeddings.append(document.embedding) docs.append(cleaned_content) ids.append(doc_id) + logger.debug(f"Inserted document: {document.id} | {document.name} | {document.meta_data}") if len(docs) > 0 and self._collection is not None: self._collection.add(ids=ids, embeddings=docs_embeddings, documents=docs) - logger.debug(f"Inserted {len(docs)} documents") + logger.debug(f"Committed {len(docs)} documents") else: logger.error("Collection does not exist") @@ -137,7 +142,7 @@ def upsert(self, documents: List[Document]) -> None: Args: documents (List[Document]): List of documents to upsert """ - logger.debug(f"Inserting {len(documents)} documents") + logger.debug(f"Upserting {len(documents)} documents") ids: List = [] docs: List = [] docs_embeddings: List = [] @@ -149,10 +154,12 @@ def upsert(self, documents: List[Document]) -> None: docs_embeddings.append(document.embedding) docs.append(cleaned_content) ids.append(doc_id) + logger.debug(f"Upserted document: {document.id} | {document.name} | {document.meta_data}") if len(docs) > 0 and self._collection is not None: self._collection.upsert(ids=ids, embeddings=docs_embeddings, documents=docs) - logger.debug(f"Inserted {len(docs)} documents") + logger.debug(f"Committed {len(docs)} documents") + else: logger.error("Collection does not exist") @@ -233,7 +240,11 @@ def get_count(self) -> int: return 0 def optimize(self) -> None: - pass + raise NotImplementedError def clear(self) -> bool: - return False + try: + self.client.delete_collection(name=self.collection) + return True + except: + return False From b7122df0c96733745f7a89dcf317c0a4809dd40a Mon Sep 17 00:00:00 2001 From: jacob weiss Date: Fri, 27 Sep 2024 08:29:53 -0600 Subject: [PATCH 2/2] clean up bare try except --- phi/vectordb/chroma/chromadb.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/phi/vectordb/chroma/chromadb.py b/phi/vectordb/chroma/chromadb.py index c5e915e25b..f5dc5f9d87 100644 --- a/phi/vectordb/chroma/chromadb.py +++ b/phi/vectordb/chroma/chromadb.py @@ -74,7 +74,7 @@ def create(self) -> None: self._collection = self.client.create_collection( name=self.collection, metadata={"hnsw:space": self.distance.value} ) - + else: logger.debug(f"Collection already exists: {self.collection}") self._collection = self.client.get_collection(name=self.collection) @@ -159,7 +159,7 @@ def upsert(self, documents: List[Document]) -> None: if len(docs) > 0 and self._collection is not None: self._collection.upsert(ids=ids, embeddings=docs_embeddings, documents=docs) logger.debug(f"Committed {len(docs)} documents") - + else: logger.error("Collection does not exist") @@ -246,5 +246,6 @@ def clear(self) -> bool: try: self.client.delete_collection(name=self.collection) return True - except: + except Exception as e: + logger.error(f"Error clearing collection: {e}") return False