src/backend/llm.py

"""JSON completion clients for the query planner.

Two wire formats cover essentially every model a user is likely to bring:

* ``openai``    -- POST {base}/chat/completions, ``Authorization: Bearer``.
  The de-facto standard: OpenAI, DeepSeek, Groq, Together, Fireworks,
  OpenRouter, vLLM, Ollama, LM Studio, and most self-hosted gateways.
* ``anthropic`` -- POST {base}/v1/messages, ``x-api-key`` + ``anthropic-version``.

Deliberately stdlib-only. The Space image is CPU-basic and already carries
torch, faiss, and the encoder; one JSON POST per planned query does not
justify adding an SDK to the runtime requirements.

The built-in provider is whatever the operator configures, over the OpenAI
wire format. The deployed Space calls OpenAI. Users may instead supply their
own provider per request; those credentials live in the caller's browser, are
used for exactly one outbound call, and are never persisted or logged here.

Configuration for the built-in provider:

    ENCODE_PLANNER_API_KEY   required; without it the built-in provider is
                             absent. OPENAI_API_KEY and DEEPSEEK_API_KEY are
                             read as fallbacks, in that order, so an existing
                             deployment keeps working after this rename.
    ENCODE_PLANNER_BASE_URL  default https://api.openai.com/v1
    ENCODE_PLANNER_MODEL     default gpt-5.6-luna
    ENCODE_PLANNER_TIMEOUT   seconds, default 24
"""

from __future__ import annotations

import ipaddress
import json
import os
import re
import socket
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

# The models this deployment offers, in the order the picker shows them: the
# first available one is the default. Each is one user-visible choice that may
# be reachable through more than one endpoint, so a laptop with a direct
# DeepSeek key and a Space with an OpenRouter key offer the same named model
# without the user seeing the difference. Endpoint, key variable and model id
# travel together: sending one provider's model name to another is a 404 that
# reads like an outage, and sending the wrong key is a 401 that reads like a
# bad secret.
#
# Labels are what the picker shows. Keep them short and free of the provider
# namespace: "DeepSeek V4 Flash", not "deepseek/deepseek-v4-flash".
BUILTINS = (
    {"id": "luna", "label": "GPT 5.6 Luna", "endpoints": (
        {"key": "OPENAI_API_KEY", "base": "https://api.openai.com/v1",
         "model": "gpt-5.6-luna"},
        {"key": "OPEN_API_KEY", "base": "https://api.openai.com/v1",
         "model": "gpt-5.6-luna"},
    )},
    {"id": "deepseek", "label": "DeepSeek V4 Flash", "endpoints": (
        {"key": "OPENROUTER_API_KEY", "base": "https://openrouter.ai/api/v1",
         "model": "deepseek/deepseek-v4-flash"},
        {"key": "DEEPSEEK_API_KEY", "base": "https://api.deepseek.com/v1",
         "model": "deepseek-v4-flash"},
    )},
)
# The operator's own endpoint, if configured. It leads the list and becomes the
# default, which is how a model co-hosted with ENCODE is served.
CUSTOM_ID = "custom"
DEFAULT_BASE_URL = BUILTINS[0]["endpoints"][0]["base"]
DEFAULT_MODEL = BUILTINS[0]["endpoints"][0]["model"]

DEFAULT_TIMEOUT = 24.0
ANTHROPIC_VERSION = "2023-06-01"
KINDS = ("openai", "anthropic")

_ENV_FILE = Path(__file__).resolve().parents[2] / ".env"


def _load_dotenv() -> None:
    """Read the repo-root .env for local development.

    Deployment does not use this: the Space injects the planner key as a
    secret, and `.dockerignore` is allowlist-style so no .env ever reaches the
    image. Real environment variables always win over the file.
    """
    if not _ENV_FILE.exists():
        return
    try:
        text = _ENV_FILE.read_text(encoding="utf-8")
    except OSError:
        return
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if line.startswith("export "):          # SAGE-core writes this form
            line = line[7:].lstrip()
        key, sep, value = line.partition("=")
        if not sep:
            continue
        key = key.strip()
        value = value.strip().strip("'\"")
        if key and key not in os.environ:
            os.environ[key] = value


_load_dotenv()


class LlmError(RuntimeError):
    """Any failure that should surface to the caller as a planner outage."""


# -- provider resolution ----------------------------------------------------

def _env(name: str) -> str:
    return os.environ.get(name, "").strip()


