"""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 ```` (phase 15, steering notes) -> the composed answer ends with `` (tuning: )`` — 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 ```` section (phase 15). _TUNING_BLOCK_RE = re.compile(r"\n(.*?)\n", 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 , 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"}, )