fix: code block syntax cannot be displayed correctly in react mode (#16904)

api/core/agent/output_parser/cot_output_parser.py

@@ -12,39 +12,45 @@ class CotAgentOutputParser:
     def handle_react_stream_output(
         cls, llm_response: Generator[LLMResultChunk, None, None], usage_dict: dict
     ) -> Generator[Union[str, AgentScratchpadUnit.Action], None, None]:
-        def parse_action(json_str):
+        def parse_action(action) -> Union[str, AgentScratchpadUnit.Action]:
-            try:
+            action_name = None
-                action = json.loads(json_str, strict=False)
+            action_input = None
-                action_name = None
+            if isinstance(action, str):
-                action_input = None
+                try:
+                    action = json.loads(action, strict=False)
+                except json.JSONDecodeError:
+                    return action or ""
 
-                # cohere always returns a list
+            # cohere always returns a list
-                if isinstance(action, list) and len(action) == 1:
+            if isinstance(action, list) and len(action) == 1:
-                    action = action[0]
+                action = action[0]
 
-                for key, value in action.items():
+            for key, value in action.items():
-                    if "input" in key.lower():
+                if "input" in key.lower():
-                        action_input = value
+                    action_input = value
-                    else:
-                        action_name = value
-
-                if action_name is not None and action_input is not None:
-                    return AgentScratchpadUnit.Action(
-                        action_name=action_name,
-                        action_input=action_input,
-                    )
                 else:
-                    return json_str or ""
+                    action_name = value
-            except:
-                return json_str or ""
 
-        def extra_json_from_code_block(code_block) -> Generator[Union[str, AgentScratchpadUnit.Action], None, None]:
+            if action_name is not None and action_input is not None:
-            code_blocks = re.findall(r"```(.*?)```", code_block, re.DOTALL)
+                return AgentScratchpadUnit.Action(
-            if not code_blocks:
+                    action_name=action_name,
-                return
+                    action_input=action_input,
-            for block in code_blocks:
+                )
-                json_text = re.sub(r"^[a-zA-Z]+\n", "", block.strip(), flags=re.MULTILINE)
+            else:
-                yield parse_action(json_text)
+                return json.dumps(action)
+
+        def extra_json_from_code_block(code_block) -> list[Union[list, dict]]:
+            blocks = re.findall(r"```[json]*\s*([\[{].*[]}])\s*```", code_block, re.DOTALL | re.IGNORECASE)
+            if not blocks:
+                return []
+            try:
+                json_blocks = []
+                for block in blocks:
+                    json_text = re.sub(r"^[a-zA-Z]+\n", "", block.strip(), flags=re.MULTILINE)
+                    json_blocks.append(json.loads(json_text, strict=False))
+                return json_blocks
+            except:
+                return []
 
         code_block_cache = ""
         code_block_delimiter_count = 0
@@ -78,7 +84,7 @@ class CotAgentOutputParser:
                 delta = response_content[index : index + steps]
                 yield_delta = False
 
-                if delta == "`":
+                if not in_json and delta == "`":
                     last_character = delta
                     code_block_cache += delta
                     code_block_delimiter_count += 1
@@ -159,8 +165,14 @@ class CotAgentOutputParser:
                 if code_block_delimiter_count == 3:
                     if in_code_block:
                         last_character = delta
-                        yield from extra_json_from_code_block(code_block_cache)
+                        action_json_list = extra_json_from_code_block(code_block_cache)
-                        code_block_cache = ""
+                        if action_json_list:
+                            for action_json in action_json_list:
+                                yield parse_action(action_json)
+                            code_block_cache = ""
+                        else:
+                            index += steps
+                            continue
 
                     in_code_block = not in_code_block
                     code_block_delimiter_count = 0

api/tests/unit_tests/core/agent/output_parser/test_cot_output_parser.py

@@ -0,0 +1,70 @@
+import json
+from collections.abc import Generator
+
+from core.agent.entities import AgentScratchpadUnit
+from core.agent.output_parser.cot_output_parser import CotAgentOutputParser
+from core.model_runtime.entities.llm_entities import AssistantPromptMessage, LLMResultChunk, LLMResultChunkDelta
+
+
+def mock_llm_response(text) -> Generator[LLMResultChunk, None, None]:
+    for i in range(len(text)):
+        yield LLMResultChunk(
+            model="model",
+            prompt_messages=[],
+            delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=text[i], tool_calls=[])),
+        )
+
+
+def test_cot_output_parser():
+    test_cases = [
+        {
+            "input": 'Through: abc\nAction: ```{"action": "Final Answer", "action_input": "```echarts\n {}\n```"}```',
+            "action": {"action": "Final Answer", "action_input": "```echarts\n {}\n```"},
+            "output": 'Through: abc\n {"action": "Final Answer", "action_input": "```echarts\\n {}\\n```"}',
+        },
+        # code block with json
+        {
+            "input": 'Through: abc\nAction: ```json\n{"action": "Final Answer", "action_input": "```echarts\n {'
+            '}\n```"}```',
+            "action": {"action": "Final Answer", "action_input": "```echarts\n {}\n```"},
+            "output": 'Through: abc\n {"action": "Final Answer", "action_input": "```echarts\\n {}\\n```"}',
+        },
+        # code block with JSON
+        {
+            "input": 'Through: abc\nAction: ```JSON\n{"action": "Final Answer", "action_input": "```echarts\n {'
+            '}\n```"}```',
+            "action": {"action": "Final Answer", "action_input": "```echarts\n {}\n```"},
+            "output": 'Through: abc\n {"action": "Final Answer", "action_input": "```echarts\\n {}\\n```"}',
+        },
+        # list
+        {
+            "input": 'Through: abc\nAction: ```[{"action": "Final Answer", "action_input": "```echarts\n {}\n```"}]```',
+            "action": {"action": "Final Answer", "action_input": "```echarts\n {}\n```"},
+            "output": 'Through: abc\n {"action": "Final Answer", "action_input": "```echarts\\n {}\\n```"}',
+        },
+        # no code block
+        {
+            "input": 'Through: abc\nAction: {"action": "Final Answer", "action_input": "```echarts\n {}\n```"}',
+            "action": {"action": "Final Answer", "action_input": "```echarts\n {}\n```"},
+            "output": 'Through: abc\n {"action": "Final Answer", "action_input": "```echarts\\n {}\\n```"}',
+        },
+        # no code block and json
+        {"input": "Through: abc\nAction: efg", "action": {}, "output": "Through: abc\n efg"},
+    ]
+
+    parser = CotAgentOutputParser()
+    usage_dict = {}
+    for test_case in test_cases:
+        # mock llm_response as a generator by text
+        llm_response: Generator[LLMResultChunk, None, None] = mock_llm_response(test_case["input"])
+        results = parser.handle_react_stream_output(llm_response, usage_dict)
+        output = ""
+        for result in results:
+            if isinstance(result, str):
+                output += result
+            elif isinstance(result, AgentScratchpadUnit.Action):
+                if test_case["action"]:
+                    assert result.to_dict() == test_case["action"]
+                output += json.dumps(result.to_dict())
+        if test_case["output"]:
+            assert output == test_case["output"]