def _custom_builtin() -> dict | None:
    """The operator's own endpoint, from ENCODE_PLANNER_*, if configured."""
    key = _env("ENCODE_PLANNER_API_KEY")
    base = _env("ENCODE_PLANNER_BASE_URL")
    model = _env("ENCODE_PLANNER_MODEL")
    if not (key and (base or model)):
        return None
    return {"id": CUSTOM_ID, "label": model or "Built-in model",
            "base_url": base or DEFAULT_BASE_URL,
            "model": model or DEFAULT_MODEL, "api_key": key}


def _resolve(entry: dict) -> dict | None:
    """The first endpoint of a catalog entry whose key is present."""
    for endpoint in entry["endpoints"]:
        key = _env(endpoint["key"])
        if key:
            return {"id": entry["id"], "label": entry["label"],
                    "base_url": endpoint["base"], "model": endpoint["model"],
                    "api_key": key}
    return None


def builtin_catalog() -> list[dict]:
    """Every built-in this deployment can actually serve, default first."""
    custom = _custom_builtin()
    served = [r for r in (_resolve(e) for e in BUILTINS) if r]
    return ([custom] if custom else []) + served


def builtin_configured() -> bool:
    return bool(builtin_catalog())


def builtin_model() -> str:
    """The default model's label, for a caller that wants one name."""
    catalog = builtin_catalog()
    return catalog[0]["label"] if catalog else DEFAULT_MODEL


def builtin_provider(choice: str | None = None) -> dict:
    """The named built-in, or the default when `choice` is absent or unknown.

    An unknown id falls back rather than failing: a picker selection saved in
    a browser outlives the deployment that offered it, and planning with the
    default beats an error the user cannot act on.
    """
    catalog = builtin_catalog()
    if not catalog:
        raise LlmError("No planner model is configured on this deployment. "
                       "Add your own model to use the Agent.")
    resolved = next((c for c in catalog if c["id"] == choice), catalog[0])
    return {"kind": "openai", "base_url": resolved["base_url"],
            "api_key": resolved["api_key"], "model": resolved["model"],
            "label": resolved["label"]}


def _timeout() -> float:
    try:
        return float(os.environ.get("ENCODE_PLANNER_TIMEOUT", "") or DEFAULT_TIMEOUT)
    except ValueError:
        return DEFAULT_TIMEOUT


def allowed_hosts() -> set[str]:
    """Endpoints the *operator* has explicitly permitted, as `host` or
    `host:port`, comma separated in ENCODE_PLANNER_ALLOW_HOSTS.

    This exists for a model co-hosted with ENCODE -- a vLLM or Ollama server on
    the same Linux box or VPC, reachable only over http on a private address.
    Whoever sets this env var already owns that machine, so permitting it is
    not a privilege escalation; an end user typing the same address into the
    UI form is exactly the request-forgery case the check below blocks.
    """
    raw = os.environ.get("ENCODE_PLANNER_ALLOW_HOSTS", "")
    return {h.strip().lower() for h in raw.split(",") if h.strip()}


# A user-supplied base_url turns this server into a request forwarder, so
# without this check a custom model is a server-side request forgery
# primitive: someone could point it at cloud instance metadata
# (169.254.169.254) or at services reachable only from inside the
# deployment's network.
#
# Note the built-in provider does NOT come through here -- it is configured
# from the environment by the operator, so it may already point anywhere,
# including an http endpoint on localhost.
def _assert_allowed_target(base_url: str) -> None:
    parts = urllib.parse.urlsplit(base_url)
    host = (parts.hostname or "").lower()
    if not host:
        raise LlmError("Model base URL has no host")
    if parts.scheme not in ("http", "https"):
        raise LlmError("Model base URL must start with http:// or https://")

    # Operator-permitted endpoints skip the public-address requirement; that
    # is the whole point of the allowlist.
    hostport = f"{host}:{parts.port}" if parts.port else host
    allow = allowed_hosts()
    if host in allow or hostport in allow:
        return

    if parts.scheme != "https":
        raise LlmError("Model base URL must start with https://")
    try:
        infos = socket.getaddrinfo(host, parts.port or 443, proto=socket.IPPROTO_TCP)
    except socket.gaierror:
        raise LlmError(f"Could not resolve model host '{host}'") from None
    for info in infos:
        ip = ipaddress.ip_address(info[4][0])
        if (ip.is_private or ip.is_loopback or ip.is_link_local
                or ip.is_reserved or ip.is_multicast or ip.is_unspecified):
            raise LlmError("Model base URL must point at a public host. An "
                           "endpoint inside this network has to be permitted "
                           "by the operator via ENCODE_PLANNER_ALLOW_HOSTS.")


