From 6cf258a809f42725bdf0cecf18789821fd0e6437 Mon Sep 17 00:00:00 2001 From: shirukai <308899573@qq.com> Date: Mon, 31 Mar 2025 16:27:29 +0800 Subject: [PATCH] fix: code block syntax cannot be displayed correctly in react mode (#16904) --- .../agent/output_parser/cot_output_parser.py | 76 +++++++++++-------- .../output_parser/test_cot_output_parser.py | 70 +++++++++++++++++ 2 files changed, 114 insertions(+), 32 deletions(-) create mode 100644 api/tests/unit_tests/core/agent/output_parser/test_cot_output_parser.py diff --git a/api/core/agent/output_parser/cot_output_parser.py b/api/core/agent/output_parser/cot_output_parser.py index 61fa774ea5..7c8f09e6b9 100644 --- a/api/core/agent/output_parser/cot_output_parser.py +++ b/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): - try: - action = json.loads(json_str, strict=False) - action_name = None - action_input = None + def parse_action(action) -> Union[str, AgentScratchpadUnit.Action]: + action_name = None + action_input = None + if isinstance(action, str): + try: + action = json.loads(action, strict=False) + except json.JSONDecodeError: + return action or "" - # cohere always returns a list - if isinstance(action, list) and len(action) == 1: - action = action[0] + # cohere always returns a list + if isinstance(action, list) and len(action) == 1: + action = action[0] - for key, value in action.items(): - if "input" in key.lower(): - 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, - ) + for key, value in action.items(): + if "input" in key.lower(): + action_input = value else: - return json_str or "" - except: - return json_str or "" + action_name = value - def extra_json_from_code_block(code_block) -> Generator[Union[str, AgentScratchpadUnit.Action], None, None]: - code_blocks = re.findall(r"```(.*?)```", code_block, re.DOTALL) - if not code_blocks: - return - for block in code_blocks: - json_text = re.sub(r"^[a-zA-Z]+\n", "", block.strip(), flags=re.MULTILINE) - yield parse_action(json_text) + 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.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) - code_block_cache = "" + action_json_list = extra_json_from_code_block(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 diff --git a/api/tests/unit_tests/core/agent/output_parser/test_cot_output_parser.py b/api/tests/unit_tests/core/agent/output_parser/test_cot_output_parser.py new file mode 100644 index 0000000000..4a613e35b0 --- /dev/null +++ b/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"]