From d939fa1df7c0ed4d29dadc4b1389c5379898b57c Mon Sep 17 00:00:00 2001 From: Yyunozor Date: Wed, 22 Jul 2026 01:54:05 +0200 Subject: [PATCH] fix(llm): pin claude-cli output to a JSON schema when supported (#2076) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The claude-cli backend delivers the extraction schema in the user turn and trusts the model to emit raw JSON. Newer Claude Code releases treat that prompt as an agentic task and report the result in prose instead ("Knowledge graph extracted — 21 nodes, 20 edges…"), so the graph parses empty, reads as truncation, and adaptive-retry bisects without ever converging. Pass --json-schema (structured output) when the CLI advertises it — probed once via `claude --help` and cached — so the object shape is constrained regardless of prompt framing. Older CLIs that predate the flag keep the user-turn prompt as a fallback. The `result` envelope still carries the JSON string, so the parse path is unchanged. --- graphify/llm.py | 66 ++++++++++++++++++++++++++++++++ tests/test_claude_cli_backend.py | 57 +++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/graphify/llm.py b/graphify/llm.py index 6c0602358..2c2dbc429 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1337,6 +1337,62 @@ def _claude_cli_envelope(stdout: str) -> dict: return envelope +# A JSON Schema pinning the top-level shape graphify consumes. Passed to +# `claude -p --json-schema` (structured output) so the CLI CONSTRAINS the model +# to emit the object directly instead of relying on it CHOOSING to honour a +# "raw JSON only" instruction in the prompt. Item internals stay loose so a +# valid extraction is never rejected; the `result` envelope field still carries +# the JSON string, so the parse path is unchanged. See #2076. +_EXTRACTION_JSON_SCHEMA = json.dumps( + { + "type": "object", + "properties": { + "nodes": {"type": "array", "items": {"type": "object"}}, + "edges": {"type": "array", "items": {"type": "object"}}, + "hyperedges": {"type": "array", "items": {"type": "object"}}, + }, + "required": ["nodes", "edges"], + } +) + +# Cache the `--json-schema` capability probe per resolved claude command so it +# runs at most once per process (extract fans a chunk out per file/slice). +_JSON_SCHEMA_SUPPORT: dict[str, bool] = {} + + +def _claude_cli_supports_json_schema(claude_cmd: str) -> bool: + """Return True if this Claude Code CLI accepts ``--json-schema``. + + Structured output (``--json-schema``) landed in newer Claude Code releases. + Probing ``claude --help`` for the flag is a direct capability check — more + reliable than guessing a version boundary — so graphify uses structured + output where it exists and falls back to the user-turn prompt on older CLIs + that predate it. Any probe failure is treated as "unsupported" (safe + fallback). Result is cached per resolved command. + """ + import subprocess + + cached = _JSON_SCHEMA_SUPPORT.get(claude_cmd) + if cached is not None: + return cached + try: + proc = subprocess.run( + [claude_cmd, "--help"], + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=30, + check=False, + **_no_window_kwargs(), + ) + supported = "--json-schema" in (proc.stdout or "") + except (OSError, subprocess.SubprocessError): + supported = False + _JSON_SCHEMA_SUPPORT[claude_cmd] = supported + return supported + + def _call_claude_cli(user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False, images: list[_ImageRef] | None = None) -> dict: """Call Claude via the locally-installed Claude Code CLI (`claude -p`). @@ -1425,6 +1481,16 @@ def _call_claude_cli(user_message: str, max_tokens: int = 8192, *, deep_mode: bo cli_model = os.environ.get("GRAPHIFY_CLAUDE_CLI_MODEL", "").strip() if cli_model: cli_args.extend(["--model", cli_model]) + # Constrain the output shape structurally where the CLI supports it. Newer + # Claude Code releases increasingly treat a bare file-dump prompt as an + # agentic task and REPORT the extraction in prose ("Knowledge graph + # extracted — 21 nodes, 20 edges…") instead of returning it; that parses to + # zero nodes, reads as truncation, and gets bisected without ever + # converging (#2076). --json-schema pins the object shape regardless of + # that framing; the user-turn prompt above stays as the fallback for older + # CLIs that predate the flag. + if _claude_cli_supports_json_schema(claude_cmd): + cli_args.extend(["--json-schema", _EXTRACTION_JSON_SCHEMA]) proc = subprocess.run( cli_args, input=combined_message, diff --git a/tests/test_claude_cli_backend.py b/tests/test_claude_cli_backend.py index c695fb62d..28c4e1915 100644 --- a/tests/test_claude_cli_backend.py +++ b/tests/test_claude_cli_backend.py @@ -154,6 +154,63 @@ def test_user_turn_preserves_untrusted_source_guardrails(fake_claude): assert "untrusted_source" in sent +# ---------- structured output via --json-schema (#2076) ---------- +# Newer Claude Code CLIs treat a bare file-dump prompt as an agentic task and +# REPORT the extraction in prose instead of returning JSON, so the graph comes +# out empty and adaptive-retry bisects forever. When the CLI supports +# `--json-schema`, graphify constrains the output shape structurally so the +# model must emit the object regardless of framing. Older CLIs that predate the +# flag fall back to the user-turn prompt, unchanged. + + +def test_json_schema_flag_added_when_cli_supports_it(monkeypatch, fake_claude): + """When the CLI advertises --json-schema, it is passed with a schema that + pins the top-level {nodes, edges} shape graphify parses.""" + monkeypatch.setattr(llm, "_claude_cli_supports_json_schema", lambda cmd: True) + llm._call_claude_cli("dummy source", max_tokens=8192) + argv = fake_claude.call_args.args[0] + assert "--json-schema" in argv + schema = json.loads(argv[argv.index("--json-schema") + 1]) + assert schema["type"] == "object" + assert set(schema["required"]) == {"nodes", "edges"} + + +def test_json_schema_flag_absent_when_cli_lacks_it(monkeypatch, fake_claude): + """Older CLIs without --json-schema must not receive the flag (it would be + an unknown option) — extraction falls back to the user-turn prompt.""" + monkeypatch.setattr(llm, "_claude_cli_supports_json_schema", lambda cmd: False) + result = llm._call_claude_cli("dummy source", max_tokens=8192) + argv = fake_claude.call_args.args[0] + assert "--json-schema" not in argv + assert len(result["nodes"]) == 2 # result envelope still carries the JSON + + +def test_supports_json_schema_detects_flag_in_help(): + llm._JSON_SCHEMA_SUPPORT.clear() + help_text = "Options:\n --json-schema JSON Schema for structured output\n" + completed = MagicMock(returncode=0, stdout=help_text, stderr="") + with patch("subprocess.run", return_value=completed): + assert llm._claude_cli_supports_json_schema("/fake/claude-new") is True + + +def test_supports_json_schema_false_when_flag_absent(): + llm._JSON_SCHEMA_SUPPORT.clear() + help_text = "Options:\n --output-format text|json|stream-json\n" + completed = MagicMock(returncode=0, stdout=help_text, stderr="") + with patch("subprocess.run", return_value=completed): + assert llm._claude_cli_supports_json_schema("/fake/claude-old") is False + + +def test_supports_json_schema_false_and_cached_on_probe_error(): + """A probe that fails to run is treated as unsupported (safe fallback) and + cached so it is not re-probed for every chunk.""" + llm._JSON_SCHEMA_SUPPORT.clear() + with patch("subprocess.run", side_effect=OSError("boom")) as run: + assert llm._claude_cli_supports_json_schema("/fake/claude-broken") is False + assert llm._claude_cli_supports_json_schema("/fake/claude-broken") is False + assert run.call_count == 1 + + # ---------- Windows path resolution (#1072) ----------