def normalize_provider(spec: object) -> dict:
    """Validate a user-supplied provider. Raises LlmError with a message meant
    to be read by whoever filled in the form."""
    if not isinstance(spec, dict):
        raise LlmError("Model settings are missing")
    kind = str(spec.get("kind") or "").strip().lower()
    if kind not in KINDS:
        raise LlmError(f"API format must be one of: {', '.join(KINDS)}")
    base_url = str(spec.get("base_url") or "").strip().rstrip("/")
    model = str(spec.get("model") or "").strip()
    api_key = str(spec.get("api_key") or "").strip()
    if not base_url:
        raise LlmError("Base URL is required")
    if not model:
        raise LlmError("Model name is required")
    if not api_key:
        raise LlmError("API key is required")
    _assert_allowed_target(base_url)
    label = str(spec.get("label") or "").strip() or model
    return {"kind": kind, "base_url": base_url, "model": model,
            "api_key": api_key, "label": label[:60]}


# -- request shaping --------------------------------------------------------

# Providers on the OpenAI wire format disagree about three fields, and the
# disagreement is only discoverable from a 400. OpenAI's newer models renamed
# the token cap to max_completion_tokens and accept only the default
# temperature; DeepSeek, Groq, vLLM and most gateways still want max_tokens.
# So: start with the spelling the host is known to want, and treat a 400 that
# names one of these fields as instructions rather than as a failure.
_OPENAI_HOSTS = ("api.openai.com",)


def _token_cap_field(base_url: str) -> str:
    host = urllib.parse.urlsplit(base_url).hostname or ""
    return ("max_completion_tokens" if host.lower() in _OPENAI_HOSTS
            else "max_tokens")


def _relax_body(body: dict, message: str) -> bool:
    """Rename or drop the one field a 400 objected to.

    Returns True when the body changed and the call is worth retrying. The
    token cap flips to whichever spelling this body is not using, so the same
    rule serves a model that wants max_completion_tokens and a gateway that
    rejects it.
    """
    m = message.lower()
    if "max_tokens" in m or "max_completion_tokens" in m:
        if "max_tokens" in body:
            body["max_completion_tokens"] = body.pop("max_tokens")
            return True
        if "max_completion_tokens" in body:
            body["max_tokens"] = body.pop("max_completion_tokens")
            return True
    if "temperature" in m and "temperature" in body:
        body.pop("temperature")          # reasoning models take the default only
        return True
    if "response_format" in m and "response_format" in body:
        body.pop("response_format")      # prompt still asks for JSON
        return True
    if "reasoning" in m and "reasoning" in body:
        body.pop("reasoning")            # costs the panel, not the plan
        return True
    return False


# What each endpoint turned out to accept, learned from its own 400s and kept
# for the life of the process. Without this every plan would pay one rejected
# request to rediscover the same fact.
_SHAPE_MEMO: dict[str, dict] = {}


def _shape_key(p: dict) -> str:
    return f"{urllib.parse.urlsplit(p['base_url']).hostname}|{p['model']}"


def _remember_shape(p: dict, body: dict) -> None:
    _SHAPE_MEMO[_shape_key(p)] = {
        "cap": ("max_completion_tokens" if "max_completion_tokens" in body
                else "max_tokens"),
        "dropped": tuple(f for f in ("temperature", "response_format")
                         if f not in body)}


def _reasoning_opt_in(base_url: str) -> dict:
    """OpenRouter returns the model's reasoning only when asked for it.

    DeepSeek sends `reasoning_content` unprompted; through OpenRouter the same
    model needs `reasoning` in the body, and the text then arrives as
    `delta.reasoning`, which the stream reader already handles. A provider
    that does not know the field answers 400 and _relax_body drops it.
    """
    host = (urllib.parse.urlsplit(base_url).hostname or "").lower()
    return {"reasoning": {"enabled": True}} if host == "openrouter.ai" else {}


def _openai_body(p: dict, system: str, user: str, max_tokens: int, **extra) -> dict:
    memo = _SHAPE_MEMO.get(_shape_key(p), {})
    dropped = memo.get("dropped", ())
    body = {"model": p["model"],
            "messages": [{"role": "system", "content": system},
                         {"role": "user", "content": user}],
            **extra}
    if "response_format" not in dropped:
        # json_object is the portable JSON mode. Gateways that don't know it
        # generally ignore it; the prompt asks for JSON anyway and the
        # response is parsed leniently below.
        body["response_format"] = {"type": "json_object"}
    if "temperature" not in dropped:
        body["temperature"] = 0.0
    body[memo.get("cap") or _token_cap_field(p["base_url"])] = max_tokens
    return body


