Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 98 additions & 4 deletions skillopt/utils/json_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)?

Expand Down Expand Up @@ -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
37 changes: 37 additions & 0 deletions tests/test_json_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from skillopt.utils.json_utils import (
_top_level_brace_objects,
_top_level_bracket_arrays,
extract_json,
extract_json_array,
)
Expand Down Expand Up @@ -85,6 +86,26 @@ 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') == []

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."""

Expand Down Expand Up @@ -186,3 +207,19 @@ 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

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]