From ee50f78d99bf3a5d463473a27f87ca52705856d5 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Thu, 21 Nov 2024 18:26:22 +0800
Subject: [PATCH] Add component 'Template' (#3562)

### What problem does this PR solve?

#3560

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
---
 agent/component/__init__.py | 1 +
 agent/component/generate.py | 2 +-
 agent/component/template.py | 2 +-
 api/apps/kb_app.py          | 4 +---
 api/apps/user_app.py        | 4 ++--
 rag/utils/es_conn.py        | 4 +++-
 6 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/agent/component/__init__.py b/agent/component/__init__.py
index a977f03f4..d8e31fb86 100644
--- a/agent/component/__init__.py
+++ b/agent/component/__init__.py
@@ -30,6 +30,7 @@ from .tushare import TuShare, TuShareParam
 from .akshare import AkShare, AkShareParam
 from .crawler import Crawler, CrawlerParam
 from .invoke import Invoke, InvokeParam
+from .template import Template, TemplateParam
 
 
 def component_class(class_name):
diff --git a/agent/component/generate.py b/agent/component/generate.py
index 53a0a4553..c6640c889 100644
--- a/agent/component/generate.py
+++ b/agent/component/generate.py
@@ -145,7 +145,7 @@ class Generate(ComponentBase):
         else: retrieval_res = pd.DataFrame([])
 
         for n, v in kwargs.items():
-            prompt = re.sub(r"\{%s\}" % re.escape(n), re.escape(str(v)), prompt)
+            prompt = re.sub(r"\{%s\}" % re.escape(n), str(v), prompt)
 
         if not self._param.inputs and prompt.find("{input}") >= 0:
             retrieval_res = self.get_input()
diff --git a/agent/component/template.py b/agent/component/template.py
index 5fe1e715c..06f84114e 100644
--- a/agent/component/template.py
+++ b/agent/component/template.py
@@ -79,7 +79,7 @@ class Template(ComponentBase):
             self._param.inputs.append({"component_id": para["component_id"], "content": kwargs[para["key"]]})
 
         for n, v in kwargs.items():
-            content = re.sub(r"\{%s\}" % re.escape(n), re.escape(str(v)), content)
+            content = re.sub(r"\{%s\}" % re.escape(n), str(v), content)
 
         return Template.be_output(content)
 
diff --git a/api/apps/kb_app.py b/api/apps/kb_app.py
index f98fe51a0..523d13876 100644
--- a/api/apps/kb_app.py
+++ b/api/apps/kb_app.py
@@ -167,9 +167,7 @@ def rm():
         if not KnowledgebaseService.delete_by_id(req["kb_id"]):
             return get_data_error_result(
                 message="Database error (Knowledgebase removal)!")
-        tenants = UserTenantService.query(user_id=current_user.id)
-        for tenant in tenants:
-            settings.docStoreConn.deleteIdx(search.index_name(tenant.tenant_id), req["kb_id"])
+        settings.docStoreConn.delete({"kb_id": req["kb_id"]}, search.index_name(kbs[0].tenant_id), req["kb_id"])
         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 a9dc620b3..7136fa737 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
+                "api_base": settings.LLM_BASE_URL
+                #"max_tokens": llm.max_tokens if llm.max_tokens else 8192
             }
         )
 
diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py
index 9de87de25..651c7a6b0 100644
--- a/rag/utils/es_conn.py
+++ b/rag/utils/es_conn.py
@@ -238,7 +238,7 @@ class ESConnection(DocStoreConnection):
         for _ in range(ATTEMPT_TIME):
             try:
                 r = self.es.bulk(index=(indexName), operations=operations,
-                                 refresh=False, timeout="600s")
+                                 refresh=False, timeout="60s")
                 if re.search(r"False", str(r["errors"]), re.IGNORECASE):
                     return res
 
@@ -249,7 +249,9 @@ class ESConnection(DocStoreConnection):
                 return res
             except Exception as e:
                 logging.warning("ESConnection.insert got exception: " + str(e))
+                res = []
                 if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE):
+                    res.append(str(e))
                     time.sleep(3)
                     continue
         return res