src/services/llm.py

"""LLM client for OpenAI-compatible APIs with circuit breaker + retry.
Supports multiple providers with model-specific tuning profiles."""
import asyncio
import json
import logging
import os
from dataclasses import dataclass, field
from typing import AsyncIterator, Optional

import httpx

from .circuit_breaker import CircuitBreaker, CircuitBreakerOpenError, retry_with_backoff
from .secrets import get_secret

logger = logging.getLogger("rga_auditor.llm")


class LLMError(RuntimeError):
    pass


# ── Provider profiles ────────────────────────────────────────────────────────
# Each profile maps the modes (black_box, white_box, analyze, verify, gaps) to
# model-specific parameters. Add new profiles here or via env overrides.

@dataclass
class _ModeParams:
    temperature: float = 0.5
    max_tokens: int = 4096
    extra: dict = field(default_factory=dict)


@dataclass
class _Profile:
    name: str
    base_url: str
    model: str
    modes: dict[str, _ModeParams] = field(default_factory=lambda: {
        "black_box": _ModeParams(temperature=0.5, max_tokens=4096,
            extra={"reasoning_effort": "medium", "reasoning_summary": False}),
        "white_box": _ModeParams(temperature=0.3, max_tokens=5120,
            extra={"reasoning_effort": "high", "reasoning_summary": True, "reasoning_summary_wait": True}),
        "analyze":  _ModeParams(temperature=0.3, max_tokens=4096,
            extra={"reasoning_effort": "high", "reasoning_summary": True, "reasoning_summary_wait": True}),
        "verify":  _ModeParams(temperature=0.5, max_tokens=4096,
            extra={"reasoning_effort": "high", "reasoning_summary": False}),
        "gaps":   _ModeParams(temperature=0.5, max_tokens=4096,
            extra={"reasoning_effort": "medium", "reasoning_summary": False}),
    })


PROFILES: dict[str, _Profile] = {
    "mercury": _Profile(
        name="Mercury 2",
        base_url="https://api.inceptionlabs.ai/v1",
        model="mercury-2",
        modes={
            "black_box": _ModeParams(temperature=0.5, max_tokens=4096,
                extra={"reasoning_effort": "medium", "reasoning_summary": False}),
            "white_box": _ModeParams(temperature=0.3, max_tokens=5120,
                extra={"reasoning_effort": "high", "reasoning_summary": True, "reasoning_summary_wait": True}),
            "analyze":  _ModeParams(temperature=0.3, max_tokens=8192,
                extra={"reasoning_effort": "high", "reasoning_summary": True, "reasoning_summary_wait": True}),
            "verify":  _ModeParams(temperature=0.5, max_tokens=4096,
                extra={"reasoning_effort": "high", "reasoning_summary": False}),
            "gaps":   _ModeParams(temperature=0.5, max_tokens=4096,
                extra={"reasoning_effort": "medium", "reasoning_summary": False}),
        },
    ),
    "minimax": _Profile(
        name="Minimax M3 (NVIDIA)",
        base_url=(os.getenv("LLM_BASE_URL") or "https://integrate.api.nvidia.com/v1").rstrip("/"),
        model="minimaxai/minimax-m3",
        modes={
            "black_box": _ModeParams(temperature=0.3, max_tokens=4096),
            "white_box": _ModeParams(temperature=0.3, max_tokens=4096),
            "analyze":  _ModeParams(temperature=0.3, max_tokens=4096),
            "verify":  _ModeParams(temperature=0.3, max_tokens=4096),
            "gaps":   _ModeParams(temperature=0.3, max_tokens=4096),
        },
    ),
}

_DEFAULT_PROFILE = os.getenv("LLM_PROVIDER", "mercury")


def _get_profile(name: Optional[str] = None) -> _Profile:
    key = name or _DEFAULT_PROFILE
    p = PROFILES.get(key)
    if p:
        return p
    logger.warning("Unknown LLM profile %r, falling back to mercury", key)
    return PROFILES["mercury"]


# ── LLM client ───────────────────────────────────────────────────────────────

