From cb5fb751e91808e057df23e7c18f10b5f91cf71f Mon Sep 17 00:00:00 2001
From: Junlin Zhou <jameszhou2108@hotmail.com>
Date: Fri, 11 Aug 2023 05:26:07 +0800
Subject: [PATCH] Enhance regex of structured_chat agents' output parser
 (#8965)

Current regex only extracts agent's action between '` ``` ``` `', this
commit will extract action between both '` ```json ``` `' and '` ``` ```
`'

This is very similar to #7511
Co-authored-by: zjl <junlinzhou@yzbigdata.com>
---
 .../agents/structured_chat/output_parser.py   |  4 +-
 .../unit_tests/agents/test_structured_chat.py | 47 +++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 libs/langchain/tests/unit_tests/agents/test_structured_chat.py

diff --git a/libs/langchain/langchain/agents/structured_chat/output_parser.py b/libs/langchain/langchain/agents/structured_chat/output_parser.py
index 6d40848e0ed..5b1b598acaa 100644
--- a/libs/langchain/langchain/agents/structured_chat/output_parser.py
+++ b/libs/langchain/langchain/agents/structured_chat/output_parser.py
@@ -19,12 +19,14 @@ logger = logging.getLogger(__name__)
 class StructuredChatOutputParser(AgentOutputParser):
     """Output parser for the structured chat agent."""
 
+    pattern = re.compile(r"```(?:json)?\n(.*?)```", re.DOTALL)
+
     def get_format_instructions(self) -> str:
         return FORMAT_INSTRUCTIONS
 
     def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         try:
-            action_match = re.search(r"```(.*?)```?", text, re.DOTALL)
+            action_match = self.pattern.search(text)
             if action_match is not None:
                 response = json.loads(action_match.group(1).strip(), strict=False)
                 if isinstance(response, list):
diff --git a/libs/langchain/tests/unit_tests/agents/test_structured_chat.py b/libs/langchain/tests/unit_tests/agents/test_structured_chat.py
new file mode 100644
index 00000000000..356c3d87e2d
--- /dev/null
+++ b/libs/langchain/tests/unit_tests/agents/test_structured_chat.py
@@ -0,0 +1,47 @@
+"""Unittests for langchain.agents.chat package."""
+from typing import Tuple
+
+from langchain.agents.structured_chat.output_parser import StructuredChatOutputParser
+from langchain.schema import AgentAction
+
+output_parser = StructuredChatOutputParser()
+
+
+def get_action_and_input(text: str) -> Tuple[str, str]:
+    output = output_parser.parse(text)
+    if isinstance(output, AgentAction):
+        return output.tool, str(output.tool_input)
+    else:
+        return "Final Answer", output.return_values["output"]
+
+
+def test_parse_with_language() -> None:
+    llm_output = """I can use the `foo` tool to achieve the goal.
+
+    Action:
+    ```json
+    {
+      "action": "foo",
+      "action_input": "bar"
+    }
+    ```
+    """
+    action, action_input = get_action_and_input(llm_output)
+    assert action == "foo"
+    assert action_input == "bar"
+
+
+def test_parse_without_language() -> None:
+    llm_output = """I can use the `foo` tool to achieve the goal.
+
+    Action:
+    ```
+    {
+      "action": "foo",
+      "action_input": "bar"
+    }
+    ```
+    """
+    action, action_input = get_action_and_input(llm_output)
+    assert action == "foo"
+    assert action_input == "bar"