Usage

@dataclass
class Usage:
    input_tokens: int = 0
    output_tokens: int = 0
    cached_tokens: int = 0
    cache_write_tokens: int = 0
    cache_miss_tokens: int = 0
    total_tokens: int = 0

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.input_tokens}")
print(f"Output tokens: {result.usage.output_tokens}")
print(f"Cached input tokens: {result.usage.cached_tokens}")
print(f"Total tokens: {result.usage.total_tokens}")

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}")

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")

# 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.input_tokens / 1000) * PRICE_PER_1K_INPUT
    output_cost = (usage.output_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}")