diff --git a/pyrit/exceptions/exceptions_helpers.py b/pyrit/exceptions/exceptions_helpers.py index 73df4e2586..3662a7cdd7 100644 --- a/pyrit/exceptions/exceptions_helpers.py +++ b/pyrit/exceptions/exceptions_helpers.py @@ -115,10 +115,16 @@ def extract_json_from_string(response_msg: str) -> str: str: The extracted JSON string if found, otherwise the original string. """ - json_pattern = re.compile(r"\{.*\}|\[.*\]") - match = json_pattern.search(response_msg) - if match: - return match.group(0) + decoder = json.JSONDecoder() + for index, char in enumerate(response_msg): + if char not in "[{": + continue + + try: + _, end_index = decoder.raw_decode(response_msg, idx=index) + return response_msg[index:end_index] + except json.JSONDecodeError: + continue return response_msg diff --git a/tests/unit/exceptions/test_exceptions_helpers.py b/tests/unit/exceptions/test_exceptions_helpers.py index c8ddccb536..54435bd746 100644 --- a/tests/unit/exceptions/test_exceptions_helpers.py +++ b/tests/unit/exceptions/test_exceptions_helpers.py @@ -59,6 +59,10 @@ def test_remove_end_md_json(input_str, expected_output): ("No JSON here", "No JSON here"), ('jsn\n{"key": "value"}\n```', '{"key": "value"}'), ('Some text before JSON {"a": [1,2,3], "b": {"c": 4}} some text after JSON', '{"a": [1,2,3], "b": {"c": 4}}'), + ( + 'Some text before JSON {\n "key": "value",\n "nested": {"enabled": true}\n} some text after JSON', + '{\n "key": "value",\n "nested": {"enabled": true}\n}', + ), ], ) def test_extract_json_from_string(input_str, expected_output):