প্রম্পট চেইনিং: ধাপ, নির্দিষ্ট স্কিমা ও promptlab হেল্পার

ticket text
   |  stage 1: extract   (model)   -> JSON record: order_id, issue, action, email
   |  stage 2: validate  (Python)  -> accepted record, or a stop with a reason
   |  branch:  missing info? unusual issue?  (Python)  -> stop or hand to a person
   |  stage 3: draft     (model)   -> reply text for a person to review
reply draft (never sent automatically)

"""promptlab: a tiny, stdlib-only helper for prompt experiments (Python 3.12).

A Client wraps any provider function, records a Response for every call and can
append one JSON line per attempt to a log file. Tutorials use toy providers, so
all numbers here are repeatable stand-ins, never measurements of a real model.
"""
from __future__ import annotations

import json
import string
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence

Message = Mapping[str, str]
Provider = Callable[..., "str | dict[str, Any]"]


@dataclass(frozen=True)
class Usage:
    """Token counts for one call; cached_input_tokens is the part of the input served from a cache."""

    input_tokens: int
    output_tokens: int
    cached_input_tokens: int = 0


@dataclass
class Response:
    """What complete() returns: the text, its Usage, a latency in milliseconds and extra metadata."""

    text: str
    usage: Usage
    latency_ms: int
    meta: dict[str, Any] = field(default_factory=dict)


class ProviderError(RuntimeError):
    """Raised by Client.complete() after the failure was logged and all retries were used."""


def estimate_tokens(text: str) -> int:
    """Roughly estimate tokens as one per four characters (at least 1). A rule of thumb, not a tokenizer."""
    return max(1, round(len(text) / 4))


class Client:
    """Call a provider(messages, **settings) function and keep a record of every attempt.

    The provider returns a str or a dict with the key "text" and optionally "cached_input_tokens".
    Latency is a deterministic fake derived from a call counter unless real_timing=True.
    """

    def __init__(self, provider: Provider, model: str = "toy", log_path: str | Path | None = None,
                 retries: int = 0, real_timing: bool = False) -> None:
        self.provider, self.model, self.retries, self.real_timing = provider, model, retries, real_timing
        self.log_path = Path(log_path) if log_path else None
        self.calls = 0  # attempts made so far; drives the fake latency and numbers the log lines

    def _log(self, record: dict[str, Any]) -> None:
        if self.log_path is None:
            return
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        with self.log_path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, sort_keys=True, ensure_ascii=False) + "\n")

    def complete(self, messages: Sequence[Message], **settings: Any) -> Response:
        """Send messages to the provider; retry on any exception; log every attempt, failures included."""
        if not messages or any("role" not in m or "content" not in m for m in messages):
            raise ValueError("messages must be a non-empty list of {'role': ..., 'content': ...} items")
        prompt = [dict(m) for m in messages]
        for attempt in range(1, self.retries + 2):
            self.calls += 1
            base = {"attempt": attempt, "call": self.calls, "messages": prompt, "model": self.model, "settings": settings}
            started = time.perf_counter()
            try:
                raw = self.provider(prompt, **settings)
                data = {"text": raw} if isinstance(raw, str) else dict(raw)
                text = str(data["text"])  # a reply without "text" counts as a provider failure
            except Exception as exc:  # the provider may fail in any way; record it, then retry or give up
                self._log({**base, "error": f"{type(exc).__name__}: {exc}", "status": "error"})
                if attempt > self.retries:
                    raise ProviderError(f"{type(exc).__name__}: {exc} (after {attempt} attempt(s))") from exc
                continue
            tokens_in = sum(estimate_tokens(m["content"]) for m in prompt)
            cached = min(int(data.get("cached_input_tokens", 0)), tokens_in)
            usage = Usage(tokens_in, estimate_tokens(text), cached)
            if self.real_timing:
                latency = round((time.perf_counter() - started) * 1000)
            else:
                latency = 40 + 5 * usage.output_tokens + 10 * (self.calls % 3)
            response = Response(text, usage, latency, {"attempt": attempt, "call": self.calls, "model": self.model})
            self._log({**base, "latency_ms": latency, "meta": response.meta, "status": "ok", "text": text,
                       "usage": asdict(usage)})
            return response
        raise AssertionError("unreachable")  # pragma: no cover


