mirror of
https://github.com/safishamsi/graphify.git
synced 2026-08-29 01:36:33 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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) ----------
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user