Project: ticket sorter

from typing import Literal

from pydantic import BaseModel, Field, ValidationError


class Triage(BaseModel):
    category: Literal["billing", "shipping", "other"]
    priority: int = Field(ge=1, le=5)


def check(answer):
    try:
        return Triage.model_validate_json(answer)
    except ValidationError:
        return None


FORMAT = """Reply with one line of JSON and nothing else, like this:
{"category": "shipping", "priority": 3}
priority is 1 (can wait) to 5 (urgent)."""

BEFORE = """You sort customer support tickets for an online shop.

Categories:
- billing: payments, charges, refunds
- shipping: parcels, delivery, damaged boxes
- other: anything else

""" + FORMAT

AFTER = """You sort customer support tickets for an online shop.

Categories:
- billing: payments, charges, refunds, and any fee a customer is asked to pay, even at the door
- shipping: where a parcel is, late or lost deliveries, damaged boxes
- other: account changes, cancelling an order, anything else

""" + FORMAT

EXAMPLES = [
    {"role": "user", "content": "You took money from my card two times"},
    {"role": "assistant", "content": '{"category": "billing", "priority": 4}'},
    {"role": "user", "content": "Where is my delivery? It is a week late"},
    {"role": "assistant", "content": '{"category": "shipping", "priority": 3}'},
    {"role": "user", "content": "Can I change the email on my account?"},
    {"role": "assistant", "content": '{"category": "other", "priority": 2}'},
]

TICKETS = [
    ("I was charged twice for one order", "billing"),
    ("My parcel has not arrived", "shipping"),
    ("Can I get a refund for the blue mug?", "billing"),
    ("How do I change my password?", "other"),
    ("The parcel arrived but the box was crushed", "shipping"),
    ("The courier asked me to pay a delivery fee at the door", "billing"),
    ("I want to cancel my order before it ships", "other"),
]

import time

from groq import Groq

from prompt import AFTER, BEFORE, EXAMPLES, TICKETS, check

MODEL = "openai/gpt-oss-120b"
client = Groq()
prices = next(model.pricing for model in client.models.list().data if model.id == MODEL)

def sort_ticket(system, text):
    messages = [{"role": "system", "content": system}, *EXAMPLES, {"role": "user", "content": text}]
    response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0)
    return response.choices[0].message.content, response.usage

def run(label, system):
    right, tokens_in, tokens_out = 0, 0, 0
    started = time.perf_counter()
    print(label)
    for text, expected in TICKETS:
        answer, usage = sort_ticket(system, text)
        tokens_in, tokens_out = tokens_in + usage.prompt_tokens, tokens_out + usage.completion_tokens
        result = check(answer)
        ok = result is not None and result.category == expected
        right += ok
        print(f"  {'ok ' if ok else 'BAD'} {expected:9} {answer}")
    seconds = time.perf_counter() - started
    cost = tokens_in * float(prices["prompt"]) + tokens_out * float(prices["completion"])
    print(f"  accuracy {right}/{len(TICKETS)}, {tokens_in} tokens in, {tokens_out} out, "
          f"${cost:.6f}, {seconds / len(TICKETS):.2f} seconds a ticket")


run("before", BEFORE)
run("after", AFTER)