feat(ml): prompt registry + per-request variant selection

Replaces the hardcoded "v1" label with a real prompt registry:

  ml/serving/prompts.py       — keyed by version: v1 (baseline),
                                v2-mentor (calm/specific persona),
                                v3-few-shot (v1 persona + curated examples)
  ml/serving/main.py          — POST /generate accepts optional prompt_version,
                                422 on unknown, echoes the version actually used
                                back in the response
  services/api/src/config.ts  — TIP_PROMPT_VERSION: empty / single / comma-list
                                (uniform random per request)
  services/api/src/routes/recommender.ts
                              — pickPromptVersion() drives selection; the
                                response's prompt_version (not a stale TS
                                constant) is what lands in tip_scores so the
                                #92 reward-analytics dashboard shows real
                                per-variant reaction rates

Closes #84.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-04-24 15:44:04 +00:00
parent aa4bdd8f09
commit 430804e9a5
9 changed files with 294 additions and 44 deletions

View File

@@ -4,7 +4,7 @@ Python. Owns models, features, training, online scoring.
| Dir | Role | Phase |
|---|---|---|
| `serving/` | FastAPI online scorer (`/score`, `/generate`) + LiteLLM gateway, called by `recommender` | 12 |
| `serving/` | FastAPI online scorer (`/score`, `/generate`) + LiteLLM gateway + prompt registry (`prompts.py`), called by `recommender` | 12 |
| `features/` | context assembler (`context.py`): signals → `PromptContext`; Feast adapter later | 2 |
| `pipelines/` | batch feature + training DAGs (Prefect/Airflow) | 4 |
| `registry/` | MLflow-backed model registry integration | 4 |
@@ -17,3 +17,7 @@ Python. Owns models, features, training, online scoring.
- Online inference must be stateless and < 50ms p99.
- Training reads from the offline feature store; serving reads from the online feature store; definitions are shared (no train/serve skew).
- Shadow deploys before any policy change that affects real users.
## Prompt registry
`serving/prompts.py` keys tip-generation prompts by stable version string. Adding a new variant means adding an entry — no caller changes. Selection precedence: `POST /generate` body's `prompt_version` field → env `DEFAULT_PROMPT_VERSION``"v1"`. The TypeScript recommender drives selection via `TIP_PROMPT_VERSION` (single value or comma-separated rotation); the version actually used flows back in the response and is persisted to `tip_scores.prompt_version` so the admin reward-analytics dashboard can bucket reactions per variant.

View File

