Merge pull request #136 from Yif-Yang/fix/json-array-scan-linear

perf(json): keep array scan linear on malformed output

skillopt/utils/json_utils.py

@@ -82,6 +82,30 @@ def _top_level_bracket_arrays(text: str) -> list[str]:
     i, n = 0, len(text)
     outer_in_str = False
     outer_esc = False
+
+    # Precompute balanced object spans once.  Looking ahead to EOF for every
+    # unmatched ``{`` makes malformed model output quadratic, while a stack
+    # keeps the scan linear and still lets a later valid array remain visible.
+    brace_ends = [-1] * n
+    brace_stack: list[int] = []
+    brace_in_str = False
+    brace_esc = False
+    for pos, current in enumerate(text):
+        if brace_in_str:
+            if brace_esc:
+                brace_esc = False
+            elif current == "\\":
+                brace_esc = True
+            elif current == '"':
+                brace_in_str = False
+            continue
+        if current == '"':
+            brace_in_str = True
+        elif current == "{":
+            brace_stack.append(pos)
+        elif current == "}" and brace_stack:
+            brace_ends[brace_stack.pop()] = pos
+
     while i < n:
         ch = text[i]
         if outer_in_str:
@@ -101,31 +125,8 @@ def _top_level_bracket_arrays(text: str) -> list[str]:
             # 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
+            object_end = brace_ends[i]
-            in_str = False
+            i = object_end + 1 if object_end >= 0 else i + 1
-            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

tests/test_json_utils.py

@@ -105,6 +105,21 @@ class TestTopLevelBracketArrays:
         text = "Explanation uses {placeholder, final answer [1, 2]"
         assert _top_level_bracket_arrays(text) == ["[1, 2]"]
 
+    def test_many_unmatched_braces_scan_linearly(self) -> None:
+        class CountingText(str):
+            def __new__(cls, value: str):
+                instance = super().__new__(cls, value)
+                instance.reads = 0
+                return instance
+
+            def __getitem__(self, key):
+                self.reads += 1
+                return super().__getitem__(key)
+
+        text = CountingText("{" * 1_000 + " final answer [1, 2]")
+        assert _top_level_bracket_arrays(text) == ["[1, 2]"]
+        assert text.reads < len(text) * 4
+
 
 class TestExtractJsonTolerantFallback:
     """extract_json — json_repair fallback for malformed non-OpenAI output."""