src/llm.py

"""Pluggable LLM provider: Ollama (local) or Anthropic (Claude).

Selected via LLM_PROVIDER. Both expose the same .complete(system, user) API so
the rest of the app never needs to know which backend is in use.
"""
from __future__ import annotations

from typing import Protocol

import requests

from .config import CONFIG


class LLMProvider(Protocol):
    name: str

    def complete(self, system: str, user: str) -> str: ...


class OllamaProvider:
    name = "ollama"

    def __init__(self, model: str | None = None, host: str | None = None):
        self.model = model or CONFIG.ollama_model
        self.host = (host or CONFIG.ollama_host).rstrip("/")

    def complete(self, system: str, user: str) -> str:
        resp = requests.post(
            f"{self.host}/api/chat",
            json={
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "stream": False,
                "options": {"temperature": 0.2},
            },
            timeout=180,
        )
        resp.raise_for_status()
        return resp.json()["message"]["content"].strip()


class AnthropicProvider:
    name = "anthropic"

    def __init__(self, model: str | None = None, api_key: str | None = None):
        from anthropic import Anthropic

        key = api_key or CONFIG.anthropic_api_key
        if not key:
            raise RuntimeError("ANTHROPIC_API_KEY is not set but LLM_PROVIDER=anthropic.")
        self.model = model or CONFIG.anthropic_model
        self.client = Anthropic(api_key=key)

    def complete(self, system: str, user: str) -> str:
        msg = self.client.messages.create(
            model=self.model,
            max_tokens=1024,
            temperature=0.2,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(block.text for block in msg.content if block.type == "text").strip()


def get_provider(name: str | None = None) -> LLMProvider:
    name = (name or CONFIG.llm_provider).lower()
    if name == "anthropic":
        return AnthropicProvider()
    if name == "ollama":
        return OllamaProvider()
    raise ValueError(f"Unknown LLM_PROVIDER: {name!r}")