diff --git a/litellm/llms/together_ai/cost_calculator.py b/litellm/llms/together_ai/cost_calculator.py index a66b9884847..4c57fc5509a 100644 --- a/litellm/llms/together_ai/cost_calculator.py +++ b/litellm/llms/together_ai/cost_calculator.py @@ -28,7 +28,7 @@ def has_together_registry_pricing(model: str, cost_map: Mapping[str, object]) -> return ( isinstance(bare_entry, Mapping) and "input_cost_per_token" in bare_entry - and bare_entry.get("litellm_provider") in (None, "together_ai") + and ("litellm_provider" not in bare_entry or bare_entry["litellm_provider"] in (None, "together_ai")) ) diff --git a/tests/unit/llms/together_ai/test_cost_calculator.py b/tests/unit/llms/together_ai/test_cost_calculator.py index 3c46175295b..19d3e6d2e60 100644 --- a/tests/unit/llms/together_ai/test_cost_calculator.py +++ b/tests/unit/llms/together_ai/test_cost_calculator.py @@ -78,11 +78,15 @@ def test_together_completion_cost_keeps_unregistered_model_size_fallback(monkeyp ) == pytest.approx(100 * input_rate + 20 * output_rate) -@pytest.mark.parametrize("bare_provider", ["openai", None, "together_ai"]) +@pytest.mark.parametrize( + ("bare_provider", "declares_provider"), + [("openai", True), (None, False), (None, True), ("together_ai", True)], +) @pytest.mark.parametrize("provider_scoped", [False, True]) def test_together_completion_cost_ignores_other_provider_bare_prices( monkeypatch: pytest.MonkeyPatch, bare_provider: str | None, + declares_provider: bool, provider_scoped: bool, ) -> None: model: Final = "example/colliding-3B" @@ -104,7 +108,7 @@ def test_together_completion_cost_ignores_other_provider_bare_prices( { "input_cost_per_token": bare_rates[0], "output_cost_per_token": bare_rates[1], - **({"litellm_provider": bare_provider} if bare_provider is not None else {}), + **({"litellm_provider": bare_provider} if declares_provider else {}), }, ) if provider_scoped: