fix(llm): recover JSON from reasoning-first model replies (#2882)
Reasoning-first models think out loud, fence the answer, or both, so the reply is not a bare JSON object and the strict parse dropped the whole chunk. Try the whole reply first (unchanged fast path), then each balanced-brace candidate, preferring braces that carry the extraction keys so narration braces do not shadow the answer; a shape restatement that carries the keys but no content is kept only as a last resort.
This commit is contained in:
rajarshidattapy
committed by
safishamsi
safishamsi
parent
2b02b7a7dc
commit
bce2c9b3d3
+138
-53
@@ -12,7 +12,7 @@ import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
@@ -986,11 +986,114 @@ def _sanitize_fragment(parsed: dict) -> dict:
|
||||
return parsed
|
||||
|
||||
|
||||
# Keys that identify an extraction fragment. Used to tell the graph object
|
||||
# apart from a brace that merely appeared in the model's narration (#2882).
|
||||
_FRAGMENT_KEYS = ("nodes", "edges", "hyperedges")
|
||||
_FRAGMENT_KEY_TOKENS = tuple(f'"{k}"' for k in _FRAGMENT_KEYS)
|
||||
# Bound on how many `{` positions are probed, so a pathological response with
|
||||
# thousands of braces cannot turn recovery into a quadratic scan. Applied to
|
||||
# the likely and the unlikely candidate lists separately, so a wall of noise
|
||||
# braces cannot crowd out an answer that comes after it.
|
||||
_MAX_OBJECT_CANDIDATES = 64
|
||||
# Reasoning models (nemotron, deepseek-r1, qwq, …) emit their chain of thought
|
||||
# in a <think> block ahead of the answer. It is prose, and it routinely
|
||||
# contains braces, so it is removed before any brace scanning.
|
||||
_THINK_BLOCK_RE = re.compile(r"<(think|thinking|reasoning)>.*?</\1>", re.S | re.I)
|
||||
_FENCE_RE = re.compile(r"```[ \t]*([A-Za-z0-9_+-]*)[ \t]*\r?\n(.*?)```", re.S)
|
||||
|
||||
|
||||
def _balanced_object(text: str, start: int) -> str | None:
|
||||
"""Return the balanced ``{...}`` substring starting at ``start``, else None."""
|
||||
depth = 0
|
||||
in_string = False
|
||||
escape = False
|
||||
for i in range(start, len(text)):
|
||||
ch = text[i]
|
||||
if escape:
|
||||
escape = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
escape = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return text[start:i + 1]
|
||||
return None
|
||||
|
||||
|
||||
def _json_object_candidates(text: str) -> list[int]:
|
||||
"""Indices of ``{`` that plausibly start an extraction fragment.
|
||||
|
||||
Braces followed shortly by one of ``_FRAGMENT_KEYS`` are tried first, so a
|
||||
model that narrates before answering — "Here's a thinking process: 1.
|
||||
**Analyze User Input:** …" with braces in the narration — does not have its
|
||||
real answer masked by the first brace in the text (#2882).
|
||||
"""
|
||||
preferred: list[int] = []
|
||||
rest: list[int] = []
|
||||
idx = text.find("{")
|
||||
while idx != -1:
|
||||
bucket = (
|
||||
preferred
|
||||
if any(k in text[idx:idx + 200] for k in _FRAGMENT_KEY_TOKENS)
|
||||
else rest
|
||||
)
|
||||
if len(bucket) < _MAX_OBJECT_CANDIDATES:
|
||||
bucket.append(idx)
|
||||
elif len(preferred) >= _MAX_OBJECT_CANDIDATES and len(rest) >= _MAX_OBJECT_CANDIDATES:
|
||||
break
|
||||
idx = text.find("{", idx + 1)
|
||||
return preferred + rest
|
||||
|
||||
|
||||
def _json_fragment_candidates(text: str) -> "Iterator[str]":
|
||||
"""Yield candidate JSON texts from a model reply, most-likely first.
|
||||
|
||||
Two sources, in order:
|
||||
|
||||
* fenced blocks — every fence, not just the first in the text, since a
|
||||
reasoning preamble often opens a ```python or ```text block of its own
|
||||
before the answer's ```json block. JSON-tagged and untagged fences come
|
||||
first; a fence in another language is still yielded, since models
|
||||
mislabel the tag.
|
||||
* balanced ``{...}`` objects lifted out of surrounding prose, at each
|
||||
plausible start rather than only the first `{` in the text.
|
||||
|
||||
Both read the ORIGINAL text. Rewriting it in place — as the old fence
|
||||
handling did, cutting from the first ``` to the last — let a fence in the
|
||||
narration truncate the real answer before it was ever parsed (#2882).
|
||||
"""
|
||||
for _lang, body in sorted(
|
||||
_FENCE_RE.findall(text), key=lambda b: b[0].strip().lower() not in ("json", "")
|
||||
):
|
||||
yield body.strip()
|
||||
for start in _json_object_candidates(text):
|
||||
blob = _balanced_object(text, start)
|
||||
if blob is not None:
|
||||
yield blob
|
||||
|
||||
|
||||
def _parse_llm_json(raw: str) -> dict:
|
||||
"""Strip optional markdown fences and parse JSON. Returns empty fragment on failure.
|
||||
|
||||
Caps the input at `_LLM_JSON_MAX_BYTES` so a hostile or runaway model
|
||||
response cannot exhaust memory inside `json.loads` (F-016).
|
||||
|
||||
Plenty of models will not return a bare JSON object no matter how the
|
||||
prompt is worded: they think out loud first, wrap the answer in a fence, or
|
||||
do both (#2882). So the whole reply is tried first, then each candidate
|
||||
:func:`_json_fragment_candidates` finds. An object carrying none of the
|
||||
extraction keys is kept only as a last resort — reasoning-first models
|
||||
routinely restate the schema (``{"description": "graph fragment"}``) before
|
||||
answering, and the narration must never shadow the answer that follows it.
|
||||
"""
|
||||
if len(raw) > _LLM_JSON_MAX_BYTES:
|
||||
print(
|
||||
@@ -999,65 +1102,47 @@ def _parse_llm_json(raw: str) -> dict:
|
||||
file=sys.stderr,
|
||||
)
|
||||
return {"nodes": [], "edges": [], "hyperedges": []}
|
||||
# Strategy 1: strip whitespace, then handle markdown fences anywhere in the
|
||||
# text (not only at offset 0 — the original code only stripped fences when
|
||||
# `raw.startswith("```")`, missing the common case where Claude prepends a
|
||||
# preamble like "Here's the extracted entities:\n\n```json\n{...}\n```").
|
||||
stripped = raw.strip()
|
||||
fence_start = stripped.find("```")
|
||||
if fence_start != -1:
|
||||
after_fence = stripped[fence_start + 3 :]
|
||||
# Optional language tag (json, JSON, javascript, etc.) up to newline.
|
||||
nl = after_fence.find("\n")
|
||||
if nl != -1 and after_fence[:nl].strip().lower() in {"json", "javascript", "js", ""}:
|
||||
after_fence = after_fence[nl + 1 :]
|
||||
fence_end = after_fence.rfind("```")
|
||||
if fence_end != -1:
|
||||
stripped = after_fence[:fence_end].strip()
|
||||
else:
|
||||
stripped = after_fence.strip()
|
||||
|
||||
stripped = _THINK_BLOCK_RE.sub(" ", raw).strip()
|
||||
|
||||
try:
|
||||
parsed = json.loads(stripped)
|
||||
if isinstance(parsed, dict):
|
||||
return _sanitize_fragment(parsed)
|
||||
# Top-level array/scalar (common LLM output) is not a usable graph
|
||||
# fragment; fall through to the next strategy rather than returning a
|
||||
# non-dict that callers will try to subscript (e.g. result["input_tokens"]).
|
||||
# fragment; fall through rather than returning a non-dict that callers
|
||||
# will try to subscript (e.g. result["input_tokens"]).
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
# Strategy 2: extract the first balanced JSON object found anywhere in
|
||||
# the text. Handles the case where Claude wraps the JSON in prose without
|
||||
# any markdown fence ("The extracted graph is { ... }. Hope this helps!").
|
||||
start = stripped.find("{")
|
||||
if start != -1:
|
||||
depth = 0
|
||||
in_string = False
|
||||
escape = False
|
||||
for i in range(start, len(stripped)):
|
||||
ch = stripped[i]
|
||||
if escape:
|
||||
escape = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
escape = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
try:
|
||||
parsed = json.loads(stripped[start : i + 1])
|
||||
if isinstance(parsed, dict):
|
||||
return _sanitize_fragment(parsed)
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
break
|
||||
|
||||
# Preference ladder, weakest last. A model that restates the required shape
|
||||
# before answering — "the schema is `{"nodes": [], "edges": []}`" — produces
|
||||
# a candidate that carries the extraction keys but no content, and taking it
|
||||
# would let the restatement shadow the answer just as surely as a prose
|
||||
# object would (#2882).
|
||||
empty_fragment: dict | None = None # right shape, nothing in it
|
||||
fallback: dict | None = None # parses, but not a fragment at all
|
||||
for candidate in _json_fragment_candidates(stripped):
|
||||
try:
|
||||
parsed = json.loads(candidate)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if not isinstance(parsed, dict):
|
||||
continue
|
||||
if any(k in parsed for k in _FRAGMENT_KEYS):
|
||||
if any(parsed.get(k) for k in _FRAGMENT_KEYS):
|
||||
return _sanitize_fragment(parsed)
|
||||
if empty_fragment is None:
|
||||
empty_fragment = parsed
|
||||
elif fallback is None:
|
||||
fallback = parsed
|
||||
|
||||
# A genuinely empty extraction is still a valid answer, and still reads as
|
||||
# hollow downstream, so it outranks an object that is not a fragment at all.
|
||||
for weaker in (empty_fragment, fallback):
|
||||
if weaker is not None:
|
||||
return _sanitize_fragment(weaker)
|
||||
|
||||
print(
|
||||
f"[graphify] LLM returned invalid JSON, skipping chunk "
|
||||
f"(first 200 chars: {raw[:200]!r})",
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
"""#2882: reasoning-first models narrate before they answer.
|
||||
|
||||
`nvidia/nemotron-*`, gemma and friends reply with a chain of thought and only
|
||||
then emit the JSON — sometimes bare, sometimes fenced, sometimes after a fence
|
||||
of their own. The narration routinely contains braces, so the first `{` in the
|
||||
text is not the answer. `_parse_llm_json` used to try exactly that one
|
||||
candidate and give up, dropping chunks whose answer was sitting right there:
|
||||
|
||||
[graphify] LLM returned invalid JSON, skipping chunk (first 200 chars:
|
||||
'Let me analyze the provided source files to extract a knowledge graph...')
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from graphify import llm
|
||||
|
||||
EMPTY = {"nodes": [], "edges": [], "hyperedges": []}
|
||||
|
||||
|
||||
def test_chain_of_thought_preamble_then_bare_json():
|
||||
raw = (
|
||||
"Let me analyze the provided source files to extract a knowledge graph "
|
||||
"fragment. I need to follow the rules carefully:\n\n"
|
||||
"1. EXTRACTED: relationship explicit in source\n"
|
||||
"2. INFERRED: reasonable inference\n"
|
||||
"3. The shape I must emit is { ... } with the keys listed above\n\n"
|
||||
'{"nodes": [{"id": "a", "label": "A"}], "edges": [], "hyperedges": []}'
|
||||
)
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "a", "label": "A"}]
|
||||
|
||||
|
||||
def test_narration_fence_precedes_the_answer_fence():
|
||||
"""A ```text block in the narration must not swallow the ```json answer.
|
||||
|
||||
The old fence handling took the FIRST ``` anywhere in the text and then cut
|
||||
at the LAST one, which mangled everything in between.
|
||||
"""
|
||||
raw = (
|
||||
"Here's a thinking process:\n\n"
|
||||
"1. **Analyze User Input:**\n"
|
||||
" - User provided several `<untrusted_source>` blocks.\n"
|
||||
" - Sketching the shape first:\n"
|
||||
"```text\n"
|
||||
"{ one node per file, plus edges }\n"
|
||||
"```\n\n"
|
||||
"Now the answer:\n\n"
|
||||
"```json\n"
|
||||
'{"nodes": [{"id": "docs_x", "label": "x.md"}], "edges": []}\n'
|
||||
"```\n"
|
||||
)
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "docs_x", "label": "x.md"}]
|
||||
|
||||
|
||||
def test_think_block_is_stripped_before_parsing():
|
||||
"""deepseek-r1 / qwq style: the chain of thought is tagged, and it is full
|
||||
of braces that would otherwise be probed as candidate objects."""
|
||||
raw = (
|
||||
"<think>\nI should emit {nodes, edges}. Let me check each file: "
|
||||
"{a.py} imports {b.py}...\n</think>\n"
|
||||
'{"nodes": [{"id": "b"}], "edges": [{"source": "a", "target": "b"}]}'
|
||||
)
|
||||
result = llm._parse_llm_json(raw)
|
||||
assert result["nodes"] == [{"id": "b"}]
|
||||
assert result["edges"] == [{"source": "a", "target": "b"}]
|
||||
|
||||
|
||||
def test_parseable_prose_object_does_not_shadow_the_answer():
|
||||
"""An earlier object that happens to be valid JSON is not the fragment."""
|
||||
raw = (
|
||||
'For reference the schema is {"description": "graph fragment"}.\n'
|
||||
'{"nodes": [{"id": "real"}], "edges": []}'
|
||||
)
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "real"}]
|
||||
|
||||
|
||||
def test_pure_prose_reply_still_degrades_to_empty_fragment():
|
||||
"""No answer anywhere: the empty fragment, never an exception, so the
|
||||
hollow detector takes over."""
|
||||
raw = (
|
||||
"Here's a thinking process:\n1. The user wants a knowledge graph.\n"
|
||||
"2. But I will describe it in words: { a depends on b }.\n"
|
||||
)
|
||||
assert llm._parse_llm_json(raw) == EMPTY
|
||||
|
||||
|
||||
def test_candidate_scan_is_bounded():
|
||||
"""A pathological reply full of braces must not turn recovery quadratic."""
|
||||
raw = "{ noise } " * 5000 + '{"nodes": [{"id": "z"}], "edges": []}'
|
||||
# The answer carries the fragment keys, so it is probed before the noise.
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "z"}]
|
||||
# Likely and unlikely candidates are capped separately, so the noise
|
||||
# cannot crowd out the answer, and neither list can grow without bound.
|
||||
assert len(llm._json_object_candidates(raw)) <= 2 * llm._MAX_OBJECT_CANDIDATES
|
||||
|
||||
|
||||
def test_empty_schema_restatement_does_not_shadow_the_answer():
|
||||
"""A restated shape carries the extraction keys but no content.
|
||||
|
||||
Preferring any candidate that merely *has* nodes/edges let the restatement
|
||||
win over the answer below it — the same shadowing this module exists to
|
||||
prevent, one level deeper than a keyless prose object.
|
||||
"""
|
||||
raw = (
|
||||
"The schema is:\n"
|
||||
"```json\n"
|
||||
'{"nodes": [], "edges": []}\n'
|
||||
"```\n"
|
||||
"Now the answer:\n"
|
||||
'{"nodes": [{"id": "real"}], "edges": []}\n'
|
||||
)
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "real"}]
|
||||
|
||||
|
||||
def test_empty_restatement_without_a_fence_does_not_shadow_either():
|
||||
raw = (
|
||||
'Shape: {"nodes": [], "edges": [], "hyperedges": []}\n'
|
||||
'Answer: {"nodes": [{"id": "real"}], "edges": []}\n'
|
||||
)
|
||||
assert llm._parse_llm_json(raw)["nodes"] == [{"id": "real"}]
|
||||
|
||||
|
||||
def test_a_genuinely_empty_extraction_is_still_returned():
|
||||
"""The model looked and found nothing: that is a valid answer, and it must
|
||||
keep reading as an empty fragment so the hollow detector takes over."""
|
||||
assert llm._parse_llm_json('{"nodes": [], "edges": [], "hyperedges": []}') == EMPTY
|
||||
assert llm._parse_llm_json(
|
||||
'Nothing to extract here.\n```json\n{"nodes": [], "edges": []}\n```'
|
||||
)["nodes"] == []
|
||||
|
||||
|
||||
def test_empty_fragment_outranks_a_non_fragment_object():
|
||||
"""Right shape but empty still beats an object that is not a fragment."""
|
||||
raw = (
|
||||
'Notes: {"description": "graph fragment", "version": 2}\n'
|
||||
'{"nodes": [], "edges": []}\n'
|
||||
)
|
||||
result = llm._parse_llm_json(raw)
|
||||
assert "description" not in result
|
||||
assert result["nodes"] == []
|
||||
Reference in New Issue
Block a user