brain-of-reese/tests/e2e/mock_llm.py

"""Deterministic OpenAI-compatible mock for E2E tests (aipi stand-in).

Implements just enough of the aipi surface:

* ``GET  /v1/models``
* ``POST /v1/embeddings`` — real bag-of-words vectors (768-dim, L2-normed).
  Because similarity is *genuine token overlap*, the relevance threshold
  behaves the same way it will in production: related questions score high,
  unrelated ones score low and trigger honest deflection.
* ``POST /v1/chat/completions`` — streaming (SSE) or not. The content keys
  off markers in the system prompt:
  - user message containing ``write a long answer`` -> a ~900-word
    deterministic numbered answer (long-answers story, phase 11)
  - ``DEFLECT_MODE``  -> honest "I haven't done anything like that" answer
  - otherwise         -> upbeat answer quoting the provided document context
  - user message containing ``pretend to think slowly`` -> 3s warm-up delay
    (used by the loading-feedback story).
  - system prompt containing ``<tuning>`` (phase 15, steering notes) ->
    the composed answer ends with `` (tuning: <first note line>)`` —
    makes prompt injection observable in the UI deterministically.

  ``max_tokens`` is honored deterministically (token ≈ whitespace word),
  like a real endpoint: an answer longer than the cap is truncated. This
  is what makes the phase-11 truncation regression observable.
"""
from __future__ import annotations

import hashlib
import math
import os
import re
import signal
import threading
import time
import uuid
from typing import Any

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

DIM = 768
TOKEN_RE = re.compile(r"[a-z0-9]+")


def embed_text(text: str) -> list[float]:
    vec = [0.0] * DIM
    for tok in TOKEN_RE.findall(text.lower()):
        idx = int(hashlib.md5(tok.encode()).hexdigest(), 16) % DIM
        vec[idx] += 1.0
    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
    return [v / norm for v in vec]


def _messages(body: dict[str, Any]) -> list[dict[str, str]]:
    return body.get("messages", [])


def _system(body: dict[str, Any]) -> str:
    return " ".join(m.get("content", "") for m in _messages(body) if m.get("role") == "system")


def _user(body: dict[str, Any]) -> str:
    parts = [m.get("content", "") for m in _messages(body) if m.get("role") == "user"]
    return parts[-1] if parts else ""


def _context(body: dict[str, Any]) -> str:
    """The document context is the longest system/user message in practice."""
    msgs = _messages(body)
    return max((m.get("content", "") for m in msgs), key=len)


LONG_ANSWER_TRIGGER = "write a long answer"
#: ~920 words — comfortably past the old hard 700-token cap (where the
#: tail would be cut) yet short enough to stream in ~8s at the mock's
#: per-chunk pacing.
LONG_ANSWER_LINES = 40
LONG_ANSWER_END = "LONG-ANSWER-END"


def long_answer() -> str:
    """~900-word deterministic walkthrough (phase 11): numbered steps plus
    a unique final line that must survive the stream untruncated."""
    lines = [
        f"{i}. Step {i}: configure node-{i} with the homelab defaults and "
        f"verify that step {i} of the long walkthrough is complete before moving on."
        for i in range(1, LONG_ANSWER_LINES + 1)
    ]
    lines.append(LONG_ANSWER_END)
    return "\n".join(lines)


#: First numbered note line of a ``<tuning>`` section (phase 15).
_TUNING_BLOCK_RE = re.compile(r"<tuning>\n(.*?)\n</tuning>", re.S)
_NOTE_LINE_RE = re.compile(r"^\d+\.\s*(.+)$")


def first_tuning_note(system: str) -> str | None:
    """The first steering note in the system prompt, or ``None``.

    The prompt numbers notes 1..N oldest-first (see
    ``app.rag.prompts.build_steering_section``); the mock echoes the first
    one into its answer so prompt injection is observable in the UI.
    """
    block = _TUNING_BLOCK_RE.search(system)
    if not block:
        return None
    for line in block.group(1).splitlines():
        m = _NOTE_LINE_RE.match(line.strip())
        if m:
            return m.group(1).strip()
    return None


