src/endopath/llm.py

"""The provider seam for every LLM call.

Every model call in the app reaches the provider through `get_client().messages.create(...)`, the subset
of the Anthropic Messages API the call sites use: a forced tool call over a system prompt, text, and
optionally rendered-page images, read back as `response.stop_reason` plus `response.content` tool_use
blocks. The extraction model (#99), the schema inducer (#92), the crosswalk drafter (#93), the two chat
co-pilots (#98, #103), and the validation judge (#97) all share that surface.

`get_client` returns a client that speaks it for the configured provider. The default is the hosted
Anthropic API. Setting `LLM_PROVIDER=openai_compatible` with `LLM_BASE_URL` points every call at a local,
OpenAI-compatible server (vLLM, Ollama, an on-premise gateway) for privacy, with no change at any call
site: the client translates the request to `/v1/chat/completions` and the response back to the Anthropic
shape. `LLM_MODEL`, when set, overrides the per-call model for either provider; the local provider
requires it, since an Anthropic model id means nothing to a local server.

Config is read lazily, at the point a client is built, so importing this module leaks no key into the
environment and never un-skips the real-API tests.
"""

from __future__ import annotations

import contextlib
import hashlib
import json
import os
import time
from contextvars import ContextVar
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Any, Callable, Optional

ANTHROPIC = "anthropic"
OPENAI_COMPATIBLE = "openai_compatible"

PROVIDER_ENV = "LLM_PROVIDER"
MODEL_ENV = "LLM_MODEL"
BASE_URL_ENV = "LLM_BASE_URL"
API_KEY_ENV = "LLM_API_KEY"  # the local provider's key; the hosted path uses ANTHROPIC_API_KEY.


def _config() -> dict:
    from dotenv import load_dotenv

    load_dotenv()  # idempotent, and it never overrides an already-set var.
    return {
        "provider": (os.environ.get(PROVIDER_ENV) or ANTHROPIC).strip() or ANTHROPIC,
        "model": (os.environ.get(MODEL_ENV) or "").strip() or None,
        "base_url": (os.environ.get(BASE_URL_ENV) or "").strip() or None,
        "api_key": (os.environ.get(API_KEY_ENV) or "").strip() or None,
    }


# --- Per-call provenance (#116) ------------------------------------------------
# The seam emits one record per call to a configured recorder (endopath.provenance persists it). The
# recorder is a callback so this module stays decoupled from storage. `call_context` tags the calls made
# inside it (a case, a mapping edge), read from a contextvar so it crosses the model seam without a
# parameter at every call site.

_recorder: Optional[Callable[[dict], None]] = None
_context: ContextVar[Optional[str]] = ContextVar("llm_call_context", default=None)


def configure_recorder(recorder: Optional[Callable[[dict], None]]) -> None:
    """Register the recorder every model call reports to, or None to record nothing (the default, so a
    test or an offline run persists no provenance)."""
    global _recorder
    _recorder = recorder


@contextlib.contextmanager
def call_context(context: Optional[str]):
    """Tag every model call made inside this block with `context` (for example `case:TCGA-AJ-A3NH`), so
    the provenance record links the call to what it was for."""
    token = _context.set(context)
    try:
        yield
    finally:
        _context.reset(token)


def _prompt_hash(system: Any, tools: Optional[list]) -> str:
    """A stable hash of the prompt template (the system prompt and the tool definitions), so a change to
    the instructions or the tool schema shows as a new prompt version. The per-call user content is not
    hashed: it is different every call and identifies nothing."""
    payload = {
        "system": _system_text(system),
        "tools": [
            {"name": tool.get("name"), "input_schema": tool.get("input_schema")} for tool in (tools or [])
        ],
    }
    return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode("utf-8")).hexdigest()


