Merge pull request #103 from zixuanguo786-ctrl/codex/robust-json-array-extraction

[codex] Harden JSON array extraction

skillopt/utils/json_utils.py

@@ -71,6 +71,95 @@ def _top_level_brace_objects(text: str) -> list[str]:
     return spans
 
 
+def _top_level_bracket_arrays(text: str) -> list[str]:
+    """Return balanced top-level ``[...]`` spans in ``text``.
+
+    This mirrors :func:`_top_level_brace_objects` for array responses. Arrays
+    nested inside a JSON object are not top-level array answers, so they are
+    ignored while the outer scanner is inside a ``{...}`` span.
+    """
+    spans: list[str] = []
+    i, n = 0, len(text)
+    outer_in_str = False
+    outer_esc = False
+    while i < n:
+        ch = text[i]
+        if outer_in_str:
+            if outer_esc:
+                outer_esc = False
+            elif ch == "\\":
+                outer_esc = True
+            elif ch == '"':
+                outer_in_str = False
+            i += 1
+            continue
+        if ch == '"':
+            outer_in_str = True
+            i += 1
+            continue
+        if ch == "{":
+            # Skip balanced objects, including arrays nested inside them. An
+            # unmatched brace may be ordinary prose, though, so do not let it
+            # poison the rest of the scan and hide a later valid array.
+            depth = 0
+            in_str = False
+            esc = False
+            j = i
+            while j < n:
+                current = text[j]
+                if in_str:
+                    if esc:
+                        esc = False
+                    elif current == "\\":
+                        esc = True
+                    elif current == '"':
+                        in_str = False
+                elif current == '"':
+                    in_str = True
+                elif current == "{":
+                    depth += 1
+                elif current == "}":
+                    depth -= 1
+                    if depth == 0:
+                        i = j + 1
+                        break
+                j += 1
+            else:
+                i += 1
+            continue
+        if ch != "[":
+            i += 1
+            continue
+
+        depth = 0
+        in_str = False
+        esc = False
+        start = i
+        while i < n:
+            ch = text[i]
+            if in_str:
+                if esc:
+                    esc = False
+                elif ch == "\\":
+                    esc = True
+                elif ch == '"':
+                    in_str = False
+            elif ch == '"':
+                in_str = True
+            elif ch == "[":
+                depth += 1
+            elif ch == "]":
+                depth -= 1
+                if depth == 0:
+                    spans.append(text[start:i + 1])
+                    i += 1
+                    break
+            i += 1
+        else:
+            break  # unterminated final array
+    return spans
+
+
 def _looks_json_like(span: str) -> bool:
     """Heuristic: does ``span`` look like an intended JSON object (vs. prose)?
 
@@ -163,10 +252,15 @@ def extract_json_array(text: str) -> list | None:
             return json.loads(m.group(1))
         except json.JSONDecodeError:
             pass
-    m = re.search(r"\[.*\]", text, re.DOTALL)
+
-    if m:
+    parsed_arrays = []
+    for candidate in _top_level_bracket_arrays(text):
         try:
-            return json.loads(m.group(0))
+            parsed = json.loads(candidate)
         except json.JSONDecodeError:
-            pass
+            continue
+        if isinstance(parsed, list):
+            parsed_arrays.append(parsed)
+    if len(parsed_arrays) == 1:
+        return parsed_arrays[0]
     return None

tests/test_json_utils.py

@@ -5,6 +5,7 @@ import pytest
 
 from skillopt.utils.json_utils import (
     _top_level_brace_objects,
+    _top_level_bracket_arrays,
     extract_json,
     extract_json_array,
 )
@@ -85,6 +86,26 @@ class TestTopLevelBraceObjects:
         assert _top_level_brace_objects(text) == ['{"edit": "right"}']
 
 
+class TestTopLevelBracketArrays:
+    """_top_level_bracket_arrays — string/object-aware top-level array scan."""
+
+    def test_single_clean_array(self) -> None:
+        assert _top_level_bracket_arrays("[1, 2]") == ["[1, 2]"]
+
+    def test_two_top_level_arrays(self) -> None:
+        assert _top_level_bracket_arrays("[1]\n[2]") == ["[1]", "[2]"]
+
+    def test_array_inside_object_is_ignored(self) -> None:
+        assert _top_level_bracket_arrays('{"items": [1, 2]}') == []
+
+    def test_bracket_inside_quoted_prose_is_ignored(self) -> None:
+        assert _top_level_bracket_arrays('label is "set it to [x]" done') == []
+
+    def test_unmatched_prose_brace_does_not_hide_later_array(self) -> None:
+        text = "Explanation uses {placeholder, final answer [1, 2]"
+        assert _top_level_bracket_arrays(text) == ["[1, 2]"]
+
+
 class TestExtractJsonTolerantFallback:
     """extract_json — json_repair fallback for malformed non-OpenAI output."""
 
@@ -186,3 +207,19 @@ class TestExtractJsonArray:
         """extract_json_array should not match a bare JSON object."""
         text = '{"this is an object": true}'
         assert extract_json_array(text) is None
+
+    def test_ignores_prose_brackets_before_valid_array(self) -> None:
+        text = "Use [not JSON] as prose. Final answer: [1, 2, 3]"
+        assert extract_json_array(text) == [1, 2, 3]
+
+    def test_multiple_valid_arrays_are_ambiguous(self) -> None:
+        text = "First candidate [1], second candidate [2]"
+        assert extract_json_array(text) is None
+
+    def test_nested_object_array_not_confused_with_array_answer(self) -> None:
+        text = '{"items": [1, 2, 3]}'
+        assert extract_json_array(text) is None
+
+    def test_unmatched_prose_brace_before_valid_array(self) -> None:
+        text = "Explanation uses {placeholder, final answer [1, 2]"
+        assert extract_json_array(text) == [1, 2]