From d92e3bd6209157e792c830ee768fccb58c1cb76a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= <hjlarry@163.com>
Date: Tue, 5 Nov 2024 18:21:41 +0800
Subject: [PATCH] fix: special prompt not work for comfyUI tool (#10307)

---
 .../builtin/comfyui/tools/comfyui_workflow.py | 26 ++++++++++++++++---
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py
index 79fe08a86b..d62772cda7 100644
--- a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py
+++ b/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: