refac

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"],
                 }

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.