From b12d1af8bb3645e3969eb53f6f57ac5d96dcf6c8 Mon Sep 17 00:00:00 2001 From: tegwick Date: Thu, 21 May 2026 03:19:27 +0200 Subject: [PATCH] Support Claude Code JSON schema execution --- llm_connect/claude_code.py | 96 ++++++++++++++++++++++++++------------ tests/test_claude_code.py | 58 +++++++++++++++++++++++ 2 files changed, 123 insertions(+), 31 deletions(-) create mode 100644 tests/test_claude_code.py diff --git a/llm_connect/claude_code.py b/llm_connect/claude_code.py index fa8c786..0b0854e 100644 --- a/llm_connect/claude_code.py +++ b/llm_connect/claude_code.py @@ -2,9 +2,12 @@ Claude Code CLI adapter — runs the ``claude`` CLI as a subprocess. """ -import asyncio -import subprocess -from typing import Optional +import asyncio +import json +import os +import subprocess +from pathlib import Path +from typing import Optional from llm_connect.adapter import LLMAdapter from llm_connect.models import RunConfig, LLMResponse @@ -23,25 +26,23 @@ class ClaudeCodeAdapter(LLMAdapter): length limits (compiled prompts can exceed 30 KB). """ - def __init__( - self, - cli_path: str = "claude", - model: Optional[str] = None, - config: Optional[LLMConfig] = None, - ): - self._config = config or LLMConfig(provider="claude-code") - self._cli_path = cli_path or self._config.claude_cli_path - self._model = model + def __init__( + self, + cli_path: Optional[str] = None, + model: Optional[str] = None, + config: Optional[LLMConfig] = None, + ): + self._config = config or LLMConfig(provider="claude-code") + self._cli_path = cli_path or self._resolve_cli_path() + self._model = model # ── LLMAdapter interface ──────────────────────────────────────── - def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: - self._preflight_budget(config) - cmd = [self._cli_path, "--print"] - if self._model: - cmd.extend(["--model", self._model]) - - timeout = config.timeout_seconds or self._config.timeout_seconds + def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: + self._preflight_budget(config) + cmd = self._build_command(config) + + timeout = config.timeout_seconds or self._config.timeout_seconds try: result = subprocess.run( @@ -85,12 +86,10 @@ class ClaudeCodeAdapter(LLMAdapter): self._consume_budget(config, response) return response - async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: - """Native async implementation using asyncio.create_subprocess_exec.""" - self._preflight_budget(config) - cmd = [self._cli_path, "--print"] - if self._model: - cmd.extend(["--model", self._model]) + async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: + """Native async implementation using asyncio.create_subprocess_exec.""" + self._preflight_budget(config) + cmd = self._build_command(config) timeout = config.timeout_seconds or self._config.timeout_seconds @@ -140,14 +139,49 @@ class ClaudeCodeAdapter(LLMAdapter): self._consume_budget(config, response) return response - def validate_config(self, config: RunConfig) -> bool: - try: - result = subprocess.run( - [self._cli_path, "--version"], + def validate_config(self, config: RunConfig) -> bool: + try: + result = subprocess.run( + [self._cli_path, "--version"], capture_output=True, text=True, timeout=10, ) return result.returncode == 0 - except (subprocess.TimeoutExpired, FileNotFoundError, OSError): - return False + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return False + + def _build_command(self, config: RunConfig) -> list[str]: + cmd = [self._cli_path, "--print"] + if self._model: + cmd.extend(["--model", self._model]) + + json_schema = _json_schema_arg(config) + if json_schema: + cmd.extend(["--json-schema", json_schema]) + return cmd + + def _resolve_cli_path(self) -> str: + configured = ( + os.environ.get("LLM_CONNECT_CLAUDE_CLI_PATH") + or os.environ.get("CLAUDE_CLI_PATH") + or self._config.claude_cli_path + ) + if configured and configured != "claude": + return configured + + local_cli = Path.home() / ".local" / "bin" / "claude" + if local_cli.exists(): + return str(local_cli) + return configured or "claude" + + +def _json_schema_arg(config: RunConfig) -> str | None: + schema = (config.model_params or {}).get("json_schema") + if not schema: + return None + if isinstance(schema, str): + return schema + if isinstance(schema, dict): + return json.dumps(schema, separators=(",", ":")) + return None diff --git a/tests/test_claude_code.py b/tests/test_claude_code.py new file mode 100644 index 0000000..d359a24 --- /dev/null +++ b/tests/test_claude_code.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from llm_connect.claude_code import ClaudeCodeAdapter +from llm_connect.config import LLMConfig +from llm_connect.models import RunConfig + + +def test_execute_prompt_passes_json_schema_to_claude_cli(monkeypatch): + calls: dict[str, object] = {} + + def fake_run(cmd, input, capture_output, text, timeout): # noqa: ANN001 + calls["cmd"] = cmd + calls["input"] = input + calls["capture_output"] = capture_output + calls["text"] = text + calls["timeout"] = timeout + return SimpleNamespace( + returncode=0, + stdout='{"summary":"ok","recommendations":[]}', + stderr="", + ) + + monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run) + adapter = ClaudeCodeAdapter(cli_path="/custom/claude") + + response = adapter.execute_prompt( + "Produce a report.", + RunConfig( + timeout_seconds=42, + model_params={"json_schema": {"type": "object"}}, + ), + ) + + assert calls == { + "cmd": [ + "/custom/claude", + "--print", + "--json-schema", + '{"type":"object"}', + ], + "input": "Produce a report.", + "capture_output": True, + "text": True, + "timeout": 42, + } + assert response.content == '{"summary":"ok","recommendations":[]}' + + +def test_claude_code_adapter_prefers_env_cli_path(monkeypatch): + monkeypatch.setenv("LLM_CONNECT_CLAUDE_CLI_PATH", "/home/me/bin/claude") + + adapter = ClaudeCodeAdapter( + config=LLMConfig(provider="claude-code", claude_cli_path="claude") + ) + + assert adapter._cli_path == "/home/me/bin/claude"