瀏覽代碼

fix: special prompt not work for comfyUI tool (#10307)

非法操作 5 月之前
父節點
當前提交
d92e3bd620
共有 1 個文件被更改,包括 22 次插入4 次删除
  1. 22 4
      api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py

+ 22 - 4
api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py

@@ -8,6 +8,20 @@ from core.tools.provider.builtin.comfyui.tools.comfyui_client import ComfyUiClie
 from core.tools.tool.builtin_tool import BuiltinTool
 
 
+def sanitize_json_string(s):
+    escape_dict = {
+        "\n": "\\n",
+        "\r": "\\r",
+        "\t": "\\t",
+        "\b": "\\b",
+        "\f": "\\f",
+    }
+    for char, escaped in escape_dict.items():
+        s = s.replace(char, escaped)
+
+    return s
+
+
 class ComfyUIWorkflowTool(BuiltinTool):
     def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]:
         comfyui = ComfyUiClient(self.runtime.credentials["base_url"])
@@ -26,13 +40,17 @@ class ComfyUIWorkflowTool(BuiltinTool):
         set_prompt_with_ksampler = True
         if "{{positive_prompt}}" in workflow:
             set_prompt_with_ksampler = False
-            workflow = workflow.replace("{{positive_prompt}}", positive_prompt)
-            workflow = workflow.replace("{{negative_prompt}}", negative_prompt)
+            workflow = workflow.replace("{{positive_prompt}}", positive_prompt.replace('"', "'"))
+            workflow = workflow.replace("{{negative_prompt}}", negative_prompt.replace('"', "'"))
 
         try:
             prompt = json.loads(workflow)
-        except:
-            return self.create_text_message("the Workflow JSON is not correct")
+        except json.JSONDecodeError:
+            cleaned_string = sanitize_json_string(workflow)
+            try:
+                prompt = json.loads(cleaned_string)
+            except:
+                return self.create_text_message("the Workflow JSON is not correct")
 
         if set_prompt_with_ksampler:
             try: