Track AI agent token cost in Python: usage, budget, stop

Your agent ran for ten minutes and the bill is bigger than the task. To see why, read the usage object in every API response, add it up per task and multiply by the model's price per 1M tokens. To cap it, compare that running total with a budget after each step and stop the loop when it is crossed. Below is a small Python class that does exactly this.

The example uses the OpenAI-compatible endpoint https://api.cheaptokens.dev/v1 and the model id deepseek-ai/DeepSeek-V4-Flash-0731 (380K context, up to 32,768 output tokens per answer). If you want to estimate a run before starting it, that is a different job, and How to estimate the cost of an AI agent covers it. This article is about measuring and capping while the run is going.

Where the numbers are: the usage object#

Every non-streaming chat completion comes back with a usage object. DeepSeek's API reference (checked 26 October 2026) lists these fields:

• prompt_tokens — everything you sent this time: system prompt, history, tool results.
• completion_tokens — what the model wrote back.
• total_tokens — the two added together.

DeepSeek also documents optional extras: prompt_cache_hit_tokens and prompt_cache_miss_tokens (they add up to prompt_tokens), and completion_tokens_details.reasoning_tokens for thinking models. The page shows no publication date, and these extras are extensions of the OpenAI format, so the code below treats them as optional and relies only on the first two fields.

Why log every step and not just the last one? Because an agent re-sends its growing context each time. Step 1 might send 2,000 input tokens, step 30 might send 9,000, and the bill is the sum of all of them. The reason is explained in why a long chat costs more, and the input/output split in input vs output tokens.

A meter that stops the agent#

Here is the whole thing: a meter that adds up tokens, prints one line per step and raises an error when the budget is crossed, plus a bare agent loop around it. Install the SDK with pip install openai and put your key in CHEAPTOKENS_API_KEY.

import os
from openai import OpenAI

client = OpenAI(
    base_url='https://api.cheaptokens.dev/v1',
    api_key=os.environ['CHEAPTOKENS_API_KEY'],
)
MODEL = 'deepseek-ai/DeepSeek-V4-Flash-0731'

# example prices, USD per 1M tokens: use the current ones from the pricing table
PRICE_IN = 0.10
PRICE_OUT = 0.40


class BudgetExceeded(Exception):
    pass


class TaskMeter:
    def __init__(self, budget_usd):
        self.budget = budget_usd
        self.steps = 0
        self.tokens_in = 0
        self.tokens_out = 0
        self.reasoning = 0

    @property
    def cost(self):
        return (self.tokens_in * PRICE_IN + self.tokens_out * PRICE_OUT) / 1_000_000

    def add(self, usage):
        self.steps += 1
        self.tokens_in += usage.prompt_tokens
        self.tokens_out += usage.completion_tokens
        details = getattr(usage, 'completion_tokens_details', None)
        self.reasoning += getattr(details, 'reasoning_tokens', 0) or 0
        print(f'step {self.steps}: in={usage.prompt_tokens} '
              f'out={usage.completion_tokens} task=${self.cost:.4f}')
        if self.cost > self.budget:
            raise BudgetExceeded(
                f'budget ${self.budget:.2f} crossed at step {self.steps}: '
                f'${self.cost:.4f}')


def run_task(goal, budget_usd=0.50, max_steps=40):
    meter = TaskMeter(budget_usd)
    messages = [
        {'role': 'system', 'content': 'Work step by step. Reply DONE when finished.'},
        {'role': 'user', 'content': goal},
    ]
    for _ in range(max_steps):
        resp = client.chat.completions.create(
            model=MODEL, messages=messages, max_tokens=1024)
        meter.add(resp.usage)  # log first, then decide
        reply = resp.choices[0].message.content or ''
        messages.append({'role': 'assistant', 'content': reply})
        if 'DONE' in reply:
            break
        messages.append({'role': 'user', 'content': 'Continue.'})
    return meter


try:
    meter = run_task('Summarise the main risks of a database migration.')
    print(f'finished: {meter.steps} steps, ${meter.cost:.4f}')
except BudgetExceeded as err:
    print('stopped:', err)

Three choices in there are worth copying.

Log before you decide. meter.add records the step first and only then checks the budget, so a stopped run still leaves a complete trace. The check happens after the call that crossed the line, so you can overshoot by at most one step. max_tokens=1024 caps the output of that last step.

Two hard stops, not one. max_steps catches an agent that loops without spending much; the dollar budget catches one that spends a lot per step. You want both.

Prices live in one place. PRICE_IN and PRICE_OUT are placeholder numbers. Copy the current values for your model from the pricing page; right now DeepSeek-V4-Flash is $0.015 per 1M input tokens and $0.10 per 1M output tokens.

What a 40-step run adds up to#

Say your agent takes 40 steps, and because the context grows, an average step sends 6,000 input tokens and gets 400 tokens back. That is 40 × 6,000 = 240,000 input tokens and 40 × 400 = 16,000 output tokens for the task.

On DeepSeek-V4-Flash that costs $0.0052. The table runs the same workload on the other models and puts OpenRouter's reference price next to ours:

In total, with input and output counted together, DeepSeek-V4-Flash comes out cheaper than OpenRouter for this task. Notice where the tokens are, too: 240,000 of the 256,000 are input. The context your agent drags along is the bill, which is why the per-step in= number in your log is the one to watch. A budget of $0.50 per task in the example above is a cap on a runaway, not a target.

Switching models only means changing MODEL. Pick the id from the CheapTokens docs: the names differ between providers, and a short id from a price label will not work in the API. For long agent loops, DeepSeek-V4-Flash has one price at any hour of the day, so a task costs the same at noon and at 3 a.m. See the DeepSeek page for its limits.

When the model is thinking#

If an answer is much longer in tokens than the text you see, look at reasoning_tokens. Thinking tokens count as output, so a reasoning-heavy step can produce far more output tokens than the visible text suggests, and they are billed as output. The meter above collects them in meter.reasoning when the API returns the field and leaves it at 0 when it does not, so it never crashes on a missing key. Print it next to out= if you want to see how much of the output was thinking. Reasoning effort and thinking tokens explains why.

Streaming: usage comes last#

If your agent streams, there is no usage in the normal chunks. In the OpenAI format you ask for it with stream_options={'include_usage': True}. DeepSeek's reference (checked 26 October 2026) says the option works only together with stream=True (otherwise the API answers 400), that usage is null in every chunk but one, and that the last chunk before data: [DONE] carries the totals for the whole request. So in a streaming loop, call meter.add once, on the chunk where usage is not empty. The non-streaming loop above is the simplest place to start; streaming in Python has the loop itself.

If the numbers look wrong#

• usage is None. You are streaming without include_usage. Add it, or switch to a non-streaming call.
• The total is lower than the balance drop. Check that every call goes through the meter: retries, a summarising call or a tool-picking call made outside run_task are easy to miss. Wrap client.chat.completions.create in one function and route everything through it.
• Cost per step keeps climbing. That is normal for a loop that appends everything to messages. Trim old tool output or summarise history; tool output out of context shows the idea.
• BudgetExceeded fires on step 2. Your budget is smaller than one step. Check that PRICE_IN is per 1M tokens, not per 1K.

To store the numbers, write meter.steps, meter.tokens_in, meter.tokens_out and meter.cost as one JSON line per task. A week of those lines tells you the real tokens per step for your agent, which is the number to put into every estimate afterwards. Models, prices and the first request are in the docs; tool calls are covered in function calling in Python.