UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-25 13:29:53 +00:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-10-03 20:51:21 -07:00
parent 2fc07fd6a2
commit 57360b7a61
2 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -68,19 +68,18 @@ class ChromaClient:
def query(
self, collection_name: str, filter: dict, limit: int = 1
) -> Optional[SearchResult]:
) -> Optional[GetResult]:
# Query the items from the collection based on the filter.
collection = self.client.get_collection(name=collection_name)
if collection:
result = collection.query(
result = collection.get(
where=filter,
n_results=limit,
limit=limit,
)
return SearchResult(
return GetResult(
**{
"ids": result["ids"],
"distances": result["distances"],
"documents": result["documents"],
"metadatas": result["metadatas"],
}

4
backend/open_webui/apps/retrieval/vector/dbs/milvus.py
View File

@ -137,7 +137,7 @@ class MilvusClient:
def query(
self, collection_name: str, filter: dict, limit: int = 1
) -> Optional[SearchResult]:
) -> Optional[GetResult]:
# Query the items from the collection based on the filter.
filter_string = " && ".join(
[
@ -152,7 +152,7 @@ class MilvusClient:
limit=limit,
)
return self._result_to_search_result([result])
return self._result_to_get_result([result])
def get(self, collection_name: str) -> Optional[GetResult]:
# Get all the items in the collection.