diff --git a/aborist/qa/verify.py b/aborist/qa/verify.py index fc3f0b9..bbde372 100644 --- a/aborist/qa/verify.py +++ b/aborist/qa/verify.py @@ -685,12 +685,138 @@ _JSON_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*\n(.*?)\n\s*```\s*$", re.DOTALL _TRAILING_COMMA_RE = re.compile(r",(\s*[}\]])") +def _repair_truncated_json(text: str) -> tuple[str, list[str]]: + """Best-effort completion of truncated JSON. + + Walks ``text`` once tracking string state and bracket/brace stack. + At end-of-input, if the parse is unbalanced (stuck mid-string, + open ``[`` / ``{`` without matching close) or has a dangling + structural artifact (trailing comma, partial key), repair so the + result is parseable. Returns ``(repaired, fixups)``. + + Targets the truncation pattern observed on Hermes-3-8B JSON-mode + output for broad-descriptive questions: the model writes one + long claim text and runs out of ``max_tokens`` mid-sentence, + yielding e.g.:: + + {"claims":[{"text":"The Apollo program was the United + States spaceflight effort which landed... + + No closing ``"``, no closing ``}``, no closing ``]``, no closing + outer ``}``. The lenient parse path can't recover any structure. + Self-healing closes the open string, balances the stack, and + drops trailing commas / partial keys so the partial content is + preserved as a single claim with whatever fields survived. + + Fixups recorded: + - ``close_string`` — appended ``"`` to close an open string + - ``drop_partial_key`` — dropped a key without value + (``,"key":`` or ``,"key`` or ``"key":``) + - ``strip_trailing_comma`` — removed a comma immediately + before stack close (separate from the regex pass which + only handles structurally-correct trailing commas) + - ``close_brace`` / ``close_bracket`` — appended ``}`` / ``]`` + per open frame on the stack + + Conservative: never inserts content (no key names, no values, + no commas), only closes / drops. Worst case the repair is a + no-op or makes parsing fail in a different way; never silently + fabricates data. + """ + if not text or not text.strip(): + return text, [] + + fixups: list[str] = [] + stack: list[str] = [] + in_string = False + escaped = False + + for ch in text: + if escaped: + escaped = False + continue + if in_string: + if ch == "\\": + escaped = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "{" or ch == "[": + stack.append(ch) + elif ch == "}" or ch == "]": + if stack: + stack.pop() + + if not stack and not in_string: + return text, fixups + + repaired = text + + if in_string: + repaired += '"' + fixups.append("close_string") + + # Drop dangling structural fragments after the last legitimate + # value, walking back from end-of-string. Order matters: handle + # partial-key (`,"foo":` or `,"foo"` or `"foo":`) before + # trailing-comma so we don't strip the wrong comma. + while True: + rstripped = repaired.rstrip() + if not rstripped: + break + # Partial key: `..., "key": ` (colon at end after trim). + m = _PARTIAL_KEY_COLON_RE.search(rstripped) + if m and m.end() == len(rstripped): + repaired = rstripped[: m.start()] + if "drop_partial_key" not in fixups: + fixups.append("drop_partial_key") + continue + # Partial key: `..., "key"` (key without colon). + m = _PARTIAL_KEY_NO_COLON_RE.search(rstripped) + if m and m.end() == len(rstripped): + repaired = rstripped[: m.start()] + if "drop_partial_key" not in fixups: + fixups.append("drop_partial_key") + continue + # Trailing comma — strip when it would otherwise stick before close. + if rstripped.endswith(","): + repaired = rstripped[:-1] + if "strip_trailing_comma" not in fixups: + fixups.append("strip_trailing_comma") + continue + break + + # Close stack in reverse, mapping `{` → `}`, `[` → `]`. + while stack: + opener = stack.pop() + if opener == "{": + repaired += "}" + fixups.append("close_brace") + else: + repaired += "]" + fixups.append("close_bracket") + + return repaired, fixups + + +# Partial-key patterns for _repair_truncated_json. ``,"foo":`` or +# ``"foo":`` at end-of-string after rstrip = a key with no value. +_PARTIAL_KEY_COLON_RE = re.compile(r',?\s*"[^"]*"\s*:\s*$') +# ``,"foo"`` or ``"foo"`` with no following colon — also dangling. +_PARTIAL_KEY_NO_COLON_RE = re.compile(r',\s*"[^"]*"\s*$') + + def _lenient_json_parse(raw: str) -> tuple[object, list[str]]: """Parse ``raw`` as JSON, defensively peeling common model drift. Returns ``(parsed_obj, fixups_applied)`` — fixups list is empty when strict parse succeeded, otherwise names what we had to peel - (``"fence"``, ``"prose_trim"``, ``"curly_quotes"``, ``"trailing_comma"``). + (``"fence"``, ``"prose_trim"``, ``"curly_quotes"``, + ``"trailing_comma"``) or repair (``"close_string"``, + ``"close_brace"``, ``"close_bracket"``, ``"drop_partial_key"``, + ``"strip_trailing_comma"``). Raises ``json.JSONDecodeError`` if the lenient pass also fails. The fixups list lands in the verify payload so an agent can @@ -741,6 +867,22 @@ def _lenient_json_parse(raw: str) -> tuple[object, list[str]]: text = new_text fixups.append("trailing_comma") + # 5. Try strict parse before invoking the truncation repair — + # peeling alone may have made it valid. + try: + return _json.loads(text), fixups + except _json.JSONDecodeError: + pass + + # 6. Self-heal truncated JSON: close open strings, drop dangling + # partial keys / trailing commas, balance the bracket stack. The + # model ran out of max_tokens mid-output; close what we can and + # parse the partial structure rather than failing the whole run. + repaired, repair_fixups = _repair_truncated_json(text) + if repair_fixups: + text = repaired + fixups.extend(repair_fixups) + return _json.loads(text), fixups # Default allowed source roles for claim_lattice_pointer mode. Roles diff --git a/bench/qa_questions.txt b/bench/qa_questions.txt index aade294..189e678 100644 --- a/bench/qa_questions.txt +++ b/bench/qa_questions.txt @@ -19,9 +19,9 @@ tell me about charles babbage tell me about the apollo program # entity list — invites lazy-anchor on a magnet chunk -list the founders of microsoft what dinosaurs were in the first jurassic park film? who are the members of the beatles? +name simpsons family members including pets? # relationship / multi-fact who is supermans girlfriend?