Avoid ambiguous JSON array extraction

LLM responses can include bracketed prose before the actual JSON array. The old greedy regex spanned from the first '[' to the last ']', causing valid single-array answers to be dropped and making multiple arrays indistinguishable.

Constraint: Keep object extraction behavior unchanged.

Rejected: Return first parseable array | silently guesses when a response has multiple valid arrays.

Confidence: high

Scope-risk: narrow

Reversibility: clean

Directive: Keep array extraction conservative when more than one valid top-level array is present.

Tested: uv run --with pytest pytest tests/test_json_utils.py -q

Tested: uv run --with ruff ruff check skillopt/utils/json_utils.py tests/test_json_utils.py

Not-tested: Full benchmark rollouts.
This commit is contained in:
Ziiii
2026-07-06 15:42:01 +08:00
parent e4ea6a6771
commit bf85781ffb
2 changed files with 106 additions and 4 deletions
+77 -4
View File
@@ -71,6 +71,74 @@ def _top_level_brace_objects(text: str) -> list[str]:
return spans 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
brace_depth = 0
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 == "{":
brace_depth += 1
i += 1
continue
if ch == "}":
brace_depth = max(0, brace_depth - 1)
i += 1
continue
if ch != "[" or brace_depth:
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: def _looks_json_like(span: str) -> bool:
"""Heuristic: does ``span`` look like an intended JSON object (vs. prose)? """Heuristic: does ``span`` look like an intended JSON object (vs. prose)?
@@ -163,10 +231,15 @@ def extract_json_array(text: str) -> list | None:
return json.loads(m.group(1)) return json.loads(m.group(1))
except json.JSONDecodeError: except json.JSONDecodeError:
pass pass
m = re.search(r"\[.*\]", text, re.DOTALL)
if m: parsed_arrays = []
for candidate in _top_level_bracket_arrays(text):
try: try:
return json.loads(m.group(0)) parsed = json.loads(candidate)
except json.JSONDecodeError: except json.JSONDecodeError:
pass continue
if isinstance(parsed, list):
parsed_arrays.append(parsed)
if len(parsed_arrays) == 1:
return parsed_arrays[0]
return None return None
+29
View File
@@ -5,6 +5,7 @@ import pytest
from skillopt.utils.json_utils import ( from skillopt.utils.json_utils import (
_top_level_brace_objects, _top_level_brace_objects,
_top_level_bracket_arrays,
extract_json, extract_json,
extract_json_array, extract_json_array,
) )
@@ -85,6 +86,22 @@ class TestTopLevelBraceObjects:
assert _top_level_brace_objects(text) == ['{"edit": "right"}'] 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') == []
class TestExtractJsonTolerantFallback: class TestExtractJsonTolerantFallback:
"""extract_json — json_repair fallback for malformed non-OpenAI output.""" """extract_json — json_repair fallback for malformed non-OpenAI output."""
@@ -186,3 +203,15 @@ class TestExtractJsonArray:
"""extract_json_array should not match a bare JSON object.""" """extract_json_array should not match a bare JSON object."""
text = '{"this is an object": true}' text = '{"this is an object": true}'
assert extract_json_array(text) is None 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