def _record_call(response: Any, *, kwargs: dict, provider: str, model: Optional[str], started: float) -> None:
    if _recorder is None:
        return
    usage = getattr(response, "usage", None)
    record = {
        "provider": provider,
        "model": model,  # the model the call site requested (or the LLM_MODEL override)
        "response_model": getattr(response, "model", None),  # the model the API says handled it
        "prompt_sha256": _prompt_hash(kwargs.get("system"), kwargs.get("tools")),
        "input_tokens": getattr(usage, "input_tokens", None),
        "output_tokens": getattr(usage, "output_tokens", None),
        "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None),
        "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None),
        "latency_ms": round((time.monotonic() - started) * 1000, 1),
        "request_id": getattr(response, "_request_id", None),  # Anthropic-issued id, correlatable
        "context": _context.get(),
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    try:
        _recorder(record)
    except Exception:  # noqa: BLE001 - recording provenance must never break the real model call
        pass


def get_client() -> Any:
    """The LLM client for the configured provider, exposing `.messages.create(...)` over the Anthropic
    Messages surface the call sites use. Raises RuntimeError when the provider is misconfigured (no key
    for the hosted path, no base URL or model for the local path)."""
    cfg = _config()
    if cfg["provider"] == ANTHROPIC:
        return _AnthropicClient(model_override=cfg["model"])
    if cfg["provider"] == OPENAI_COMPATIBLE:
        if not cfg["base_url"]:
            raise RuntimeError(
                f"{PROVIDER_ENV}={OPENAI_COMPATIBLE} needs {BASE_URL_ENV} set to the server's base URL"
            )
        if not cfg["model"]:
            raise RuntimeError(
                f"{PROVIDER_ENV}={OPENAI_COMPATIBLE} needs {MODEL_ENV} set to the served model name"
            )
        return _OpenAICompatibleClient(
            base_url=cfg["base_url"], model=cfg["model"], api_key=cfg["api_key"]
        )
    raise RuntimeError(
        f"unknown {PROVIDER_ENV} {cfg['provider']!r} (known: {ANTHROPIC}, {OPENAI_COMPATIBLE})"
    )


class _AnthropicClient:
    """The hosted path: forwards to the Anthropic SDK, overriding the model when LLM_MODEL is set."""

    def __init__(self, *, model_override: Optional[str]) -> None:
        import anthropic

        api_key = os.environ.get("ANTHROPIC_API_KEY")
        if not api_key:
            raise RuntimeError(
                "ANTHROPIC_API_KEY is not set. Copy .env.example to .env and fill it in, or set "
                f"{PROVIDER_ENV}={OPENAI_COMPATIBLE} with {BASE_URL_ENV} to use a local model."
            )
        self._client = anthropic.Anthropic(api_key=api_key)
        self._model_override = model_override
        self.messages = _AnthropicMessages(self)


class _AnthropicMessages:
    def __init__(self, outer: _AnthropicClient) -> None:
        self._outer = outer

    def create(self, **kwargs: Any) -> Any:
        if self._outer._model_override:
            kwargs["model"] = self._outer._model_override
        started = time.monotonic()
        response = self._outer._client.messages.create(**kwargs)
        _record_call(response, kwargs=kwargs, provider=ANTHROPIC, model=kwargs.get("model"), started=started)
        return response


class _OpenAICompatibleClient:
    """The local path: translates an Anthropic Messages call to an OpenAI `/v1/chat/completions` request
    and the response back to the Anthropic shape (`stop_reason` plus `content` tool_use blocks), so a
    call site reads the result exactly as it reads a hosted one."""

    def __init__(self, *, base_url: str, model: str, api_key: Optional[str]) -> None:
        self._base_url = base_url.rstrip("/")
        self._model = model
        self._api_key = api_key
        self.messages = _OpenAICompatibleMessages(self)

    def _create(
        self,
        *,
        max_tokens: Optional[int] = None,
        system: Any = None,
        tools: Optional[list] = None,
        tool_choice: Any = None,
        messages: Optional[list] = None,
        **_ignored: Any,
    ) -> Any:
        import httpx

        body: dict = {
            "model": self._model,
            "messages": _to_openai_messages(system, messages or []),
        }
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        if tools:
            body["tools"] = [_to_openai_tool(tool) for tool in tools]
        if tool_choice:
            body["tool_choice"] = _to_openai_tool_choice(tool_choice)
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        started = time.monotonic()
        response = httpx.post(
            f"{self._base_url}/v1/chat/completions", json=body, headers=headers, timeout=600.0
        )
        response.raise_for_status()
        result = _to_anthropic_response(response.json())
        _record_call(
            result, kwargs={"system": system, "tools": tools}, provider=OPENAI_COMPATIBLE, model=self._model, started=started
        )
        return result


class _OpenAICompatibleMessages:
    def __init__(self, outer: _OpenAICompatibleClient) -> None:
        self._outer = outer

    def create(self, **kwargs: Any) -> Any:
        return self._outer._create(**kwargs)


# --- Request translation: Anthropic Messages -> OpenAI chat completions --------


def _system_text(system: Any) -> str:
    if system is None:
        return ""
    if isinstance(system, str):
        return system
    # A list of {type: text, text, cache_control?}; the local server may not cache, so drop the hint.
    return "\n\n".join(block.get("text", "") for block in system if block.get("type") == "text")


def _to_openai_content(content: Any) -> Any:
    if isinstance(content, str):
        return content
    parts: list[dict] = []
    for block in content:
        kind = block.get("type")
        if kind == "text":
            parts.append({"type": "text", "text": block.get("text", "")})
        elif kind == "image":
            source = block.get("source", {})
            data_uri = f"data:{source.get('media_type')};base64,{source.get('data')}"
            parts.append({"type": "image_url", "image_url": {"url": data_uri}})
    return parts


def _to_openai_messages(system: Any, messages: list) -> list:
    out: list[dict] = []
    system_text = _system_text(system)
    if system_text:
        out.append({"role": "system", "content": system_text})
    for message in messages:
        out.append({"role": message["role"], "content": _to_openai_content(message["content"])})
    return out


def _to_openai_tool(tool: dict) -> dict:
    return {
        "type": "function",
        "function": {
            "name": tool["name"],
            "description": tool.get("description", ""),
            "parameters": tool["input_schema"],
        },
    }


def _to_openai_tool_choice(tool_choice: Any) -> Any:
    if isinstance(tool_choice, dict) and tool_choice.get("type") == "tool":
        return {"type": "function", "function": {"name": tool_choice["name"]}}
    return tool_choice


# --- Response translation: OpenAI chat completions -> Anthropic Message ---------

# OpenAI finish_reason -> the Anthropic stop_reason the call sites check ("max_tokens" fails loudly).
_STOP_REASON = {"length": "max_tokens", "tool_calls": "tool_use", "stop": "end_turn"}


def _to_anthropic_response(data: dict) -> SimpleNamespace:
    choice = (data.get("choices") or [{}])[0]
    message = choice.get("message") or {}
    stop_reason = _STOP_REASON.get(choice.get("finish_reason"), choice.get("finish_reason"))

    content: list[SimpleNamespace] = []
    for call in message.get("tool_calls") or []:
        function = call.get("function") or {}
        try:
            parsed = json.loads(function.get("arguments") or "{}")
        except json.JSONDecodeError:
            parsed = {}
        content.append(SimpleNamespace(type="tool_use", name=function.get("name"), input=parsed))
    text = message.get("content")
    if text:
        content.append(SimpleNamespace(type="text", text=text))

    usage_data = data.get("usage") or {}
    usage = SimpleNamespace(
        input_tokens=usage_data.get("prompt_tokens"),
        output_tokens=usage_data.get("completion_tokens"),
    )
    return SimpleNamespace(stop_reason=stop_reason, content=content, usage=usage)