llm-connect/llm_connect/config.py

"""
LLM configuration and API key resolution.
"""

import os
import re
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class LLMConfig:
    """Configuration for an LLM adapter.

    Attributes:
        provider: Backend identifier (``"openrouter"`` or ``"claude-code"``).
        model: Model name / path sent to the provider.
        api_key: Resolved API key (may be ``None`` for CLI backends).
        api_base: Base URL for HTTP-based providers.
        claude_cli_path: Path to the ``claude`` CLI binary.
        timeout_seconds: Per-request timeout.
        max_retries: Number of retry attempts on transient errors.
        extra: Arbitrary provider-specific overrides.
    """

    provider: str = "openrouter"
    model: str = "moonshotai/kimi-k3"
    api_key: str | None = None
    api_base: str = "https://openrouter.ai/api/v1"
    claude_cli_path: str = "claude"
    timeout_seconds: int = 300
    max_retries: int = 3
    extra: dict[str, Any] = field(default_factory=dict)


# Future supplier custody, then the live llm-connect workload lane.
_OPENROUTER_BAO_CANDIDATES: tuple[tuple[str, str], ...] = (
    ("platform/suppliers/openrouter/llm-connect", "OPENROUTER_API_KEY"),
    (
        "platform/workloads/activity-core/llm-connect/llm-connect-provider-secrets",
        "OPENROUTER_API_KEY",
    ),
)
_BAO_PATH_RE = re.compile(r"[A-Za-z0-9_./-]+")
_BAO_FIELD_RE = re.compile(r"[A-Za-z0-9_]+")
_bao_misses: set[tuple[str, str]] = set()


def clear_openbao_key_cache() -> None:
    """Forget OpenBao misses. Tests use this so a later stub can succeed."""
    _bao_misses.clear()


def resolve_api_key(
    explicit: str | None = None,
    env_var: str = "OPENROUTER_API_KEY",
    key_file_paths: list[Path] | None = None,
) -> str | None:
    """Return an API key from the first available source.

    An explicit argument always wins. Otherwise ``LLM_CONNECT_KEY_SOURCE``
    selects the order. ``auto`` (the default) tries OpenBao, then the
    environment, then a key file. ``env`` skips OpenBao. ``openbao`` uses
    OpenBao only.

    OpenBao is read with ``bao kv get -field=...``. A missing binary, a
    denied token, or an empty field is skipped. The value is not logged.
    """
    if explicit:
        return explicit

    source = os.environ.get("LLM_CONNECT_KEY_SOURCE", "auto").strip().lower()
    if source not in {"auto", "env", "openbao"}:
        source = "auto"
    if source == "openbao":
        return _openbao_api_key(env_var)
    if source == "env":
        return _env_or_file_api_key(env_var, key_file_paths)
    return _openbao_api_key(env_var) or _env_or_file_api_key(env_var, key_file_paths)


def _env_or_file_api_key(env_var: str, key_file_paths: list[Path] | None) -> str | None:
    from_env = os.environ.get(env_var)
    if from_env:
        return from_env.strip()
    for path in key_file_paths or []:
        try:
            text = path.read_text().strip()
            if text:
                return text
        except OSError:
            continue
    return None


def _openbao_api_key(env_var: str) -> str | None:
    """Read one field from OpenBao. Return ``None`` when nothing is readable."""
    if env_var != "OPENROUTER_API_KEY":
        return None
    candidates: list[tuple[str, str]] = []
    custom = os.environ.get("LLM_CONNECT_OPENROUTER_BAO_PATH", "").strip()
    if custom:
        field = os.environ.get("LLM_CONNECT_OPENROUTER_BAO_FIELD", "OPENROUTER_API_KEY").strip()
        candidates.append((custom, field))
    candidates.extend(_OPENROUTER_BAO_CANDIDATES)
    for path, field in candidates:
        value = _bao_kv_field(path, field)
        if value:
            return value
    return None


def _bao_kv_field(path: str, field: str) -> str | None:
    if not _BAO_PATH_RE.fullmatch(path) or not _BAO_FIELD_RE.fullmatch(field):
        return None
    key = (path, field)
    if key in _bao_misses:
        return None
    try:
        proc = subprocess.run(
            ["bao", "kv", "get", f"-field={field}", path],
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        _bao_misses.add(key)
        return None
    if proc.returncode != 0:
        _bao_misses.add(key)
        return None
    value = proc.stdout.strip()
    if not value:
        _bao_misses.add(key)
        return None
    return value


def find_project_root(start: Path | None = None) -> Path | None:
    """Walk up from *start* (default CWD) looking for ``pyproject.toml``.

    Returns the directory containing the marker file, or ``None``.
    """
    current = (start or Path.cwd()).resolve()
    for directory in [current, *current.parents]:
        if (directory / "pyproject.toml").is_file():
            return directory
    return None


def load_config(
    provider: str = "openrouter",
    model: str | None = None,
    api_key: str | None = None,
    **overrides: Any,
) -> LLMConfig:
    """Build an :class:`LLMConfig` with sensible defaults.

    For the ``openrouter`` provider the API key is resolved via
    :func:`resolve_api_key` (OpenBao, then env var, then project-root key file).
    """
    root = find_project_root()
    key_file_paths = [root / "apikey-openrouter.txt"] if root else []

    resolved_key = api_key
    if provider == "openrouter" and not resolved_key:
        resolved_key = resolve_api_key(
            explicit=None,
            env_var="OPENROUTER_API_KEY",
            key_file_paths=key_file_paths,
        )

    defaults: dict[str, Any] = {
        "provider": provider,
        "model": model or "moonshotai/kimi-k3",
        "api_key": resolved_key,
    }
    defaults.update(overrides)
    return LLMConfig(**defaults)