Optimise Prompts with DSPy: Signatures, Metrics, Few-Shot Compilers and Held-Out Tests

"""Tidewell's ticket triage, written as a DSPy program instead of a hand-tuned prompt."""
import json
import re
from collections import Counter
from typing import Literal

import dspy
import tiktoken

import harness as H

TEAMS = ["billing", "cancellation", "account_access", "security", "bug", "feature_request"]
PRIORITIES = ["P1", "P2", "P3"]


def load(split):
    """The labelled tickets of one split ("train", "dev" or "test") as dspy.Examples whose input is the ticket."""
    return [dspy.Example(**r).with_inputs("ticket") for r in json.load(open("tickets.json"))[split]]


def plan_of(ticket):
    """The customer's plan, from the ticket's signature line."""
    m = re.search(r"(Starter|Pro|Enterprise)", ticket.strip().split("\n")[-1])
    return m.group(1) if m else "unknown"


# ---------- Step 1: declare the task ----------

# TODO (Step 1): add the two outputs, team and priority, as dspy.OutputField()s typed with Literal[...] so DSPy
# only accepts one of TEAMS and one of PRIORITIES. The docstring becomes the model's instructions.
class Triage(dspy.Signature):
    """Route a Tidewell support ticket to a team and set its priority."""
    ticket: str = dspy.InputField()


def metric(example, pred, trace=None):
    """How right a prediction is. When scoring (trace is None): 1.0 for team and priority right, 0.5 for one of them,
    0.0 for neither. When an optimiser asks (trace is not None): True only if both are right, because a demo that is
    half right teaches the model half a mistake."""
    # TODO (Step 1): compare pred.team and pred.priority with the example's; a score, or a bool when trace is set.
    raise NotImplementedError("Step 1: write metric()")


# ---------- Step 2: measure, then read the misses ----------

def evaluate(program, data):
    """(score, rows): the program's mean metric on data as a percentage, and one (example, prediction, score) row per
    example. Use dspy.Evaluate with H.THREADS threads."""
    raise NotImplementedError("evaluate() arrives in Step 2")


def breakdown(rows, field, group):
    """{group value: (right, total)}: how often `field` ("team" or "priority") was right, per value of
    group(example), e.g. group=lambda e: plan_of(e.ticket)."""
    raise NotImplementedError("breakdown() arrives in Step 2")


def confusion(rows, field):
    """Counter of (right answer, predicted answer) for the rows where `field` was wrong."""
    raise NotImplementedError("confusion() arrives in Step 2")


# ---------- Step 3: compile with examples ----------

def few_shot(k, signature=None):
    """A Predict(signature or Triage) program compiled with LabeledFewShot: k labelled train tickets as demos."""
    raise NotImplementedError("few_shot() arrives in Step 3")


def bootstrap():
    """A Predict(Triage) program compiled with BootstrapFewShot: up to 4 demos the model itself got right on train
    (by metric with a trace) plus labelled ones, 8 demos at most."""
    raise NotImplementedError("bootstrap() arrives in Step 3")


def prompt_tokens(program, ticket):
    """Tokens in the prompt the program sends for one ticket: every message DSPy's ChatAdapter formats from the
    program's signature, demos and the ticket, counted with tiktoken's cl100k_base."""
    raise NotImplementedError("prompt_tokens() arrives in Step 3")


# ---------- Step 4: search on dev, check on test ----------

def search():
    """A Predict(Triage) program compiled with BootstrapFewShotWithRandomSearch: 1 random candidate demo set (plus
    the ones it always tries), each scored on dev, the best kept. Same demo limits as bootstrap()."""
    raise NotImplementedError("search() arrives in Step 4")


def held_out(programs):
    """{name: {"dev", "test", "gap"}} for {name: program}: the dev and test scores and dev minus test."""
    raise NotImplementedError("held_out() arrives in Step 4")


# ---------- Step 5: say what the labels cannot ----------

def with_notes(notes):
    """Triage with the notes added after its instructions (Signature.with_instructions)."""
    raise NotImplementedError("with_notes() arrives in Step 5")


def choose(dev_scores):
    """The name with the best dev score; on a tie, the first one listed (list cheaper programs first)."""
    raise NotImplementedError("choose() arrives in Step 5")