UsageTracker

Track token usage across requests.

Quick Example

from mamba_agents.tokens import UsageTracker

tracker = UsageTracker()

# Record usage
tracker.record_usage(
    prompt_tokens=100,
    completion_tokens=50,
    model="gpt-4o",
)

# Get aggregate
usage = tracker.get_total_usage()
print(f"Total: {usage.total_tokens}")

# Get history
for record in tracker.get_usage_history():
    print(f"{record.timestamp}: {record.total_tokens}")

# Reset
tracker.reset()

Data Classes

TokenUsage

Aggregate usage statistics.

UsageRecord

Single usage record.

API Reference

UsageTracker

UsageTracker(cost_rates: dict[str, float] | None = None)

Track token usage across requests.

Provides per-request tracking, session aggregates, and cost estimation.

Initialize the usage tracker.

PARAMETER	DESCRIPTION
cost_rates	Optional cost per 1000 tokens for different models. TYPE: dict[str, float] | None DEFAULT: None

Source code in src/mamba_agents/tokens/tracker.py

def __init__(
    self,
    cost_rates: dict[str, float] | None = None,
) -> None:
    """Initialize the usage tracker.

    Args:
        cost_rates: Optional cost per 1000 tokens for different models.
    """
    self._records: list[UsageRecord] = []
    self._totals = TokenUsage()
    self._cost_rates = cost_rates or {}

record_usage

record_usage(
    usage: Usage,
    model: str | None = None,
    tool_name: str | None = None,
) -> None

Record usage from a pydantic-ai result.

PARAMETER	DESCRIPTION
usage	Usage object from pydantic-ai. TYPE: Usage
model	Optional model name. TYPE: str | None DEFAULT: None
tool_name	Optional tool name for tool calls. TYPE: str | None DEFAULT: None

Source code in src/mamba_agents/tokens/tracker.py

def record_usage(
    self,
    usage: Usage,
    model: str | None = None,
    tool_name: str | None = None,
) -> None:
    """Record usage from a pydantic-ai result.

    Args:
        usage: Usage object from pydantic-ai.
        model: Optional model name.
        tool_name: Optional tool name for tool calls.
    """
    # Extract token counts from pydantic-ai Usage
    # Use new API (input_tokens/output_tokens) with fallback to deprecated names
    prompt_tokens = (
        getattr(usage, "input_tokens", None) or getattr(usage, "request_tokens", None) or 0
    )
    completion_tokens = (
        getattr(usage, "output_tokens", None) or getattr(usage, "response_tokens", None) or 0
    )
    total_tokens = usage.total_tokens or (prompt_tokens + completion_tokens)

    record = UsageRecord(
        timestamp=datetime.now(),
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        model=model,
        tool_name=tool_name,
    )

    self._records.append(record)

    # Update totals
    self._totals.prompt_tokens += prompt_tokens
    self._totals.completion_tokens += completion_tokens
    self._totals.total_tokens += total_tokens
    self._totals.request_count += 1

record_raw

record_raw(
    prompt_tokens: int,
    completion_tokens: int,
    model: str | None = None,
    tool_name: str | None = None,
) -> None

Record raw token counts.

PARAMETER	DESCRIPTION
prompt_tokens	Number of prompt tokens. TYPE: int
completion_tokens	Number of completion tokens. TYPE: int
model	Optional model name. TYPE: str | None DEFAULT: None
tool_name	Optional tool name. TYPE: str | None DEFAULT: None

Source code in src/mamba_agents/tokens/tracker.py

def record_raw(
    self,
    prompt_tokens: int,
    completion_tokens: int,
    model: str | None = None,
    tool_name: str | None = None,
) -> None:
    """Record raw token counts.

    Args:
        prompt_tokens: Number of prompt tokens.
        completion_tokens: Number of completion tokens.
        model: Optional model name.
        tool_name: Optional tool name.
    """
    total = prompt_tokens + completion_tokens

    record = UsageRecord(
        timestamp=datetime.now(),
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total,
        model=model,
        tool_name=tool_name,
    )

    self._records.append(record)

    self._totals.prompt_tokens += prompt_tokens
    self._totals.completion_tokens += completion_tokens
    self._totals.total_tokens += total
    self._totals.request_count += 1

