BaseLM: a readable keyword model for DSPy

import json
import re

import dspy
from dspy.clients.openai_format import to_openai_chat_request

KEYWORDS = {
    "billing": ["charged", "refund", "invoice", "payment"],
    "shipping": ["parcel", "delivery", "courier", "arrived"],
    "account": ["password", "login", "email", "account"],
}


def read_fields(text):
    """The [[ ## name ## ]] blocks of a message, as a dictionary."""
    parts = re.split(r"\[\[ ## (\w+) ## \]\]\n", text)
    return {name: value.split("\n\nRespond with")[0].strip() for name, value in zip(parts[1::2], parts[2::2])}


def words(text):
    return {w for w in re.findall(r"[a-z]+", text.lower()) if len(w) > 3}


def sort_ticket(ticket, demos, attempt=0):
    # With examples in the prompt: copy the label of the one that shares most words.
    scored = [(len(words(ticket) & words(d["ticket"])), d["category"]) for d in demos if "category" in d]
    best = max(scored, default=(0, None))
    if best[0] > 0:
        return best[1], f"It shares the most words with an example labelled {best[1]}."
    # Without: keyword rules.
    for category, keys in KEYWORDS.items():
        found = [k for k in keys if k in ticket.lower()]
        if found:
            return category, f"The ticket mentions '{found[0]}'."
    guess = ["account", "billing", "shipping"][attempt % 3]
    return guess, "No keyword matched, so I guessed."


class ShopLM(dspy.BaseLM):
    """A stand-in language model: keyword rules instead of a neural network."""

    forward_contract = "typed_lm"

    def forward(self, request):
        messages = to_openai_chat_request(request)["messages"]
        wanted = re.findall(r"^\d+\. `(\w+)`", messages[0]["content"].split("Your output fields are:")[1], re.M)
        inputs = read_fields(messages[-1]["content"])
        demos = [read_fields(u["content"]) | read_fields(a["content"]) for u, a in zip(messages[1:-1:2], messages[2:-1:2])]

        ticket = inputs.get("ticket", "")
        # DSPy numbers repeated tries of the same request (rollout_id); each try guesses differently.
        cache = request.config.cache
        attempt = cache.rollout_id if cache and cache.rollout_id else 0
        category, why = sort_ticket(ticket, demos, attempt)
        order = re.search(r"A-\d{4}", ticket)
        # A ReAct trajectory arrives as more [[ ## ]] blocks: observation_0 is the first tool's result.
        looked_up = inputs.get("observation_0") if inputs.get("tool_name_0") == "lookup_order" else None
        tool_used = "tool_name_0" in inputs
        call_tool = bool(order) and not tool_used
        answers = {
            "category": category,
            "reasoning": why,
            "reply": f"Your order is {looked_up}." if looked_up else f"Thanks for your message. Our {category} team will help.",
            "next_thought": "I should look the order up." if call_tool else "I have what I need.",
            "next_tool_name": "lookup_order" if call_tool else "finish",
            "next_tool_args": json.dumps({"order_id": order.group()} if call_tool else {}),
        }
        text = "\n\n".join(f"[[ ## {name} ## ]]\n{answers.get(name, '')}" for name in wanted)
        return dspy.LMResponse.from_text(text + "\n\n[[ ## completed ## ]]", model=self.model)

import dspy

from shop_lm import ShopLM

dspy.configure(lm=ShopLM(model="shop/keywords"))
sort = dspy.Predict("ticket -> category")
for ticket in ["I was charged twice", "My parcel never arrived", "I want my money back"]:
    print(ticket, "->", sort(ticket=ticket).category)