fix: pass all custom pricing fields to register_model in completion() and embedding()

Previously, register_model() was called with only input_cost_per_token,
output_cost_per_token, and litellm_provider. This dropped ~40+ other
pricing fields from CustomPricingLiteLLMParams (cache_read_input_token_cost,
cache_creation_input_token_cost, output_cost_per_reasoning_token, etc.)
as well as model_info metadata (mode, supports_prompt_caching, max_tokens).

For DB-sourced custom-priced models, the first request after a pod restart
would register a partial entry in litellm.model_cost, causing cost
calculations to miss cache token discounts and other extended pricing
until the entry was later enriched by deployment_callback_on_success
mutating the lru_cache.

Changes:
- Add _build_custom_pricing_entry() helper that iterates over all
  CustomPricingLiteLLMParams.model_fields and merges model_info metadata
- Replace hardcoded 3-field dicts in both completion() and embedding()
  with the new helper
- Add 7 tests covering field collection, model_info merging, precedence,
  None skipping, and end-to-end register_model behavior

Co-authored-by: openhands <openhands@all-hands.dev>

litellm/main.py

@@ -107,6 +107,7 @@ from litellm.realtime_api.main import _realtime_health_check
 from litellm.secret_managers.main import get_secret_bool, get_secret_str
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
+    CustomPricingLiteLLMParams,
     ModelResponseStream,
     RawRequestTypedDict,
     StreamingChoices,