def _openai_request(p: dict, system: str, user: str, max_tokens: int):
    return (f"{p['base_url']}/chat/completions",
            _openai_body(p, system, user, max_tokens, stream=False),
            {"Content-Type": "application/json",
             "Authorization": f"Bearer {p['api_key']}"})


def _anthropic_request(p: dict, system: str, user: str, max_tokens: int):
    body = {"model": p["model"], "max_tokens": max_tokens, "system": system,
            "messages": [{"role": "user", "content": user}]}
    return (f"{p['base_url']}/v1/messages", body,
            {"Content-Type": "application/json", "x-api-key": p["api_key"],
             "anthropic-version": ANTHROPIC_VERSION})


def _openai_text(payload: dict) -> tuple[str, str]:
    try:
        choice = payload["choices"][0]
        return choice["message"]["content"] or "", choice.get("finish_reason") or ""
    except (KeyError, IndexError, TypeError):
        raise LlmError("Model response had no message content") from None


def _anthropic_text(payload: dict) -> tuple[str, str]:
    try:
        parts = [b.get("text", "") for b in payload["content"] if b.get("type") == "text"]
    except (KeyError, TypeError):
        raise LlmError("Model response had no message content") from None
    if not parts:
        raise LlmError("Model response had no message content")
    stop = payload.get("stop_reason") or ""
    return "".join(parts), "length" if stop == "max_tokens" else stop


def _http_detail(err: urllib.error.HTTPError) -> str:
    """The provider's own `error.message`, bounded.

    Never echo the raw body: it can quote the request, and the request carries
    the user's query. This one field is a fixed diagnostic about request shape,
    which is what makes a 4xx debuggable and what _relax_body reads.
    """
    try:
        msg = json.loads(err.read().decode()).get("error", {}).get("message", "")
    except Exception:
        return ""
    return msg if isinstance(msg, str) and 0 < len(msg) <= 200 else ""


# One retry per adjustable field, and there are three of them.
MAX_SHAPE_RETRIES = 3


_FENCE = re.compile(r"```(?:json)?\s*(.+?)\s*```", re.S)


def parse_json_object(text: str) -> dict:
    """Lenient parse: providers without a JSON mode often wrap the object in a
    markdown fence or add a sentence around it. Strict parse first, then peel."""
    for candidate in (text, ):
        try:
            value = json.loads(candidate)
            if isinstance(value, dict):
                return value
        except json.JSONDecodeError:
            pass
    fenced = _FENCE.search(text)
    if fenced:
        try:
            value = json.loads(fenced.group(1))
            if isinstance(value, dict):
                return value
        except json.JSONDecodeError:
            pass
    start, end = text.find("{"), text.rfind("}")
    if 0 <= start < end:
        try:
            value = json.loads(text[start:end + 1])
            if isinstance(value, dict):
                return value
        except json.JSONDecodeError:
            pass
    raise LlmError("Model did not return valid JSON")


