322 lines
12 KiB
Python
322 lines
12 KiB
Python
"""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).
|
||
- user message containing ``think out loud`` -> the answer is preceded by
|
||
~800 chars of deterministic ``reasoning_content`` chunks (the
|
||
thinking-display story, phase 17).
|
||
- 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"
|
||
|
||
#: Phase 17 (thinking-display story): a user message containing this
|
||
#: substring (case-insensitive) is answered with a deterministic
|
||
#: ``reasoning_content`` stream ahead of the content — same convention as
|
||
#: the other user-message triggers above. Existing E2E questions do not
|
||
#: contain the substring, so every other suite is unaffected.
|
||
THINKING_TRIGGER = "think out loud"
|
||
|
||
|
||
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
|
||
|
||
|
||
def compose_thinking(body: dict[str, Any]) -> str:
|
||
"""Deterministic reasoning scratchpad (thinking-display story, phase 17).
|
||
|
||
A fixed 4-line "Step 1… Step 4" template quoting the first ~60 chars
|
||
of the user question: unique per question, byte-stable across runs,
|
||
~700–900 chars total (≈ 60–75 frames at the mock's 12-char/0.02s
|
||
pacing). The ``Step 2: Check my notes`` line fragment is what the E2E
|
||
assertions key off.
|
||
"""
|
||
q = _user(body).strip()[:60]
|
||
return (
|
||
f"Step 1: Read the question carefully — “{q}” — and figure out what kind of "
|
||
"answer it wants (a how-to, a lookup, or a design decision) before touching "
|
||
"the docs, so I don't over- or under-answer.\n"
|
||
"Step 2: Check my notes for the closest match. The homelab kubernetes file "
|
||
"is the obvious candidate, but I should also consider whether a deployments "
|
||
"note covers the same ground better.\n"
|
||
"Step 3: Re-read the relevant sections top to bottom so every specific — "
|
||
"hosts, versions, ports, schedules — is exact as written rather than "
|
||
"remembered, and note which document each fact comes from.\n"
|
||
"Step 4: Draft the answer around those specifics, keep it tight with short "
|
||
"paragraphs and bullets where it helps, cite the documents by path, and "
|
||
"double-check that nothing is invented."
|
||
)
|
||
|
||
|
||
@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, thinking: str = "") -> Any:
|
||
"""SSE frames for one chat completion (phase 17: + reasoning).
|
||
|
||
When ``thinking`` is non-empty its 12-char slices go out FIRST as
|
||
``delta.reasoning_content`` frames — same 0.02s cadence and envelope
|
||
as the content frames, the aipi wire convention (reasoning before
|
||
content). Without ``thinking`` the output is byte-identical to the
|
||
content-only stream, so the other story suites are unaffected.
|
||
"""
|
||
model = "turbo"
|
||
chunk_id = f"chatcmpl-{uuid.uuid4()}"
|
||
if delay:
|
||
time.sleep(delay)
|
||
for piece in re.findall(r".{1,12}", thinking, re.S):
|
||
payload = {
|
||
"id": chunk_id,
|
||
"object": "chat.completion.chunk",
|
||
"created": int(time.time()),
|
||
"model": model,
|
||
"choices": [
|
||
{"index": 0, "delta": {"reasoning_content": piece}, "finish_reason": None}
|
||
],
|
||
}
|
||
yield f"data: {json_dumps(payload)}\n\n"
|
||
time.sleep(0.02)
|
||
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
|
||
thinking = compose_thinking(body) if THINKING_TRIGGER in _user(body).lower() else ""
|
||
|
||
if not body.get("stream"):
|
||
message: dict[str, Any] = {"role": "assistant", "content": answer}
|
||
if thinking:
|
||
# Harmless future-proofing: the app only uses streaming, but a
|
||
# non-streaming client that reads the field gets the reasoning.
|
||
message["reasoning_content"] = thinking
|
||
return {
|
||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||
"object": "chat.completion",
|
||
"created": int(time.time()),
|
||
"model": body.get("model", "turbo"),
|
||
"choices": [
|
||
{"index": 0, "message": message, "finish_reason": "stop"}
|
||
],
|
||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||
}
|
||
|
||
return StreamingResponse(
|
||
_sse_stream(answer, delay, thinking=thinking),
|
||
media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||
)
|