Usage

Token usage statistics

The Usage class tracks token consumption for AI requests.

Definition

@dataclass
class Usage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

Fields

prompt_tokens · int

Number of tokens in the input (prompt + messages).

completion_tokens · int

Number of tokens in the AI’s response.

total_tokens · int

Total tokens used (prompt_tokens + completion_tokens).

Usage

From generate_text

from ai_query import generate_text
from ai_query.providers import openai

result = await generate_text(
    model=openai("gpt-4o"),
    prompt="Explain Python in one sentence."
)

print(f"Input tokens: {result.usage.prompt_tokens}")
print(f"Output tokens: {result.usage.completion_tokens}")
print(f"Total tokens: {result.usage.total_tokens}")

From stream_text

from ai_query import stream_text
from ai_query.providers import google

result = stream_text(
    model=google("gemini-2.0-flash"),
    prompt="Write a haiku."
)

async for chunk in result.text_stream:
    print(chunk, end="")

# Usage is available after streaming
usage = await result.usage
print(f"\nTotal tokens: {usage.total_tokens}")

Multi-Step Usage

When using tools, result.usage contains the final model step’s usage. Each model call’s usage is available on its corresponding step:

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="Research this topic.",
    tools=tools
)

# Usage reported by the final model call
print(f"Final step tokens: {result.usage.total_tokens}")

# Provider-reported usage for every model call
for i, step in enumerate(result.steps):
    print(f"Step {i+1}: {step.usage.total_tokens} tokens")

Cost Estimation

Use token counts to estimate costs:

# Example pricing (check provider for current rates)
PRICE_PER_1K_INPUT = 0.01   # $0.01 per 1K input tokens
PRICE_PER_1K_OUTPUT = 0.03  # $0.03 per 1K output tokens

def estimate_cost(usage: Usage) -> float:
    input_cost = (usage.prompt_tokens / 1000) * PRICE_PER_1K_INPUT
    output_cost = (usage.completion_tokens / 1000) * PRICE_PER_1K_OUTPUT
    return input_cost + output_cost

result = await generate_text(model=openai("gpt-4o"), prompt="Hello")
cost = estimate_cost(result.usage)
print(f"Estimated cost: ${cost:.4f}")