fix xai cost calculation (#14127)

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:

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

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)