litellm/scripts/adaptive_router_demo/traffic.py

"""
Synthetic traffic generator for the adaptive_router demo dashboard.

What it does:
  - Sends labeled multi-turn chat requests to the proxy's adaptive router.
  - For each turn, peeks at the `x-litellm-adaptive-router-model` response
    header to learn which underlying model was picked.
  - Draws a Bernoulli outcome from a hard-coded ORACLE table that says
    "model M succeeds at request type T with probability p".
  - Sends a final follow-up turn whose user message is engineered to
    BOTH classify into the same RequestType AND match the
    satisfaction regex on success (so the bandit's `(type, model)` cell
    gets +alpha). On failure we send a neutral follow-up so no signal
    fires — over time, models the oracle favors accumulate alpha faster.

Why this shape:
  - The post-call hook gates signal recording on len(messages) >= 4.
    A single 5-message request passes the gate in one round-trip, which
    keeps the demo cheap.
  - Mock responses (`mock_response=...`) skip the real LLM call but still
    flow through routing + post-call hooks, so no API keys / no spend.

Run:
  uv run python scripts/adaptive_router_demo/traffic.py \\
      --proxy-url http://localhost:4000 \\
      --api-key   sk-1234 \\
      --router    smart-cheap-router \\
      --rounds    100 \\
      --rate      0.5

Open `dashboard.html` in a browser alongside this and watch the bars move.
"""

from __future__ import annotations

import argparse
import asyncio
import random
import sys
import uuid
from typing import Dict, List, Tuple

import httpx

# ---- prompts (paired with the RequestType the classifier will assign) ----
# Each prompt is engineered to (a) classify into the listed type and (b) make
# sense as a user request. Keep prompts short to limit token cost.
PROMPTS: Dict[str, List[str]] = {
    "code_generation": [
        "Write a Python function that flattens a nested list",
        "Create a TypeScript function that debounces another function",
        "Build a Rust function that parses a CSV string",
        "Generate a SQL function that returns running totals",
    ],
    "factual_lookup": [
        "What is the capital of New Zealand?",
        "When was the Treaty of Westphalia signed?",
        "Who is the current Secretary General of the UN?",
        "Where is Mount Kilimanjaro located?",
    ],
    "writing": [
        "Write an email declining a meeting politely",
        "Draft a paragraph introducing a product launch",
        "Compose a short blog post about morning routines",
        "Rewrite this sentence to be more concise: ...",
    ],
}

# Engineered satisfaction follow-ups — each one is designed to:
#   (1) match the satisfaction regex (thanks/great/works/perfect/etc.), AND
#   (2) re-classify into the SAME RequestType as the first prompt
# so that signals attribute to the right (type, model) bandit cell.
SATISFY: Dict[str, str] = {
    "code_generation": "thanks, that works! now write me a python function that does the inverse",
    "factual_lookup":  "perfect, thanks! who is the current prime minister?",
    "writing":         "great, thanks! now write a follow-up email confirming attendance",
}

# Neutral follow-up — does not match any signal regex, does not move the bandit.
NEUTRAL_FOLLOWUP = "ok, noted"

# Oracle: P(success | request_type, model). Tunable.
# Defaults: smart dominates code/writing; both are fine for factual_lookup.
ORACLE: Dict[str, Dict[str, float]] = {
    "code_generation": {"smart": 0.92, "fast": 0.35},
    "factual_lookup":  {"smart": 0.90, "fast": 0.85},
    "writing":         {"smart": 0.85, "fast": 0.55},
}

# Fabricated assistant turn — content doesn't matter for the hook, only the role.
FAB_ASSISTANT = "Got it. Working on that now."


def _build_messages(prompt: str, last_user: str) -> List[Dict[str, str]]:
    """5-message conversation that passes the SIGNAL_GATE_MIN_MESSAGES=4 gate."""
    return [
        {"role": "user",      "content": prompt},
        {"role": "assistant", "content": FAB_ASSISTANT},
        {"role": "user",      "content": "ok continue"},
        {"role": "assistant", "content": FAB_ASSISTANT},
        {"role": "user",      "content": last_user},
    ]