def render(template: str, **variables: Any) -> str:
    """Fill {name} fields in template; raise ValueError if a variable is missing or not used. To print a literal brace, write it twice."""
    names = {name for _, name, _, _ in string.Formatter().parse(template) if name}
    if names - variables.keys():
        raise ValueError(f"missing variables: {sorted(names - variables.keys())}")
    if variables.keys() - names:
        raise ValueError(f"unused variables: {sorted(variables.keys() - names)}")
    return template.format(**variables)

from promptlab import Client, ProviderError, estimate_tokens, render

def toy_echo(messages, **settings):
    """Stand-in provider: repeats the last message in capitals."""
    return messages[-1]["content"].upper()

client = Client(toy_echo, model="toy-echo", log_path="logs/demo.jsonl")
reply = client.complete([{"role": "user", "content": render("Say hi to {name}.", name="Aiko")}])
print(reply.text)
print(reply.usage)
print(reply.latency_ms, reply.meta)
print(open("logs/demo.jsonl", encoding="utf-8").read().count("\n"), "log line")

SAY HI TO AIKO.
Usage(input_tokens=4, output_tokens=4, cached_input_tokens=0)
70 {'attempt': 1, 'call': 1, 'model': 'toy-echo'}
1 log line

import json, re

EXTRACT = """Read the support ticket. Reply with JSON only, with the keys order_id, issue, action, email.
issue is one of: damaged, late, wrong_item, other. action is one of: refund, replace, none.
Use null for anything the ticket does not state. Do not guess.

### TICKET
{ticket}"""

DRAFT = """Write a polite reply of at most 60 words to a customer.
You may use only these facts:
order: {order_id}
problem: {issue}
our next step: {step}
Do not promise anything that is not listed."""

def toy_extract(messages, **settings):
    """Stand-in: finds fields with simple rules and answers in JSON. A real model reads more loosely."""
    text = messages[-1]["content"].split("### TICKET", 1)[1].lower()
    order = re.search(r"order\s*#?\s*(\d{4,})", text)
    mail = re.search(r"[\w.]+@[\w.]+\w", text)
    issue = next((k for w, k in [("cracked", "damaged"), ("late", "late"), ("wrong", "wrong_item")] if w in text), "other")
    action = next((k for w, k in [("refund", "refund"), ("replace", "replace")] if w in text), "none")
    return json.dumps({"order_id": order and order.group(1), "issue": issue, "action": action,
                       "email": mail and mail.group(0)})

def toy_draft(messages, **settings):
    """Stand-in: fills a sentence from the facts it is given."""
    facts = dict(line.split(": ", 1) for line in messages[-1]["content"].splitlines() if ": " in line)
    return f"Hello, thank you for writing about order {facts['order']}. We are sorry about the {facts['problem']} problem. Our next step: {facts['our next step']}."

SCHEMA = {  # the fixed contract between stage 1 and the rest of the chain
    "order_id": (str, None),
    "issue": (str, {"damaged", "late", "wrong_item", "other"}),
    "action": (str, {"refund", "replace", "none"}),
    "email": (str, None),
}
REQUIRED = ["order_id", "email"]
STEP = {"refund": "we will start your refund", "replace": "we will send a replacement",
        "none": "a colleague will review your case"}

def check_record(text):
    """Stage 2 (plain Python): return (record, problems). No model is involved."""
    try:
        record = json.loads(text)
    except json.JSONDecodeError as exc:
        return None, [f"not valid JSON ({exc.msg})"]
    if not isinstance(record, dict) or set(record) != set(SCHEMA):
        return None, [f"keys must be exactly {sorted(SCHEMA)}"]
    problems = []
    for key, (kind, allowed) in SCHEMA.items():
        value = record[key]
        if value is not None and not isinstance(value, kind):
            problems.append(f"{key}: expected {kind.__name__}")
        elif allowed and value not in allowed:
            problems.append(f"{key}: {value!r} is not one of {sorted(allowed)}")
    return record, problems

