Building Cost-Aware LLM Pipelines with Token Budgets
Large language models are powerful, but they are also metered resources. Every token that enters or leaves a model has a price, and when you build pipelines that chain multiple prompts, retrievers, and agents together, costs can spiral out of control faster than you might expect. A single misbehaving retrieval step that stuffs 30,000 tokens of context into a prompt can turn a $0.02 call into a $0.90 call — and at scale, that difference is the line between a sustainable product and a money pit.
Cost-aware LLM pipelines are designed to treat token usage as a first-class engineering concern. Instead of blindly sending requests and hoping the bill stays reasonable, you set explicit token budgets, track consumption at each stage, and enforce limits that prevent runaway costs. This tutorial walks through the concepts, the architecture, and a working implementation you can adapt to your own applications.
What Is a Token Budget?
A token budget is a hard or soft limit on the number of tokens that a pipeline, a sub-pipeline, or an individual request is allowed to consume. Tokens are counted in two directions: input tokens (what you send to the model) and output tokens (what the model generates). Most providers price these differently, with output tokens typically costing 3–5x more than input tokens.
A budget can be expressed in several ways:
- Per-request budget: A cap on tokens for a single user request, including all intermediate LLM calls.
- Per-session budget: A cap that accumulates across multiple turns of a conversation.
- Per-tenant budget: A cap tied to a customer or organization, often used in multi-tenant SaaS.
- Per-pipeline-stage budget: A cap on a specific component, like a retrieval-augmentation step or a summarization step.
The key insight is that budgets are composable. A per-tenant budget might be the sum of many per-request budgets, and each request might allocate sub-budgets to retrieval, reasoning, and response generation stages.
Why Token Budgets Matter
Without budgets, LLM applications are vulnerable to several common failure modes. A user might submit an unusually long input that blows up your context window. A retrieval step might return far more documents than expected. An agent might get stuck in a loop, making dozens of redundant calls. A summarization chain might recursively expand rather than compress. In each case, the cost is real and often unpredictable.
Budgets give you three things: predictability, protection, and observability. Predictability means you can forecast spend based on expected traffic. Protection means a single pathological request cannot drain your budget for the day. Observability means you can see exactly where tokens are being spent and optimize accordingly.
This matters most in production settings where you are serving many users. In development, a $5 experiment is fine. In production, a thousand users each triggering a $5 experiment is a $5,000 daily bill. Budgets turn that risk into a controlled, measurable parameter.
Architecture of a Cost-Aware Pipeline
A cost-aware pipeline has three main components: a token counter, a budget allocator, and a budget enforcer. The counter measures tokens before and after each LLM call. The allocator distributes a total budget across pipeline stages. The enforcer checks remaining budget before allowing a call to proceed and raises an exception or triggers a fallback when the budget is exhausted.
The pipeline itself is typically structured as a sequence of stages, where each stage receives a context object that carries the remaining budget. Each stage is responsible for updating the budget after it consumes tokens. This makes the budget a shared, mutable resource that flows through the pipeline alongside the actual data.
Here is a simplified diagram of the flow:
User Request
│
▼
┌─────────────────┐
│ Budget Allocator │ ── assigns 4000 tokens total
└─────────────────┘
│
▼
┌─────────────────┐
│ Retrieval Stage │ ── uses 1500 input tokens
└─────────────────┘
│
▼
┌─────────────────┐
│ Reasoning Stage │ ── uses 1200 input + 800 output tokens
└─────────────────┘
│
▼
┌─────────────────┐
│ Response Stage │ ── uses 300 output tokens
└─────────────────┘
│
▼
Final Answer (200 tokens remaining)
Implementing a Token Counter
The first building block is an accurate token counter. The best approach is to use the same tokenizer your model provider uses. For OpenAI models, that is the tiktoken library. For other providers, you may need to use their SDK or approximate with a character-based heuristic (roughly 4 characters per token for English text).
import tiktoken
class TokenCounter:
def __init__(self, model: str = "gpt-4o"):
try:
self.encoding = tiktoken.encoding_for_model(model)
except KeyError:
self.encoding = tiktoken.get_encoding("cl100k_base")
def count(self, text: str) -> int:
return len(self.encoding.encode(text))
def count_messages(self, messages: list[dict]) -> int:
"""Count tokens for a chat-formatted message list."""
total = 0
for msg in messages:
total += self.count(msg.get("content", ""))
# overhead per message (role tags, separators)
total += 4
total += 2 # priming tokens
return total
The count_messages method accounts for the overhead that chat APIs add for role tags and message separators. This overhead is small per message but adds up when you have long conversation histories.
Implementing a Budget Tracker
The budget tracker holds the remaining token budget and provides methods for checking and consuming tokens. It is designed to be passed through the pipeline as a shared resource.
from dataclasses import dataclass, field
@dataclass
class BudgetExceededError(Exception):
requested: int
remaining: int
message: str = ""
def __str__(self):
return (
f"Budget exceeded: requested {self.requested} tokens "
f"but only {self.remaining} remain. {self.message}"
)
@dataclass
class BudgetTracker:
total: int
used: int = 0
input_used: int = 0
output_used: int = 0
breakdown: dict = field(default_factory=dict)
@property
def remaining(self) -> int:
return self.total - self.used
def check(self, amount: int):
"""Raise an error if the requested amount exceeds remaining budget."""
if amount > self.remaining:
raise BudgetExceededError(
requested=amount,
remaining=self.remaining,
message="Consider reducing context or using a cheaper model."
)
def consume(self, input_tokens: int, output_tokens: int, stage: str = "default"):
"""Record token usage for a completed LLM call."""
total = input_tokens + output_tokens
self.used += total
self.input_used += input_tokens
self.output_used += output_tokens
if stage not in self.breakdown:
self.breakdown[stage] = {"input": 0, "output": 0}
self.breakdown[stage]["input"] += input_tokens
self.breakdown[stage]["output"] += output_tokens
def estimate_input(self, messages: list[dict], counter: TokenCounter) -> int:
"""Estimate input tokens for a message list before sending."""
return counter.count_messages(messages)
def report(self) -> dict:
return {
"total_budget": self.total,
"total_used": self.used,
"remaining": self.remaining,
"input_used": self.input_used,
"output_used": self.output_used,
"breakdown_by_stage": self.breakdown,
}
Notice that the tracker separates input and output token accounting. This is important because output tokens are typically more expensive, and you may want to apply different budget weights to them. We will cover weighted budgets in the best practices section.
Building the Pipeline Stages
Each pipeline stage follows a consistent pattern: estimate the input tokens it will need, check the budget, make the LLM call, then record actual usage. The estimate-check-call-record cycle ensures that a stage never starts work it cannot afford to finish.
import openai
class LLMStage:
"""Base class for pipeline stages that call an LLM."""
def __init__(self, name: str, model: str, counter: TokenCounter):
self.name = name
self.model = model
self.counter = counter
self.client = openai.OpenAI()
def call_llm(
self,
messages: list[dict],
budget: BudgetTracker,
max_output: int = 500,
temperature: float = 0.7,
) -> str:
# 1. Estimate input tokens
estimated_input = budget.estimate_input(messages, self.counter)
# 2. Check budget for input + max output
budget.check(estimated_input + max_output)
# 3. Make the call
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
max_tokens=max_output,
temperature=temperature,
)
# 4. Record actual usage from the API response
usage = response.usage
budget.consume(
input_tokens=usage.prompt_tokens,
output_tokens=usage.completion_tokens,
stage=self.name,
)
return response.choices[0].message.content
The call_llm method uses the API response's usage object to record actual token counts rather than relying solely on estimates. This is critical because the provider's count is the source of truth for billing. Your estimate is only used for the pre-flight budget check.
Putting It Together: A Retrieval-Augmented Pipeline
Now let us assemble a complete retrieval-augmented generation (RAG) pipeline with budget enforcement. The pipeline retrieves relevant documents, compresses them if needed, and generates an answer.
class RetrievalStage:
"""Retrieves documents and trims them to fit the budget."""
def __init__(self, counter: TokenCounter, max_context_tokens: int = 2000):
self.counter = counter
self.max_context_tokens = max_context_tokens
def retrieve(self, query: str, corpus: list[str], budget: BudgetTracker) -> str:
# Simple keyword-based retrieval for illustration
scored = []
query_terms = set(query.lower().split())
for doc in corpus:
doc_terms = set(doc.lower().split())
score = len(query_terms & doc_terms)
scored.append((score, doc))
scored.sort(key=lambda x: x[0], reverse=True)
# Accumulate documents until we hit the context limit
context_parts = []
token_count = 0
for score, doc in scored:
if score == 0:
break
doc_tokens = self.counter.count(doc)
if token_count + doc_tokens > self.max_context_tokens:
# Trim the last document to fit
remaining_space = self.max_context_tokens - token_count
if remaining_space > 100:
trimmed = doc[: remaining_space * 4] # approx 4 chars/token
context_parts.append(trimmed)
token_count += self.counter.count(trimmed)
break
context_parts.append(doc)
token_count += doc_tokens
budget.consume(
input_tokens=token_count,
output_tokens=0,
stage="retrieval",
)
return "\n\n".join(context_parts)
class RAGPipeline:
def __init__(self, model: str = "gpt-4o", total_budget: int = 4000):
self.counter = TokenCounter(model)
self.budget = BudgetTracker(total=total_budget)
self.retrieval = RetrievalStage(self.counter, max_context_tokens=2000)
self.generation = LLMStage("generation", model, self.counter)
def run(self, query: str, corpus: list[str]) -> dict:
try:
# Stage 1: Retrieve context
context = self.retrieval.retrieve(query, corpus, self.budget)
# Stage 2: Generate answer
messages = [
{
"role": "system",
"content": (
"Answer the user's question using only the provided "
"context. If the context does not contain the answer, "
"say you don't know."
),
},
{
"role": "user",
"content": f"Context:\n{context}\n\nQuestion: {query}",
},
]
answer = self.generation.call_llm(
messages,
self.budget,
max_output=500,
)
return {
"answer": answer,
"budget_report": self.budget.report(),
}
except BudgetExceededError as e:
return {
"answer": None,
"error": str(e),
"budget_report": self.budget.report(),
}
# Example usage
if __name__ == "__main__":
corpus = [
"The Eiffel Tower was completed in 1889 for the World's Fair in Paris.",
"Python is a programming language known for its readability.",
"Photosynthesis converts sunlight into chemical energy in plants.",
"The Great Wall of China spans over 13,000 miles.",
]
pipeline = RAGPipeline(model="gpt-4o", total_budget=3000)
result = pipeline.run("When was the Eiffel Tower built?", corpus)
print("Answer:", result["answer"])
print("Budget Report:", result["budget_report"])
Handling Budget Exhaustion Gracefully
When a budget is exceeded, you have several options beyond simply failing. The right choice depends on your application's requirements.
- Fallback to a cheaper model: If the primary model is too expensive for the remaining budget, switch to a smaller, cheaper model for the final response.
- Reduce context: Trim retrieved documents more aggressively, summarize conversation history, or drop less relevant messages.
- Return partial results: If the retrieval stage succeeded but generation would exceed the budget, return the retrieved context with a note that full generation was skipped.
- Queue for later processing: For non-real-time workloads, defer the request to a batch job with a separate budget pool.
Here is an example of a fallback strategy that switches to a cheaper model when the budget is tight:
class FallbackGenerationStage(LLMStage):
def __init__(self, counter: TokenCounter, primary_model: str, fallback_model: str):
self.primary = LLMStage("generation_primary", primary_model, counter)
self.fallback = LLMStage("generation_fallback", fallback_model, counter)
def call_llm(self, messages, budget, max_output=500, temperature=0.7):
estimated_input = budget.estimate_input(messages, self.counter)
try:
budget.check(estimated_input + max_output)
return self.primary.call_llm(messages, budget, max_output, temperature)
except BudgetExceededError:
# Try the cheaper model with a smaller output cap
budget.check(estimated_input + max_output // 2)
return self.fallback.call_llm(
messages, budget, max_output // 2, temperature
)
Weighted Budgets for Cost Accuracy
Token counts alone do not capture cost differences between input and output tokens, or between different models. A more sophisticated approach uses weighted budgets where each token type has a cost weight. This lets you express budgets in terms of actual dollars or credits rather than raw token counts.
@dataclass
class WeightedBudgetTracker:
total_credits: float
used_credits: float = 0
# Cost per 1K tokens for different categories
pricing: dict = field(default_factory=lambda: {
"gpt-4o_input": 0.0025,
"gpt-4o_output": 0.01,
"gpt-4o-mini_input": 0.00015,
"gpt-4o-mini_output": 0.0006,
})
@property
def remaining_credits(self) -> float:
return self.total_credits - self.used_credits
def consume(
self,
input_tokens: int,
output_tokens: int,
model: str,
stage: str = "default",
):
input_cost = (input_tokens / 1000) * self.pricing[f"{model}_input"]
output_cost = (output_tokens / 1000) * self.pricing[f"{model}_output"]
self.used_credits += input_cost + output_cost
def check(self, estimated_input: int, max_output: int, model: str):
est_cost = (
(estimated_input / 1000) * self.pricing[f"{model}_input"]
+ (max_output / 1000) * self.pricing[f"{model}_output"]
)
if est_cost > self.remaining_credits:
raise BudgetExceededError(
requested=int(est_cost * 1000),
remaining=int(self.remaining_credits * 1000),
message=f"Estimated cost ${est_cost:.4f} exceeds remaining ${self.remaining_credits:.4f}",
)
With weighted budgets, you can set a budget like "no more than $0.05 per request" and the tracker will enforce it regardless of which models or token ratios are involved. This is the approach you should use in production.
Best Practices
Always use the provider's usage data for final accounting. Your local token estimates are good for pre-flight checks, but the API response's usage field is what you will be billed for. Small discrepancies between your tokenizer and the provider's can compound over millions of requests.
Set budgets at multiple levels. A per-request budget prevents individual requests from running wild, but you also need per-tenant and per-day budgets to protect against aggregate cost spikes. Implement these as nested trackers that are checked in sequence.
Log budget reports for every request. The breakdown by stage is invaluable for debugging cost issues. If you see that 80% of your tokens are going to retrieval, you know where to optimize. Store these reports alongside your request logs for later analysis.
Pre-compute and cache token counts for static content. System prompts, few-shot examples, and other fixed text do not change between requests. Count them once at startup and reuse the count rather than re-tokenizing on every call.
Use streaming to enforce output budgets in real time. When you use streaming responses, you can count output tokens as they arrive and abort the stream if the budget is exceeded. This prevents a model that rambles from consuming more output tokens than you allocated.
def stream_with_budget(
messages: list[dict],
budget: BudgetTracker,
counter: TokenCounter,
model: str,
max_output: int,
):
estimated_input = budget.estimate_input(messages, counter)
budget.check(estimated_input + max_output)
stream = openai.OpenAI().chat.completions.create(
model=model,
messages=messages,
max_tokens=max_output,
stream=True,
stream_options={"include_usage": True},
)
output_text = ""
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
output_text += chunk.choices[0].delta.content
yield chunk.choices[0].delta.content
# The final chunk includes usage data
if chunk.usage:
budget.consume(
input_tokens=chunk.usage.prompt_tokens,
output_tokens=chunk.usage.completion_tokens,
stage="streaming_generation",
)
Test with pathological inputs. Deliberately feed your pipeline extremely long inputs, empty inputs, and inputs designed to trigger maximum retrieval. Verify that the budget enforcement catches these cases before they reach the API. This is the equivalent of load testing, but for cost.
Review and adjust budgets regularly. As your prompts evolve and your user base grows, the appropriate budget will change. Set up dashboards that show actual token usage distributions across requests, and adjust your budgets to cover the 99th percentile without being wasteful.
Conclusion
Building cost-aware LLM pipelines with token budgets is not optional for any application that runs in production. By treating token consumption as a measurable, enforceable resource — with counters, trackers, and stage-level allocation — you gain predictability, protection, and the observability needed to optimize over time. The implementation patterns in this tutorial give you a foundation: start with a simple per-request budget, add weighted cost tracking as you move toward production, and layer in multi-level budgets and fallback strategies as your application grows. The result is a pipeline that delivers reliable answers without delivering surprise bills.