From f3c4c2b8e35d5e93539a84a915ad83362bbfa896 Mon Sep 17 00:00:00 2001
From: David Tippett <dtaivpp@github.com>
Date: Sat, 1 Mar 2025 13:26:19 -0500
Subject: [PATCH] Changed to use colleciton name and fixed bulk indexing
 missing index.

---
 .../retrieval/vector/dbs/opensearch.py        | 43 ++++++++++---------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/backend/open_webui/retrieval/vector/dbs/opensearch.py b/backend/open_webui/retrieval/vector/dbs/opensearch.py
index b8186b3f9..bfae4d9db 100644
--- a/backend/open_webui/retrieval/vector/dbs/opensearch.py
+++ b/backend/open_webui/retrieval/vector/dbs/opensearch.py
@@ -49,7 +49,7 @@ class OpenSearchClient:
             ids=ids, distances=distances, documents=documents, metadatas=metadatas
         )
 
-    def _create_index(self, index_name: str, dimension: int):
+    def _create_index(self, collection_name: str, dimension: int):
         body = {
             "mappings": {
                 "properties": {
@@ -72,24 +72,24 @@ class OpenSearchClient:
                 }
             }
         }
-        self.client.indices.create(index=f"{self.index_prefix}_{index_name}", body=body)
+        self.client.indices.create(index=f"{self.index_prefix}_{collection_name}", body=body)
 
     def _create_batches(self, items: list[VectorItem], batch_size=100):
         for i in range(0, len(items), batch_size):
             yield items[i : i + batch_size]
 
-    def has_collection(self, index_name: str) -> bool:
+    def has_collection(self, collection_name: str) -> bool:
         # has_collection here means has index.
         # We are simply adapting to the norms of the other DBs.
-        return self.client.indices.exists(index=f"{self.index_prefix}_{index_name}")
+        return self.client.indices.exists(index=f"{self.index_prefix}_{collection_name}")
 
-    def delete_colleciton(self, index_name: str):
+    def delete_colleciton(self, collection_name: str):
         # delete_collection here means delete index.
         # We are simply adapting to the norms of the other DBs.
-        self.client.indices.delete(index=f"{self.index_prefix}_{index_name}")
+        self.client.indices.delete(index=f"{self.index_prefix}_{collection_name}")
 
     def search(
-        self, index_name: str, vectors: list[list[float]], limit: int
+        self, collection_name: str, vectors: list[list[float]], limit: int
     ) -> Optional[SearchResult]:
         query = {
             "size": limit,
@@ -108,7 +108,7 @@ class OpenSearchClient:
         }
 
         result = self.client.search(
-            index=f"{self.index_prefix}_{index_name}", body=query
+            index=f"{self.index_prefix}_{collection_name}", body=query
         )
 
         return self._result_to_search_result(result)
@@ -141,21 +141,21 @@ class OpenSearchClient:
         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)
+    def _create_index_if_not_exists(self, collection_name: str, dimension: int):
+        if not self.has_index(collection_name):
+            self._create_index(collection_name, dimension)
 
-    def get(self, index_name: str) -> Optional[GetResult]:
+    def get(self, collection_name: str) -> Optional[GetResult]:
         query = {"query": {"match_all": {}}, "_source": ["text", "metadata"]}
 
         result = self.client.search(
-            index=f"{self.index_prefix}_{index_name}", body=query
+            index=f"{self.index_prefix}_{collection_name}", body=query
         )
         return self._result_to_get_result(result)
 
-    def insert(self, index_name: str, items: list[VectorItem]):
-        if not self.has_index(index_name):
-            self._create_index(index_name, dimension=len(items[0]["vector"]))
+    def insert(self, collection_name: str, items: list[VectorItem]):
+        self._create_index_if_not_exists(collection_name=collection_name,
+                                         dimension=len(items[0]["vector"]))
 
         for batch in self._create_batches(items):
             actions = [
@@ -173,15 +173,16 @@ class OpenSearchClient:
             ]
             self.client.bulk(actions)
 
-    def upsert(self, index_name: str, items: list[VectorItem]):
-        if not self.has_index(index_name):
-            self._create_index(index_name, dimension=len(items[0]["vector"]))
+    def upsert(self, collection_name: str, items: list[VectorItem]):
+        self._create_index_if_not_exists(collection_name=collection_name,
+                                         dimension=len(items[0]["vector"]))
 
         for batch in self._create_batches(items):
             actions = [
                 {
                     "index": {
                         "_id": item["id"],
+                        "_index": f"{self.index_prefix}_{collection_name}",
                         "_source": {
                             "vector": item["vector"],
                             "text": item["text"],
@@ -193,9 +194,9 @@ class OpenSearchClient:
             ]
             self.client.bulk(actions)
 
-    def delete(self, index_name: str, ids: list[str]):
+    def delete(self, collection_name: str, ids: list[str]):
         actions = [
-            {"delete": {"_index": f"{self.index_prefix}_{index_name}", "_id": id}}
+            {"delete": {"_index": f"{self.index_prefix}_{collection_name}", "_id": id}}
             for id in ids
         ]
         self.client.bulk(body=actions)