diff --git a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py index 02d253bf6..6e52503a1 100644 --- a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py +++ b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py @@ -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"], } diff --git a/backend/open_webui/apps/retrieval/vector/dbs/milvus.py b/backend/open_webui/apps/retrieval/vector/dbs/milvus.py index 61b3c23c8..b5bbb24b3 100644 --- a/backend/open_webui/apps/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/apps/retrieval/vector/dbs/milvus.py @@ -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.