From 430804e9a56f408dad05d3100f9c9fa759ab1d32 Mon Sep 17 00:00:00 2001 From: alvis Date: Fri, 24 Apr 2026 15:44:04 +0000 Subject: [PATCH] feat(ml): prompt registry + per-request variant selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .env.example | 8 ++ ml/README.md | 6 +- ml/serving/main.py | 38 ++----- ml/serving/prompts.py | 105 ++++++++++++++++++ ml/serving/tests/test_generate.py | 67 ++++++++++- services/api/src/config.ts | 8 ++ .../src/routes/__tests__/recommender.test.ts | 1 + .../routes/__tests__/recommender.unit.test.ts | 47 +++++++- services/api/src/routes/recommender.ts | 58 ++++++++-- 9 files changed, 294 insertions(+), 44 deletions(-) create mode 100644 ml/serving/prompts.py diff --git a/.env.example b/.env.example index 9af2780..b603d3a 100644 --- a/.env.example +++ b/.env.example @@ -37,3 +37,11 @@ TODOIST_CLIENT_SECRET= NATS_URL= # How often the background scheduler refreshes Todoist tasks per active user (ms). TODOIST_SYNC_INTERVAL_MS=900000 + +# Tip prompt selection — empty = use ml/serving default (v1). +# Pin a single variant: "v2-mentor" +# Rotate uniformly across variants: "v1,v2-mentor,v3-few-shot" +# Buckets show up in the admin reward-analytics dashboard (#92). +TIP_PROMPT_VERSION= +# Default version on the Python side when the API doesn't specify one. +DEFAULT_PROMPT_VERSION=v1 diff --git a/ml/README.md b/ml/README.md index 5bd1747..e62659d 100644 --- a/ml/README.md +++ b/ml/README.md @@ -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` | 1–2 | +| `serving/` | FastAPI online scorer (`/score`, `/generate`) + LiteLLM gateway + prompt registry (`prompts.py`), called by `recommender` | 1–2 | | `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. diff --git a/ml/serving/main.py b/ml/serving/main.py index 607fe5c..2bd9225 100644 --- a/ml/serving/main.py +++ b/ml/serving/main.py @@ -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"], ) diff --git a/ml/serving/prompts.py b/ml/serving/prompts.py new file mode 100644 index 0000000..2eec45d --- /dev/null +++ b/ml/serving/prompts.py @@ -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] diff --git a/ml/serving/tests/test_generate.py b/ml/serving/tests/test_generate.py index 14a3d8d..180875c 100644 --- a/ml/serving/tests/test_generate.py +++ b/ml/serving/tests/test_generate.py @@ -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"] diff --git a/services/api/src/config.ts b/services/api/src/config.ts index c8561f9..3ba6e57 100644 --- a/services/api/src/config.ts +++ b/services/api/src/config.ts @@ -43,4 +43,12 @@ export const config = { /** How often to proactively sync Todoist tasks in the background (ms) */ TODOIST_SYNC_INTERVAL_MS: parseInt(optional('TODOIST_SYNC_INTERVAL_MS', String(15 * 60 * 1000)), 10), + + /** + * Tip prompt version selection. Single value (e.g. "v2-mentor") pins one + * variant; comma-separated list (e.g. "v1,v2-mentor,v3-few-shot") rotates + * uniformly per request so #92's reward-analytics dashboard accumulates + * comparable buckets. Empty → ml/serving's own default ("v1"). + */ + TIP_PROMPT_VERSION: optional('TIP_PROMPT_VERSION', ''), }; diff --git a/services/api/src/routes/__tests__/recommender.test.ts b/services/api/src/routes/__tests__/recommender.test.ts index 962d65d..3882d89 100644 --- a/services/api/src/routes/__tests__/recommender.test.ts +++ b/services/api/src/routes/__tests__/recommender.test.ts @@ -134,6 +134,7 @@ describe('POST /recommend integration', () => { json: async () => ({ candidates: [{ id: 'adv-1', content: 'Take a break.', rationale: 'You deserve it.' }], model: 'tip-generator', + prompt_version: 'v1', }), } as any); } diff --git a/services/api/src/routes/__tests__/recommender.unit.test.ts b/services/api/src/routes/__tests__/recommender.unit.test.ts index 44aa7bf..d8d73ae 100644 --- a/services/api/src/routes/__tests__/recommender.unit.test.ts +++ b/services/api/src/routes/__tests__/recommender.unit.test.ts @@ -2,8 +2,9 @@ * Pure-function unit tests for recommender logic — no DB, no HTTP. * These can import directly from the module without any mocking. */ -import { describe, it, expect } from 'vitest'; -import { inferReward, dueAgeDays } from '../recommender.js'; +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { inferReward, dueAgeDays, pickPromptVersion } from '../recommender.js'; +import { config } from '../../config.js'; describe('inferReward', () => { it('dismiss → -1', () => expect(inferReward('dismiss', null)).toBe(-1.0)); @@ -37,3 +38,45 @@ describe('dueAgeDays', () => { expect(dueAgeDays({ date: yesterday })).toBeGreaterThan(0); }); }); + +describe('pickPromptVersion', () => { + // Save + restore the original env-driven config field across tests. + let original: string; + beforeEach(() => { original = config.TIP_PROMPT_VERSION; }); + afterEach(() => { (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = original; }); + + it('empty config → null (let ml/serving pick its default)', () => { + (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = ''; + expect(pickPromptVersion()).toBeNull(); + }); + + it('whitespace-only config → null', () => { + (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = ' '; + expect(pickPromptVersion()).toBeNull(); + }); + + it('single value → that value', () => { + (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = 'v2-mentor'; + expect(pickPromptVersion()).toBe('v2-mentor'); + }); + + it('comma-separated → uniformly samples from the set', () => { + (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = 'v1,v2-mentor,v3-few-shot'; + const seen = new Set(); + // With 100 trials, the chance of missing any of 3 buckets is (2/3)^100 ≈ 0 — test is reliable. + for (let i = 0; i < 100; i++) { + const picked = pickPromptVersion(); + expect(picked).not.toBeNull(); + seen.add(picked!); + } + expect(seen).toEqual(new Set(['v1', 'v2-mentor', 'v3-few-shot'])); + }); + + it('trims whitespace around comma-separated entries', () => { + (config as { TIP_PROMPT_VERSION: string }).TIP_PROMPT_VERSION = ' v1 , v2-mentor '; + for (let i = 0; i < 20; i++) { + const picked = pickPromptVersion()!; + expect(['v1', 'v2-mentor']).toContain(picked); + } + }); +}); diff --git a/services/api/src/routes/recommender.ts b/services/api/src/routes/recommender.ts index 77bb743..1f5ee57 100644 --- a/services/api/src/routes/recommender.ts +++ b/services/api/src/routes/recommender.ts @@ -13,7 +13,19 @@ import { SignalAggregator } from '../signals/aggregator.js'; const router: ExpressRouter = Router(); -const PROMPT_VERSION = 'v1'; +/** + * Pick a prompt version for this request. `config.TIP_PROMPT_VERSION` is either + * empty (let ml/serving pick its default), a single version, or a comma-separated + * list to rotate uniformly across requests so the #92 dashboard accumulates + * comparable buckets per variant. Exported for testing. + */ +export function pickPromptVersion(): string | null { + const raw = config.TIP_PROMPT_VERSION.trim(); + if (!raw) return null; + const versions = raw.split(',').map((v) => v.trim()).filter(Boolean); + if (!versions.length) return null; + return versions[Math.floor(Math.random() * versions.length)] ?? null; +} // --------------------------------------------------------------------------- // Signal aggregator — register sources here as new integrations are added @@ -117,12 +129,19 @@ interface LlmCandidate { rationale?: string; } +interface LlmGenerateResult { + candidates: TipCandidate[]; + promptVersion: string | null; + model: string | null; +} + async function fetchLlmCandidates( userId: string, signals: Signal[], hour: number, dayOfWeek: number, -): Promise { + promptVersion: string | null, +): Promise { try { const tasks = signals.slice(0, 10).map((s) => ({ content: s.content, @@ -137,13 +156,18 @@ async function fetchLlmCandidates( user_id: userId, context: { tasks, hour_of_day: hour, day_of_week: dayOfWeek }, n: 3, + ...(promptVersion ? { prompt_version: promptVersion } : {}), }), signal: AbortSignal.timeout(15_000), }); - if (!res.ok) return []; - const data = (await res.json()) as { candidates: LlmCandidate[]; model?: string }; + if (!res.ok) return { candidates: [], promptVersion: null, model: null }; + const data = (await res.json()) as { + candidates: LlmCandidate[]; + model?: string; + prompt_version?: string; + }; const now = new Date().toISOString(); - return data.candidates.map((c) => ({ + const candidates: TipCandidate[] = data.candidates.map((c) => ({ id: `llm:${c.id}`, content: c.content, source: 'llm' as const, @@ -152,8 +176,13 @@ async function fetchLlmCandidates( createdAt: now, features: { is_overdue: false, task_age_days: 0, priority: 1 }, })); + return { + candidates, + promptVersion: data.prompt_version ?? null, + model: data.model ?? null, + }; } catch { - return []; + return { candidates: [], promptVersion: null, model: null }; } } @@ -181,9 +210,16 @@ router.post('/recommend', requireAuth, async (req: AuthenticatedRequest, res: Re const signals = await aggregator.fetchAll(req.userId!); const signalCandidates = signals.map(signalToCandidate); - const llmCandidates = await fetchLlmCandidates(req.userId!, signals, hour, dayOfWeek); + const requestedPromptVersion = pickPromptVersion(); + const llmResult = await fetchLlmCandidates( + req.userId!, + signals, + hour, + dayOfWeek, + requestedPromptVersion, + ); - const allCandidates: TipCandidate[] = [...signalCandidates, ...llmCandidates]; + const allCandidates: TipCandidate[] = [...signalCandidates, ...llmResult.candidates]; if (!allCandidates.length) { res.status(204).end(); return; @@ -227,8 +263,10 @@ router.post('/recommend', requireAuth, async (req: AuthenticatedRequest, res: Re candidateCount: allCandidates.length, latencyMs, servedAt, - promptVersion: isLlmTip ? PROMPT_VERSION : null, - llmModel: isLlmTip ? 'tip-generator' : null, + // Trust the version/model the generator reports; falls back to whatever + // we asked for so the bucket isn't mislabeled if /generate omits it. + promptVersion: isLlmTip ? (llmResult.promptVersion ?? requestedPromptVersion ?? null) : null, + llmModel: isLlmTip ? (llmResult.model ?? 'tip-generator') : null, tipKind: tip.kind ?? null, });