OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-09 16:07:54 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-10-03 20:58:56 -07:00
parent 57360b7a61
commit 124a17e826
1 changed files with 36 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
backend/open_webui/apps/retrieval/vector/dbs/chroma.py
View File

@ -49,42 +49,49 @@ class ChromaClient:
self, collection_name: str, vectors: list[list[float | int]], limit: int self, collection_name: str, vectors: list[list[float | int]], limit: int
) -> Optional[SearchResult]: ) -> Optional[SearchResult]:
# Search for the nearest neighbor items based on the vectors and return 'limit' number of results. # Search for the nearest neighbor items based on the vectors and return 'limit' number of results.
collection = self.client.get_collection(name=collection_name) try:
if collection: collection = self.client.get_collection(name=collection_name)
result = collection.query( if collection:
query_embeddings=vectors, result = collection.query(
n_results=limit, query_embeddings=vectors,
) n_results=limit,
)
return SearchResult( return SearchResult(
**{ **{
"ids": result["ids"], "ids": result["ids"],
"distances": result["distances"], "distances": result["distances"],
"documents": result["documents"], "documents": result["documents"],
"metadatas": result["metadatas"], "metadatas": result["metadatas"],
} }
) )
return None return None
except Exception as e:
return None
def query( def query(
self, collection_name: str, filter: dict, limit: int = 1 self, collection_name: str, filter: dict, limit: int = 1
) -> Optional[GetResult]: ) -> Optional[GetResult]:
# Query the items from the collection based on the filter. # 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( try:
**{ collection = self.client.get_collection(name=collection_name)
"ids": result["ids"], if collection:
"documents": result["documents"], result = collection.get(
"metadatas": result["metadatas"], where=filter,
} limit=limit,
) )
return None
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]: def get(self, collection_name: str) -> Optional[GetResult]:
# Get all the items in the collection. # Get all the items in the collection.