Move json utility to the custom json parsing; apply it to the monologue-like agents (#1740)

agenthub/micro/agent.py

@@ -1,10 +1,7 @@
-from json import JSONDecodeError
-
 from jinja2 import BaseLoader, Environment
 
 from opendevin.controller.agent import Agent
 from opendevin.controller.state.state import State
-from opendevin.core.exceptions import LLMOutputError
 from opendevin.core.utils import json
 from opendevin.events.action import Action, action_from_dict
 from opendevin.llm.llm import LLM
@@ -14,26 +11,11 @@ from .registry import all_microagents
 
 
 def parse_response(orig_response: str) -> Action:
-    depth = 0
-    start = -1
-    for i, char in enumerate(orig_response):
-        if char == '{':
-            if depth == 0:
-                start = i
-            depth += 1
-        elif char == '}':
-            depth -= 1
-            if depth == 0 and start != -1:
-                response = orig_response[start : i + 1]
-                try:
-                    action_dict = json.loads(response)
-                    action = action_from_dict(action_dict)
-                    return action
-                except JSONDecodeError as e:
-                    raise LLMOutputError(
-                        'Invalid JSON in response. Please make sure the response is a valid JSON object.'
-                    ) from e
-    raise LLMOutputError('No valid JSON object found in response.')
+    # attempt to load the JSON dict from the response
+    action_dict = json.loads(orig_response)
+
+    # load the action from the dict
+    return action_from_dict(action_dict)
 
 
 def to_json(obj, **kwargs):

agenthub/monologue_agent/utils/prompts.py

@@ -1,8 +1,4 @@
-import re
-from json import JSONDecodeError
-
 from opendevin.core.config import config
-from opendevin.core.exceptions import LLMOutputError
 from opendevin.core.utils import json
 from opendevin.events.action import (
     Action,
@@ -159,7 +155,7 @@ def get_request_action_prompt(
     }
 
 
-def parse_action_response(response: str) -> Action:
+def parse_action_response(orig_response: str) -> Action:
     """
     Parses a string to find an action within it
 
@@ -169,35 +165,13 @@ def parse_action_response(response: str) -> Action:
     Returns:
     - Action: The action that was found in the response string
     """
-    try:
-        action_dict = json.loads(response)
-    except JSONDecodeError:
-        # Find response-looking json in the output and use the more promising one. Helps with weak llms
-        response_json_matches = re.finditer(
-            r"""{\s*\"action\":\s?\"(\w+)\"(?:,?|,\s*\"args\":\s?{((?:.|\s)*?)})\s*}""",
-            response,
-        )  # Find all response-looking strings
-
-        def rank(match):
-            return (
-                len(match[2]) if match[1] == 'message' else 130
-            )  # Crudely rank multiple responses by length
-
-        try:
-            action_dict = json.loads(
-                max(response_json_matches, key=rank)[0]
-            )  # Use the highest ranked response
-        except (ValueError, JSONDecodeError):
-            raise LLMOutputError(
-                'Invalid JSON, the response must be well-formed JSON as specified in the prompt.'
-            )
-    except (ValueError, TypeError):
-        raise LLMOutputError(
-            'Invalid JSON, the response must be well-formed JSON as specified in the prompt.'
-        )
+    # attempt to load the JSON dict from the response
+    action_dict = json.loads(orig_response)
+
     if 'content' in action_dict:
         # The LLM gets confused here. Might as well be robust
         action_dict['contents'] = action_dict.pop('content')
+
     return action_from_dict(action_dict)
 
 

agenthub/planner_agent/prompt.py

@@ -175,9 +175,12 @@ def parse_response(response: str) -> Action:
     Returns:
     - Action: A valid next action to perform from model output
     """
+    # attempt to load the JSON dict from the response
     action_dict = json.loads(response)
+
     if 'contents' in action_dict:
         # The LLM gets confused here. Might as well be robust
         action_dict['content'] = action_dict.pop('contents')
+
     action = action_from_dict(action_dict)
     return action

opendevin/core/utils/json.py

@@ -2,6 +2,8 @@ import json
 
 from json_repair import repair_json
 
+from opendevin.core.exceptions import LLMOutputError
+
 
 def my_encoder(obj):
     """
@@ -25,14 +27,26 @@ def dumps(obj, **kwargs):
     return json.dumps(obj, default=my_encoder, **kwargs)
 
 
-def loads(s, **kwargs):
+def loads(json_str, **kwargs):
     """
     Create a JSON object from str
     """
-    json_start = s.find('{')
-    json_end = s.rfind('}') + 1
-    if json_start == -1 or json_end == -1:
-        raise ValueError('Invalid response: no JSON found')
-    s = s[json_start:json_end]
-    s = repair_json(s)
-    return json.loads(s, **kwargs)
+    depth = 0
+    start = -1
+    for i, char in enumerate(json_str):
+        if char == '{':
+            if depth == 0:
+                start = i
+            depth += 1
+        elif char == '}':
+            depth -= 1
+            if depth == 0 and start != -1:
+                response = json_str[start : i + 1]
+                try:
+                    json_str = repair_json(response)
+                    return json.loads(json_str, **kwargs)
+                except (json.JSONDecodeError, ValueError, TypeError) as e:
+                    raise LLMOutputError(
+                        'Invalid JSON in response. Please make sure the response is a valid JSON object.'
+                    ) from e
+    raise LLMOutputError('No valid JSON object found in response.')

tests/unit/test_response_parsing.py

@@ -1,13 +1,23 @@
 import pytest
 
-from agenthub.micro.agent import LLMOutputError, parse_response
+from agenthub.micro.agent import parse_response as parse_response_micro
+from agenthub.monologue_agent.utils.prompts import (
+    parse_action_response as parse_response_monologue,
+)
+from agenthub.planner_agent.prompt import parse_response as parse_response_planner
+from opendevin.core.exceptions import LLMOutputError
+from opendevin.core.utils.json import loads as custom_loads
 from opendevin.events.action import (
     FileWriteAction,
     MessageAction,
 )
 
 
-def test_parse_single_complete_json():
+@pytest.mark.parametrize(
+    'parse_response_module',
+    [parse_response_micro, parse_response_planner, parse_response_monologue],
+)
+def test_parse_single_complete_json(parse_response_module):
     input_response = """
     {
         "action": "message",
@@ -19,10 +29,15 @@ def test_parse_single_complete_json():
     expected = MessageAction(
         "The following typos were fixed:\n* 'futur' -> 'future'\n* 'imagin' -> 'imagine'\n* 'techological' -> 'technological'\n* 'responsability' -> 'responsibility'\nThe corrected file is ./short_essay.txt."
     )
-    assert parse_response(input_response) == expected
+    result = parse_response_module(input_response)
+    assert result == expected
 
 
-def test_parse_json_with_surrounding_text():
+@pytest.mark.parametrize(
+    'parse_response_module',
+    [parse_response_micro, parse_response_planner, parse_response_monologue],
+)
+def test_parse_json_with_surrounding_text(parse_response_module):
     input_response = """
     Some initial text that is not JSON formatted.
     {
@@ -37,10 +52,15 @@ def test_parse_json_with_surrounding_text():
     expected = FileWriteAction(
         path='./updated_file.txt', content='Updated text content here...'
     )
-    assert parse_response(input_response) == expected
+    result = parse_response_module(input_response)
+    assert result == expected
 
 
-def test_parse_first_of_multiple_jsons():
+@pytest.mark.parametrize(
+    'parse_response_module',
+    [parse_response_micro, parse_response_planner, parse_response_monologue],
+)
+def test_parse_first_of_multiple_jsons(parse_response_module):
     input_response = """
     I will firstly do
     {
@@ -59,10 +79,18 @@ def test_parse_first_of_multiple_jsons():
     }
     """
     expected = FileWriteAction(path='./short_essay.txt', content='Text content here...')
-    assert parse_response(input_response) == expected
+    result = parse_response_module(input_response)
+    assert result == expected
 
 
 def test_invalid_json_raises_error():
-    input_response = '{"action": "write", "args": { "path": "./short_essay.txt", "content": "Missing closing brace"'
+    # This should fail if repair_json is able to fix this faulty JSON
+    input_response = '{"action": "write", "args": { "path": "./short_essay.txt", "content": "Missing closing brace" }'
+    with pytest.raises(LLMOutputError):
+        custom_loads(input_response)
+
+
+def test_no_json_found():
+    input_response = 'This is just a string with no JSON object.'
     with pytest.raises(LLMOutputError):
-        parse_response(input_response)
+        custom_loads(input_response)