class LLM:
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        profile: Optional[str] = None,
    ) -> None:
        # Resolve profile (may come from LLM_PROVIDER env or LLM_MODEL → "custom")
        self.profile = _get_profile(profile)
        # Env var overrides take priority over profile values
        self.model = model or self.profile.model
        self.base_url = self.profile.base_url.rstrip("/")
        # Allow LLM_BASE_URL to contain the full path; strip /chat/completions
        # since callers append it.
        suffix = "/chat/completions"
        if self.base_url.endswith(suffix):
            self.base_url = self.base_url[:-len(suffix)]
        if self.profile.model == "mercury-2":
            self.api_key = api_key or get_secret("INCEPTION_API_KEY") or os.getenv("INCEPTION_API_KEY")
        else:
            self.api_key = api_key or get_secret("LLM_API_KEY") or os.getenv("LLM_API_KEY")
        if not self.api_key:
            raise RuntimeError(f"No API key found for {self.profile.name}. Set INCEPTION_API_KEY or LLM_API_KEY accordingly.")
            raise RuntimeError("No LLM API key found. Set LLM_API_KEY or INCEPTION_API_KEY in .env")
        self._cb = CircuitBreaker(name="llm", failure_threshold=5, recovery_timeout_s=30.0)

    @property
    def _headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}

    def is_available(self) -> bool:
        return self._cb.is_available()

    def _params_for(self, mode: str, override_temperature: Optional[float] = None, override_max_tokens: Optional[int] = None) -> dict:
        p = self.profile.modes.get(mode, self.profile.modes["black_box"])
        params = {
            "temperature": override_temperature if override_temperature is not None else p.temperature,
            "max_tokens": override_max_tokens or p.max_tokens,
        }
        params.update(p.extra)
        return params

    def _build_payload(self, prompt: str, system: Optional[str], mode: str,
                       temperature: Optional[float], max_tokens: Optional[int],
                       **overrides) -> dict:
        messages: list[dict] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        params = self._params_for(mode, temperature, max_tokens)
        # Apply any per-call overrides
        # Supported overrides: model (switch model name for this call)
        call_model = overrides.pop("model", self.model)
        params.update(overrides)

        payload = {
            "model": call_model,
            "messages": messages,
            **params,
        }
        logger.info("LLM.%s: model=%s mode=%s temp=%s max_tokens=%s extra=%s prompt_len=%d",
                     mode, self.model, mode,
                     params.get("temperature"), params.get("max_tokens"),
                     {k: v for k, v in params.items() if k not in ("temperature", "max_tokens")},
                     len(prompt))
        return payload

    async def chat(self, prompt: str, system: Optional[str] = None, mode: str = "black_box",
                   temperature: Optional[float] = None, max_tokens: Optional[int] = None,
                   **overrides) -> str:
        try:
            return await self._cb.call(
                self._do_chat_with_retry, prompt, system=system, mode=mode,
                temperature=temperature, max_tokens=max_tokens, **overrides
            )
        except (CircuitBreakerOpenError, LLMError) as e:
            if self.profile.model != "mercury-2":
                logger.warning("Falling back to mercury after %s error: %s", type(e).__name__, e)
                fallback = LLM(profile="mercury")
                return await fallback.chat(prompt, system=system, mode=mode,
                                           temperature=temperature, max_tokens=max_tokens, **overrides)
            raise

    @retry_with_backoff(max_retries=2, base_delay_s=0.5, retryable_exceptions=(httpx.ConnectError, httpx.TimeoutException, httpx.RemoteProtocolError))
    async def _do_chat_with_retry(self, prompt: str, system: Optional[str] = None, mode: str = "black_box",
                                   temperature: Optional[float] = None, max_tokens: Optional[int] = None,
                                   **overrides) -> str:
        return await self._do_chat(prompt, system=system, mode=mode,
                                    temperature=temperature, max_tokens=max_tokens, **overrides)

    async def _do_chat(self, prompt: str, system: Optional[str] = None, mode: str = "black_box",
                        temperature: Optional[float] = None, max_tokens: Optional[int] = None,
                        **overrides) -> str:
        payload = self._build_payload(prompt, system, mode, temperature, max_tokens, **overrides)
        timeout = 300.0
        async with httpx.AsyncClient(timeout=timeout) as client:
            r = await client.post(f"{self.base_url}/chat/completions", json=payload, headers=self._headers)
            if r.status_code != 200:
                body = (await r.aread())[:500].decode("utf-8", "ignore")
                logger.error("LLM.chat: HTTP %d: %s", r.status_code, body)
                raise LLMError(f"{self.profile.name} {r.status_code}: {body}")
            data = r.json()
            content = data["choices"][0]["message"].get("content")
            if content is None:
                logger.warning("LLM.chat: null content in response (reasoning summary mode?) — returning empty")
                content = ""
            logger.info("LLM.chat: response length=%d, preview=%.200s", len(content), content)
            return content

    async def astream(self, prompt: str, system: Optional[str] = None, mode: str = "black_box",
                       temperature: Optional[float] = None, max_tokens: Optional[int] = None,
                       **overrides) -> AsyncIterator[str]:
        if not self._cb.is_available():
            if self.profile.model != "mercury-2":
                logger.warning("Circuit open for %s, falling back to mercury stream", self.profile.name)
                fallback = LLM(profile="mercury")
                async for chunk in fallback.astream(prompt, system=system, mode=mode,
                                                     temperature=temperature, max_tokens=max_tokens, **overrides):
                    yield chunk
                return
            logger.warning("LLM circuit is OPEN — streaming unavailable")
            yield "The LLM service is temporarily unavailable. Please try again shortly."
            return
        payload = self._build_payload(prompt, system, mode, temperature, max_tokens, **overrides)
        payload["stream"] = True
        timeout = 600.0
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                async with client.stream("POST", f"{self.base_url}/chat/completions", json=payload, headers=self._headers) as r:
                    if r.status_code != 200:
                        body = await r.aread()
                        raise LLMError(f"{self.profile.name} {r.status_code}: {body[:500].decode('utf-8', 'ignore')}")
                    async for line in r.aiter_lines():
                        if not line or not line.startswith("data: "):
                            continue
                        chunk = line[6:].strip()
                        if chunk == "[DONE]":
                            break
                        try:
                            data = json.loads(chunk)
                            delta = data["choices"][0].get("delta", {}).get("content")
                            if delta:
                                yield delta
                        except (json.JSONDecodeError, KeyError, IndexError):
                            continue
        except Exception as e:
            if self.profile.model != "mercury-2":
                logger.warning("Falling back to mercury stream after error: %s", e)
                fallback = LLM(profile="mercury")
                async for chunk in fallback.astream(prompt, system=system, mode=mode,
                                                     temperature=temperature, max_tokens=max_tokens, **overrides):
                    yield chunk
                return
            logger.error("LLM.astream failed: %s", e)
            await self._cb.call(lambda: (_ for _ in ()).throw(e))
            raise


_llm: Optional[LLM] = None


def get_llm(profile: Optional[str] = None) -> LLM:
    global _llm
    if _llm is None or profile:
        _llm = LLM(profile=profile)
    return _llm


def list_profiles() -> list[dict]:
    return [{"key": k, "name": p.name, "model": p.model, "base_url": p.base_url} for k, p in PROFILES.items()]