Prompt Engineering for Production Code: Context First, Tagged Examples, Positive Constraints, and an Acceptance Check

"""Build a code-generation prompt: your code first, examples tagged, positive constraints, the task last."""
import tiktoken

enc = tiktoken.get_encoding("o200k_base")


def build(task: str, context_files: dict[str, str], examples: list[tuple[str, str]],
          constraints: list[str], budget: int = 6_000) -> str:
    docs, used = [], 0
    for i, (path, text) in enumerate(context_files.items(), 1):           # most relevant first
        block = f'<document index="{i}">\n<source>{path}</source>\n<document_content>\n{text}\n</document_content>\n</document>'
        n = len(enc.encode(block))
        if used + n > budget:
            break                                                          # drop whole files, never half a function
        docs.append(block); used += n
    ex = "\n".join(f"<example>\n<request>{q}</request>\n<code>\n{a}\n</code>\n</example>" for q, a in examples)
    rules = "\n".join(f"- {c}" for c in constraints)
    return (f"<documents>\n{chr(10).join(docs)}\n</documents>\n\n<examples>\n{ex}\n</examples>\n\n"
            f"<constraints>\n{rules}\n</constraints>\n\n<task>\n{task}\n</task>\n"
            "Return one Python module in a single fenced python code block and nothing else.")


CONSTRAINTS = [
    "Raise app.errors.RateLimitError(key, retry_after) when a request is refused.",
    "Log through the logger argument passed to the constructor.",
    "Take the current time as a parameter (now: float) so tests can control it.",
    "Annotate every function and method with types.",
    "Catch only the specific exceptions you expect, by name.",
]

"""Accept generated code only if it compiles, follows the house rules (checked by AST) and passes behavior tests."""
import ast
import re
import types


def extract(reply: str) -> str:
    m = re.search(r"`{3}python\n(.*?)`{3}", reply, re.S)
    if not m:
        raise ValueError("no fenced python block in the reply")
    return m.group(1)


def house_rules(tree: ast.AST) -> list[str]:
    problems = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "print":
            problems.append(f"line {node.lineno}: print() instead of the logger")
        if isinstance(node, ast.ExceptHandler) and (node.type is None or getattr(node.type, "id", "") == "Exception"):
            problems.append(f"line {node.lineno}: broad except")
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if node.returns is None or any(a.annotation is None for a in node.args.args if a.arg != "self"):
                problems.append(f"line {node.lineno}: {node.name} is missing type hints")
        if isinstance(node, ast.Call) and getattr(getattr(node.func, "value", None), "id", "") == "time":
            problems.append(f"line {node.lineno}: reads the clock directly; take now as a parameter")
    return problems


def behavior(mod) -> list[str]:
    class Log:
        def __init__(self): self.lines = []
        def info(self, *a, **k): self.lines.append(a)
        warning = info
    b = mod.TokenBucket(capacity=2, refill_per_sec=1.0, logger=Log())
    errs = []
    if not (b.take("u1", now=0.0) and b.take("u1", now=0.0)):
        errs.append("first two requests should pass")
    try:
        b.take("u1", now=0.0); errs.append("third request should raise RateLimitError")
    except mod.RateLimitError as e:
        if not 0 < e.retry_after <= 1.0:
            errs.append(f"retry_after {e.retry_after} should be within (0, 1]")
    if not b.take("u1", now=1.0):
        errs.append("a token should refill after 1 second")
    if not b.take("u2", now=0.0):
        errs.append("keys must not share a bucket")
    return errs


def accept(reply: str, errors_module_src: str) -> dict:
    try:
        src = extract(reply)
        tree = ast.parse(src)
    except (ValueError, SyntaxError) as e:
        return {"accepted": False, "problems": [str(e)]}
    problems = house_rules(tree)
    mod = types.ModuleType("generated")
    exec(compile(errors_module_src, "app/errors.py", "exec"), mod.__dict__)    # provides RateLimitError
    try:
        exec(compile(src.replace("from app.errors import RateLimitError", ""), "<generated>", "exec"), mod.__dict__)
        problems += behavior(mod)
    except Exception as e:                                                     # the harness may catch broadly
        problems.append(f"runtime: {type(e).__name__}: {e}")
    return {"accepted": not problems, "problems": problems}

from prompt_builder import build, CONSTRAINTS
from acceptance import accept

FENCE = "`" * 3   # a code fence, spelled out so this listing doesn't contain one

ERRORS = '''class RateLimitError(Exception):
    def __init__(self, key: str, retry_after: float):
        super().__init__(f"{key} rate limited; retry after {retry_after:.2f}s")
        self.key, self.retry_after = key, retry_after
'''
prompt = build("Write TokenBucket(capacity, refill_per_sec, logger) with take(key, now) -> bool, one bucket per key.",
               {"app/errors.py": ERRORS, "app/middleware/auth.py": "def require_user(req, logger): ..."},
               [("a cache with expiry", "class TTLCache:\n    def get(self, key: str, now: float) -> str | None: ...")],
               CONSTRAINTS)
print(prompt[:120].replace("\n", " | "), "...", prompt[-160:].replace("\n", " | "))

tutorial = FENCE + '''python
import time
class TokenBucket:
    def __init__(self, capacity, refill_per_sec, logger):
        self.capacity, self.rate, self.buckets = capacity, refill_per_sec, {}
    def take(self, key, now=None):
        try:
            now = now or time.time()
            tokens, last = self.buckets.get(key, (self.capacity, now))
            tokens = min(self.capacity, tokens + (now - last) * self.rate)
            if tokens < 1:
                print("rate limited", key)
                return False
            self.buckets[key] = (tokens - 1, now)
            return True
        except:
            return True
''' + FENCE
constrained = FENCE + '''python
import logging

from app.errors import RateLimitError


class TokenBucket:
    def __init__(self, capacity: int, refill_per_sec: float, logger: logging.Logger) -> None:
        self.capacity, self.rate, self.log = capacity, refill_per_sec, logger
        self.buckets: dict[str, tuple[float, float]] = {}

    def take(self, key: str, now: float) -> bool:
        tokens, last = self.buckets.get(key, (float(self.capacity), now))
        tokens = min(float(self.capacity), tokens + (now - last) * self.rate)
        if tokens < 1.0:
            retry_after = (1.0 - tokens) / self.rate
            self.log.warning("rate limited", extra={"key": key, "retry_after": retry_after})
            raise RateLimitError(key, retry_after)
        self.buckets[key] = (tokens - 1.0, now)
        return True
''' + FENCE
for name, reply in (("tutorial-style", tutorial), ("constrained", constrained)):
    r = accept(reply, ERRORS)
    print(f"{name}: accepted={r['accepted']}")
    for p in r["problems"]:
        print("   -", p)

<documents> | <document index="1"> | <source>app/errors.py</source> | <document_content> | class RateLimitError(Exception): |      ... refill_per_sec, logger) with take(key, now) -> bool, one bucket per key. | </task> | Return one Python module in a single fenced python code block and nothing else.
rate limited u1
rate limited u1
tutorial-style: accepted=False
   - line 3: __init__ is missing type hints
   - line 5: take is missing type hints
   - line 15: broad except
   - line 7: reads the clock directly; take now as a parameter
   - line 11: print() instead of the logger
   - third request should raise RateLimitError
   - a token should refill after 1 second
constrained: accepted=True