Designing Multi-Step Prompt Chains

From a Working Chain to a Well-Designed Chain Lesson 61 introduced prompt chaining and built a complete chain for support tickets. A chain that works on a few examples, however, is not yet a chain you can rely on in production. Real chains must handle failures, be easy to change, report what happened at every step, and stay within cost and time limits.

This lesson covers the design of multi-step chains:

Specifying each step as a contract: inputs, outputs, validation, and failure behavior Managing state that flows through the chain Choosing error-handling strategies for each step Recording a trace for debugging and monitoring Building a small, reusable chain runner in Python Lesson 63 then covers how to decide which steps a complex task needs.

Principle 1: Treat Each Step as a Contract A contract is a precise description of what a step needs, what it produces, and how its result is judged. Writing the contract before writing the prompt prevents most chain problems.

The Step Contract Part Question Example (Extraction Step) Purpose What is this step's single job? Turn the ticket text into structured fields Inputs Which values from the state does it read? ticket Output What exactly does it produce, in which format? JSON: issue_type, course, days_since_purchase Validation How do we know the output is acceptable? Allowed issue types; days is a non-negative integer or null On failure What happens if validation fails? Retry once; then send to a human Performed by Model, code, or person? Model Why Contracts Matter Each step can be built and tested independently, because its inputs and outputs are defined. Steps can be replaced without changing the rest of the chain, as long as the contract is kept. For example, you could replace a model-based extraction step with a cheaper model, or with a rule-based parser, and nothing else needs to change. Failures are local: when validation fails, you know exactly which step broke its contract. Principle 2: Make the State Explicit The state is the collection of all values available as the chain runs: the original input plus every step's output.

Initial state: {"ticket": "..."} After extract: {"ticket": "...", "extracted": {...}} After decide: {"ticket": "...", "extracted": {...}, "decision": "refund_not_eligible"} After draft: {"ticket": "...", "extracted": {...}, "decision": "...", "reply": "..."} Why a Single State Object Helps Every step reads what it needs from the same place, which makes dependencies visible. The full state at any point shows exactly what the chain knew at that moment, which is invaluable for debugging. The state can be saved, so a chain can resume after a failure without repeating expensive steps. Rules for State Rule Reason Steps add new keys; they do not modify earlier values Earlier results remain available for later steps and for debugging Use clear, stable key names Steps and logs refer to the same names Store structured values, not raw model text Later steps and code can use them directly Keep the original input Later steps often need it (Lesson 61) Principle 3: Choose a Failure Strategy for Each Step Every step can fail: the model may return invalid JSON, a value may be outside the allowed set, an API call may time out. Decide in advance what each step does on failure.

Strategy When to Use Example Retry Failures are random and a second attempt often succeeds Invalid JSON from a model Retry with feedback The model can fix its output if told what was wrong "issue_type must be one of..." (similar to Reflexion, Lesson 60) Fallback A simpler or safer alternative exists Use a default category, a template reply, or a different model Skip The step is optional An optional "improve tone" step Stop and escalate Continuing would be risky Policy check failed: send to a human Limit Retries Every retry costs time and money. One or two retries are usually enough to recover from random failures. If a step keeps failing, the problem is systematic (Lesson 46), and more retries will not help.

Escalation Is a Valid Outcome A chain that returns "needs human review" for an uncertain case is working correctly. A chain that always produces an answer, even when steps failed, is not.

Principle 4: Choose the Right Executor for Each Step Executor Best For Example Code Exact rules, calculations, lookups, formatting Refund eligibility, totals, database queries Small or fast model Simple, well-defined language tasks Extraction, classification Larger model Complex reasoning or high-quality writing Drafting a nuanced reply Person High-stakes or ambiguous decisions Approving an exception to policy Choosing the executor per step is one of the main advantages of chains: you can use expensive resources only where they matter.

Principle 5: Record a Trace A trace is a record of every step execution: which step ran, which attempt it was, whether it succeeded, what errors occurred, and how long it took.

extract attempt 1 failed invalid issue_type 0.84s extract attempt 2 ok 0.79s decision attempt 1 ok 0.00s draft attempt 1 ok 1.92s check attempt 1 ok 0.65s Why Traces Matter Debugging: you can see exactly where a chain failed and what the step received. Monitoring: if a step's failure rate increases after a model update, the trace shows it. Cost and latency: you can see which steps are slow or retried often. Debugging multi-step workflows with traces is covered in Lesson 74.

A Reusable Chain Runner in Python The following code implements the principles above: steps with contracts, an explicit state, retries, validation, and a trace. It runs without the API, using simple functions as steps, so you can test the chain logic first and plug in model calls later.

The Step and the Runner import time from dataclasses import dataclass, field from typing import Any, Callable

@dataclass class Step: """One link in a chain: what it does, how it is validated, how often it may retry.""" name: str run: Callable[[dict], Any] validate: Callable[[Any], list[str]] = lambda output: [] retries: int = 1 output_key: str = ""

def __post_init__(self): self.output_key = self.output_key or self.name

@dataclass class ChainResult: status: str state: dict trace: list[dict] = field(default_factory=list)

def run_chain(steps: list[Step], initial_state: dict) -> ChainResult: state = dict(initial_state) trace = [] for step in steps: for attempt in range(1, step.retries + 2): start = time.perf_counter() try: output = step.run(state) errors = step.validate(output) except Exception as error: output, errors = None, [f"{type(error).name}: {error}"] trace.append({ "step": step.name, "attempt": attempt, "ok": not errors, "errors": errors, "seconds": round(time.perf_counter() - start, 3), }) if not errors: state[step.output_key] = output break else: return ChainResult("failed", state, trace) return ChainResult("completed", state, trace) Code Explanation Step

run is any function that receives the current state and returns an output. It can call a model, run code, or query a database. validate returns a list of error messages; an empty list means the output is valid. The default accepts everything. retries is the number of additional attempts after the first one. output_key is the name under which the output is stored in the state. __post_init__ sets it to the step's name if none is given. run_chain()

Copies the initial state, so the caller's dictionary is not modified. For each step, tries up to retries + 1 times. Catches exceptions from run or validate and turns them into error messages, so one failing step does not crash the program. Records every attempt in the trace, including its duration. On success, stores the output and moves to the next step (break). The for ... else construction: the else block runs only if the loop finished without break, meaning every attempt failed. In that case, the chain stops and returns "failed" with the state and trace collected so far. Testing the Runner With Simple Steps calls = {"count": 0}

def flaky_extract(state: dict) -> dict: """Simulates a model that returns an invalid value on the first attempt.""" calls["count"] += 1 if calls["count"] == 1: return {"issue_type": "refund"} return {"issue_type": "refund_request", "days_since_purchase": 10}

def validate_extract(output: dict) -> list[str]: return [] if output.get("issue_type") in {"refund_request", "other"} else ["invalid issue_type"]

def decide(state: dict) -> str: days = state["extracted"]["days_since_purchase"] return "refund_eligible" if days <= 7 else "refund_not_eligible"

steps = [ Step("extract", flaky_extract, validate_extract, retries=1, output_key="extracted"), Step("decision", decide), ]

result = run_chain(steps, {"ticket": "Refund please, I bought it 10 days ago."}) print(result.status, result.state["decision"]) for record in result.trace: print(record["step"], record["attempt"], record["ok"], record["errors"]) Output:

completed refund_not_eligible extract 1 False ['invalid issue_type'] extract 2 True [] decision 1 True [] What the Test Shows The first extraction attempt returned an invalid issue type. Validation caught it, and the runner retried. The second attempt succeeded, and its output was stored as extracted. The decision step, implemented in code, read extracted from the state and applied the policy. The trace records the failed attempt, which you would otherwise never see. Plugging In Model Calls To use a real model in a step, write a run function that builds a prompt from the state and calls the API:

import json

from openai import OpenAI

client = OpenAI() MODEL = "gpt-5.4" # Replace with a model available in your account

EXTRACT_INSTRUCTIONS = """Extract fields from the ticket in as JSON with keys issue_type ("refund_request" or "other") and days_since_purchase (integer or null). Return only the JSON object."""

