Skip to content

Commit

Permalink
fix: special prompt not work for comfyUI tool (langgenius#10307)
Browse files Browse the repository at this point in the history
  • Loading branch information
hjlarry authored and JunXu01 committed Nov 9, 2024
1 parent 77558d0 commit aed03be
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,20 @@
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"])
Expand All @@ -26,13 +40,17 @@ def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMe
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:
Expand Down

0 comments on commit aed03be

Please sign in to comment.