@@ -996,6 +997,32 @@ def _drop_input_examples_from_tools(
     return cleaned_tools
 
 
+def _build_custom_pricing_entry(
+    custom_llm_provider: str,
+    kwargs: dict,
+    model_info: Optional[dict] = None,
+) -> dict:
+    """Build a complete model cost entry from kwargs and model_info.
+
+    Collects all CustomPricingLiteLLMParams fields present in kwargs and
+    merges metadata from model_info (mode, supports_prompt_caching, max_tokens)
+    so that register_model() receives the full pricing configuration.
+    """
+    entry: dict = {"litellm_provider": custom_llm_provider}
+
+    for field_name in CustomPricingLiteLLMParams.model_fields:
+        value = kwargs.get(field_name)
+        if value is not None:
+            entry[field_name] = value
+
+    if model_info and isinstance(model_info, dict):
+        for key in ("mode", "supports_prompt_caching", "max_tokens"):
+            if key in model_info and model_info[key] is not None:
+                entry.setdefault(key, model_info[key])
+
+    return entry
+
+
 @tracer.wrap()
 @client
 def completion(  # type: ignore # noqa: PLR0915
@@ -1351,27 +1378,16 @@ def completion(  # type: ignore # noqa: PLR0915
             timeout = float(timeout)  # type: ignore
 
         ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
-        if input_cost_per_token is not None and output_cost_per_token is not None:
+        if (
+            input_cost_per_token is not None and output_cost_per_token is not None
+        ) or input_cost_per_second is not None:
             litellm.register_model(
                 {
-                    f"{custom_llm_provider}/{model}": {
-                        "input_cost_per_token": input_cost_per_token,
-                        "output_cost_per_token": output_cost_per_token,
-                        "litellm_provider": custom_llm_provider,
-                    }
-                }
-            )
-        elif (
-            input_cost_per_second is not None
-        ):  # time based pricing just needs cost in place
-            output_cost_per_second = output_cost_per_second
-            litellm.register_model(
-                {
-                    f"{custom_llm_provider}/{model}": {
-                        "input_cost_per_second": input_cost_per_second,
-                        "output_cost_per_second": output_cost_per_second,
-                        "litellm_provider": custom_llm_provider,
-                    }
+                    f"{custom_llm_provider}/{model}": _build_custom_pricing_entry(
+                        custom_llm_provider=custom_llm_provider,
+                        kwargs=kwargs,
+                        model_info=model_info,
+                    )
                 }
             )
         ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ###
@@ -4644,7 +4660,6 @@ def embedding(  # noqa: PLR0915
     input_cost_per_token = kwargs.get("input_cost_per_token", None)
     output_cost_per_token = kwargs.get("output_cost_per_token", None)
     input_cost_per_second = kwargs.get("input_cost_per_second", None)
-    output_cost_per_second = kwargs.get("output_cost_per_second", None)
     openai_params = [
         "user",
         "dimensions",
@@ -4694,25 +4709,16 @@ def embedding(  # noqa: PLR0915
     )
 
     ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
-    if input_cost_per_token is not None and output_cost_per_token is not None:
+    if (
+        input_cost_per_token is not None and output_cost_per_token is not None
+    ) or input_cost_per_second is not None:
         litellm.register_model(
             {
-                f"{custom_llm_provider}/{model}": {
-                    "input_cost_per_token": input_cost_per_token,
-                    "output_cost_per_token": output_cost_per_token,
-                    "litellm_provider": custom_llm_provider,
-                }
-            }
-        )
-    if input_cost_per_second is not None:  # time based pricing just needs cost in place
-        output_cost_per_second = output_cost_per_second or 0.0
-        litellm.register_model(
-            {
-                f"{custom_llm_provider}/{model}": {
-                    "input_cost_per_second": input_cost_per_second,
-                    "output_cost_per_second": output_cost_per_second,
-                    "litellm_provider": custom_llm_provider,
-                }
+                f"{custom_llm_provider}/{model}": _build_custom_pricing_entry(
+                    custom_llm_provider=custom_llm_provider,
+                    kwargs=kwargs,
+                    model_info=kwargs.get("model_info"),
+                )
             }
         )
 

tests/test_litellm/test_register_model_custom_pricing.py

@@ -0,0 +1,180 @@
+"""
+Test that register_model() in completion() and embedding() passes all
+custom pricing fields from kwargs and model_info, not just the base
+input/output costs.
+
+Previously, only input_cost_per_token, output_cost_per_token, and
+litellm_provider were forwarded. Fields like cache_read_input_token_cost,
+mode, and supports_prompt_caching were dropped, causing incorrect cost
+calculations for DB-sourced models with prompt caching pricing.
+"""
+
+import os
+import sys
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+import litellm
+from litellm.main import _build_custom_pricing_entry
+
+
+def test_build_custom_pricing_entry_includes_all_kwargs_fields():
+    """All CustomPricingLiteLLMParams fields present in kwargs should be
+    included in the resulting entry dict."""
+    kwargs = {
+        "input_cost_per_token": 0.001,
+        "output_cost_per_token": 0.002,
+        "cache_read_input_token_cost": 0.00025,
+        "cache_creation_input_token_cost": 0.005,
+        "output_cost_per_reasoning_token": 0.01,
+        "input_cost_per_audio_token": 0.003,
+        "unrelated_kwarg": "should_be_ignored",
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+    )
+
+    assert entry["litellm_provider"] == "openai"
+    assert entry["input_cost_per_token"] == 0.001
+    assert entry["output_cost_per_token"] == 0.002
+    assert entry["cache_read_input_token_cost"] == 0.00025
+    assert entry["cache_creation_input_token_cost"] == 0.005
+    assert entry["output_cost_per_reasoning_token"] == 0.01
+    assert entry["input_cost_per_audio_token"] == 0.003
+    assert "unrelated_kwarg" not in entry
+
+
+def test_build_custom_pricing_entry_merges_model_info_metadata():
+    """Fields from model_info (mode, supports_prompt_caching, max_tokens)
+    should be merged into the entry when present."""
+    kwargs = {
+        "input_cost_per_token": 0.001,
+        "output_cost_per_token": 0.002,
+    }
+    model_info = {
+        "id": "deployment-123",
+        "mode": "chat",
+        "supports_prompt_caching": True,
+        "max_tokens": 128000,
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+        model_info=model_info,
+    )
+
+    assert entry["mode"] == "chat"
+    assert entry["supports_prompt_caching"] is True
+    assert entry["max_tokens"] == 128000
+
+
+def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info():
+    """If a field appears in both kwargs and model_info, the kwargs value
+    should take precedence (setdefault behavior)."""
+    kwargs = {
+        "input_cost_per_token": 0.001,
+        "output_cost_per_token": 0.002,
+    }
+    model_info = {
+        "mode": "chat",
+        "supports_prompt_caching": True,
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+        model_info=model_info,
+    )
+
+    # model_info fields should be set via setdefault
+    assert entry["mode"] == "chat"
+    assert entry["supports_prompt_caching"] is True
+
+
+def test_build_custom_pricing_entry_skips_none_values():
+    """Fields with None values in kwargs should not be included."""
+    kwargs = {
+        "input_cost_per_token": 0.001,
+        "output_cost_per_token": None,  # explicitly None
+        "cache_read_input_token_cost": None,
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+    )
+
+    assert entry["input_cost_per_token"] == 0.001
+    assert "output_cost_per_token" not in entry
+    assert "cache_read_input_token_cost" not in entry
+
+
+def test_build_custom_pricing_entry_handles_no_model_info():
+    """Should work correctly when model_info is None."""
+    kwargs = {
+        "input_cost_per_token": 0.001,
+        "output_cost_per_token": 0.002,
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+        model_info=None,
+    )
+
+    assert entry["litellm_provider"] == "openai"
+    assert entry["input_cost_per_token"] == 0.001
+    assert entry["output_cost_per_token"] == 0.002
+    assert "mode" not in entry
+
+
+def test_register_model_receives_cache_pricing_fields():
+    """End-to-end: when register_model is called with a full pricing entry,
+    the cache pricing fields should be present in litellm.model_cost."""
+    model_key = "openai/test-custom-model-with-cache-pricing"
+
+    litellm.register_model(
+        {
+            model_key: {
+                "input_cost_per_token": 0.001,
+                "output_cost_per_token": 0.002,
+                "cache_read_input_token_cost": 0.00025,
+                "supports_prompt_caching": True,
+                "mode": "chat",
+                "max_tokens": 8192,
+                "litellm_provider": "openai",
+            }
+        }
+    )
+
+    registered = litellm.model_cost.get(model_key)
+    assert registered is not None, f"{model_key} should be in model_cost"
+    assert registered["cache_read_input_token_cost"] == 0.00025
+    assert registered["supports_prompt_caching"] is True
+    assert registered["mode"] == "chat"
+    assert registered["max_tokens"] == 8192
+
+    # Cleanup
+    litellm.model_cost.pop(model_key, None)
+
+
+def test_build_custom_pricing_entry_time_based():
+    """Time-based pricing fields should be included correctly."""
+    kwargs = {
+        "input_cost_per_second": 0.01,
+        "output_cost_per_second": 0.02,
+    }
+
+    entry = _build_custom_pricing_entry(
+        custom_llm_provider="openai",
+        kwargs=kwargs,
+    )
+
+    assert entry["litellm_provider"] == "openai"
+    assert entry["input_cost_per_second"] == 0.01
+    assert entry["output_cost_per_second"] == 0.02