llm-connect/llm_connect/claude_code.py

"""
Claude Code CLI adapter — runs the ``claude`` CLI as a subprocess.
"""

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
from llm_connect.config import LLMConfig
from llm_connect._token_estimator import estimate_tokens
from llm_connect.exceptions import (
    LLMSubprocessError,
    LLMTimeoutError,
)


class ClaudeCodeAdapter(LLMAdapter):
    """LLM adapter that shells out to the ``claude`` CLI with ``--print``.

    The compiled prompt is piped via **stdin** to avoid shell argument
    length limits (compiled prompts can exceed 30 KB).
    """

    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._build_command(config)

        timeout = config.timeout_seconds or self._config.timeout_seconds

        try:
            result = subprocess.run(
                cmd,
                input=prompt,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired as exc:
            raise LLMTimeoutError(
                f"claude CLI timed out after {timeout}s",
                cause=exc,
            ) from exc

        if result.returncode != 0:
            raise LLMSubprocessError(
                f"claude CLI exited with code {result.returncode}",
                return_code=result.returncode,
                stderr=result.stderr,
            )

        content = result.stdout
        prompt_tokens = estimate_tokens(prompt)
        completion_tokens = estimate_tokens(content)

        response = LLMResponse(
            content=content,
            model=self._model or "claude-code-cli",
            usage={
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            finish_reason="stop",
            metadata={
                "provider": "claude-code",
                "cli_path": self._cli_path,
            },
        )
        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._build_command(config)

        timeout = config.timeout_seconds or self._config.timeout_seconds

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                proc.communicate(input=prompt.encode()),
                timeout=timeout,
            )
        except asyncio.TimeoutError as exc:
            raise LLMTimeoutError(
                f"claude CLI timed out after {timeout}s",
                cause=exc,
            ) from exc

        if proc.returncode != 0:
            raise LLMSubprocessError(
                f"claude CLI exited with code {proc.returncode}",
                return_code=proc.returncode,
                stderr=stderr_bytes.decode(),
            )

        content = stdout_bytes.decode()
        prompt_tokens = estimate_tokens(prompt)
        completion_tokens = estimate_tokens(content)

        response = LLMResponse(
            content=content,
            model=self._model or "claude-code-cli",
            usage={
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            finish_reason="stop",
            metadata={
                "provider": "claude-code",
                "cli_path": self._cli_path,
                "async": True,
            },
        )
        self._consume_budget(config, response)
        return response

    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

    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