def compose_answer(body: dict[str, Any]) -> str:
    system = _system(body)
    user = _user(body)
    if LONG_ANSWER_TRIGGER in user.lower():
        answer = long_answer()
    elif "DEFLECT_MODE" in system:
        answer = (
            "Ah — I haven't done anything like that, so I don't want to make stuff up! "
            "You're thinking bigger than my notes for a second. Try asking about "
            "kubernetes, backups, or deploying a new service — I know those inside out. "
            "You've got this!"
        )
    else:
        ctx = _context(body)
        snippet = ctx[:220].replace("\n", " ").strip()
        answer = (
            f"Great question — you've absolutely got this! Here's what my notes say about "
            f"“{user.strip()[:80]}”: {snippet}… That's the gist from the docs; happy to "
            "dig into any of it. (Deterministic mock answer for E2E.)"
        )
    # Steering (phase 15): when the system prompt carries <tuning>, the
    # answer ends with the first note — deterministically observable.
    note = first_tuning_note(system)
    if note:
        answer = f"{answer} (tuning: {note})"
    return answer


@app.post("/__shutdown__")
def shutdown() -> dict[str, Any]:
    """Test hook (loading-feedback story): terminate this mock process to
    simulate an LLM outage. The E2E fixture restores a fresh instance on
    the same port afterwards, so the rest of the session keeps working."""

    def _die() -> None:
        time.sleep(0.1)  # let the HTTP response flush before we exit
        os.kill(os.getpid(), signal.SIGTERM)

    threading.Thread(target=_die, daemon=True).start()
    return {"status": "shutting down"}


@app.get("/v1/models")
def models() -> dict[str, Any]:
    return {
        "object": "list",
        "data": [
            {"id": "turbo", "object": "model"},
            {"id": "embed", "object": "model"},
            {"id": "lite", "object": "model"},
        ],
    }


@app.post("/v1/embeddings")
def embeddings(body: dict[str, Any]) -> dict[str, Any]:
    raw = body.get("input")
    if isinstance(raw, str):
        raw = [raw]
    inputs: list[Any] = list(raw) if isinstance(raw, list) else []
    data = [
        {"object": "embedding", "index": i, "embedding": embed_text(t)}
        for i, t in enumerate(inputs)
    ]
    return {
        "object": "list",
        "data": data,
        "model": body.get("model", "embed"),
        "usage": {"prompt_tokens": 8, "total_tokens": 8},
    }


def _sse_stream(answer: str, delay: float) -> Any:
    model = "turbo"
    chunk_id = f"chatcmpl-{uuid.uuid4()}"
    if delay:
        time.sleep(delay)
    for piece in re.findall(r".{1,12}", answer, re.S):
        payload = {
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model,
            "choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
        }
        yield f"data: {json_dumps(payload)}\n\n"
        time.sleep(0.02)
    yield (
        "data: "
        + json_dumps(
            {
                "id": chunk_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
        )
        + "\n\n"
    )
    yield "data: [DONE]\n\n"


def json_dumps(obj: dict[str, Any]) -> str:
    import json

    return json.dumps(obj)


def _apply_max_tokens(answer: str, max_tokens: Any) -> str:
    """Deterministic stand-in for the endpoint's output cap: one token ≈
    one whitespace-separated word. Answers within the cap pass through
    byte-identical, so existing (short) answers are unaffected."""
    if not isinstance(max_tokens, int) or max_tokens <= 0:
        return answer
    words = answer.split()
    if len(words) <= max_tokens:
        return answer
    return " ".join(words[:max_tokens])


@app.post("/v1/chat/completions")
def chat_completions(body: dict[str, Any]) -> Any:
    answer = _apply_max_tokens(compose_answer(body), body.get("max_tokens"))
    delay = 3.0 if "pretend to think slowly" in _user(body) else 0.0

    if not body.get("stream"):
        return {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.get("model", "turbo"),
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": answer},
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
        }

    return StreamingResponse(
        _sse_stream(answer, delay),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )