diff --git a/graphify/llm.py b/graphify/llm.py index a05c9d4632..97e199b339 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -12,7 +12,7 @@ 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 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)>.*?", 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})", diff --git a/tests/test_llm_parser_reasoning.py b/tests/test_llm_parser_reasoning.py new file mode 100644 index 0000000000..5c16c7ba25 --- /dev/null +++ b/tests/test_llm_parser_reasoning.py @@ -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 `` 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 = ( + "\nI should emit {nodes, edges}. Let me check each file: " + "{a.py} imports {b.py}...\n\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"] == []