diff --git a/skillopt/utils/json_utils.py b/skillopt/utils/json_utils.py index f1fab6e9..408cb230 100644 --- a/skillopt/utils/json_utils.py +++ b/skillopt/utils/json_utils.py @@ -71,6 +71,74 @@ 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 + 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: """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)) 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 diff --git a/tests/test_json_utils.py b/tests/test_json_utils.py index 286efd7c..0728dc6b 100644 --- a/tests/test_json_utils.py +++ b/tests/test_json_utils.py @@ -5,6 +5,7 @@ from skillopt.utils.json_utils import ( _top_level_brace_objects, + _top_level_bracket_arrays, extract_json, extract_json_array, ) @@ -85,6 +86,22 @@ def test_real_object_after_quoted_brace(self) -> None: 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: """extract_json — json_repair fallback for malformed non-OpenAI output.""" @@ -186,3 +203,15 @@ def test_object_not_confused_with_array(self) -> None: """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