fix(llm): pin claude-cli output to a JSON schema when supported (#2076)

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.
This commit is contained in:
Yyunozor
2026-07-22 15:51:03 +01:00
committed by safishamsi
parent 8c5f585587
commit 7116a6f5f5
2 changed files with 123 additions and 0 deletions
+66
View File
@@ -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,
+57
View File
@@ -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 <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 <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) ----------