← AI Agents Production Engineering
🤖 AI Agents

Agent Cost Attribution and Token Tracking Guide

Source: mortalapps.com
TL;DR
  • Cost attribution is the process of mapping LLM token consumption to specific business units, users, or projects using metadata.
  • It solves the visibility gap where a single API key masks multi-departmental spend, preventing accurate ROI analysis.
  • Production systems require this for financial governance, internal chargebacks, and identifying runaway agent loops.
  • Implementation relies on OpenTelemetry span attributes to inject business context into the telemetry pipeline.
  • A key tradeoff involves the overhead of high-cardinality metadata versus the granularity of financial reporting.

Why This Matters

In the experimental phase, AI agent costs are often treated as a centralized research expense. However, as agentic workflows move into production, organizations face the 'black box' billing problem. LLM providers typically bill at the API key or account level, providing no native mechanism to distinguish between a customer support agent and a sales automation agent. Without granular ai agent token cost attribution, engineering leaders cannot calculate the unit economics of their features or implement internal chargeback models. This approach was developed to bridge the gap between technical execution and financial accountability. If ignored, production systems risk unmonitored budget depletion, as a single recursive agent loop or a high-volume batch process can incur thousands of dollars in costs before a monthly bill arrives. Furthermore, without attribution, it is impossible to identify which specific prompts or models are driving the highest costs, making optimization efforts a matter of guesswork. Implementing a robust attribution framework using OpenTelemetry ensures that every token is accounted for at the point of generation, allowing teams to scale AI operations with the same financial rigor applied to cloud infrastructure. This is superior to using multiple API keys, which introduces significant management overhead and complicates rate limit handling across a shared model pool.

Core Concepts

To build a cost attribution system, engineers must master several architectural concepts that link runtime execution to financial data.

  • Span Attributes: Key-value pairs attached to OpenTelemetry spans. For cost tracking, these include gen_ai.usage.input_tokens, gen_ai.usage.output_tokens, and custom business keys like app.business_unit.
  • Context Propagation: The mechanism by which metadata (like a Team ID) is passed through a distributed system, ensuring that even downstream tool calls are attributed to the original requester.
  • Cardinality: The number of unique values in a dataset. High-cardinality attributes (like user_id) can significantly increase the storage requirements and cost of observability platforms.
  • Chargeback vs. Showback: Chargeback involves actual internal billing between departments; showback provides visibility without the financial transfer. Both require the same attribution data.
  • Token Normalization: The process of converting raw token counts from different providers (OpenAI, Anthropic, Google) into a standardized cost metric based on current pricing tables.

How It Works

The lifecycle of a cost-attributed agent request follows a structured path from the application layer to the financial dashboard.

1. Context Injection

When a request enters the agent system, the application middleware extracts business metadata (e.g., org_id, project_code) from the incoming JWT or session context. This metadata is injected into the active OpenTelemetry context, ensuring it is available to all subsequent spans in the trace.

2. Instrumented LLM Interaction

The agent performs its reasoning loop. Each call to an LLM is wrapped in an instrumented client. When the LLM returns a response, the client extracts the usage statistics (prompt, completion, and total tokens) and attaches them as attributes to the current span, alongside the business metadata injected in step one.

3. Telemetry Export and Collection

The OpenTelemetry SDK batches these spans and exports them to an OTel Collector. The collector can be configured to perform 'processor' logic, such as dropping sensitive prompt text while retaining the token counts and attribution tags to minimize storage costs.

4. Aggregation and Cost Mapping

Telemetry data flows into a time-series database (e.g., Prometheus) or an OLAP store (e.g., ClickHouse). A separate 'Pricing Service' maintains a lookup table of current model prices per 1k tokens. A scheduled job or real-time query joins the token counts with the pricing table, grouped by the attribution tags.

5. Failure Paths and Edge Cases

If an LLM call fails or times out, usage data may be missing. The system must handle 'partial attribution' where the trace exists but token counts are null. Similarly, if the OTel Collector is unreachable, the agent should continue to function (non-blocking), but a local buffer or retry logic is required to prevent 'dark spend' - untracked costs that occur during telemetry outages.

Architecture

The architecture consists of four primary layers. First, the Application Layer hosts the AI agent and uses an OTel-instrumented SDK to manage context. Second, the Telemetry Pipeline utilizes an OpenTelemetry Collector as a gateway to receive, process, and route spans. Third, the Storage Layer uses a high-performance database like ClickHouse or a managed observability platform to store span attributes. Finally, the Analytics Layer consists of a pricing engine and a visualization tool (like Grafana). Data flows from the agent to the collector via OTLP (OpenTelemetry Protocol). The collector enriches the data before sending it to storage. The analytics layer then queries the storage, applying model-specific pricing logic to the aggregated token counts to generate cost-per-team reports.

Implementing Custom Span Attributes for Attribution

Effective attribution starts with standardizing the attributes used across all agents. While OpenTelemetry provides the gen_ai semantic conventions, business-specific tags must be added manually. Engineers should use a consistent prefix, such as attr., to distinguish attribution metadata from operational metadata. For example, attr.cost_center, attr.environment, and attr.feature_id. These should be set at the start of the trace and propagated using the baggage header if the agent calls external microservices.

Handling Streaming Token Counts

One of the most significant technical hurdles in agent cost tracking is streaming responses. Many LLM providers do not return token counts in the final chunk of a stream by default (e.g., OpenAI requires stream_options={'include_usage': True}). Without this configuration, the telemetry layer will record zero tokens for every streaming request. The instrumentation logic must be designed to intercept the final usage object in the stream and update the span attributes retrospectively before the span is closed.

