From fd6b0398591aa59b31879bef1f701ec18cb93fd8 Mon Sep 17 00:00:00 2001
From: Vineeth B V <37930821+vinsdragonis@users.noreply.github.com>
Date: Thu, 6 Feb 2025 12:04:14 +0530
Subject: [PATCH] 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.
---
 .../retrieval/vector/dbs/opensearch.py        | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/backend/open_webui/retrieval/vector/dbs/opensearch.py b/backend/open_webui/retrieval/vector/dbs/opensearch.py
index b3d8b5eb8a..c4732e1bc0 100644
--- a/backend/open_webui/retrieval/vector/dbs/opensearch.py
+++ b/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)