Add component 'Template' (infiniflow#3562)

### What problem does this PR solve? infiniflow#3560 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
KevinHuSh · Nov 21, 2024 · ee50f78 · ee50f78
1 parent 193b08a
commit ee50f78
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 8 deletions.
diff --git a/agent/component/__init__.py b/agent/component/__init__.py
@@ -30,6 +30,7 @@
 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
@@ -145,7 +145,7 @@ def _run(self, history, **kwargs):
         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
@@ -79,7 +79,7 @@ def _run(self, history, **kwargs):
             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
@@ -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
@@ -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
@@ -238,7 +238,7 @@ def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str) ->
         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 @@ def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str) ->
                 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