get_total_usage

get_total_usage() -> TokenUsage

Get total usage statistics.

RETURNS	DESCRIPTION
TokenUsage	TokenUsage with aggregate statistics.

Source code in src/mamba_agents/tokens/tracker.py

def get_total_usage(self) -> TokenUsage:
    """Get total usage statistics.

    Returns:
        TokenUsage with aggregate statistics.
    """
    return self._totals

get_usage_history

get_usage_history() -> list[UsageRecord]

Get usage history.

RETURNS	DESCRIPTION
list[UsageRecord]	List of all usage records.

Source code in src/mamba_agents/tokens/tracker.py

def get_usage_history(self) -> list[UsageRecord]:
    """Get usage history.

    Returns:
        List of all usage records.
    """
    return self._records.copy()

get_cost_estimate

get_cost_estimate(model: str | None = None) -> float

Estimate cost based on usage.

PARAMETER	DESCRIPTION
model	Model to use for cost calculation. If None, uses default rate. TYPE: str | None DEFAULT: None

RETURNS	DESCRIPTION
float	Estimated cost in USD.

Source code in src/mamba_agents/tokens/tracker.py

def get_cost_estimate(self, model: str | None = None) -> float:
    """Estimate cost based on usage.

    Args:
        model: Model to use for cost calculation.
               If None, uses default rate.

    Returns:
        Estimated cost in USD.
    """
    # Get rate per 1000 tokens
    rate = 0.0
    if model and model in self._cost_rates:
        rate = self._cost_rates[model]
    elif "default" in self._cost_rates:
        rate = self._cost_rates["default"]

    # Calculate cost
    return (self._totals.total_tokens / 1000) * rate

get_breakdown_by_tool

get_breakdown_by_tool() -> dict[str, TokenUsage]

Get token usage broken down by tool.

RETURNS	DESCRIPTION
dict[str, TokenUsage]	Dictionary mapping tool names to usage.

Source code in src/mamba_agents/tokens/tracker.py

def get_breakdown_by_tool(self) -> dict[str, TokenUsage]:
    """Get token usage broken down by tool.

    Returns:
        Dictionary mapping tool names to usage.
    """
    breakdown: dict[str, TokenUsage] = {}

    for record in self._records:
        key = record.tool_name or "_agent"

        if key not in breakdown:
            breakdown[key] = TokenUsage()

        breakdown[key].prompt_tokens += record.prompt_tokens
        breakdown[key].completion_tokens += record.completion_tokens
        breakdown[key].total_tokens += record.total_tokens
        breakdown[key].request_count += 1

    return breakdown

reset

reset() -> None

Reset all tracking data.

Source code in src/mamba_agents/tokens/tracker.py

def reset(self) -> None:
    """Reset all tracking data."""
    self._records.clear()
    self._totals = TokenUsage()

TokenUsage dataclass

TokenUsage(
    prompt_tokens: int = 0,
    completion_tokens: int = 0,
    total_tokens: int = 0,
    request_count: int = 0,
    cached_tokens: int = 0,
)

Aggregate token usage statistics.

ATTRIBUTE	DESCRIPTION
prompt_tokens	Total prompt tokens. TYPE: int
completion_tokens	Total completion tokens. TYPE: int
total_tokens	Total tokens. TYPE: int
request_count	Number of requests. TYPE: int
cached_tokens	Tokens served from cache. TYPE: int

UsageRecord dataclass

UsageRecord(
    timestamp: datetime,
    prompt_tokens: int,
    completion_tokens: int,
    total_tokens: int,
    model: str | None = None,
    tool_name: str | None = None,
)

A single usage record.

ATTRIBUTE	DESCRIPTION
timestamp	When the usage was recorded. TYPE: datetime
prompt_tokens	Tokens in the prompt. TYPE: int
completion_tokens	Tokens in the completion. TYPE: int
total_tokens	Total tokens used. TYPE: int
model	Model used for this request. TYPE: str | None
tool_name	Optional tool name if tool call. TYPE: str | None