The Pricing Engine: Decoupling Cost from Telemetry

It is a common mistake to calculate the dollar cost at the time of the LLM call and store it as a span attribute. This is an anti-pattern because LLM pricing changes frequently (e.g., price drops for GPT-4o). Instead, store the raw token counts and the model string (e.g., gpt-4o-2024-05-13). The pricing engine should be a separate service or a SQL view that applies the current price to the historical token data. This allows for 'back-billing' or recalculating costs if a provider offers a retroactive discount or if a pricing error is discovered.

Aggregation Strategies: Prometheus vs. OLAP

Choosing the right storage for attribution data depends on the required granularity. Prometheus is excellent for real-time alerting (e.g., 'Team A has spent $50 in the last hour'), but it struggles with high-cardinality data like user_id. For detailed monthly chargeback reports involving thousands of users, an OLAP database like ClickHouse or BigQuery is preferred. These systems can ingest millions of spans and perform complex joins with organizational metadata (e.g., mapping cost_center IDs to department names) in seconds.

Multi-Model Normalization

In a production environment, agents often use a mix of models (e.g., Haiku for fast classification, Opus for complex reasoning). The attribution system must normalize these into a common currency. This involves tracking not just 'tokens' but 'input_tokens' and 'output_tokens' separately, as output tokens are typically 3-4x more expensive. Furthermore, cached tokens (like those in Anthropic or OpenAI) must be tracked using specific attributes like gen_ai.usage.input_cache_hit_tokens to ensure the attribution reflects the actual discounted price paid to the provider.

Code Example

Instrumenting an OpenAI call with custom attribution metadata using OpenTelemetry.
Python
import os
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from openai import OpenAI

# Initialize OpenTelemetry
trace.set_tracer_provider(TracerProvider())
tracer = trace.get_tracer(__name__)
trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

def run_attributed_agent(prompt, team_id, project_id):
    with tracer.start_as_current_span("agent_execution") as span:
        # Inject attribution metadata
        span.set_attribute("attr.team_id", team_id)
        span.set_attribute("attr.project_id", project_id)
        
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            # Ensure usage is tracked even if streaming (if applicable)
            stream=False
        )
        
        # Record token usage on the span
        usage = response.usage
        span.set_attribute("gen_ai.usage.input_tokens", usage.input_tokens)
        span.set_attribute("gen_ai.usage.output_tokens", usage.output_tokens)
        span.set_attribute("gen_ai.response.model", "gpt-4o")
        
        return response.choices[0].message.content

# Example usage
output = run_attributed_agent("Analyze this budget.", "finance-dept", "q4-audit")
Expected Output
The code executes the LLM call and prints an OpenTelemetry span to the console containing attributes: attr.team_id='finance-dept', attr.project_id='q4-audit', and the specific token counts returned by OpenAI.

Key Takeaways

Token attribution is essential for moving AI agents from research to a sustainable production business model.
OpenTelemetry is the industry-standard framework for capturing and propagating cost-related metadata.
Always separate raw token metrics from the volatile pricing logic of LLM providers.
Streaming responses require specific configuration to ensure usage data is not lost.
High-cardinality attribution data should be stored in OLAP databases rather than standard time-series metrics stores.
Context propagation ensures that costs are correctly attributed even in complex, multi-service agent workflows.
Effective cost tracking enables proactive alerting for runaway agent loops, preventing financial surprises.

Frequently Asked Questions

What is the difference between token tracking and cost attribution? +
Token tracking measures the raw volume of input and output data processed by an LLM. Cost attribution maps those volumes to specific business entities (like a team or project) using metadata to determine who is responsible for the spend.
When should I avoid using high-cardinality tags like user_id in my spans? +
Avoid using user_id in Prometheus-based metrics, as it can cause a 'label cardinality explosion' that degrades performance. Instead, store user-level attribution in an OLAP database like ClickHouse or a logging system.
How do I track costs for agents using local models like Llama 3? +
For local models, the 'cost' is usually compute (GPU/vCPU time) rather than tokens. You should still track tokens to measure efficiency, but your attribution model will map usage to infrastructure depreciation or cloud instance hourly rates.
What happens if the OpenTelemetry collector goes down? +
If the collector is down, the agent should continue to function to avoid downtime. However, unless you have local buffering or a retry mechanism in your SDK, you will lose the token data for that period, resulting in 'dark spend'.
Can I use multiple API keys instead of OTel attribution? +
While possible for 2-3 teams, it doesn't scale. Multiple keys complicate rate limit management, require complex secret rotation, and don't provide the granular 'per-feature' or 'per-user' visibility that OTel attributes offer.
How does this work with LangChain or LangGraph? +
Most frameworks support OpenTelemetry via integration libraries (like OpenInference). You can wrap the framework's provider calls and use context managers to inject your attribution tags into the underlying OTel spans.
What is the best way to handle LLM price changes? +
Store the model name and version (e.g., 'gpt-4o-2024-08-06') in your spans. Maintain a separate pricing table. Join these at query time so you can update prices retroactively without changing your telemetry data.
Does token tracking include the cost of vector database queries? +
Standard LLM token tracking does not. To get a full 'Agent Cost,' you must also instrument your vector DB calls (e.g., Pinecone, Weaviate) with similar attribution tags to capture read/write unit costs.