def model_extract(state: dict) -> dict: response = client.responses.create( model=MODEL, instructions=EXTRACT_INSTRUCTIONS, input=f"\n{state['ticket']}\n", ) return json.loads(response.output_text)

steps[0] = Step("extract", model_extract, validate_extract, retries=1, output_key="extracted") If the model returns invalid JSON, json.loads raises an exception, which the runner records as an error and retries. No other part of the chain changes, because the step's contract is the same.

Designing Validation for Each Step Validation is the most important design decision after the step's purpose. Different steps need different checks.

Step Type Useful Validation Extraction Required keys present, types correct, values in allowed sets (Lesson 22) Classification Label in the allowed list (Lesson 18) Generation Length limits, required phrases present, forbidden content absent (Lesson 35) Summarization Numbers and names match the source (Lesson 20) Code generation Code parses; tests pass (Lesson 60) Decision by code Usually no validation needed, but test the function itself Validate Meaning, Not Only Format A JSON object can have the right keys and still be wrong. Where possible, check values against the input: an extracted course name should appear in the ticket, and an extracted number should appear in the source text.

Controlling Cost and Latency Technique Effect Use code for exact steps No model cost, no randomness Use smaller models for simple steps Lower cost per step Run independent steps in parallel Lower total latency (Lesson 63) Limit retries per step Bounded worst-case cost Stop early on failure No wasted calls after a step fails Cache results of deterministic steps Avoid repeating identical work Set a budget for the whole chain, such as a maximum number of model calls, and measure actual usage from the trace.

Common Mistakes Mistake 1: Steps Without Contracts What it is: Writing prompts for each step without defining their inputs, outputs, and validation.

Why it happens: Prompts are written first and connected later.

What it causes: Mismatched formats between steps and failures that are hard to locate.

How to fix it: Write the contract table for each step before writing its prompt.

Mistake 2: Hidden State What it is: Passing values between steps through scattered variables.

Why it happens: The chain grows step by step in a script.

What it causes: It becomes unclear what each step depends on, and debugging requires reading all the code.

How to fix it: Use a single state object with clear keys.

Mistake 3: The Same Failure Strategy Everywhere What it is: Retrying every step three times regardless of its type.

Why it happens: A global retry setting is easy to implement.

What it causes: Wasted calls on systematic failures and risky continuation after critical failures.

How to fix it: Choose retry, fallback, skip, or escalate per step.

Mistake 4: No Trace What it is: Logging only the final output.

Why it happens: The final output is what users see.

What it causes: When something goes wrong, there is no way to know which step caused it.

How to fix it: Record every attempt with its status, errors, and timing.

Mistake 5: Testing Only the Whole Chain What it is: Testing the chain end to end only.

Why it happens: It seems simpler.

What it causes: Failures are hard to attribute, and every test costs several model calls.

How to fix it: Test each step separately, and test the runner logic with simple functions, as shown above.

Best Practices Write a contract for every step. Purpose, inputs, output, validation, failure behavior, executor.

Keep one explicit state. Steps add to it; nothing is silently overwritten.

Validate after every step. Check format and, where possible, meaning.

Choose failure strategies deliberately. Retry for random failures; escalate when continuing is risky.

Use the cheapest reliable executor for each step. Code first, then smaller models, then larger ones.

Record a trace. It is your main tool for debugging and monitoring.

Test steps and runner independently. Then test the full chain.

Try It Yourself Exercise 1: Write Step Contracts Write the contract table for each step of this chain: "Generate a weekly progress email for a student from their activity log." Include at least four steps, and for each one decide its executor, validation, and failure strategy.

Exercise 2: Extend the Runner Add a fallback field to Step: an optional function that is called when all attempts fail. If it returns a value, the chain stores it and continues instead of failing. Test it with a step that always fails and a fallback that returns a default category.

Exercise 3: Measure From the Trace Run the chain on ten inputs (using simulated or real model steps). From the traces, calculate the failure rate of each step on its first attempt, the average number of attempts per step, and the total time per step. Decide which step most needs improvement.