scripts/llm_backends.py
A supporting file of the persona skill, shown as source.
#!/usr/bin/env python3
"""
Shared LLM backend adapters for Claude and Codex CLIs.
"""
from __future__ import annotations
import asyncio
import json
import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path
BACKEND_CHOICES = ("auto", "claude-cli", "codex-cli")
REPORT_BACKEND_CHOICES = ("same", "python", "claude-cli", "codex-cli")
CODEX_ENV_MARKERS = ("CODEX_THREAD_ID", "CODEX_CI")
CLAUDE_ENV_MARKERS = (
"CLAUDECODE",
"CLAUDE_CODE",
"CLAUDE_CODE_ENTRYPOINT",
"CLAUDE_PROJECT_DIR",
"CLAUDE_SESSION_ID",
"CLAUDE_CODE_SIMPLE",
)
def _backend_command(backend: str) -> str:
if backend == "claude-cli":
return "claude"
if backend == "codex-cli":
return "codex"
raise ValueError(f"Unsupported backend: {backend}")
def backend_available(backend: str) -> bool:
"""Return True when the corresponding CLI is available on PATH."""
return shutil.which(_backend_command(backend)) is not None
def infer_backend_from_env(env: dict[str, str] | None = None) -> str | None:
"""Infer runtime backend from environment markers when possible."""
env = env or os.environ
if any(env.get(key) for key in CODEX_ENV_MARKERS):
return "codex-cli"
if any(env.get(key) for key in CLAUDE_ENV_MARKERS):
return "claude-cli"
return None
def resolve_backend(requested: str | None = None, env: dict[str, str] | None = None) -> str:
"""Resolve auto/explicit backend selection to a concrete CLI."""
backend = (requested or "auto").strip()
if backend not in BACKEND_CHOICES:
raise ValueError(f"Unknown backend: {backend}. Valid: {BACKEND_CHOICES}")
if backend != "auto":
if not backend_available(backend):
raise FileNotFoundError(
f"Requested backend is not available on PATH: {_backend_command(backend)}"
)
return backend
inferred = infer_backend_from_env(env=env)
if inferred and backend_available(inferred):
return inferred
for candidate in ("claude-cli", "codex-cli"):
if backend_available(candidate):
return candidate
raise FileNotFoundError("Neither `claude` nor `codex` CLI is available on PATH")
def resolve_report_backend(requested: str | None, simulation_backend: str) -> str:
"""Resolve the report backend, allowing `same` and deterministic Python fallback."""
backend = (requested or "same").strip()
if backend not in REPORT_BACKEND_CHOICES:
raise ValueError(
f"Unknown report backend: {backend}. Valid: {REPORT_BACKEND_CHOICES}"
)
if backend == "same":
return simulation_backend
return backend
def default_model_for_backend(backend: str) -> str | None:
"""Return backend-specific default model alias, if any."""
if backend == "claude-cli":
return "sonnet"
return None
def resolve_model(model: str | None, backend: str) -> str | None:
"""Resolve model override or backend default."""
if model:
return model
return default_model_for_backend(backend)
def format_model_label(model: str | None) -> str:
"""Human-readable model label for logs."""
return model or "(CLI default)"
def describe_backend(backend: str) -> str:
"""Human-readable backend description."""
if backend == "claude-cli":
return "claude CLI (independent context per persona)"
if backend == "codex-cli":
return "codex CLI (independent context per persona)"
return backend
def build_json_only_schema(required_keys: list[str] | None = None) -> dict:
"""Build a minimal JSON schema that locks the top-level contract."""
response_properties: dict[str, object] = {}
if required_keys:
for key in required_keys:
response_properties[key] = {}
schema: dict[str, object] = {
"type": "object",
"properties": {
"responses": {
"type": "object",
"properties": response_properties,
"additionalProperties": True,
}
},
"required": ["responses"],
"additionalProperties": False,
}
if required_keys:
schema["properties"]["responses"]["required"] = required_keys
return schema
def extract_json_from_text(text: str):
"""Extract a JSON value from raw text or fenced code blocks."""
text = (text or "").strip()
if not text:
raise ValueError("Empty response text")
try:
return json.loads(text)
except json.JSONDecodeError:
pass
patterns = [
r"```json\s*\n(.*?)```",
r"```\s*\n(.*?)```",
r"(\{.*\})",
r"(\[.*\])",
]
for pattern in patterns:
match = re.search(pattern, text, re.DOTALL)
if not match:
continue
candidate = match.group(1).strip()
try:
return json.loads(candidate)
except json.JSONDecodeError:
continue
raise ValueError(f"Could not extract valid JSON from response:\n{text[:500]}")
def build_codex_prompt(system_prompt: str, user_message: str, output_mode: str) -> str:
"""Combine system instructions and user message for Codex exec."""
mode_instruction = {
"json": "Return only a JSON object that matches the provided schema exactly.",
"text": "Return only the requested final text. Do not wrap it in code fences.",
}.get(output_mode, "Return only the requested final output.")
return (
"You are being called through a non-interactive CLI adapter.\n"
"Treat the SYSTEM PROMPT section as higher priority than the USER MESSAGE section.\n"
"Do not ask clarifying questions. Do not inspect local files unless absolutely necessary.\n"
f"{mode_instruction}\n\n"
"<SYSTEM_PROMPT>\n"
f"{system_prompt}\n"
"</SYSTEM_PROMPT>\n\n"
"<USER_MESSAGE>\n"
f"{user_message}\n"
"</USER_MESSAGE>\n"
)
# Optional claude CLI flags introduced after the original release. Each is
# emitted only when the installed CLI advertises it, so older CLIs keep the
# legacy command line unchanged.
CLAUDE_CAPABILITY_FLAGS = (
"--json-schema",
"--safe-mode",
"--fallback-model",
"--effort",
"--max-budget-usd",
)
_claude_capabilities_cache: dict[str, bool] | None = None
def detect_claude_capabilities(help_text: str | None = None) -> dict[str, bool]:
"""Detect which optional claude CLI flags this installation supports.
Runs `claude --help` once per process and caches the result. Pass
`help_text` to parse a known help output instead (used by tests).
"""
global _claude_capabilities_cache
if help_text is not None:
return {flag: flag in help_text for flag in CLAUDE_CAPABILITY_FLAGS}
if _claude_capabilities_cache is None:
try:
proc = subprocess.run(
["claude", "--help"],
capture_output=True,
text=True,
timeout=30,
)
detected = proc.stdout or ""
except (OSError, subprocess.TimeoutExpired, subprocess.SubprocessError):
detected = ""
_claude_capabilities_cache = {
flag: flag in detected for flag in CLAUDE_CAPABILITY_FLAGS
}
return _claude_capabilities_cache
def _build_claude_print_command(
system_prompt: str,
model: str | None = None,
json_schema: dict | None = None,
isolation: bool = True,
fallback_model: str | None = None,
effort: str | None = None,
max_budget_usd: float | None = None,
capabilities: dict[str, bool] | None = None,
) -> list[str]:
"""Build the claude -p command, gating newer flags on CLI capabilities."""
caps = capabilities if capabilities is not None else detect_claude_capabilities()
cmd = [
"claude",
"-p",
"--system-prompt",
system_prompt,
"--output-format",
"json",
"--tools",
"",
"--no-session-persistence",
]
if isolation and caps.get("--safe-mode"):
cmd.append("--safe-mode")
if json_schema is not None and caps.get("--json-schema"):
cmd.extend(["--json-schema", json.dumps(json_schema)])
if fallback_model and caps.get("--fallback-model"):
cmd.extend(["--fallback-model", fallback_model])
if effort and caps.get("--effort"):
cmd.extend(["--effort", effort])
if max_budget_usd is not None and caps.get("--max-budget-usd"):
cmd.extend(["--max-budget-usd", str(max_budget_usd)])
if model:
cmd.extend(["--model", model])
return cmd
def _parse_claude_envelope(stdout: str) -> dict:
"""Parse the claude -p --output-format json envelope defensively.
With --json-schema the CLI leaves `result` empty and returns the parsed
object in a separate top-level `structured_output` key, so both fields
are surfaced and the caller picks whichever is populated.
"""
try:
payload = json.loads(stdout)
except json.JSONDecodeError as error:
raise RuntimeError(
f"claude CLI returned non-JSON output: {stdout.strip()[:1000]}"
) from error
if not isinstance(payload, dict):
raise RuntimeError(f"Unexpected claude CLI payload: {str(payload)[:1000]}")
if payload.get("is_error"):
detail = payload.get("result") or payload.get("subtype") or "unknown error"
raise RuntimeError(f"claude CLI reported an error: {str(detail)[:1000]}")
usage = payload.get("usage") or {}
return {
"text": payload.get("result") or "",
"structured": payload.get("structured_output"),
"input_tokens": usage.get("input_tokens", 0),
"output_tokens": usage.get("output_tokens", 0),
"cache_creation_input_tokens": usage.get("cache_creation_input_tokens", 0),
"cache_read_input_tokens": usage.get("cache_read_input_tokens", 0),
"total_cost_usd": payload.get("total_cost_usd"),
"model_usage": payload.get("modelUsage") or {},
}
def preview_command(
backend: str,
model: str | None = None,
structured_output: bool = False,
isolation: bool = True,
fallback_model: str | None = None,
effort: str | None = None,
capabilities: dict[str, bool] | None = None,
) -> str:
"""Render a backend command preview for dry-run output."""
if backend == "claude-cli":
caps = capabilities if capabilities is not None else detect_claude_capabilities()
lines = [
"claude -p \\",
' --system-prompt "..." \\',
" --output-format json \\",
' --tools "" \\',
" --no-session-persistence \\",
]
if isolation and caps.get("--safe-mode"):
lines.append(" --safe-mode \\")
if structured_output and caps.get("--json-schema"):
lines.append(" --json-schema '{...}' \\")
if fallback_model and caps.get("--fallback-model"):
lines.append(f" --fallback-model {fallback_model} \\")
if effort and caps.get("--effort"):
lines.append(f" --effort {effort} \\")
if model:
lines.append(f" --model {model}")
return "\n".join(lines)
lines = [
"codex -a never -s read-only \\",
]
if model:
lines.append(f" -m {model} \\")
lines.extend(
[
" exec \\",
" --skip-git-repo-check \\",
" --ephemeral \\",
]
)
if structured_output:
lines.extend(
[
" --output-schema /tmp/schema.json \\",
" -o /tmp/response.json \\",
]
)
else:
lines.append(" -o /tmp/response.txt \\")
lines.append(" -")
return "\n".join(lines)
async def _communicate_async(
cmd: list[str],
stdin_text: str,
cwd: Path | None = None,
timeout_s: int = 300,
) -> tuple[int, str, str]:
proc = await asyncio.create_subprocess_exec(
*cmd,
cwd=str(cwd) if cwd else None,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
try:
stdout, stderr = await asyncio.wait_for(
proc.communicate(input=stdin_text.encode()),
timeout=timeout_s,
)
except asyncio.TimeoutError as error:
proc.kill()
await proc.communicate()
raise TimeoutError(f"Command timed out after {timeout_s}s: {' '.join(cmd[:4])}") from error
return proc.returncode, stdout.decode(), stderr.decode()
def _communicate_sync(
cmd: list[str],
stdin_text: str,
cwd: Path | None = None,
timeout_s: int = 300,
) -> tuple[int, str, str]:
try:
proc = subprocess.run(
cmd,
input=stdin_text,
capture_output=True,
text=True,
timeout=timeout_s,
cwd=str(cwd) if cwd else None,
)
except subprocess.TimeoutExpired as error:
raise TimeoutError(f"Command timed out after {timeout_s}s: {' '.join(cmd[:4])}") from error
return proc.returncode, proc.stdout, proc.stderr
def _build_codex_exec_command(
model: str | None,
output_path: Path,
schema_path: Path | None = None,
) -> list[str]:
cmd = ["codex", "-a", "never", "-s", "read-only"]
if model:
cmd.extend(["-m", model])
cmd.extend(["exec", "--skip-git-repo-check", "--ephemeral"])
if schema_path is not None:
cmd.extend(["--output-schema", str(schema_path)])
cmd.extend(["-o", str(output_path), "-"])
return cmd
async def run_json_completion_async(
*,
backend: str,
system_prompt: str,
user_message: str,
model: str | None = None,
json_schema: dict | None = None,
cwd: Path | None = None,
timeout_s: int = 300,
isolation: bool = True,
structured_output: bool = True,
fallback_model: str | None = None,
effort: str | None = None,
max_budget_usd: float | None = None,
) -> dict:
"""Run a single JSON completion against the resolved backend."""
resolved_backend = resolve_backend(backend)
resolved_model = resolve_model(model, resolved_backend)
if resolved_backend == "claude-cli":
claude_schema = json_schema if structured_output else None
cmd = _build_claude_print_command(
system_prompt,
model=resolved_model,
json_schema=claude_schema,
isolation=isolation,
fallback_model=fallback_model,
effort=effort,
max_budget_usd=max_budget_usd,
)
returncode, stdout, stderr = await _communicate_async(
cmd, user_message, cwd=cwd, timeout_s=timeout_s
)
if returncode != 0 and claude_schema is not None and "schema" in (
(stderr or stdout).lower()
):
# Per-run degradation: if the CLI rejected this particular schema,
# retry once without --json-schema and fall back to text extraction.
cmd = _build_claude_print_command(
system_prompt,
model=resolved_model,
isolation=isolation,
fallback_model=fallback_model,
effort=effort,
max_budget_usd=max_budget_usd,
)
returncode, stdout, stderr = await _communicate_async(
cmd, user_message, cwd=cwd, timeout_s=timeout_s
)
if returncode != 0:
raise RuntimeError((stderr or stdout).strip()[:1000])
envelope = _parse_claude_envelope(stdout)
if envelope["structured"] is not None:
data = envelope["structured"]
else:
data = extract_json_from_text(envelope["text"])
return {
"backend": resolved_backend,
"resolved_model": resolved_model,
"text": envelope["text"],
"data": data,
"input_tokens": envelope["input_tokens"],
"output_tokens": envelope["output_tokens"],
"cache_creation_input_tokens": envelope["cache_creation_input_tokens"],
"cache_read_input_tokens": envelope["cache_read_input_tokens"],
"total_cost_usd": envelope["total_cost_usd"],
"model_usage": envelope["model_usage"],
"usage_supported": True,
}
prompt = build_codex_prompt(system_prompt, user_message, output_mode="json")
with tempfile.TemporaryDirectory() as tmpdir:
tmpdir_path = Path(tmpdir)
output_path = tmpdir_path / "response.json"
schema_path = None
if json_schema is not None:
schema_path = tmpdir_path / "schema.json"
schema_path.write_text(json.dumps(json_schema), encoding="utf-8")
cmd = _build_codex_exec_command(
model=resolved_model,
output_path=output_path,
schema_path=schema_path,
)
returncode, stdout, stderr = await _communicate_async(
cmd, prompt, cwd=cwd, timeout_s=timeout_s
)
if returncode != 0:
raise RuntimeError((stderr or stdout).strip()[:1000])
text = output_path.read_text(encoding="utf-8").strip()
return {
"backend": resolved_backend,
"resolved_model": resolved_model,
"text": text,
"data": extract_json_from_text(text),
"input_tokens": 0,
"output_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"total_cost_usd": None,
"model_usage": {},
"usage_supported": False,
}
def run_text_completion(
*,
backend: str,
system_prompt: str,
user_message: str,
model: str | None = None,
cwd: Path | None = None,
timeout_s: int = 300,
isolation: bool = True,
fallback_model: str | None = None,
effort: str | None = None,
max_budget_usd: float | None = None,
) -> dict:
"""Run a single text completion against the resolved backend."""
resolved_backend = resolve_backend(backend)
resolved_model = resolve_model(model, resolved_backend)
if resolved_backend == "claude-cli":
cmd = _build_claude_print_command(
system_prompt,
model=resolved_model,
isolation=isolation,
fallback_model=fallback_model,
effort=effort,
max_budget_usd=max_budget_usd,
)
returncode, stdout, stderr = _communicate_sync(
cmd, user_message, cwd=cwd, timeout_s=timeout_s
)
if returncode != 0:
raise RuntimeError((stderr or stdout).strip()[:1000])
envelope = _parse_claude_envelope(stdout)
return {
"backend": resolved_backend,
"resolved_model": resolved_model,
"text": envelope["text"],
"input_tokens": envelope["input_tokens"],
"output_tokens": envelope["output_tokens"],
"cache_creation_input_tokens": envelope["cache_creation_input_tokens"],
"cache_read_input_tokens": envelope["cache_read_input_tokens"],
"total_cost_usd": envelope["total_cost_usd"],
"model_usage": envelope["model_usage"],
"usage_supported": True,
}
prompt = build_codex_prompt(system_prompt, user_message, output_mode="text")
with tempfile.TemporaryDirectory() as tmpdir:
output_path = Path(tmpdir) / "response.txt"
cmd = _build_codex_exec_command(model=resolved_model, output_path=output_path)
returncode, stdout, stderr = _communicate_sync(
cmd, prompt, cwd=cwd, timeout_s=timeout_s
)
if returncode != 0:
raise RuntimeError((stderr or stdout).strip()[:1000])
return {
"backend": resolved_backend,
"resolved_model": resolved_model,
"text": output_path.read_text(encoding="utf-8"),
"input_tokens": 0,
"output_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"total_cost_usd": None,
"model_usage": {},
"usage_supported": False,
}