async def _send(
    client: httpx.AsyncClient,
    proxy_url: str,
    api_key: str,
    router: str,
    session_id: str,
    messages: List[Dict[str, str]],
    mock_response: str,
) -> Tuple[bool, str]:
    """Returns (ok, chosen_model)."""
    body = {
        "model": router,
        "messages": messages,
        "metadata": {"litellm_session_id": session_id},
        "mock_response": mock_response,
    }
    try:
        r = await client.post(
            f"{proxy_url}/v1/chat/completions",
            json=body,
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=15.0,
        )
        r.raise_for_status()
    except Exception as e:  # noqa: BLE001
        print(f"  request failed: {e}", file=sys.stderr)
        return False, ""
    chosen = r.headers.get("x-litellm-adaptive-router-model", "")
    return True, chosen


async def _drive_one_session(
    client: httpx.AsyncClient,
    proxy_url: str,
    api_key: str,
    router: str,
    request_type: str,
    prompt: str,
) -> str:
    """Run one labeled session. Returns the chosen model (for logging)."""
    session_id = f"demo-{uuid.uuid4()}"

    # Send the engineered 5-message conversation. The follow-up is chosen
    # AFTER we observe what model the router would pick — but since the
    # router is sticky-per-session, the model on this single round-trip
    # IS the model we're crediting.
    #
    # Pre-decide success based on the oracle for whichever model gets picked.
    # We can't know the pick before sending, so: send a neutral follow-up
    # first to learn the pick, then send a second round with credit attached.
    #
    # Round 1: neutral follow-up → no signal fires, but we learn the pick.
    ok, chosen = await _send(
        client, proxy_url, api_key, router, session_id,
        _build_messages(prompt, NEUTRAL_FOLLOWUP),
        mock_response=FAB_ASSISTANT,
    )
    if not ok or not chosen:
        return ""

    # Decide outcome from oracle.
    p = ORACLE.get(request_type, {}).get(chosen, 0.5)
    success = random.random() < p
    follow_up = SATISFY[request_type] if success else NEUTRAL_FOLLOWUP

    # Round 2: include the round-1 turns + a new follow-up. On success the
    # follow-up matches satisfaction → +alpha for (request_type, chosen).
    history = _build_messages(prompt, NEUTRAL_FOLLOWUP) + [
        {"role": "assistant", "content": FAB_ASSISTANT},
        {"role": "user",      "content": follow_up},
    ]
    await _send(
        client, proxy_url, api_key, router, session_id, history,
        mock_response=FAB_ASSISTANT,
    )
    return chosen


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--proxy-url", default="http://localhost:4000")
    ap.add_argument("--api-key",   required=True, help="proxy key with /v1/chat/completions perms")
    ap.add_argument("--router",    default="smart-cheap-router")
    ap.add_argument("--rounds",    type=int, default=100)
    ap.add_argument("--rate",      type=float, default=0.5,
                    help="seconds between sessions; lower = faster")
    ap.add_argument("--types",     default="code_generation,factual_lookup,writing",
                    help="comma-separated subset of request types to drive")
    args = ap.parse_args()

    types = [t.strip() for t in args.types.split(",") if t.strip() in PROMPTS]
    if not types:
        print(f"ERROR: no valid types. Choose from: {list(PROMPTS)}", file=sys.stderr)
        sys.exit(2)

    print(f"driving {args.rounds} sessions across types: {types}")
    print(f"oracle: {ORACLE}")
    print(f"proxy: {args.proxy_url}  router: {args.router}\n")

    counts: Dict[Tuple[str, str], int] = {}
    async with httpx.AsyncClient() as client:
        for i in range(args.rounds):
            rt = random.choice(types)
            prompt = random.choice(PROMPTS[rt])
            chosen = await _drive_one_session(
                client, args.proxy_url, args.api_key, args.router, rt, prompt,
            )
            if chosen:
                counts[(rt, chosen)] = counts.get((rt, chosen), 0) + 1
            if (i + 1) % 10 == 0:
                summary = ", ".join(
                    f"{rt}/{m}={n}" for (rt, m), n in sorted(counts.items())
                )
                print(f"  round {i + 1}/{args.rounds}  picks: {summary}")
            await asyncio.sleep(args.rate)

    print("\nfinal pick distribution:")
    for (rt, m), n in sorted(counts.items()):
        print(f"  {rt:22s} → {m:8s}  {n}")


if __name__ == "__main__":
    asyncio.run(main())