Merge pull request #554 from guardrails-ai/mory91/check-json-text

Additional JSON parsing

Author: zsimjee

guardrails/utils/json_utils.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Type, Union
 
+import regex
+
 from guardrails.datatypes import (
     URL,
     Boolean,
@@ -370,6 +372,14 @@ def extract_json_from_ouput(output: str) -> Tuple[Optional[Dict], Optional[Excep
         has_block, block_start, block_end = has_code_block(output)
         if has_block and block_start is not None and block_end is not None:
             extracted_code_block = get_code_block(output, block_start, block_end)
+        else:
+            json_pattern = regex.compile(r"\{(?:[^{}]+|\{(?:(?R)|[^{}]+)*\})*\}")
+            json_groups = json_pattern.findall(output)
+            json_start, json_end = output.find("{"), output.rfind("}")
+            if len(json_groups) > 0 and len(json_groups[0]) == (
+                json_end - json_start + 1
+            ):
+                extracted_code_block = json_groups[0]
 
     # Treat the output as a JSON string, and load it into a dict.
     error = None

tests/unit_tests/utils/test_json_utils.py

@@ -32,14 +32,46 @@
 ```
 """
 
-invalid_json_code_block = """
+invalid_json_code_block__quotes = """
 ```json
 {
     a: 1
 }
 ```
 """
 
+invalid_json_code_block__braces = """
+```json
+{"choice": {"action": "flight", "random_key": "random_value"}
+```
+"""
+
+invalid_json_no_block__quotes = """
+{
+    a: 1
+}
+"""
+
+invalid_json_no_block__braces = """
+{"choice": {"action": "flight", "random_key": "random_value"}
+"""
+
+text_with_no_code_block = """
+Here is the data you requested
+
+{ "a": 1 , "b": { "c": [{"d": 2}, {"e": 3}]}}
+"""
+
+text_with_json_code_block = """
+Here is the data you requested
+
+```json
+{
+    "a": 1
+}
+```
+"""
+
 not_even_json = "This isn't even json..."
 
 
@@ -49,12 +81,39 @@
         (json_code_block, {"a": 1}, None),
         (anonymous_code_block, {"a": 1}, None),
         (no_code_block, {"a": 1}, None),
+        (text_with_no_code_block, {"a": 1, "b": {"c": [{"d": 2}, {"e": 3}]}}, None),
+        (text_with_json_code_block, {"a": 1}, None),
         (js_code_block, None, "Expecting value: line 1 column 1 (char 0)"),
         (
-            invalid_json_code_block,
+            invalid_json_code_block__quotes,
+            None,
+            "Expecting property name enclosed in double quotes: line 2 column 5 (char 6)",  # noqa
+        ),
+        (
+            invalid_json_code_block__braces,
+            None,
+            "Expecting ',' delimiter: line 1 column 62 (char 61)",  # noqa
+        ),
+        (
+            invalid_json_no_block__quotes,
             None,
             "Expecting property name enclosed in double quotes: line 2 column 5 (char 6)",  # noqa
         ),
+        (
+            invalid_json_no_block__braces,
+            None,
+            "Expecting ',' delimiter: line 3 column 1 (char 63)",  # noqa
+        ),
+        (
+            invalid_json_code_block__quotes,
+            None,
+            "Expecting property name enclosed in double quotes: line 2 column 5 (char 6)",  # noqa
+        ),
+        (
+            invalid_json_code_block__braces,
+            None,
+            "Expecting ',' delimiter: line 1 column 62 (char 61)",  # noqa
+        ),
         (not_even_json, None, "Expecting value: line 1 column 1 (char 0)"),
     ],
 )