def run_chain(ticket, extractor, writer):
    """Three stages: extract (model) -> validate and branch (Python) -> draft (model)."""
    try:
        raw = extractor.complete([{"role": "user", "content": render(EXTRACT, ticket=ticket)}]).text
    except ProviderError as exc:
        return {"status": "failed", "stage": "extract", "why": str(exc)}
    record, problems = check_record(raw)
    if problems:
        return {"status": "failed", "stage": "validate", "why": "; ".join(problems)}
    missing = [key for key in REQUIRED if record[key] is None]
    if missing:                                    # stop: do not let a model invent the gap
        return {"status": "needs_info", "ask_customer_for": missing}
    if record["issue"] == "other":                 # branch: this case is not ours to automate
        return {"status": "human", "record": record}
    prompt = render(DRAFT, order_id=record["order_id"], issue=record["issue"].replace("_", " "),
                    step=STEP[record["action"]])
    try:
        reply = writer.complete([{"role": "user", "content": prompt}]).text
    except ProviderError as exc:
        return {"status": "failed", "stage": "draft", "why": str(exc)}
    return {"status": "drafted", "to": record["email"], "reply": reply}

from pprint import pprint

TICKETS = {
    "A": "Hello, order 48213 arrived with a cracked screen. Please send a refund. Kwame, kwame@example.com",
    "B": "My parcel never came and nobody answers. Please help. Lena, lena@example.org",
    "C": "Order #51877 is fine, but how do I change the language of the app? Mateo, mateo@example.net",
}
extractor, writer = Client(toy_extract), Client(toy_draft, log_path="logs/chain.jsonl")
for name, ticket in TICKETS.items():
    print(name)
    pprint(run_chain(ticket, extractor, writer), width=76)

A
{'reply': 'Hello, thank you for writing about order 48213. We are sorry '
          'about the damaged problem. Our next step: we will start your '
          'refund.',
 'status': 'drafted',
 'to': 'kwame@example.com'}
B
{'ask_customer_for': ['order_id'], 'status': 'needs_info'}
C
{'record': {'action': 'none',
            'email': 'mateo@example.net',
            'issue': 'other',
            'order_id': '51877'},
 'status': 'human'}

def chatty(messages, **settings):           # breaks the contract: prose around the JSON
    return "Sure! Here is the data: " + toy_extract(messages)

calls = {"n": 0}
def flaky(messages, **settings):            # fails once, then works
    calls["n"] += 1
    if calls["n"] == 1:
        raise TimeoutError("provider timed out")
    return toy_extract(messages)

def always_down(messages, **settings):
    raise ConnectionError("provider unreachable")

writer = Client(toy_draft)
for label, provider, retries in [("chatty", chatty, 0), ("flaky, 1 retry", flaky, 1), ("down, 1 retry", always_down, 1)]:
    result = run_chain(TICKETS["A"], Client(provider, retries=retries), writer)
    print(f"{label}: {result['status']} {result.get('stage', '')} {result.get('why', '')[:70]}".rstrip())

chatty: failed validate not valid JSON (Expecting value)
flaky, 1 retry: drafted
down, 1 retry: failed extract ConnectionError: provider unreachable (after 2 attempt(s))

history = "\n".join(f"> message {i}: thanks, still waiting for an update" for i in range(60))
long_ticket = TICKETS["A"] + "\n" + history
extract_prompt = render(EXTRACT, ticket=long_ticket)
draft_prompt = render(DRAFT, order_id="48213", issue="damaged", step=STEP["refund"])
print("extract stage reads :", estimate_tokens(extract_prompt), "tokens (rough)")
print("draft, facts only   :", estimate_tokens(draft_prompt), "tokens (rough)")
print("draft, ticket again :", estimate_tokens(draft_prompt + long_ticket), "tokens (rough)")

extract stage reads : 836 tokens (rough)
draft, facts only   : 50 tokens (rough)
draft, ticket again : 822 tokens (rough)

ticket_d = "Order 60412 arrived late. Please replace it. Priya"
print(run_chain(ticket_d, Client(toy_extract), Client(toy_draft)))

{'status': 'needs_info', 'ask_customer_for': ['email']}

def needs_human(record):
    return record["issue"] == "other" or (record["issue"] == "late" and record["action"] == "refund")

late_refund = {"issue": "late", "action": "refund"}
late_replace = {"issue": "late", "action": "replace"}
print(needs_human(late_refund), needs_human(late_replace))

True False

def wrong_enum(messages, **settings):
    return json.dumps({"order_id": "1", "issue": "broken", "action": "refund", "email": "a@b.co"})

result = run_chain(TICKETS["A"], Client(wrong_enum), Client(toy_draft))
assert result["stage"] == "validate"
print(result["why"])

issue: 'broken' is not one of ['damaged', 'late', 'other', 'wrong_item']