From 4fffee6695611121f6e2be97006914fca5257923 Mon Sep 17 00:00:00 2001
From: Zhichang Yu <yuzhichang@gmail.com>
Date: Fri, 14 Mar 2025 17:34:02 +0800
Subject: [PATCH] Regards kb_id at ElasticSearch insert, update, delete.
 (#6105)

### What problem does this PR solve?

Regards kb_id at ElasticSearch insert, update, delete. Close #6066

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 rag/utils/es_conn.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py
index 89687e1bd..a8dd65317 100644
--- a/rag/utils/es_conn.py
+++ b/rag/utils/es_conn.py
@@ -284,6 +284,7 @@ class ESConnection(DocStoreConnection):
             assert "_id" not in d
             assert "id" in d
             d_copy = copy.deepcopy(d)
+            d_copy["kb_id"] = knowledgebaseId
             meta_id = d_copy.pop("id", "")
             operations.append(
                 {"index": {"_index": indexName, "_id": meta_id}})
@@ -316,6 +317,7 @@ class ESConnection(DocStoreConnection):
     def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
         doc = copy.deepcopy(newValue)
         doc.pop("id", None)
+        condition["kb_id"] = knowledgebaseId
         if "id" in condition and isinstance(condition["id"], str):
             # update specific single document
             chunkId = condition["id"]
@@ -399,6 +401,7 @@ class ESConnection(DocStoreConnection):
     def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
         qry = None
         assert "_id" not in condition
+        condition["kb_id"] = knowledgebaseId
         if "id" in condition:
             chunk_ids = condition["id"]
             if not isinstance(chunk_ids, list):