Selaa lähdekoodia

Merge pull request #9693 from vinsdragonis/main

fix: Fixed error occurring when using OpenSearch as a vector db
Timothy Jaeryang Baek 2 kuukautta sitten
vanhempi
commit
d5a815b19c

+ 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, collection_name: str, filter: dict, limit: Optional[int] = None
+    ) -> Optional[GetResult]:
+        if not self.has_collection(collection_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}_{collection_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)

+ 1 - 1
backend/requirements.txt

@@ -47,7 +47,7 @@ fake-useragent==1.5.1
 chromadb==0.6.2
 pymilvus==2.5.0
 qdrant-client~=1.12.0
-opensearch-py==2.7.1
+opensearch-py==2.8.0
 
 
 transformers