interactive-story/backend/app/contextwindow.py

"""M11: what the inference server will *actually* accept, as opposed to what we budgeted.

M8 found the failure this module exists to prevent. The application budgets a
prompt up to `Settings.context_token_budget` — 16,384 by default — while Ollama
enforces a window of its own, and on a machine with no VRAM that window defaults
to **4,096**. The request still returns HTTP 200. Nothing warns anybody. What
actually happens is worse than an error: `llama.cpp` drops the **oldest** tokens,
and the oldest tokens in this application are the system block — the narrator
rules and the campaign canon. The symptom is a narrator that forgets canon deep
into a long session, with nothing on screen explaining why, and every acceptance
test that reads a returned 200 as success passing throughout.

The invariant M11 requires:

    The application must not silently budget more narrator input
    than the configured Ollama runtime will actually accept.

Note the word *silently*. There are two honest outcomes and this module produces
both: either the window is **verified**, in which case the budget is capped to it
so the prompt physically cannot overflow; or it is **unverified**, in which case
the assembly says so, in the context report, on the connection test, and in the
turn's stored provenance. What must not happen is the third thing — assembling
16,384 tokens against a 4,096-token server and calling the result a turn.

## Why this is not solved by sending `num_ctx`

It was tried, and it is documented in `DEVELOPMENT.md`. Ollama's
OpenAI-compatible endpoint accepts `num_ctx` — nested in `options` or at the top
level — returns 200, and ignores it. Worse, it reloads the model at its own
default, so priming the server through the native API first does not help
either: the next request resets the window. The window is a property of how the
model is loaded, not of the request, so the only things that change it are a
model with `num_ctx` baked in (`/api/create`) or `OLLAMA_CONTEXT_LENGTH` on the
server. Both are operator actions. This module's job is not to change the
window; it is to find out what it is and refuse to lie about it.

## How the window is found

Ollama's native API sits beside the OpenAI-compatible one on the same host, so
this asks the server the application is already talking to, and nothing else. No
new destination, the same endpoint policy, the same TLS trust store.

    /api/ps    a loaded model reports `context_length`: the window the runtime
               is enforcing *right now*. This is the truth when it is available.
    /api/show  an unloaded model may carry `num_ctx` in its baked parameters,
               which is the window it will load with; `model_info` carries the
               architecture's own ceiling, which caps everything else.

`/api/ps` is asked first because a model that is loaded has already settled the
question. `/api/show` answers it for a model that is not loaded yet, which is the
ordinary case at the start of a session.

## What it deliberately does not do

It does not hard-code 4,096, which would cripple a correctly configured
deployment; it does not raise the budget, which is the operator's decision; it
does not fall back to a cloud probe, a bundled table of model sizes, or a guess
from the model's name. An unknown window is reported as unknown.
"""

from __future__ import annotations

import logging
import re
import time
from dataclasses import dataclass

import httpx

from . import endpoints, tlstrust

log = logging.getLogger(__name__)

#: Short, because this sits in the turn path. A server that does not answer in
#: two seconds has told us what we need to know: we cannot verify the window
#: right now, and the turn should proceed unverified rather than stall.
PROBE_TIMEOUT = 2.0
CONNECT_TIMEOUT = 1.5

#: A verified window is stable — it changes when an operator reloads a model —
#: so it is worth keeping. A failure is cached too, and for much less time,
#: because the commonest cause is a server that is starting up.
POSITIVE_TTL = 600.0
NEGATIVE_TTL = 60.0

#: Sources, in the order of how much they prove.
LOADED = "loaded"        # /api/ps: what the runtime is enforcing now
PARAMETERS = "parameters"  # /api/show: what the model will load with
UNKNOWN = "unknown"


@dataclass(frozen=True)
class Window:
    """What was learned about the server's input window, and how."""

    #: The total context in tokens — input *and* output share it — or None when
    #: it could not be determined.
    tokens: int | None
    #: One of LOADED, PARAMETERS, UNKNOWN.
    source: str
    #: The architecture's own ceiling, when the server reported one. Useful to a
    #: reader deciding whether raising the window is even possible.
    model_max: int | None = None
    #: Why the window is unknown, or how it was found. Shown to the user.
    detail: str = ""

    @property
    def verified(self) -> bool:
        return self.tokens is not None


UNVERIFIED = Window(tokens=None, source=UNKNOWN, detail="not checked")

_cache: dict[tuple[str, str], tuple[float, Window]] = {}


