app/billing/pricing.py

"""
Pricing table for LLM providers.
Used to calculate estimated costs from token usage.
"""
from typing import Dict, Optional

# Pricing per 1M tokens (as of 2024, update as needed)
PRICING_TABLE: Dict[str, Dict[str, float]] = {
    "gemini": {
        "gemini-pro": {"input": 0.50, "output": 1.50},  # $0.50/$1.50 per 1M tokens
        "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
        "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
        "gemini-1.0-pro": {"input": 0.50, "output": 1.50},
        "default": {"input": 0.50, "output": 1.50}
    },
    "openai": {
        "gpt-4": {"input": 30.00, "output": 60.00},
        "gpt-4-turbo": {"input": 10.00, "output": 30.00},
        "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
        "default": {"input": 0.50, "output": 1.50}
    }
}


def calculate_cost(
    provider: str,
    model: str,
    prompt_tokens: int,
    completion_tokens: int
) -> float:
    """
    Calculate estimated cost in USD based on token usage.
    
    Args:
        provider: "gemini" or "openai"
        model: Model name (e.g., "gemini-pro", "gpt-3.5-turbo")
        prompt_tokens: Number of input tokens
        completion_tokens: Number of output tokens
        
    Returns:
        Estimated cost in USD
    """
    provider_pricing = PRICING_TABLE.get(provider.lower(), {})
    model_pricing = provider_pricing.get(model.lower(), provider_pricing.get("default", {"input": 0.50, "output": 1.50}))
    
    # Calculate cost: (tokens / 1M) * price_per_1M
    input_cost = (prompt_tokens / 1_000_000) * model_pricing["input"]
    output_cost = (completion_tokens / 1_000_000) * model_pricing["output"]
    
    return input_cost + output_cost


def get_model_pricing(provider: str, model: str) -> Dict[str, float]:
    """Get pricing for a specific model."""
    provider_pricing = PRICING_TABLE.get(provider.lower(), {})
    return provider_pricing.get(model.lower(), provider_pricing.get("default", {"input": 0.50, "output": 1.50}))