diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index 6b39f37921d..629f8ac9fdb 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -1,9 +1,8 @@ import os -from collections.abc import Mapping from datetime import date -from typing import Any import pytest +from pydantic import BaseModel, ConfigDict def _skip_live_prompt_caching_test(): @@ -13,34 +12,53 @@ def _skip_live_prompt_caching_test(): pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay") -def cheapest_together_chat_model(*capability_flags: str) -> str: + +class TogetherCostEntry(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str | None = None + mode: str | None = None + deprecation_date: str | None = None + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + supports_function_calling: bool | None = None + supports_response_schema: bool | None = None + + +def cheapest_together_chat_model( + *, function_calling: bool = False, response_schema: bool = False +) -> str: import litellm today = date.today().isoformat() - def qualifies(name: str, entry: Mapping[str, Any]) -> bool: - deprecation_date = entry.get("deprecation_date") + def qualifies(name: str, entry: TogetherCostEntry) -> bool: return ( name.startswith("together_ai/") - and entry.get("litellm_provider") == "together_ai" - and entry.get("mode") == "chat" - and (deprecation_date is None or deprecation_date > today) - and (entry.get("input_cost_per_token") or 0.0) > 0 - and (entry.get("output_cost_per_token") or 0.0) > 0 - and all(bool(entry.get(flag)) for flag in capability_flags) + and entry.litellm_provider == "together_ai" + and entry.mode == "chat" + and (entry.deprecation_date is None or entry.deprecation_date > today) + and (entry.input_cost_per_token or 0.0) > 0 + and (entry.output_cost_per_token or 0.0) > 0 + and (not function_calling or bool(entry.supports_function_calling)) + and (not response_schema or bool(entry.supports_response_schema)) ) + registry: dict[str, TogetherCostEntry] = { + name: TogetherCostEntry.model_validate(raw) + for name, raw in litellm.model_cost.items() + if isinstance(raw, dict) and name.startswith("together_ai/") + } candidates = sorted( - ( - name - for name, entry in litellm.model_cost.items() - if isinstance(entry, Mapping) and qualifies(name, entry) - ), + (name for name, entry in registry.items() if qualifies(name, entry)), key=lambda name: ( - litellm.model_cost[name].get("input_cost_per_token") or 0.0, - litellm.model_cost[name].get("output_cost_per_token") or 0.0, + registry[name].input_cost_per_token or 0.0, + registry[name].output_cost_per_token or 0.0, name, ), ) - assert candidates, f"no live together_ai chat model in the cost map satisfies {capability_flags}" + assert candidates, ( + "no live together_ai chat model in the cost map satisfies " + f"function_calling={function_calling} response_schema={response_schema}" + ) return candidates[0] diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index 7a49d46b528..0b4e9d3952c 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -19,7 +19,7 @@ class TestTogetherAI(BaseLLMChatTest): litellm.set_verbose = True return { "model": cheapest_together_chat_model( - "supports_function_calling", "supports_response_schema" + function_calling=True, response_schema=True ) }