def stream_json(system: str, user: str, *, provider: dict | None = None,
                max_tokens: int = 4000):
    """Yield ("thinking", delta) as the model reasons, then ("text", full).

    Measured on deepseek-v4-flash: the first reasoning token arrives ~1.3s in
    while the first answer token takes ~27s, and reasoning outweighs the answer
    roughly 16:1 by volume. Streaming therefore turns almost the whole wait
    into something the user can watch, and it keeps the socket busy so a read
    timeout never fires on a slow call.

    Only the OpenAI-compatible shape is streamed here; that is the built-in
    provider's shape. A user's own model streams in their browser instead.
    """
    p = provider or builtin_provider()
    if p["kind"] != "openai":
        raise LlmError("Streaming is only implemented for OpenAI-compatible models")
    if "json" not in f"{system} {user}".lower():
        raise LlmError("prompt must mention 'json' to use json_object mode")

    body = _openai_body(p, system, user, max_tokens, stream=True,
                        **_reasoning_opt_in(p["base_url"]),
                        # DeepSeek honours this and returns
                        # completion_tokens_details.reasoning_tokens in the
                        # final frame, so the UI can report what the model
                        # actually spent rather than a character-count proxy.
                        stream_options={"include_usage": True})

    parts: list[str] = []
    finish = ""
    usage: dict = {}
    for attempt in range(MAX_SHAPE_RETRIES + 1):
        req = urllib.request.Request(
            f"{p['base_url']}/chat/completions", data=json.dumps(body).encode(),
            method="POST", headers={"Content-Type": "application/json",
                                    "Authorization": f"Bearer {p['api_key']}"})
        try:
            with urllib.request.urlopen(req, timeout=_timeout()) as resp:
                for raw in resp:
                    line = raw.decode("utf-8", "replace").strip()
                    if not line.startswith("data: "):
                        continue
                    chunk = line[6:]
                    if chunk == "[DONE]":
                        break
                    try:
                        frame = json.loads(chunk)
                    except json.JSONDecodeError:
                        continue
                    if frame.get("usage"):
                        usage = frame["usage"]      # final frame carries totals
                    try:
                        choice = frame["choices"][0]
                    except (KeyError, IndexError):
                        continue
                    finish = choice.get("finish_reason") or finish
                    delta = choice.get("delta") or {}
                    think = delta.get("reasoning_content") or delta.get("reasoning")
                    if think:
                        yield ("thinking", think)
                    if delta.get("content"):
                        parts.append(delta["content"])
            _remember_shape(p, body)
            break
        except urllib.error.HTTPError as err:
            detail = _http_detail(err)
            # Same shape negotiation as the non-streaming path. Nothing has
            # been yielded yet at this point, so the retry is invisible.
            if (err.code == 400 and attempt < MAX_SHAPE_RETRIES
                    and not parts and _relax_body(body, detail)):
                continue
            raise LlmError(f"Model returned HTTP {err.code}"
                           f"{': ' + detail if detail else ''}") from None
        except urllib.error.URLError as err:
            raise LlmError(f"Could not reach the model: {err.reason}") from None
        except TimeoutError:
            raise LlmError("The model stopped responding") from None

    if finish == "length":
        raise LlmError("Model response was cut off by the token limit")
    yield ("usage", {"reasoning_tokens":
                     (usage.get("completion_tokens_details") or {}).get("reasoning_tokens")})
    yield ("text", "".join(parts))


def complete_json(system: str, user: str, *, provider: dict | None = None,
                  max_tokens: int = 4000) -> dict:
    """One non-streaming completion, parsed as a JSON object.

    `provider` defaults to the deployment's built-in model. Raises LlmError
    for missing config, transport failure, and any unusable response; callers
    validate the object's contents.
    """
    p = provider or builtin_provider()

    # OpenAI-compatible json_object mode is rejected outright unless the word
    # "json" appears in the prompt. Trimming a prompt can silently remove it,
    # so fail here with the actual cause rather than shipping a 400.
    if p["kind"] == "openai" and "json" not in f"{system} {user}".lower():
        raise LlmError("prompt must mention 'json' to use json_object mode")

    shape = _openai_request if p["kind"] == "openai" else _anthropic_request
    url, body, headers = shape(p, system, user, max_tokens)
    for attempt in range(MAX_SHAPE_RETRIES + 1):
        req = urllib.request.Request(url, data=json.dumps(body).encode(),
                                     method="POST", headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=_timeout()) as resp:
                payload = json.loads(resp.read().decode())
            if p["kind"] == "openai":
                _remember_shape(p, body)
            break
        except urllib.error.HTTPError as err:
            detail = _http_detail(err)
            # A 400 naming a field this provider spells differently is a fact
            # about the endpoint, not a failed plan. Fix the field and resend;
            # a rejected request costs no tokens.
            if (err.code == 400 and attempt < MAX_SHAPE_RETRIES
                    and _relax_body(body, detail)):
                continue
            raise LlmError(f"Model returned HTTP {err.code}"
                           f"{': ' + detail if detail else ''}") from None
        except urllib.error.URLError as err:
            raise LlmError(f"Could not reach the model: {err.reason}") from None
        except (TimeoutError, json.JSONDecodeError) as err:
            raise LlmError(f"Model response was unusable: {type(err).__name__}") from None

    text, finish = (_openai_text if p["kind"] == "openai" else _anthropic_text)(payload)

    # Reasoning models bill their thinking against max_tokens alongside the
    # visible content. A tight cap truncates the JSON mid-object, which would
    # otherwise surface as a baffling "did not return valid JSON".
    if finish == "length":
        raise LlmError("Model response was cut off by the token limit")
    return parse_json_object(text)