diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6c6a09cd73e..01f3e2472f8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -57,6 +57,7 @@ from litellm.llms.vertex_ai.cost_calculator import ( cost_per_token as google_cost_per_token, ) from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_router +from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import ( HttpxBinaryResponseContent, @@ -341,6 +342,8 @@ def cost_per_token( # noqa: PLR0915 return deepseek_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "perplexity": return perplexity_cost_per_token(model=model, usage=usage_block) + elif custom_llm_provider == "xai": + return xai_cost_per_token(model=model, usage=usage_block) else: model_info = _cached_get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider @@ -675,9 +678,9 @@ def completion_cost( # noqa: PLR0915 or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = ( - completion_response.get("usage", {}) - ) + usage_obj: Optional[ + Union[dict, Usage] + ] = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects( @@ -1279,7 +1282,9 @@ class BaseTokenUsageProcessor: not hasattr(combined, "completion_tokens_details") or not combined.completion_tokens_details ): - combined.completion_tokens_details = CompletionTokensDetailsWrapper() + combined.completion_tokens_details = ( + CompletionTokensDetailsWrapper() + ) # Check what keys exist in the model's completion_tokens_details for attr in usage.completion_tokens_details.model_fields: diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py new file mode 100644 index 00000000000..62a48080d1c --- /dev/null +++ b/litellm/llms/xai/cost_calculator.py @@ -0,0 +1,54 @@ +""" +Helper util for handling XAI-specific cost calculation +- e.g.: reasoning tokens for grok models +""" + +from typing import Tuple, Union + +from litellm.types.utils import Usage +from litellm.utils import get_model_info + + +def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: + """ + Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. + + Input: + - model: str, the model name without provider prefix + - usage: LiteLLM Usage block, containing XAI-specific usage information + + Returns: + Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + """ + ## GET MODEL INFO + model_info = get_model_info(model=model, custom_llm_provider="xai") + + def _safe_float_cast( + value: Union[str, int, float, None, object], default: float = 0.0 + ) -> float: + """Safely cast a value to float with proper type handling for mypy.""" + if value is None: + return default + try: + return float(value) # type: ignore + except (ValueError, TypeError): + return default + + ## CALCULATE INPUT COST + input_cost_per_token = _safe_float_cast(model_info.get("input_cost_per_token")) + prompt_cost: float = (usage.prompt_tokens or 0) * input_cost_per_token + + ## CALCULATE OUTPUT COST + output_cost_per_token = _safe_float_cast(model_info.get("output_cost_per_token")) + + # For XAI models, completion is billed as (visible completion tokens + reasoning tokens) + completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) + reasoning_tokens = 0 + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + reasoning_tokens = int( + getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 + ) + + completion_cost = (completion_tokens + reasoning_tokens) * output_cost_per_token + + return prompt_cost, completion_cost diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py new file mode 100644 index 00000000000..5a26d808c5d --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -0,0 +1,189 @@ +""" +Test suite for XAI cost calculation functionality. +""" + +import math +import os +import sys + +import litellm +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + Usage, +) + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.llms.xai.cost_calculator import cost_per_token + + +class TestXAICostCalculator: + """Test suite for XAI cost calculation functionality.""" + + def setup_method(self): + """Set up test environment.""" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + def test_basic_cost_calculation(self): + """Test basic cost calculation without reasoning tokens.""" + usage = Usage(prompt_tokens=12, completion_tokens=125, total_tokens=137) + + prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) + + # Expected costs for grok-3-mini: + # Input: 12 tokens * $3e-7 = $0.0000036 + # Output: 125 tokens * $5e-7 = $0.0000625 + expected_prompt_cost = 12 * 3e-7 + expected_completion_cost = 125 * 5e-7 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_reasoning_tokens_cost_calculation(self): + """Test cost calculation with reasoning tokens from completion_tokens_details.""" + usage = Usage( + prompt_tokens=12, + completion_tokens=125, + total_tokens=1086, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=949, + rejected_prediction_tokens=0, + text_tokens=None, # Not set, but doesn't matter for XAI billing + ), + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) + + # Expected costs for grok-3-mini: + # Input: 12 tokens * $3e-7 = $0.0000036 + # Completion: (125 + 949) tokens * $5e-7 = $0.000537 + expected_prompt_cost = 12 * 3e-7 + expected_completion_cost = (125 + 949) * 5e-7 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_reasoning_and_text_tokens_cost_calculation(self): + """Test cost calculation with both reasoning and text tokens.""" + usage = Usage( + prompt_tokens=12, + completion_tokens=125, + total_tokens=1086, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=949, + rejected_prediction_tokens=0, + text_tokens=76, # Explicitly set (but ignored in XAI billing) + ), + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) + + # Expected costs for grok-3-mini: + # Input: 12 tokens * $3e-7 = $0.0000036 + # Completion: (125 + 949) tokens * $5e-7 = $0.000537 + # Note: text_tokens field is ignored, only completion_tokens + reasoning_tokens matters + expected_prompt_cost = 12 * 3e-7 + expected_completion_cost = (125 + 949) * 5e-7 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_grok_4_cost_calculation(self): + """Test cost calculation for grok-4 model.""" + usage = Usage( + prompt_tokens=10, + completion_tokens=200, + total_tokens=210, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=150, + rejected_prediction_tokens=0, + text_tokens=50, # Ignored in XAI billing + ), + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-4", usage=usage) + + # Expected costs for grok-4: + # Input: 10 tokens * $3e-6 = $0.00003 + # Completion: (200 + 150) tokens * $1.5e-5 = $0.00525 + expected_prompt_cost = 10 * 3e-6 + expected_completion_cost = (200 + 150) * 1.5e-5 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_grok_3_fast_beta_cost_calculation(self): + """Test cost calculation for grok-3-fast-beta model.""" + usage = Usage( + prompt_tokens=20, + completion_tokens=300, + total_tokens=320, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=200, + rejected_prediction_tokens=0, + text_tokens=100, # Ignored in XAI billing + ), + ) + + prompt_cost, completion_cost = cost_per_token( + model="grok-3-fast-beta", usage=usage + ) + + # Expected costs for grok-3-fast-beta: + # Input: 20 tokens * $5e-6 = $0.0001 + # Completion: (300 + 200) tokens * $2.5e-5 = $0.0125 + expected_prompt_cost = 20 * 5e-6 + expected_completion_cost = (300 + 200) * 2.5e-5 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_edge_case_no_completion_tokens_details(self): + """Test cost calculation when completion_tokens_details is not present.""" + usage = Usage(prompt_tokens=12, completion_tokens=125, total_tokens=137) + + prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) + + # Should fall back to basic calculation + expected_prompt_cost = 12 * 3e-7 + expected_completion_cost = 125 * 5e-7 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_edge_case_large_reasoning_tokens(self): + """Test cost calculation when reasoning_tokens is larger than completion_tokens.""" + usage = Usage( + prompt_tokens=12, + completion_tokens=50, # Less than reasoning_tokens + total_tokens=62, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=100, # More than completion_tokens + rejected_prediction_tokens=0, + text_tokens=None, + ), + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) + + # Expected costs: + # Input: 12 tokens * $3e-7 = $0.0000036 + # Completion: (50 + 100) tokens * $5e-7 = $0.000075 + expected_prompt_cost = 12 * 3e-7 + expected_completion_cost = (50 + 100) * 5e-7 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)