Skip to content
Notifications
Clear all

Walkthrough: Adding observability (tracing, costs) to a production LangGraph agent

1 Posts
1 Users
0 Reactions
2 Views
(@brianw)
Estimable Member
Joined: 1 week ago
Posts: 72
Topic starter   [#20972]

Integrating comprehensive observability into a LangGraph agent deployed beyond a proof-of-concept stage presents a significant challenge, primarily due to the multi-step, non-linear nature of the workflows. While the built-in tracing offers a foundational view, it lacks the granularity required for true operational oversight and financial accountability. This walkthrough details a method to instrument a production LangGraph agent for enhanced tracing and, critically, accurate cost attribution across its various components, which may include multiple LLM calls, tool executions, and conditional pathways.

The core strategy involves intercepting the graph's execution stream to capture events, annotate them with business context, and forward them to a unified observability backend. This requires a custom `Checkpointer` and careful use of `Thread` and `Node` metadata. Below is a foundational configuration for a custom checkpointer that logs events and embeds cost data.

```python
from langgraph.checkpoint.base import BaseCheckpointSaver
from typing import Any, Dict, List
import uuid
from datetime import datetime

class ObservabilityCheckpointer(BaseCheckpointSaver):
def __init__(self, observability_client):
self.client = observability_client
self.run_id = str(uuid.uuid4())

def save_checkpoint(self, config: Dict[str, Any], step: int, metadata: Dict[str, Any]):
# Capture node start/end events from config
current_node = config.get("configurable", {}).get("node_name")
if current_node:
event = {
"run_id": self.run_id,
"timestamp": datetime.utcnow().isoformat(),
"node": current_node,
"step": step,
"thread_id": config.get("configurable", {}).get("thread_id"),
"metadata": metadata,
"event_type": "node_end"
}
# Enrich with cost if LLM call is present in metadata
if "lm_provider" in metadata:
event["cost"] = self._calculate_cost(
metadata["lm_provider"],
metadata.get("token_usage", {})
)
self.client.emit(event)
return {}

def _calculate_cost(self, provider: str, token_usage: Dict) -> float:
# Example cost mapping (per 1M tokens)
cost_map = {
"openai/gpt-4o-mini": {"input": 0.15, "output": 0.60},
"anthropic/claude-3-haiku": {"input": 0.80, "output": 2.40},
}
if provider not in cost_map:
return 0.0
rates = cost_map[provider]
input_cost = (token_usage.get("input_tokens", 0) / 1_000_000) * rates["input"]
output_cost = (token_usage.get("output_tokens", 0) / 1_000_000) * rates["output"]
return round(input_cost + output_cost, 6)
```

The implementation focuses on three key observability pillars:

* **Distributed Tracing:** The custom checkpointer generates a unique `run_id` for each graph invocation, tagging all subsequent events (node entries/exits, tool calls, LLM invocations). This allows reconstruction of the full execution graph, including conditional branches and loops, in a tool like Jaeger, Tempo, or a commercial APM.
* **Cost Attribution:** By intercepting metadata at each node—specifically after LLM calls—we can apply a cost calculation function. The critical step is ensuring token counts and model identifiers are propagated into the checkpointer's metadata. This enables cost breakdowns by:
* User or tenant ID (stored in `thread_id` or custom metadata)
* Specific graph path or branch taken
* Individual node or tool responsible for the expense
* **Performance Metrics:** Logging the timestamp at each node entry and exit allows for the calculation of node latency, which is essential for identifying bottlenecks in complex, multi-agent workflows.

To integrate this, you must bind the checkpointer to your graph and ensure your node functions publish the necessary metadata. For example, a node using a ChatModel should wrap the call to extract token usage.

```python
from langchain_core.messages import AIMessage
from langgraph.graph import StateGraph

def node_with_observability(state, config):
# Your node logic
messages = state["messages"]
# Assume chat_model is bound to the graph's context
response: AIMessage = config["context"]["chat_model"].invoke(messages)
# Critical: Attach cost-related metadata
metadata = {
"lm_provider": "openai/gpt-4o-mini",
"token_usage": {
"input_tokens": response.usage_metadata["input_tokens"],
"output_tokens": response.usage_metadata["output_tokens"]
} if hasattr(response, 'usage_metadata') else {}
}
return {"messages": [response]}, metadata

# During graph compilation
builder = StateGraph(...)
builder.add_node("analyst_node", node_with_observability)
...
graph = builder.compile(checkpointer=ObservabilityCheckpointer(observability_client=my_client))
```

The primary pitfalls of this approach are the increased complexity in state management and the potential performance overhead from synchronous logging. For high-volume applications, I recommend emitting events asynchronously to a queue (e.g., Redis or Kafka) to be processed by the observability backend. Furthermore, this model requires consistent metadata propagation across all nodes; any node that fails to report its token usage will create a gap in the cost analysis. The resulting data should be stored in a time-series database (like Prometheus) for real-time dashboards and in a columnar data store (like BigQuery) for detailed historical analysis and showback reporting by department or project.


Spreadsheets or it didn't happen.


   
Quote