oO/ml/serving/prompts.py

"""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
    profile_features: "dict | None"


@dataclass(frozen=True)
class Prompt:
    version: str
    system: str
    build_user: Callable[["_Ctx", int], str]


def _base_user_lines(ctx: "_Ctx") -> list[str]:
    # Overdue tasks first, then high-priority, then oldest — most actionable context at top
    tasks = sorted(
        ctx.tasks,
        key=lambda t: (not t.get("is_overdue", False), -t.get("priority", 1), -t.get("task_age_days", 0.0)),
    )
    lines = [f"Time: {ctx.hour_of_day:02d}:00, day_of_week={ctx.day_of_week}"]
    if tasks:
        overdue = [t for t in tasks if t.get("is_overdue")]
        lines.append(f"Tasks: {len(tasks)} total, {len(overdue)} overdue")
        for t in tasks[:5]:
            due = t.get("due_date", "no due date")
            lines.append(f"  - [{t.get('priority','?')}] {t.get('content','?')} (due: {due})")
    p = getattr(ctx, "profile_features", None) or {}
    if p:
        parts: list[str] = []
        if (v := p.get("completion_rate_30d")) is not None:
            parts.append(f"completion_rate={float(v):.0%}")
        if (v := p.get("dismiss_rate_30d")) is not None:
            parts.append(f"dismiss_rate={float(v):.0%}")
        if (v := p.get("preferred_hour")) is not None:
            parts.append(f"preferred_hour={int(v):02d}:00")
        if parts:
            lines.append(f"User profile: {', '.join(parts)}")
    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),
}


# ── v4-orchestrator ────────────────────────────────────────────────────────
# Not a Prompt entry — takes pre-computed agent snippets, not a _Ctx.

_SYS_V4_ORCHESTRATOR = (
    "You are a personal advisor generating a single, perfectly-timed tip. "
    "Multiple specialized agents have analyzed the user's current context and provided "
    "their insights below. Synthesize their combined perspective to generate exactly ONE "
    "tip that is specific, actionable, and relevant right now. "
    "Respond ONLY with a JSON object with keys: "
    '"id" (short slug), "content" (the tip, ≤2 sentences), '
    '"rationale" (why now, ≤1 sentence). '
    "No markdown, no prose outside the JSON object."
)


def build_orchestrator_messages(
    agent_outputs: list[dict],
    tasks: list[dict],
    hour_of_day: int,
    day_of_week: int,
) -> list[dict]:
    """Build the [system, user] message list for the orchestrator LLM call.

    agent_outputs: list of {agent_id, prompt_text} dicts.
    Falls back to raw task summary when agent_outputs is empty.
    """
    lines = [f"Current time: {hour_of_day:02d}:00, day_of_week={day_of_week}", ""]
    if agent_outputs:
        lines.append("Context from analysis agents:")
        for s in agent_outputs:
            lines.append(f"[{s['agent_id']}] {s['prompt_text']}")
    else:
        overdue = [t for t in tasks if t.get("is_overdue")]
        lines.append(
            f"No pre-computed agent context available. "
            f"Tasks: {len(tasks)} total, {len(overdue)} overdue."
        )
        for t in tasks[:3]:
            lines.append(f"  - {t.get('content', '?')}")
    lines.append("\nGenerate one tip as a JSON object.")
    return [
        {"role": "system", "content": _SYS_V4_ORCHESTRATOR},
        {"role": "user", "content": "\n".join(lines)},
    ]


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]