@@ -31,6 +31,8 @@ import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from prompts import get_prompt
app = FastAPI(title="oO ML Serving", version="1.0.0")
LITELLM_URL = os.getenv("LITELLM_URL", "http://localhost:4000")
@@ -181,6 +183,7 @@ class GenerateRequest(BaseModel):
user_id: str
context: PromptContext = PromptContext()
n: int = 3
prompt_version: Optional[str] = None # None → server default (env DEFAULT_PROMPT_VERSION)
class TipCandidate(BaseModel):
@@ -193,33 +196,11 @@ class TipCandidate(BaseModel):
class GenerateResponse(BaseModel):
candidates: list[TipCandidate]
model: str
prompt_version: str
prompt_tokens: int = 0
completion_tokens: int = 0
_GENERATE_SYSTEM = (
"You are a personal productivity coach. "
"Given the user's current context, generate actionable, specific tips. "
"Respond ONLY with a JSON array of objects, each with keys: "
'"id" (short slug), "content" (the tip, ≤2 sentences), "rationale" (why now, ≤1 sentence). '
"No markdown, no prose outside the JSON array."
)
def _build_prompt(ctx: PromptContext, n: int) -> str:
lines = [f"Time: {ctx.hour_of_day:02d}:00, day_of_week={ctx.day_of_week}"]
if ctx.tasks:
overdue = [t for t in ctx.tasks if t.get("is_overdue")]
lines.append(f"Tasks: {len(ctx.tasks)} total, {len(overdue)} overdue")
for t in ctx.tasks[:5]:
due = t.get("due_date", "no due date")
lines.append(f" - [{t.get('priority','?')}] {t.get('content','?')} (due: {due})")
for k, v in ctx.extra.items():
lines.append(f"{k}: {v}")
lines.append(f"\nGenerate {n} tips as a JSON array.")
return "\n".join(lines)
# ── Endpoints ──────────────────────────────────────────────────────────────
@app.get("/health")
@@ -253,10 +234,14 @@ async def generate(req: GenerateRequest) -> GenerateResponse:
Retries up to _MAX_GENERATE_RETRIES times on malformed JSON, appending
a correction hint to the conversation so the model can self-correct.
"""
prompt = _build_prompt(req.context, req.n)
try:
prompt_template = get_prompt(req.prompt_version)
except KeyError as e:
raise HTTPException(status_code=422, detail=f"Unknown prompt_version: {e.args[0]}")
user_msg = prompt_template.build_user(req.context, req.n)
messages: list[dict] = [
{"role": "system", "content": _GENERATE_SYSTEM},
{"role": "user", "content": prompt},
{"role": "system", "content": prompt_template.system},
{"role": "user", "content": user_msg},
]
headers = {"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}
last_parse_error: str = ""
@@ -313,6 +298,7 @@ async def generate(req: GenerateRequest) -> GenerateResponse:
return GenerateResponse(
candidates=candidates,
model=model_used,
prompt_version=prompt_template.version,
prompt_tokens=total_usage["prompt_tokens"],
completion_tokens=total_usage["completion_tokens"],
)

105
ml/serving/prompts.py Normal file
View File

@@ -0,0 +1,105 @@
"""Prompt registry for tip generation (#84).
Each entry is an immutable (system, build_user) pair keyed by a stable version
string. Adding a new version here makes it selectable via the ``prompt_version``
field on ``POST /generate`` — the selected version flows back in the response
and is persisted to ``tip_scores.prompt_version`` so the admin reward-analytics
dashboard can bucket reactions per variant.
Versions:
v1 — neutral "productivity coach" baseline (unchanged from ffdf707).
v2-mentor — calm/specific mentor persona; same structural prompt as v1.
v3-few-shot — v1 persona plus two curated example tips inside the system prompt.
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Callable, Protocol
class _Ctx(Protocol):
tasks: list[dict]
hour_of_day: int
day_of_week: int
extra: dict
@dataclass(frozen=True)
class Prompt:
version: str
system: str
build_user: Callable[["_Ctx", int], str]
def _base_user_lines(ctx: "_Ctx") -> list[str]:
lines = [f"Time: {ctx.hour_of_day:02d}:00, day_of_week={ctx.day_of_week}"]
if ctx.tasks:
overdue = [t for t in ctx.tasks if t.get("is_overdue")]
lines.append(f"Tasks: {len(ctx.tasks)} total, {len(overdue)} overdue")
for t in ctx.tasks[:5]:
due = t.get("due_date", "no due date")
lines.append(f" - [{t.get('priority','?')}] {t.get('content','?')} (due: {due})")
for k, v in ctx.extra.items():
lines.append(f"{k}: {v}")
return lines
def _build_user_v1(ctx: "_Ctx", n: int) -> str:
return "\n".join([*_base_user_lines(ctx), f"\nGenerate {n} tips as a JSON array."])
_SYS_V1 = (
"You are a personal productivity coach. "
"Given the user's current context, generate actionable, specific tips. "
"Respond ONLY with a JSON array of objects, each with keys: "
'"id" (short slug), "content" (the tip, ≤2 sentences), "rationale" (why now, ≤1 sentence). '
"No markdown, no prose outside the JSON array."
)
_SYS_V2_MENTOR = (
"You are a calm, wise mentor — the kind who has seen a thousand people get stuck on "
"the same thing and knows when a small, concrete step unblocks them. Your tips are "
"earned, never generic; they reference the user's specific context and respect that "
"their time is short. Speak plainly. Prefer one precise action over vague encouragement. "
"Respond ONLY with a JSON array of objects, each with keys: "
'"id" (short slug), "content" (the tip, ≤2 sentences), "rationale" (why now, ≤1 sentence). '
"No markdown, no prose outside the JSON array."
)
# Two curated examples illustrate the shape we want: (1) a precise micro-action
# for an overdue item, and (2) a time-aware tip that trades tiny effort now for
# reduced friction later. Kept inside the system prompt so token cost is paid
# once per conversation and not per user turn.
_SYS_V3_FEW_SHOT = _SYS_V1 + (
"\n\nExamples of the shape and tone to aim for:\n"
'[{"id":"overdue-anchor",'
'"content":"Spend the next 12 minutes on \\"Call dentist\\" — set a timer and stop '
'when it rings, done or not.",'
'"rationale":"Overdue 6 days; a fixed micro-session breaks the avoidance loop."},'
'{"id":"evening-wind-down",'
'"content":"Pick one task from tomorrow\'s list and write its first line now while '
'context is fresh.",'
'"rationale":"It is 21:00; tomorrow-you will thank present-you for not starting cold."}]'
)
PROMPTS: dict[str, Prompt] = {
"v1": Prompt("v1", _SYS_V1, _build_user_v1),
"v2-mentor": Prompt("v2-mentor", _SYS_V2_MENTOR, _build_user_v1),
"v3-few-shot": Prompt("v3-few-shot", _SYS_V3_FEW_SHOT, _build_user_v1),
}
def default_version() -> str:
return os.getenv("DEFAULT_PROMPT_VERSION", "v1")
def get_prompt(version: str | None) -> Prompt:
"""Look up a prompt by version. Falls back to ``DEFAULT_PROMPT_VERSION`` when
``version`` is ``None``; raises :class:`KeyError` for unknown versions so
callers can surface a 422 to clients."""
v = version or default_version()
if v not in PROMPTS:
raise KeyError(v)
return PROMPTS[v]

View File

@@ -8,7 +8,10 @@ import httpx
from unittest.mock import AsyncMock, patch
from httpx import AsyncClient, ASGITransport, Response
from main import app, _build_prompt, PromptContext
from main import app, PromptContext
from prompts import PROMPTS, get_prompt
_build_user_v1 = PROMPTS["v1"].build_user
def _litellm_response(candidates: list[dict]) -> Response:
@@ -96,7 +99,7 @@ def test_build_prompt_includes_tasks():
hour_of_day=9,
day_of_week=2,
)
prompt = _build_prompt(ctx, n=3)
prompt = _build_user_v1(ctx, n=3)
assert "Write report" in prompt
assert "09:00" in prompt
assert "Generate 3 tips" in prompt
@@ -105,21 +108,21 @@ def test_build_prompt_includes_tasks():
def test_build_prompt_truncates_at_five():
tasks = [{"content": f"Task {i}", "priority": 1, "is_overdue": False, "due_date": None} for i in range(8)]
ctx = PromptContext(tasks=tasks, hour_of_day=12)
prompt = _build_prompt(ctx, n=2)
prompt = _build_user_v1(ctx, n=2)
assert "Task 4" in prompt
assert "Task 5" not in prompt
def test_build_prompt_extra_fields():
ctx = PromptContext(tasks=[], hour_of_day=8, extra={"mood": "focused", "energy": "high"})
prompt = _build_prompt(ctx, n=1)
prompt = _build_user_v1(ctx, n=1)
assert "mood: focused" in prompt
assert "energy: high" in prompt
def test_build_prompt_empty_tasks_no_task_line():
ctx = PromptContext(tasks=[], hour_of_day=10)
prompt = _build_prompt(ctx, n=2)
prompt = _build_user_v1(ctx, n=2)
assert "Tasks:" not in prompt
assert "Generate 2 tips" in prompt
@@ -223,3 +226,57 @@ def test_parse_llm_json_raises_on_invalid():
from main import _parse_llm_json
with pytest.raises((ValueError, Exception)):
_parse_llm_json("this is not json")
# ── Prompt registry / selection (#84) ──────────────────────────────────────
def test_prompt_registry_contains_expected_versions():
assert set(PROMPTS.keys()) >= {"v1", "v2-mentor", "v3-few-shot"}
# v2-mentor must differ from v1 in tone — easiest assertion: different system prompt.
assert PROMPTS["v1"].system != PROMPTS["v2-mentor"].system
# v3-few-shot must include curated example content in its system prompt.
assert "Examples" in PROMPTS["v3-few-shot"].system
def test_get_prompt_unknown_raises_keyerror():
with pytest.raises(KeyError):
get_prompt("does-not-exist")
def test_get_prompt_default_when_none():
p = get_prompt(None)
assert p.version == "v1" # current DEFAULT_PROMPT_VERSION
@pytest.mark.anyio
async def test_generate_echoes_selected_prompt_version():
"""Server should report back which prompt_version it actually used."""
fake_items = [{"id": "tip-1", "content": "x", "rationale": "y"}]
mock_resp = _litellm_response(fake_items)
with patch("main.httpx.AsyncClient") as MockClient:
instance = AsyncMock()
instance.post = AsyncMock(return_value=mock_resp)
instance.__aenter__ = AsyncMock(return_value=instance)
instance.__aexit__ = AsyncMock(return_value=False)
MockClient.return_value = instance
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
resp = await client.post(
"/generate",
json={"user_id": "u1", "n": 1, "prompt_version": "v2-mentor"},
)
assert resp.status_code == 200
assert resp.json()["prompt_version"] == "v2-mentor"
@pytest.mark.anyio
async def test_generate_422_on_unknown_prompt_version():
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
resp = await client.post(
"/generate",
json={"user_id": "u1", "n": 1, "prompt_version": "nonsense"},
)
assert resp.status_code == 422
assert "Unknown prompt_version" in resp.json()["detail"]