Merge eb334bcf63 into c80f1c2ec8

litellm/llms/together_ai/cost_calculator.py

@@ -21,8 +21,15 @@ from litellm.types.utils import CallTypes
 
 def has_together_registry_pricing(model: str, cost_map: Mapping[str, object]) -> bool:
     stripped: Final = model.removeprefix("together_ai/")
-    entry: Final = cost_map.get(f"together_ai/{stripped}")
-    return isinstance(entry, Mapping) and "input_cost_per_token" in entry
+    provider_entry: Final = cost_map.get(f"together_ai/{stripped}")
+    if isinstance(provider_entry, Mapping) and "input_cost_per_token" in provider_entry:
+        return True
+    bare_entry: Final = cost_map.get(stripped)
+    return (
+        isinstance(bare_entry, Mapping)
+        and "input_cost_per_token" in bare_entry
+        and ("litellm_provider" not in bare_entry or bare_entry["litellm_provider"] in (None, "together_ai"))
+    )
 
 
 # Extract the number of billion parameters from the model name

tests/unit/llms/together_ai/test_cost_calculator.py

@@ -0,0 +1,132 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.cost_calculator import completion_cost
+from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
+
+
+@pytest.mark.parametrize("cache_rate", [0.0, 3e-8])
+@pytest.mark.parametrize("cached_tokens", [0, 75200])
+@pytest.mark.parametrize("custom_pricing", [False, True])
+def test_together_completion_cost_preserves_registered_rates(
+    monkeypatch: pytest.MonkeyPatch,
+    cache_rate: float,
+    cached_tokens: int,
+    custom_pricing: bool,
+) -> None:
+    deployment_id: Final = "00000000-3b00-4000-8000-000000000001"
+    model: Final = "example/pricing-regression-3B"
+    registry_key: Final = deployment_id if custom_pricing else f"together_ai/{model}"
+    input_rate: Final = 1.5e-7
+    output_rate: Final = 5e-7
+    prompt_tokens: Final = 89536
+    completion_tokens: Final = 2117
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "together-ai-up-to-4b",
+        {"input_cost_per_token": 9e-6, "output_cost_per_token": 9e-6, "litellm_provider": "together_ai"},
+    )
+    litellm.register_model(
+        model_cost={
+            registry_key: {
+                "input_cost_per_token": input_rate,
+                "output_cost_per_token": output_rate,
+                "cache_read_input_token_cost": cache_rate,
+                "litellm_provider": "together_ai",
+            }
+        }
+    )
+    response: Final = ModelResponse(
+        model=model,
+        usage=Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
+        ),
+    )
+
+    cost: Final = completion_cost(
+        completion_response=response,
+        custom_llm_provider="together_ai",
+        custom_pricing=custom_pricing,
+        router_model_id=deployment_id,
+    )
+
+    assert cost == pytest.approx(
+        (prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_rate + completion_tokens * output_rate
+    )
+
+
+def test_together_completion_cost_keeps_unregistered_model_size_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    input_rate: Final = 2e-6
+    output_rate: Final = 7e-6
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "together-ai-up-to-4b",
+        {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate, "litellm_provider": "together_ai"},
+    )
+    response: Final = ModelResponse(
+        model="example/unregistered-3B",
+        usage=Usage(prompt_tokens=100, completion_tokens=20),
+    )
+
+    assert completion_cost(
+        completion_response=response,
+        custom_llm_provider="together_ai",
+    ) == pytest.approx(100 * input_rate + 20 * output_rate)
+
+
+@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"
+    fallback_rates: Final = (2e-6, 7e-6)
+    bare_rates: Final = (3e-6, 11e-6)
+    scoped_rates: Final = (5e-6, 13e-6)
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "together-ai-up-to-4b",
+        {
+            "input_cost_per_token": fallback_rates[0],
+            "output_cost_per_token": fallback_rates[1],
+            "litellm_provider": "together_ai",
+        },
+    )
+    monkeypatch.setitem(
+        litellm.model_cost,
+        model,
+        {
+            "input_cost_per_token": bare_rates[0],
+            "output_cost_per_token": bare_rates[1],
+            **({"litellm_provider": bare_provider} if declares_provider else {}),
+        },
+    )
+    if provider_scoped:
+        monkeypatch.setitem(
+            litellm.model_cost,
+            f"together_ai/{model}",
+            {
+                "input_cost_per_token": scoped_rates[0],
+                "output_cost_per_token": scoped_rates[1],
+                "litellm_provider": "together_ai",
+            },
+        )
+    response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=100, completion_tokens=20))
+    input_rate, output_rate = (
+        scoped_rates if provider_scoped else bare_rates if bare_provider in (None, "together_ai") else fallback_rates
+    )
+
+    assert completion_cost(
+        completion_response=response,
+        custom_llm_provider="together_ai",
+    ) == pytest.approx(100 * input_rate + 20 * output_rate)