CostEstimator

Estimate costs based on token usage.

Quick Example

from mamba_agents.tokens import CostEstimator

estimator = CostEstimator()

# Estimate cost
breakdown = estimator.estimate(
    prompt_tokens=1000,
    completion_tokens=500,
    model="gpt-4o",
)
print(f"Total: ${breakdown.total_cost:.4f}")

# Custom rates
estimator.set_rate("my-model", 0.002)

# Get rates
all_rates = estimator.get_all_rates()

Default Rates

Model	Input (per 1M)	Output (per 1M)
gpt-4o	$2.50	$10.00
gpt-4o-mini	$0.15	$0.60
claude-3-5-sonnet	$3.00	$15.00
Local models	$0.00	$0.00

API Reference

CostEstimator

CostEstimator(custom_rates: dict[str, float] | None = None)

Estimate costs based on token usage.

Provides cost estimation for various models with configurable rates.

Initialize the cost estimator.

PARAMETER	DESCRIPTION
custom_rates	Custom cost rates per 1000 tokens. Overrides defaults for specified models. TYPE: dict[str, float] | None DEFAULT: None

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

def __init__(
    self,
    custom_rates: dict[str, float] | None = None,
) -> None:
    """Initialize the cost estimator.

    Args:
        custom_rates: Custom cost rates per 1000 tokens.
                      Overrides defaults for specified models.
    """
    self._rates = DEFAULT_COST_RATES.copy()
    if custom_rates:
        self._rates.update(custom_rates)

get_rate

get_rate(model: str) -> float

Get the cost rate for a model.

PARAMETER	DESCRIPTION
model	Model name or identifier. TYPE: str

RETURNS	DESCRIPTION
float	Cost per 1000 tokens.

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

def get_rate(self, model: str) -> float:
    """Get the cost rate for a model.

    Args:
        model: Model name or identifier.

    Returns:
        Cost per 1000 tokens.
    """
    # Check exact match
    if model in self._rates:
        return self._rates[model]

    # Check prefix match
    model_lower = model.lower()
    for key, rate in self._rates.items():
        if key.lower() in model_lower:
            return rate

    # Return default
    return self._rates.get("default", 0.0)

estimate

estimate(
    usage: TokenUsage, model: str | None = None
) -> CostBreakdown

Estimate cost for token usage.

PARAMETER	DESCRIPTION
usage	Token usage to estimate cost for. TYPE: TokenUsage
model	Model to use for rate lookup. TYPE: str | None DEFAULT: None

RETURNS	DESCRIPTION
CostBreakdown	CostBreakdown with detailed costs.

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

def estimate(
    self,
    usage: TokenUsage,
    model: str | None = None,
) -> CostBreakdown:
    """Estimate cost for token usage.

    Args:
        usage: Token usage to estimate cost for.
        model: Model to use for rate lookup.

    Returns:
        CostBreakdown with detailed costs.
    """
    rate = self.get_rate(model or "default")

    prompt_cost = (usage.prompt_tokens / 1000) * rate
    completion_cost = (usage.completion_tokens / 1000) * rate
    total_cost = prompt_cost + completion_cost

    return CostBreakdown(
        prompt_cost=prompt_cost,
        completion_cost=completion_cost,
        total_cost=total_cost,
        model=model,
    )

estimate_tokens

estimate_tokens(
    token_count: int, model: str | None = None
) -> float

Estimate cost for a token count.

PARAMETER	DESCRIPTION
token_count	Number of tokens. TYPE: int
model	Model to use for rate lookup. TYPE: str | None DEFAULT: None

RETURNS	DESCRIPTION
float	Estimated cost in USD.

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

def estimate_tokens(
    self,
    token_count: int,
    model: str | None = None,
) -> float:
    """Estimate cost for a token count.

    Args:
        token_count: Number of tokens.
        model: Model to use for rate lookup.

    Returns:
        Estimated cost in USD.
    """
    rate = self.get_rate(model or "default")
    return (token_count / 1000) * rate

set_rate

set_rate(model: str, rate: float) -> None

Set a custom rate for a model.

PARAMETER	DESCRIPTION
model	Model name. TYPE: str
rate	Cost per 1000 tokens. TYPE: float

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

def set_rate(self, model: str, rate: float) -> None:
    """Set a custom rate for a model.

    Args:
        model: Model name.
        rate: Cost per 1000 tokens.
    """
    self._rates[model] = rate

get_all_rates

get_all_rates() -> dict[str, float]

Get all configured rates.

RETURNS	DESCRIPTION
dict[str, float]	Dictionary of model names to rates.

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

def get_all_rates(self) -> dict[str, float]:
    """Get all configured rates.

    Returns:
        Dictionary of model names to rates.
    """
    return self._rates.copy()

CostBreakdown dataclass

CostBreakdown(
    prompt_cost: float,
    completion_cost: float,
    total_cost: float,
    model: str | None = None,
)

Cost breakdown by category.

ATTRIBUTE	DESCRIPTION
prompt_cost	Cost for prompt tokens. TYPE: float
completion_cost	Cost for completion tokens. TYPE: float
total_cost	Total cost. TYPE: float
model	Model used for calculation. TYPE: str | None