232 lines
7.6 KiB
Python
232 lines
7.6 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).
|
|
|
|
``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)
|
|
|
|
|
|
def compose_answer(body: dict[str, Any]) -> str:
|
|
system = _system(body)
|
|
user = _user(body)
|
|
if LONG_ANSWER_TRIGGER in user.lower():
|
|
return long_answer()
|
|
if "DEFLECT_MODE" in system:
|
|
return (
|
|
"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!"
|
|
)
|
|
ctx = _context(body)
|
|
snippet = ctx[:220].replace("\n", " ").strip()
|
|
return (
|
|
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.)"
|
|
)
|
|
|
|
|
|
@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"},
|
|
)
|