From 124a17e826b797b36f89e70d0546b761762f112b Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Thu, 3 Oct 2024 20:58:56 -0700 Subject: [PATCH] refac --- .../apps/retrieval/vector/dbs/chroma.py | 65 ++++++++++--------- 1 file changed, 36 insertions(+), 29 deletions(-) diff --git a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py index 6e52503a1..4a85c1251 100644 --- a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py +++ b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py @@ -49,42 +49,49 @@ class ChromaClient: self, collection_name: str, vectors: list[list[float | int]], limit: int ) -> Optional[SearchResult]: # Search for the nearest neighbor items based on the vectors and return 'limit' number of results. - collection = self.client.get_collection(name=collection_name) - if collection: - result = collection.query( - query_embeddings=vectors, - n_results=limit, - ) + try: + collection = self.client.get_collection(name=collection_name) + if collection: + result = collection.query( + query_embeddings=vectors, + n_results=limit, + ) - return SearchResult( - **{ - "ids": result["ids"], - "distances": result["distances"], - "documents": result["documents"], - "metadatas": result["metadatas"], - } - ) - return None + return SearchResult( + **{ + "ids": result["ids"], + "distances": result["distances"], + "documents": result["documents"], + "metadatas": result["metadatas"], + } + ) + return None + except Exception as e: + return None def query( self, collection_name: str, filter: dict, limit: int = 1 ) -> Optional[GetResult]: # Query the items from the collection based on the filter. - collection = self.client.get_collection(name=collection_name) - if collection: - result = collection.get( - where=filter, - limit=limit, - ) - return GetResult( - **{ - "ids": result["ids"], - "documents": result["documents"], - "metadatas": result["metadatas"], - } - ) - return None + try: + collection = self.client.get_collection(name=collection_name) + if collection: + result = collection.get( + where=filter, + limit=limit, + ) + + return GetResult( + **{ + "ids": result["ids"], + "documents": result["documents"], + "metadatas": result["metadatas"], + } + ) + return None + except Exception as e: + return None def get(self, collection_name: str) -> Optional[GetResult]: # Get all the items in the collection.