Просмотр исходного кода

Added a query method for OpenSearch vector db.

- This PR aims to address the error 400: "**'OpenSearchClient' object has no attribute 'query'**".
- With the implemented query() method, this issue should be resolved and allow uploaded documents to be vectorized and retrieved based on the given query.
Vineeth B V 2 месяцев назад
Родитель
Сommit
fd6b039859
1 измененных файлов с 34 добавлено и 0 удалено
  1. 34 0
      backend/open_webui/retrieval/vector/dbs/opensearch.py

+ 34 - 0
backend/open_webui/retrieval/vector/dbs/opensearch.py

@@ -113,6 +113,40 @@ class OpenSearchClient:
 
         return self._result_to_search_result(result)
 
+    def query(
+        self, index_name: str, filter: dict, limit: Optional[int] = None
+    ) -> Optional[GetResult]:
+        if not self.has_collection(index_name):
+            return None
+
+        query_body = {
+            "query": {
+                "bool": {
+                    "filter": []
+                }
+            },
+            "_source": ["text", "metadata"],
+        }
+
+        for field, value in filter.items():
+            query_body["query"]["bool"]["filter"].append({
+                "term": {field: value}
+            })
+
+        size = limit if limit else 10
+
+        try:
+            result = self.client.search(
+                index=f"{self.index_prefix}_{index_name}",
+                body=query_body,
+                size=size
+            )
+
+            return self._result_to_get_result(result)
+
+        except Exception as e:
+            return None
+
     def get_or_create_index(self, index_name: str, dimension: int):
         if not self.has_index(index_name):
             self._create_index(index_name, dimension)