outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
非法操作 2024-11-05 18:21:41 +08:00 committed by GitHub
parent 7f583ec1ac
commit d92e3bd620
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 22 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py
View File

@ -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,11 +40,15 @@ 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 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")