Merge branch 'main' into feature/support-gunicorn-test3

Author: hzy@singhand.com

agent/component/categorize.py

@@ -99,9 +99,13 @@ def _run(self, history, **kwargs):
         # If a category is found, return the category with the highest count.
         if any(category_counts.values()):
             max_category = max(category_counts.items(), key=lambda x: x[1])
-            return Categorize.be_output(self._param.category_description[max_category[0]]["to"])
+            res = Categorize.be_output(self._param.category_description[max_category[0]]["to"])
+            self.set_output(res)
+            return res
 
-        return Categorize.be_output(list(self._param.category_description.items())[-1][1]["to"])
+        res = Categorize.be_output(list(self._param.category_description.items())[-1][1]["to"])
+        self.set_output(res)
+        return res
 
     def debug(self, **kwargs):
         df = self._run([], **kwargs)

api/apps/conversation_app.py

@@ -42,6 +42,7 @@ def set_conversation():
     conv_id = req.get("conversation_id")
     is_new = req.get("is_new")
     name = req.get("name", "New conversation")
+    req["user_id"] = current_user.id
 
     if len(name) > 255:
         name = name[0:255]
@@ -64,7 +65,7 @@ def set_conversation():
         e, dia = DialogService.get_by_id(req["dialog_id"])
         if not e:
             return get_data_error_result(message="Dialog not found")
-        conv = {"id": conv_id, "dialog_id": req["dialog_id"], "name": name, "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]}
+        conv = {"id": conv_id, "dialog_id": req["dialog_id"], "name": name, "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}],"user_id": current_user.id}
         ConversationService.save(**conv)
         return get_json_result(data=conv)
     except Exception as e:
@@ -248,7 +249,7 @@ def stream():
         else:
             answer = None
             for ans in chat(dia, msg, **req):
-                answer = structure_answer(conv, ans, message_id, req["conversation_id"])
+                answer = structure_answer(conv, ans, message_id, conv.id)
                 ConversationService.update_by_id(conv.id, conv.to_dict())
                 break
             return get_json_result(data=answer)

api/apps/dialog_app.py

@@ -28,6 +28,7 @@
 
 
 @manager.route('/set', methods=['POST'])  # noqa: F821
+@validate_request("prompt_config")
 @login_required
 def set_dialog():
     req = request.json
@@ -43,33 +44,10 @@ def set_dialog():
     similarity_threshold = req.get("similarity_threshold", 0.1)
     vector_similarity_weight = req.get("vector_similarity_weight", 0.3)
     llm_setting = req.get("llm_setting", {})
-    default_prompt_with_dataset = {
-        "system": """你是一个智能助手，请总结知识库的内容来回答问题，请列举知识库中的数据详细回答。当所有知识库内容都与问题无关时，你的回答必须包括“知识库中未找到您要的答案！”这句话。回答需要考虑聊天历史。
-以下是知识库：
-{knowledge}
-以上是知识库。""",
-        "prologue": "您好，我是您的助手小樱，长得可爱又善良，can I help you?",
-        "parameters": [
-            {"key": "knowledge", "optional": False}
-        ],
-        "empty_response": "Sorry! 知识库中未找到相关内容！"
-    }
-    default_prompt_no_dataset = {
-        "system": """You are a helpful assistant.""",
-        "prologue": "您好，我是您的助手小樱，长得可爱又善良，can I help you?",
-        "parameters": [
-           
-        ],
-        "empty_response": ""
-    }
-    prompt_config = req.get("prompt_config", default_prompt_with_dataset)
-
-    if not prompt_config["system"]:
-        prompt_config["system"] = default_prompt_with_dataset["system"]
+    prompt_config = req["prompt_config"]
 
-    if not req.get("kb_ids", []):
-        if prompt_config['system'] == default_prompt_with_dataset['system'] or "{knowledge}" in prompt_config['system']:
-            prompt_config = default_prompt_no_dataset
+    if not req.get("kb_ids", []) and not prompt_config.get("tavily_api_key") and "{knowledge}" in prompt_config['system']:
+        return get_data_error_result(message="Please remove `{knowledge}` in system prompt since no knowledge base/Tavily used here.")
 
     for p in prompt_config["parameters"]:
         if p["optional"]:

api/apps/kb_app.py

@@ -47,7 +47,7 @@ def create():
         return get_data_error_result(message="Dataset name must be string.")
     if dataset_name == "":
         return get_data_error_result(message="Dataset name can't be empty.")
-    if len(dataset_name) >= DATASET_NAME_LIMIT:
+    if len(dataset_name.encode("utf-8")) >= DATASET_NAME_LIMIT:
         return get_data_error_result(
             message=f"Dataset name length is {len(dataset_name)} which is large than {DATASET_NAME_LIMIT}")
 

api/apps/sdk/chat.py

@@ -173,8 +173,10 @@ def update(tenant_id, chat_id):
     if llm:
         if "model_name" in llm:
             req["llm_id"] = llm.pop("model_name")
-            if not TenantLLMService.query(tenant_id=tenant_id, llm_name=req["llm_id"], model_type="chat"):
-                return get_error_data_result(f"`model_name` {req.get('llm_id')} doesn't exist")
+            if req.get("llm_id") is not None:
+                llm_name, llm_factory = TenantLLMService.split_model_name_and_factory(req["llm_id"])
+                if not TenantLLMService.query(tenant_id=tenant_id, llm_name=llm_name, llm_factory=llm_factory, model_type="chat"):
+                    return get_error_data_result(f"`model_name` {req.get('llm_id')} doesn't exist")
         req["llm_setting"] = req.pop("llm")
     e, tenant = TenantService.get_by_id(tenant_id)
     if not e:

api/apps/sdk/dify_retrieval.py

@@ -16,6 +16,7 @@
 from flask import request, jsonify
 
 from api.db import LLMType
+from api.db.services.document_service import DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api import settings
@@ -70,12 +71,13 @@ def retrieval(tenant_id):
 
         records = []
         for c in ranks["chunks"]:
+            e, doc = DocumentService.get_by_id( c["doc_id"])
             c.pop("vector", None)
             records.append({
                 "content": c["content_with_weight"],
                 "score": c["similarity"],
                 "title": c["docnm_kwd"],
-                "metadata": {}
+                "metadata": doc.meta_fields
             })
 
         return jsonify({"records": records})

api/constants.py

@@ -25,3 +25,4 @@
 REQUEST_MAX_WAIT_SEC = 300
 
 DATASET_NAME_LIMIT = 128
+FILE_NAME_LEN_LIMIT = 256

api/db/services/canvas_service.py

@@ -173,6 +173,19 @@ def completion(tenant_id, agent_id, question, session_id=None, stream=True, **kw
             conv.reference = []
         conv.reference.append({"chunks": [], "doc_aggs": []})
 
+        kwargs_changed = False
+        if kwargs:
+            query = canvas.get_preset_param()
+            if query:
+                for ele in query:
+                    if ele["key"] in kwargs:
+                        if ele["value"] != kwargs[ele["key"]]:
+                            ele["value"] = kwargs[ele["key"]]
+                            kwargs_changed = True
+        if kwargs_changed:
+            conv.dsl = json.loads(str(canvas))
+            API4ConversationService.update_by_id(session_id, {"dsl": conv.dsl})
+
     final_ans = {"reference": [], "content": ""}
     if stream:
         try:

api/db/services/common_service.py

@@ -254,7 +254,7 @@ def delete_by_id(cls, pid):
         # Returns:
         #     Number of records deleted
         return cls.model.delete().where(cls.model.id == pid).execute()
-
+    
     @classmethod
     @DB.connection_context()
     def delete_by_ids(cls, pids):

api/db/services/document_service.py

@@ -27,6 +27,7 @@
 from peewee import fn
 
 from api import settings
+from api.constants import IMG_BASE64_PREFIX
 from api.db import FileType, LLMType, ParserType, StatusEnum, TaskStatus, UserTenantRole
 from api.db.db_models import DB, Document, Knowledgebase, Task, Tenant, UserTenant
 from api.db.db_utils import bulk_insert_into_db
@@ -147,7 +148,26 @@ def insert(cls, doc):
     def remove_document(cls, doc, tenant_id):
         cls.clear_chunk_num(doc.id)
         try:
+            page = 0
+            page_size = 1000
+            all_chunk_ids = []
+            while True:
+                chunks = settings.docStoreConn.search(["img_id"], [], {"doc_id": doc.id}, [], OrderByExpr(),
+                                                      page * page_size, page_size, search.index_name(tenant_id),
+                                                      [doc.kb_id])
+                chunk_ids = settings.docStoreConn.getChunkIds(chunks)
+                if not chunk_ids:
+                    break
+                all_chunk_ids.extend(chunk_ids)
+                page += 1
+            for cid in all_chunk_ids:
+                if STORAGE_IMPL.obj_exist(doc.kb_id, cid):
+                    STORAGE_IMPL.rm(doc.kb_id, cid)
+            if doc.thumbnail and not doc.thumbnail.startswith(IMG_BASE64_PREFIX):
+                if STORAGE_IMPL.obj_exist(doc.kb_id, doc.thumbnail):
+                    STORAGE_IMPL.rm(doc.kb_id, doc.thumbnail)
             settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id)
+
             graph_source = settings.docStoreConn.getFields(
                 settings.docStoreConn.search(["source_id"], [], {"kb_id": doc.kb_id, "knowledge_graph_kwd": ["graph"]}, [], OrderByExpr(), 0, 1, search.index_name(tenant_id), [doc.kb_id]), ["source_id"]
             )