# How ML Intern Tracks Token Usage Across the Agent Loop

> Discover how ML Intern tracks token usage across the agent loop using LLM call site instrumentation and Litellm responses for cost analysis.

- Repository: [Hugging Face/ml-intern](https://github.com/huggingface/ml-intern)
- Tags: internals
- Published: 2026-04-24

---

**ML Intern tracks token usage by instrumenting LLM call sites with a telemetry helper that extracts usage fields from Litellm responses, emits structured `llm_call` events to the session store, and aggregates them via post-processing scripts for cost analysis.**

The `huggingface/ml-intern` repository implements a granular telemetry system to monitor LLM consumption in real time. By intercepting model responses at the agent loop level, the system captures prompt tokens, completion tokens, and cache-related metrics, enabling precise cost attribution and performance analysis across autonomous agent executions.

## Normalizing Token Usage from Litellm Responses

Token tracking begins with normalization. In [`agent/core/telemetry.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/telemetry.py), the `extract_usage` function (lines 29-69) sanitizes the raw `usage` field from Litellm responses or plain dictionaries. It returns a stable dictionary containing `prompt_tokens`, `completion_tokens`, `total_tokens`, `cache_read_tokens`, and `cache_creation_tokens`, ensuring consistent field names regardless of the underlying model provider.

## Emitting Telemetry from the Agent Loop

The `record_llm_call` async function in [`agent/core/telemetry.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/telemetry.py) serves as the central instrumentation hook. It is invoked from two specific locations in [`agent/core/agent_loop.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/agent_loop.py): lines 390-401 for **streaming** responses and lines 474-482 for **non-streaming** responses. This dual instrumentation ensures that every LLM invocation, whether delivered via chunked streaming or complete response objects, triggers telemetry capture.

When called, `record_llm_call` executes the following steps:

1. Invokes `extract_usage` on the raw response to normalize token counts.
2. Calculates USD cost using `litellm.completion_cost`, defaulting to `0.0` if the calculation fails.
3. Constructs an `llm_call` event containing the model name, latency, finish reason, cost, and all token fields.
4. Sends the event to the current `Session` via `session.send_event`.

```python

# agent/core/telemetry.py

async def record_llm_call(
    session: Any,
    *,
    model: str,
    response: Any = None,
    latency_ms: int,
    finish_reason: str | None,
) -> dict:
    """Emit an ``llm_call`` event and return the extracted usage dict."""
    usage = extract_usage(response) if response is not None else {}
    cost_usd = 0.0
    if response is not None:
        try:
            from litellm import completion_cost
            cost_usd = float(completion_cost(completion_response=response) or 0.0)
        except Exception:
            cost_usd = 0.0
    from agent.core.session import Event
    try:
        await session.send_event(Event(
            event_type="llm_call",
            data={
                "model": model,
                "latency_ms": latency_ms,
                "finish_reason": finish_reason,
                "cost_usd": cost_usd,
                **usage,
            },
        ))
    except Exception as e:
        logger.debug("record_llm_call failed (non-fatal): %s", e)
    return usage

```

## Persisting Events to Session Storage

Once emitted, events are captured by `Session.send_event` in [`agent/core/session.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/session.py). The method appends each `llm_call` event to an in-memory `logged_events` list and forwards it to connected clients. When the session terminates, `Session.save_trajectory_local` persists the entire event stream—including all token usage records—to a JSON trajectory file, creating an immutable audit trail suitable for downstream analysis.

## Aggregating Tokens for KPI Reporting

Post-execution analysis is handled by [`scripts/build_kpis.py`](https://github.com/huggingface/ml-intern/blob/main/scripts/build_kpis.py). This script reads persisted session logs and iterates over every `llm_call` event (lines 42-48), extracting token fields from `event["data"]` to compute aggregate statistics. The aggregation loop sums `prompt_tokens`, `completion_tokens`, `cache_read_tokens`, and `cache_creation_tokens`, producing per-session and per-bucket KPIs for cost control and performance monitoring.

```python

# scripts/build_kpis.py (inside the per‑event loop)

if et == "llm_call":
    out["llm_calls"] += 1
    out["tokens_prompt"]      += int(data.get("prompt_tokens") or 0)
    out["tokens_completion"] += int(data.get("completion_tokens") or 0)
    out["tokens_cache_read"] += int(data.get("cache_read_tokens") or 0)
    out["tokens_cache_creation"] += int(data.get("cache_creation_tokens") or 0)
    out["cost_usd"]          += float(data.get("cost_usd") or 0.0)

```

### Accessing Token Totals from Saved Sessions

You can programmatically audit token consumption from any saved session file:

```python
import json, pathlib

session_file = pathlib.Path("session_logs/session_abc123_20240424_101530.json")
with session_file.open() as f:
    trajectory = json.load(f)

# Sum tokens across all LLM calls in this session

total_prompt = sum(ev["data"].get("prompt_tokens", 0)
                   for ev in trajectory["events"]
                   if ev["event_type"] == "llm_call")
print(f"Prompt tokens used: {total_prompt}")

```

## Summary

- **Normalization**: The `extract_usage` helper in [`agent/core/telemetry.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/telemetry.py) standardizes Litellm responses into a consistent schema that includes cache-related token fields.
- **Instrumentation**: Call sites in [`agent/core/agent_loop.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/agent_loop.py) (streaming at lines 390-401 and non-streaming at lines 474-482) invoke `record_llm_call` to capture usage at the moment of response receipt.
- **Cost Calculation**: Each `llm_call` event includes USD cost computed via `litellm.completion_cost`, with a safe fallback to `0.0` on failure.
- **Persistence**: Events are stored in-memory via `Session.send_event` and flushed to JSON via `Session.save_trajectory_local` in [`agent/core/session.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/session.py).
- **Aggregation**: The [`scripts/build_kpis.py`](https://github.com/huggingface/ml-intern/blob/main/scripts/build_kpis.py) utility sums token fields across sessions, enabling precise cost attribution and performance reporting.

## Frequently Asked Questions

### How does ML Intern handle token tracking for streaming responses versus non-streaming responses?

Both paths use identical instrumentation. The [`agent/core/agent_loop.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/agent_loop.py) file contains separate call sites for streaming (lines 390-401) and non-streaming (lines 474-482), but both invoke the same `telemetry.record_llm_call` method with the final response chunk or complete response object. This ensures token extraction and event emission remain consistent regardless of how the model delivers content.

### What specific token fields does ML Intern track beyond standard prompt and completion counts?

According to the implementation in [`agent/core/telemetry.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/telemetry.py), the system captures five distinct fields: `prompt_tokens`, `completion_tokens`, `total_tokens`, `cache_read_tokens`, and `cache_creation_tokens`. The latter two are particularly important for providers like Anthropic that support prompt caching, allowing the system to distinguish between novel and cached prompt content.

### Where is token usage physically stored in ML Intern?

Token usage is stored as structured events within the `logged_events` list of a `Session` object defined in [`agent/core/session.py`](https://github.com/huggingface/ml-intern/blob/main/agent/core/session.py). When the session ends, `save_trajectory_local` persists these events to a JSON file on disk, typically located in a `session_logs/` directory, creating a permanent record that [`scripts/build_kpis.py`](https://github.com/huggingface/ml-intern/blob/main/scripts/build_kpis.py) can later ingest for analysis.

### How does the telemetry system handle cost calculation failures?

The `record_llm_call` function wraps the `litellm.completion_cost` invocation in a try-except block. If the cost calculation raises an exception or returns `None`, the function defaults the `cost_usd` field to `0.0` and continues emitting the event. This defensive design ensures that token tracking remains robust and non-blocking even when cost metadata is unavailable or the Litellm integration encounters errors.