LLM-infa/api/lib/llm.py

"""Generic OpenAI-compatible chat-completions client.

The LLM is *not* deployed in this container — point this client at any external
endpoint that speaks the OpenAI Chat Completions wire format:
    - https://api.openai.com/v1
    - https://openrouter.ai/api/v1
    - https://api.anthropic.com (via OpenAI-compat proxies)
    - http://host.docker.internal:11434/v1   (local Ollama)
    - http://vllm-host:8000/v1                (vLLM / llama.cpp server / LM Studio)

Configure once via environment (LLM_BASE_URL, LLM_API_KEY, LLM_MODEL),
or override per-request through /pipeline/process by passing an `llm` block.
"""
from __future__ import annotations

import json
import logging
import urllib.error
import urllib.request
from typing import Optional

from pydantic import BaseModel, Field

from api.config import settings


logger = logging.getLogger("llm")


class LLMConfig(BaseModel):
    """Wire-level config for an OpenAI-compatible chat-completions endpoint."""

    base_url: str = Field(..., description="Base URL ending in /v1 (or equivalent)")
    api_key: str = ""
    # `model` is optional so the same config can be used for /v1/models discovery
    # before the user has picked one.
    model: str = ""
    timeout: int = 120
    max_tokens: int = 1000
    temperature: float = 0.3
    extra_headers: dict[str, str] = Field(default_factory=dict)


class LLMError(RuntimeError):
    pass


def from_settings() -> Optional[LLMConfig]:
    """Build LLMConfig from environment variables, or return None if not configured."""
    if not settings.LLM_BASE_URL or not settings.LLM_MODEL:
        return None
    return LLMConfig(
        base_url=settings.LLM_BASE_URL,
        api_key=settings.LLM_API_KEY,
        model=settings.LLM_MODEL,
        timeout=settings.LLM_TIMEOUT,
        max_tokens=settings.LLM_MAX_TOKENS,
        temperature=settings.LLM_TEMPERATURE,
    )


def chat_complete(cfg: LLMConfig, messages: list[dict]) -> str:
    """Synchronous Chat Completions call. Returns assistant message content."""
    url = cfg.base_url.rstrip("/") + "/chat/completions"
    payload: dict = {
        "model": cfg.model,
        "messages": messages,
        "max_tokens": cfg.max_tokens,
        "temperature": cfg.temperature,
    }
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
        # urllib's default "Python-urllib/x.y" UA is blocked by some CDNs (Cloudflare 1010)
        "User-Agent": "LLM-infa/0.3",
    }
    if cfg.api_key:
        headers["Authorization"] = f"Bearer {cfg.api_key}"
    headers.update(cfg.extra_headers or {})

    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=body, headers=headers, method="POST")

    try:
        with urllib.request.urlopen(req, timeout=cfg.timeout) as resp:
            raw = resp.read()
    except urllib.error.HTTPError as e:
        try:
            err_body = e.read().decode("utf-8", errors="ignore")
        except Exception:
            err_body = ""
        raise LLMError(f"LLM HTTP {e.code} {e.reason}: {err_body[:500]}")
    except Exception as e:
        raise LLMError(f"LLM request failed: {e}")

    try:
        data = json.loads(raw.decode("utf-8"))
        choices = data.get("choices") or []
        if not choices:
            raise LLMError(f"LLM returned no choices: {data}")
        msg = choices[0].get("message") or {}
        content = msg.get("content")
        if content is None:
            raise LLMError(f"LLM choice has no content: {choices[0]}")
        return content.strip()
    except LLMError:
        raise
    except Exception as e:
        raise LLMError(f"Failed to parse LLM response: {e}")


def ping(cfg: LLMConfig) -> str:
    """Tiny round-trip to verify connectivity. Returns the assistant reply text."""
    return chat_complete(
        cfg,
        [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Reply with the single word OK."},
        ],
    )


def list_models(cfg: LLMConfig) -> list[dict]:
    """GET {base_url}/models. Returns a normalized list of {id, name, owned_by?}.

    Compatible with OpenAI, OpenRouter, Ollama OpenAI-compat, vLLM, LM Studio.
    """
    url = cfg.base_url.rstrip("/") + "/models"
    headers = {"Accept": "application/json", "User-Agent": "LLM-infa/0.3"}
    if cfg.api_key:
        headers["Authorization"] = f"Bearer {cfg.api_key}"
    headers.update(cfg.extra_headers or {})

    req = urllib.request.Request(url, headers=headers, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=cfg.timeout) as resp:
            raw = resp.read()
    except urllib.error.HTTPError as e:
        try:
            err_body = e.read().decode("utf-8", errors="ignore")
        except Exception:
            err_body = ""
        raise LLMError(f"LLM HTTP {e.code} {e.reason}: {err_body[:500]}")
    except Exception as e:
        raise LLMError(f"List models failed: {e}")

    try:
        data = json.loads(raw.decode("utf-8"))
    except Exception as e:
        raise LLMError(f"Failed to parse models response: {e}")

    items = data.get("data") or data.get("models") or data.get("results") or []
    out: list[dict] = []
    for item in items:
        if isinstance(item, str):
            out.append({"id": item, "name": item})
            continue
        if not isinstance(item, dict):
            continue
        mid = item.get("id") or item.get("name") or item.get("model")
        if not mid:
            continue
        entry = {"id": mid, "name": item.get("name") or mid}
        owned = item.get("owned_by") or item.get("organization")
        if owned:
            entry["owned_by"] = owned
        out.append(entry)
    out.sort(key=lambda x: x["id"].lower())
    return out