From 9f3141804f861e21a7ab54093dc145ca931b9830 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Fri, 22 Nov 2024 12:25:42 +0800
Subject: [PATCH] Fix chunk enable/disable issue (#3579)

### What problem does this PR solve?

#3576

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 api/apps/chunk_app.py | 9 ++++++---
 api/apps/user_app.py  | 4 ++--
 rag/nlp/search.py     | 2 +-
 rag/utils/es_conn.py  | 8 +++++++-
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/api/apps/chunk_app.py b/api/apps/chunk_app.py
index 0b237ace49..4606c8b1d4 100644
--- a/api/apps/chunk_app.py
+++ b/api/apps/chunk_app.py
@@ -170,9 +170,12 @@ def switch():
         e, doc = DocumentService.get_by_id(req["doc_id"])
         if not e:
             return get_data_error_result(message="Document not found!")
-        if not settings.docStoreConn.update({"id": req["chunk_ids"]}, {"available_int": int(req["available_int"])},
-                                            search.index_name(doc.tenant_id), doc.kb_id):
-            return get_data_error_result(message="Index updating failure")
+        for cid in req["chunk_ids"]:
+            if not settings.docStoreConn.update({"id": cid},
+                                                {"available_int": int(req["available_int"])},
+                                                search.index_name(DocumentService.get_tenant_id(req["doc_id"])),
+                                                doc.kb_id):
+                return get_data_error_result(message="Index updating failure")
         return get_json_result(data=True)
     except Exception as e:
         return server_error_response(e)
diff --git a/api/apps/user_app.py b/api/apps/user_app.py
index 7136fa7371..66cae41519 100644
--- a/api/apps/user_app.py
+++ b/api/apps/user_app.py
@@ -517,8 +517,8 @@ def user_register(user_id, user):
                 "llm_name": llm.llm_name,
                 "model_type": llm.model_type,
                 "api_key": settings.API_KEY,
-                "api_base": settings.LLM_BASE_URL
-                #"max_tokens": llm.max_tokens if llm.max_tokens else 8192
+                "api_base": settings.LLM_BASE_URL,
+                "max_tokens": llm.max_tokens if llm.max_tokens else 8192
             }
         )
 
diff --git a/rag/nlp/search.py b/rag/nlp/search.py
index e7b7009030..eb389bdd85 100644
--- a/rag/nlp/search.py
+++ b/rag/nlp/search.py
@@ -56,7 +56,7 @@ def get_filters(self, req):
             if key in req and req[key] is not None:
                 condition[field] = req[key]
         # TODO(yzc): `available_int` is nullable however infinity doesn't support nullable columns.
-        for key in ["knowledge_graph_kwd"]:
+        for key in ["knowledge_graph_kwd", "available_int"]:
             if key in req and req[key] is not None:
                 condition[key] = req[key]
         return condition
diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py
index 6354defc1e..0630f84136 100644
--- a/rag/utils/es_conn.py
+++ b/rag/utils/es_conn.py
@@ -121,8 +121,14 @@ def search(self, selectFields: list[str], highlightFields: list[str], condition:
         bqry = Q("bool", must=[])
         condition["kb_id"] = knowledgebaseIds
         for k, v in condition.items():
-            if not isinstance(k, str) or not v:
+            if k == "available_int":
+                if v == 0:
+                    bqry.filter.append(Q("range", available_int={"lt": 1}))
+                else:
+                    bqry.filter.append(
+                        Q("bool", must_not=Q("range", available_int={"lt": 1})))
                 continue
+            if not v: continue
             if isinstance(v, list):
                 bqry.filter.append(Q("terms", **{k: v}))
             elif isinstance(v, str) or isinstance(v, int):