def native_base(endpoint_url: str) -> str:
    """The Ollama-native base beside an OpenAI-compatible endpoint.

    `https://host:1234/v1` -> `https://host:1234`. Anything else is used as
    given, because an endpoint that is not shaped like Ollama's is one this
    cannot interrogate and should not guess about.
    """
    trimmed = (endpoint_url or "").rstrip("/")
    return re.sub(r"/v1$", "", trimmed)


def effective_budget(configured: int, window: Window | int | None) -> int:
    """The budget the prompt may actually use.

    The whole enforcement, in one line: a verified window is a ceiling. The
    configured budget still wins when it is *smaller*, because a reader who has
    deliberately asked for a shorter prompt should get one.
    """
    tokens = window.tokens if isinstance(window, Window) else window
    if tokens is None or tokens <= 0:
        return configured
    return min(configured, tokens)


def cache_clear() -> None:
    """Forgets what was learned. Called when the endpoint or model changes."""
    _cache.clear()


async def probe(endpoint_url: str, model: str, *, use_cache: bool = True) -> Window:
    """Asks the server what window `model` gets. Never raises.

    Returns `UNVERIFIED` for every failure — refused endpoint, unreachable
    server, TLS failure, a non-Ollama endpoint, an unparseable answer. The caller
    cannot act differently on those and the reader is told the same thing either
    way: the window could not be checked.
    """
    if not endpoint_url or not model:
        return Window(None, UNKNOWN, detail="no endpoint or model configured")

    key = (endpoint_url, model)
    now = time.monotonic()
    if use_cache:
        hit = _cache.get(key)
        if hit is not None and hit[0] > now:
            return hit[1]

    window = await _ask(endpoint_url, model)
    ttl = POSITIVE_TTL if window.verified else NEGATIVE_TTL
    _cache[key] = (now + ttl, window)
    return window


async def _ask(endpoint_url: str, model: str) -> Window:
    # The same policy the turn itself is held to. A window probe must not be a
    # way to reach an address inference may not (ADR 011, H12).
    reason = endpoints.rejection_reason(endpoint_url)
    if reason is not None:
        return Window(None, UNKNOWN, detail=f"endpoint not allowed — {reason}")

    base = native_base(endpoint_url)
    try:
        async with httpx.AsyncClient(
            timeout=httpx.Timeout(PROBE_TIMEOUT, connect=CONNECT_TIMEOUT),
            verify=tlstrust.ssl_context(),
        ) as client:
            loaded = await _loaded_window(client, base, model)
            if loaded is not None:
                tokens, ceiling = loaded
                return Window(
                    tokens, LOADED, ceiling,
                    f"{tokens:,} tokens, reported by the running model",
                )
            return await _declared_window(client, base, model)
    except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
        log.debug("context window probe failed for %s: %s", base, exc)
        return Window(None, UNKNOWN, detail=f"could not ask the server ({type(exc).__name__})")


async def _loaded_window(client, base: str, model: str):
    """`/api/ps`: the window a resident model is actually being served with."""
    resp = await client.get(f"{base}/api/ps")
    if resp.status_code != 200:
        return None
    for entry in (resp.json() or {}).get("models") or []:
        if entry.get("name") == model or entry.get("model") == model:
            tokens = entry.get("context_length")
            if isinstance(tokens, int) and tokens > 0:
                return tokens, None
    return None


async def _declared_window(client, base: str, model: str) -> Window:
    """`/api/show`: what the model will load with, and its architectural cap."""
    resp = await client.post(f"{base}/api/show", json={"model": model})
    if resp.status_code != 200:
        return Window(
            None, UNKNOWN,
            detail=f"the server did not describe the model (HTTP {resp.status_code})",
        )
    body = resp.json() or {}
    ceiling = _architecture_ceiling(body.get("model_info") or {})
    declared = _num_ctx(body.get("parameters"))
    if declared is None:
        return Window(
            None, UNKNOWN, ceiling,
            detail=(
                "the model sets no num_ctx, so the server will load it at its own "
                "default — which is 4,096 where there is no VRAM"
            ),
        )
    tokens = min(declared, ceiling) if ceiling else declared
    return Window(
        tokens, PARAMETERS, ceiling,
        f"{tokens:,} tokens, from the model's own num_ctx",
    )


def _num_ctx(parameters) -> int | None:
    """Reads `num_ctx` out of the plain-text parameter block Ollama returns."""
    if not isinstance(parameters, str):
        return None
    match = re.search(r"^\s*num_ctx\s+(\d+)\s*$", parameters, re.MULTILINE)
    return int(match.group(1)) if match else None


def _architecture_ceiling(model_info: dict) -> int | None:
    """`<arch>.context_length` — the largest window this model can have."""
    for key, value in model_info.items():
        if key.endswith(".context_length") and isinstance(value, int) and value > 0:
            return value
    return None