Token Usage Tracking

Why Token Usage Matters

Every call to a language model or agent consumes tokens — the units of text the model reads (input tokens) and generates (output tokens). AI Platform tracks token usage per inference call so you can attribute the cost of running evaluations and generating datasets to specific models and runs.

Model detail page showing token usage summary and a breakdown by evaluation, dataset, and purpose.

How AI Platform Tracks Usage

AI Platform reads token usage directly from the model output. Every model output type supports an optional usage field (ModelUsage) — once you populate it, AI Platform records it automatically with no additional configuration.

AI Platform does not read token counts from the provider’s raw response directly, because each provider uses a different response shape. You are responsible for mapping the provider’s token fields to AI Platform’s schema in your adapter or inference snippet.

How to Report Usage

How you populate usage depends on how your model is integrated: via an output adapter, a custom inference snippet with the identity adapter, or an agent inference snippet returning an Open Responses trace.

Output Adapter

If your model uses a model adapter, add a usage block to the output Jinja template. The template has access to the provider’s raw response — map the provider’s token fields to num_prompt_tokens and num_completion_tokens. This applies whether the adapter is used alone or alongside a custom inference snippet.

The following example targets an OpenAI-compatible provider. Field names vary by provider — inspect the raw response to find the correct keys for yours:

model_adapters/openai_chat_completion_output.jinja
{% if response.usage is defined and response.usage is not none %}
    ,"usage": {
        "num_prompt_tokens": {{ response.usage.prompt_tokens }},
        "num_completion_tokens": {{ response.usage.completion_tokens }}
    }
{% endif %}

Custom Inference Snippet with the Identity Adapter

If your model uses a custom inference snippet with the identity adapter (latticeflow$identity), no Jinja template runs — the snippet’s return value is the final model output. Populate usage directly in the returned object.

The following example uses ChatCompletionModelOutput — replace the token counts with values extracted from the provider’s response:

run_inference.py
from typing import Any
from latticeflow.core.dtypes import ChatCompletionModelOutput, ModelUsage

def run_inference(body: str, environment: dict[str, Any]) -> str:
    output = ChatCompletionModelOutput(
        choices=[...],
        usage=ModelUsage(num_prompt_tokens=42, num_completion_tokens=17),
    )
    return output.model_dump_json()

Agents and Custom AI Systems (Open Responses)

For agents and custom AI systems that return an OpenResponsesModelOutput trace, pass usage alongside items when constructing the output object. Extract the token counts from the provider’s response before building ModelUsage — for example from a message_end event in a streaming response.

The following example uses OpenResponsesModelOutput — replace the token counts with values extracted from the provider’s response:

run_inference.py
from typing import Any
from latticeflow.core.dtypes import ModelUsage, OpenResponsesModelOutput

def run_inference(body: str, environment: dict[str, Any]) -> str:
    output = OpenResponsesModelOutput(
        items=[...],
        usage=ModelUsage(num_prompt_tokens=42, num_completion_tokens=17),
    )
    return output.model_dump_json()

Validate Integration

After adding usage, verify it is reported correctly before running evaluations. Run lf test model with your model’s key — in the output, look for the usage field under Model output in the format expected by LatticeFlow AI:

lf test model openai-gpt-4-1-nano
4. Transforming model output.
...
Model output in the format expected by LatticeFlow AI:
output='{"choices":[...],"usage":{"num_completion_tokens":9,"num_prompt_tokens":9}}'
==============================================================
Successfully tested configuration of model with key 'openai-gpt-4-1-nano'.

Alternatively, open the model’s detail page and click Test Connection. Click Run, then open the Transform Output stage and check the Model Output section — usage should appear with non-zero num_prompt_tokens and num_completion_tokens.

Transform Output stage showing usage in the Model Output section

If usage is absent, verify that the provider’s raw response includes token counts — if it does not, token tracking is not supported by that provider endpoint. If it does, check that your integration code is correctly mapping to the field names expected by AI Platform.