Build a Cost-Aware Routing Layer for LLM Agents

Most teams discover their agent's token spend the same way: the cloud bill arrives and someone asks what happened. A single LLM call is cheap and predictable. A multi-step agent is neither — it can burn an order of magnitude more tokens than the call you prototyped, because every tool call, every observation, and every retry adds another round-trip through the model. Nobody plans for that. They find out at invoicing time.

The fix is not to use smaller models everywhere. That just trades your cost problem for a quality problem. The fix is a routing layer: score each task, send it to the cheapest model tier that can handle it, escalate when confidence is low, log every cent, and degrade gracefully when the budget runs out.

This tutorial builds exactly that in Python. It runs as-is with no API keys — a mock model stands in for real calls so you can see the routing, escalation, receipts, and guardrails work end to end. Swap the mock for your provider's SDK when you're ready.

What you'll build

1. Model tiers with per-tier pricing, plus rough token math.
2. A task router that scores difficulty and picks the cheapest capable tier.
3. Cascade routing: try the small model first, escalate on low confidence.
4. Cost attribution: a receipt for every run — tokens and dollars per task.
5. Budget guardrails: cached answers and simpler models instead of hard failures.

Step 1: Tiers and token math

Everything starts with knowing what a call costs. Define your tiers with illustrative per-token prices (replace these with your provider's actual pricing), and estimate tokens with a rough heuristic until you wire in a real tokenizer:

def estimate_tokens(text: str) -> int:
    """Rough heuristic: ~4 characters per token for English text."""
    return max(1, len(text) // 4)

@dataclass(frozen=True)
class ModelTier:
    name: str
    input_per_mtok: float    # $ per 1M input tokens (illustrative)
    output_per_mtok: float   # $ per 1M output tokens (illustrative)
    quality: float           # illustrative capability score, 0..1

SMALL  = ModelTier("small",  0.15,  0.60, 0.70)
MEDIUM = ModelTier("medium", 1.25,  5.00, 0.88)
LARGE  = ModelTier("large",  5.00, 15.00, 0.97)

def cost_of(tier, in_tok, out_tok) -> float:
    return in_tok / 1e6 * tier.input_per_mtok + out_tok / 1e6 * tier.output_per_mtok

Two things worth saying plainly. First, the 4-characters-per-token rule is a stand-in — fine for a tutorial, not for invoicing. In production, use your provider's tokenizer or usage API. Second, the prices above are illustrative placeholders so the math is visible; plug in real numbers before making decisions with them.

Step 2: Score the task, pick the cheapest capable tier

Routing needs one number per task: how hard is this? A perfect difficulty classifier is a research project. A decent heuristic is an afternoon. The router only needs a number, not a perfect one — start simple and improve it later:

def score_difficulty(task: str) -> float:
    t = task.lower()
    score = 0.15
    score += min(0.25, len(task) / 2000)
    for kw in ("compare", "analyz", "design", "trade-off", "tradeoff"):
        if kw in t: score += 0.08
    for kw in ("code", "sql", "debug", "prove", "architect"):
        if kw in t: score += 0.10
    score += 0.05 * t.count("?")
    if any(w in t for w in ("step by step", " first ", " then ")):
        score += 0.07
    return min(0.95, score)

def pick_tier(difficulty: float, margin: float = 0.05) -> ModelTier:
    """Cheapest tier whose capability clears the difficulty plus a margin."""
    for tier in TIERS:
        if tier.quality >= difficulty + margin:
            return tier
    return LARGE

"What is the capital of France?" scores around 0.17 and lands on the small tier. "Debug this SQL query step by step, then prove the fix handles nulls" scores around 0.62. The margin keeps you from assigning a task to a tier that's exactly at its limit.

Step 3: Cascade — try small, escalate on low confidence

Picking a tier up front is a guess. The cascade makes it a guess you can correct: run the small model first, and only pay for a bigger one when confidence comes back low.

current = pick_tier(difficulty)
while True:
    resp = model.complete(task, difficulty)   # mock here, real SDK in prod
    record_step(current, resp)                # tokens + cost, see Step 4
    if resp.confidence >= THRESHOLD or current is LARGE:
        break
    nxt = next_tier(current)
    if spent + projected_cost(nxt) > budget:
        degraded = True                       # keep what we have; see Step 5
        break
    current = nxt

The confidence number can come from whatever signal you trust: logprobs, a self-rating prompt, a verifier model, or agreement between two samples. The mechanism doesn't care — it just needs a number and a threshold. In the example code, the mock model reports lower confidence when the task is harder than the tier, which is exactly the situation where you'd want a human-designed escalation anyway.

Step 4: Every run gets a receipt

This is the part most teams skip, and it's the part that ends bill-shock. Log tokens and dollars per step, per task, per agent. When someone asks "what did that run cost?", you answer with a receipt, not a shrug:

@dataclass
class Receipt:
    task: str
    steps: list = field(default_factory=list)  # (tier, in_tok, out_tok, cost, confidence)
    degraded: bool = False
    cached: bool = False

    @property
    def total_cost(self) -> float:
        return sum(s[3] for s in self.steps)

Running five tasks through the example — once naively (everything to the large tier) and once through the router — produces this (simulated run, illustrative pricing):

=== Naive: every task goes to the large tier ===
Total cost for 5 tasks: $0.0157

=== Routed: difficulty-scored tiers + escalation + $0.05/task budget ===
- What is the capital of France?                        small(conf 0.62)  $0.0001
- Summarize this error log in two sentences: ...        small(conf 0.60) -> medium(conf 0.83)  $0.0008
- Compare three caching strategies ...                  small(conf 0.58) -> medium(conf 0.77)  $0.0012
- Design a retry and backoff strategy ...               small(conf 0.56) -> medium(conf 0.76)  $0.0013
- Debug this SQL query step by step ...                 small(conf 0.52) -> medium(conf 0.70)  $0.0017

Total cost for 5 tasks: $0.0051
Illustrative savings vs. naive: 67%

The trivial question stays on the small tier. The hard ones escalate to medium — and stop there, because medium clears the confidence threshold. Nothing in this run needed the large tier at all, which is itself useful information: it tells you where your workload actually lives.

Step 5: Guardrails — degrade, don't die

Budgets fail in two familiar ways: the task errors out mid-run when money runs out (bad), or the agent keeps spending because nobody told it to stop (worse). The routing layer handles both before they happen:

1. Before spending: if even the cheapest possible first step exceeds the remaining budget, serve a cached answer from an earlier identical task instead of spending. Mark the receipt degraded so it's honest about what happened.
2. Before escalating: if the next tier up would break the budget, keep the current answer and mark the run degraded.

# Guardrail 1: broke before spending -> cached answer, $0
floor = cost_of(SMALL, estimate_tokens(task), 100)
if floor > self.budget and task_hash in self.cache:
    return cached_receipt(degraded=True)

# Guardrail 2: escalation would break the budget -> keep current answer
if spent + projected_cost(next_tier) > self.budget:
    receipt.degraded = True
    break

The same task run twice under a near-zero budget shows the behavior:

=== Budget guardrail: same task twice, tiny $0.00001 budget ===
First run : small  degraded=True cached=False  $0.0001
Second run: cache  degraded=True cached=True   $0.0000
The second run served a cached answer for $0 instead of failing.

The first run couldn't afford to escalate, so it kept the small-tier answer and said so. The second run cost nothing at all. Neither one threw an error at the user. That is the whole point: a budget overrun becomes a slightly worse answer, not an outage.

Where this fits in a real architecture

A routing layer is one component, not a system. In the reference architecture I maintain (github.com/karmendra8386/agentic-ai-aws-reference), cost shows up in two places: the observability section calls for structured per-turn logs that include tokens and cost, and the governance pattern calls for cost dashboards next to quality dashboards. This routing layer is the control plane that acts on those numbers — the logs tell you what you spent, the router decides what you'll spend next.

The repo itself is a starter skeleton (CDK stacks for agent IAM roles and encrypted storage, Playwright regression tests), not production code — treat the routing layer here the same way. Harden the pieces for your environment: real tokenizers, real pricing, a real confidence signal, and a real cache backend.

What this gives you

Five things, in the order you'll notice them. First, the bill stops being a surprise — every run has a receipt. Second, easy work stops subsidizing hard work — the small tier handles what it can, and you can see that it does. Third, escalation becomes a deliberate decision with a paper trail instead of a silent default to the biggest model. Fourth, budgets become graceful: degraded answers instead of failed tasks. Fifth, and most useful over time, the receipts become data — you'll learn which task types actually need the large tier, and that tells you where to invest in prompts, fine-tuning, or caching instead of raw model spend.

Start with the mock, run it, read the receipts. Then swap in one real model call and watch your own numbers appear. That's the whole tutorial — and the end of bill-shock.

---

Karmendra Pandey is Practice Architect for AI/ML at TEKsystems (19 years across data engineering, ML, and cloud), currently architecting AI/ML frameworks for production AI agents on AWS.

Comments

[Paul Irolla]

Your receipts price each call, but on a long agent run the input side keeps growing: every tool call replays the conversation so far, so later steps pay again for earlier ones. A router picks the tier; it doesn't shrink what each call carries. I built Tokenade for that part. It's open source and sits between the agent and the model. On THOL (Claude Code with Sonnet, sessions over 200k tokens, cost in USD, measured 2026-09-14), it cut session cost by 39% on average.