diff --git a/llama_index/agent/react/output_parser.py b/llama_index/agent/react/output_parser.py index 7db41634c6a7b0d78ddd027835d89e46dd379255..eb66e02e388e8b8ef3c1d4397fa0abe5210411d4 100644 --- a/llama_index/agent/react/output_parser.py +++ b/llama_index/agent/react/output_parser.py @@ -1,7 +1,6 @@ """ReAct output parser.""" -import json import re from typing import Tuple @@ -16,7 +15,7 @@ from llama_index.types import BaseOutputParser def extract_tool_use(input_text: str) -> Tuple[str, str, str]: pattern = ( - r"\s*Thought: (.*?)\nAction: ([a-zA-Z0-9_]+).*?\nAction Input: .*?(\{.*?\})" + r"\s*Thought: (.*?)\nAction: ([a-zA-Z0-9_]+).*?\nAction Input: .*?(\{.*\})" ) match = re.search(pattern, input_text, re.DOTALL) @@ -50,6 +49,26 @@ def extract_final_response(input_text: str) -> Tuple[str, str]: return thought, answer +def parse_action_reasoning_step(output: str) -> ActionReasoningStep: + """ + Parse an action reasoning step from the LLM output. + """ + # Weaker LLMs may generate ReActAgent steps whose Action Input are horrible JSON strings. + # `dirtyjson` is more lenient than `json` in parsing JSON strings. + import dirtyjson as json + + thought, action, action_input = extract_tool_use(output) + json_str = extract_json_str(action_input) + # First we try json, if this fails we use ast + try: + action_input_dict = json.loads(json_str) + except Exception: + action_input_dict = action_input_parser(json_str) + return ActionReasoningStep( + thought=thought, action=action, action_input=action_input_dict + ) + + class ReActOutputParser(BaseOutputParser): """ReAct Output parser.""" @@ -85,18 +104,7 @@ class ReActOutputParser(BaseOutputParser): ) if "Action:" in output: - thought, action, action_input = extract_tool_use(output) - json_str = extract_json_str(action_input) - - # First we try json, if this fails we use ast - try: - action_input_dict = json.loads(json_str) - except json.JSONDecodeError: - action_input_dict = action_input_parser(json_str) - - return ActionReasoningStep( - thought=thought, action=action, action_input=action_input_dict - ) + return parse_action_reasoning_step(output) raise ValueError(f"Could not parse output: {output}") diff --git a/poetry.lock b/poetry.lock index f8385d6077aefa1ce2d94bd155a86d1a42023bc0..d906ea9feca65f83ff06ce647e56468b803b49a4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1177,6 +1177,17 @@ files = [ [package.extras] graph = ["objgraph (>=1.7.2)"] +[[package]] +name = "dirtyjson" +version = "1.0.8" +description = "JSON decoder for Python that can extract data from the muck" +optional = false +python-versions = "*" +files = [ + {file = "dirtyjson-1.0.8-py3-none-any.whl", hash = "sha256:125e27248435a58acace26d5c2c4c11a1c0de0a9c5124c5a94ba78e517d74f53"}, + {file = "dirtyjson-1.0.8.tar.gz", hash = "sha256:90ca4a18f3ff30ce849d100dcf4a003953c79d3a2348ef056f1d9c22231a25fd"}, +] + [[package]] name = "diskcache" version = "5.6.3" @@ -7812,4 +7823,4 @@ query-tools = ["guidance", "jsonpath-ng", "lm-format-enforcer", "rank-bm25", "sc [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "1ccf014ec22186ffcbec1325d4876089a318f0bec886beadd05b39a145e6a86a" +content-hash = "6f230242ed4fe799ae5b98f5b64e5f388ad54c304198c6de55bc056be873397c" diff --git a/pyproject.toml b/pyproject.toml index 1d30bc327fdff64888bdf2879bd19fae15f0296a..7ce3394a074bd5e003edb92d74d71241dd9ab00d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,7 @@ spacy = {optional = true, version = "^3.7.1"} aiohttp = "^3.8.6" networkx = ">=3.0" psycopg2-binary = {optional = true, version = "^2.9.9"} +dirtyjson = "^1.0.8" [tool.poetry.extras] gradientai = [ diff --git a/tests/agent/react/test_react_output_parser.py b/tests/agent/react/test_react_output_parser.py index 8d95218ea3552e7b72b7d14b6c268e929f2c5a8c..c9534048aea18b0b7ed78b522b72a8beaa72dcb4 100644 --- a/tests/agent/react/test_react_output_parser.py +++ b/tests/agent/react/test_react_output_parser.py @@ -1,9 +1,23 @@ from llama_index.agent.react.output_parser import ( extract_final_response, extract_tool_use, + parse_action_reasoning_step, ) +def test_parse_action_reasoning_step() -> None: + mock_input_text = """\ +Thought: Gotta use a tool. +Action: tool +Action Input: {'pages': ['coffee'] /* comment */, 'load_kwargs': {}, 'query_str': ''}, along those lines. +""" + assert parse_action_reasoning_step(mock_input_text).action_input == { + "pages": ["coffee"], + "load_kwargs": {}, + "query_str": "", + } + + def test_extract_tool_use() -> None: mock_input_text = """\ Thought: I need to use a tool to help me answer the question. @@ -16,6 +30,18 @@ Action Input: {"a": 1, "b": 1} assert action_input == '{"a": 1, "b": 1}' +def test_extract_tool_use_with_nested_dicts() -> None: + mock_input_text = """\ +Thought: Gotta use a tool. +Action: tool +Action Input: {"a": 1, "b": {}} +""" + thought, action, action_input = extract_tool_use(mock_input_text) + assert thought == "Gotta use a tool." + assert action == "tool" + assert action_input == '{"a": 1, "b": {}}' + + def test_extract_tool_use_() -> None: mock_input_text = """\ Thought: I need to use a tool to help me answer the question.