diff --git a/litellm/provider_capabilities.py b/litellm/provider_capabilities.py new file mode 100644 index 00000000000..c9c6538805c --- /dev/null +++ b/litellm/provider_capabilities.py @@ -0,0 +1,200 @@ +""" +Structured provider capabilities declaration for model routing and feature detection. + +Replaces ad-hoc flags scattered across model_prices_and_context_window.json +with a single typed dataclass that each provider declares. + +Design mirrors zeshim's `ProviderCapabilities` interface, adapted to LiteLLM's +Python dataclass conventions and existing model registry. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ProviderSupports: + """Boolean feature flags for a provider/model. + + Each field defaults to False — providers opt IN to capabilities. + This avoids the "silently false" problem where a missing key in a + JSON dict is indistinguishable from a capability being absent. + """ + + vision: bool = False + streaming: bool = False + multi_image: bool = False + function_calling: bool = False + structured_output: bool = False + audio_input: bool = False + audio_output: bool = False + document_input: bool = False + prompt_caching: bool = False + reasoning: bool = False + + +@dataclass(frozen=True) +class ProviderStrengths: + """Scored provider strengths for intelligent routing decisions. + + Each dimension is a float in [0.0, 1.0] where higher is better. + Used by Router strategies (latency-based, cost-based, quality-based) + to make data-driven deployment selection. + + Not all providers need to declare all dimensions — missing scores + default to 0.0 (neutral). + """ + + speed: float = 0.0 # Low latency relative to peer models + cost_efficiency: float = 0.0 # Tokens-per-dollar relative to peer models + quality: float = 0.0 # Output quality / benchmark performance + reliability: float = 0.0 # Uptime / error-rate track record + context_fidelity: float = 0.0 # Long-context adherence + + def __post_init__(self): + for name in ( + "speed", + "cost_efficiency", + "quality", + "reliability", + "context_fidelity", + ): + val = getattr(self, name) + if not (0.0 <= val <= 1.0): + raise ValueError(f"strengths.{name} must be in [0.0, 1.0], got {val}") + + +@dataclass(frozen=True) +class ProviderCapabilities: + """Complete capabilities declaration for a provider/model. + + Composed of: + supports — boolean feature flags (what the model CAN do) + strengths — scored dimensions (how WELL it does them), + used by the Router for intelligent deployment selection + + Immutable by design — capabilities don't change at runtime. + Models that gain capabilities (e.g., a vision update) should + declare a new capabilities object. + """ + + supports: ProviderSupports = field(default_factory=ProviderSupports) + strengths: ProviderStrengths = field(default_factory=ProviderStrengths) + + # ── Convenience accessors (backward-compatible with existing API) ── + + @property + def supports_vision(self) -> bool: + return self.supports.vision + + @property + def supports_streaming(self) -> bool: + return self.supports.streaming + + @property + def supports_function_calling(self) -> bool: + return self.supports.function_calling + + @property + def supports_structured_output(self) -> bool: + return self.supports.structured_output + + @property + def supports_audio(self) -> bool: + return self.supports.audio_input or self.supports.audio_output + + +# ── Pre-built capability profiles for common model families ── + +# These serve as templates that individual providers can extend or override. +# The goal: reduce duplication — OpenAI-compatible providers don't each need +# to re-declare the same capability set. + +CAPABILITIES_OPENAI_GPT4O = ProviderCapabilities( + supports=ProviderSupports( + vision=True, + streaming=True, + multi_image=True, + function_calling=True, + structured_output=True, + reasoning=True, + ), + strengths=ProviderStrengths( + speed=0.7, + cost_efficiency=0.5, + quality=0.9, + reliability=0.95, + context_fidelity=0.85, + ), +) + +CAPABILITIES_ANTHROPIC_CLAUDE = ProviderCapabilities( + supports=ProviderSupports( + vision=True, + streaming=True, + multi_image=True, + function_calling=True, + structured_output=True, + prompt_caching=True, + document_input=True, + reasoning=True, + ), + strengths=ProviderStrengths( + speed=0.6, + cost_efficiency=0.4, + quality=0.95, + reliability=0.95, + context_fidelity=0.90, + ), +) + +CAPABILITIES_GEMINI_FLASH = ProviderCapabilities( + supports=ProviderSupports( + vision=True, + streaming=True, + function_calling=True, + audio_input=True, + prompt_caching=True, + ), + strengths=ProviderStrengths( + speed=0.9, + cost_efficiency=0.85, + quality=0.7, + reliability=0.85, + context_fidelity=0.95, + ), +) + +# ── Integration helper: build from existing model_cost_map entry ── + + +def capabilities_from_model_info( + supports_vision: bool = False, + supports_function_calling: bool = False, + supports_streaming: bool = False, + supports_audio_input: bool = False, + supports_audio_output: bool = False, + supports_prompt_caching: bool = False, + supports_response_schema: bool = False, + supports_reasoning: bool = False, + **_: object, # ignore unknown keys from model_cost_map +) -> ProviderCapabilities: + """Bridge: construct ProviderCapabilities from existing model_cost_map flags. + + This allows gradual migration — existing JSON entries can populate + capabilities via this helper without rewriting the registry. + """ + return ProviderCapabilities( + supports=ProviderSupports( + vision=supports_vision, + streaming=supports_streaming, + function_calling=supports_function_calling, + structured_output=supports_response_schema, + audio_input=supports_audio_input, + audio_output=supports_audio_output, + prompt_caching=supports_prompt_caching, + reasoning=supports_reasoning, + ), + # strengths comes from model_cost_map scoring or defaults to neutral + ) diff --git a/tests/test_provider_capabilities.py b/tests/test_provider_capabilities.py new file mode 100644 index 00000000000..8fd7250b73b --- /dev/null +++ b/tests/test_provider_capabilities.py @@ -0,0 +1,177 @@ +"""Tests for litellm.provider_capabilities — structured capability declarations.""" + +import pytest +from litellm.provider_capabilities import ( + CAPABILITIES_ANTHROPIC_CLAUDE, + CAPABILITIES_GEMINI_FLASH, + CAPABILITIES_OPENAI_GPT4O, + ProviderCapabilities, + ProviderStrengths, + ProviderSupports, + capabilities_from_model_info, +) + + +class TestProviderSupports: + """Boolean feature flags — opt-in, all default to False.""" + + def test_default_all_false(self): + s = ProviderSupports() + assert s.vision is False + assert s.streaming is False + assert s.function_calling is False + assert s.structured_output is False + + def test_explicit_true(self): + s = ProviderSupports(vision=True, streaming=True) + assert s.vision is True + assert s.streaming is True + assert s.function_calling is False # still default + + def test_frozen(self): + s = ProviderSupports(vision=True) + with pytest.raises(Exception): + s.vision = False # type: ignore[misc] + + +class TestProviderStrengths: + """Scored dimensions — float in [0.0, 1.0], defaults to 0.0.""" + + def test_default_all_zero(self): + s = ProviderStrengths() + assert s.speed == 0.0 + assert s.quality == 0.0 + + def test_valid_scores(self): + s = ProviderStrengths(speed=0.8, quality=0.9, reliability=0.95) + assert s.speed == 0.8 + assert s.quality == 0.9 + + def test_score_out_of_range_raises(self): + with pytest.raises(ValueError, match="must be in"): + ProviderStrengths(speed=1.5) + + def test_score_negative_raises(self): + with pytest.raises(ValueError, match="must be in"): + ProviderStrengths(quality=-0.1) + + def test_frozen(self): + s = ProviderStrengths(speed=0.5) + with pytest.raises(Exception): + s.speed = 0.8 # type: ignore[misc] + + +class TestProviderCapabilities: + """Composed capabilities declaration.""" + + def test_default_empty(self): + cap = ProviderCapabilities() + assert cap.supports_vision is False + assert cap.supports_streaming is False + assert cap.supports_function_calling is False + + def test_backward_compat_accessors(self): + cap = ProviderCapabilities( + supports=ProviderSupports( + vision=True, streaming=True, function_calling=True + ), + ) + assert cap.supports_vision is True + assert cap.supports_streaming is True + assert cap.supports_function_calling is True + assert cap.supports_structured_output is False + + def test_supports_audio_combines_input_and_output(self): + cap = ProviderCapabilities( + supports=ProviderSupports(audio_input=True, audio_output=False), + ) + assert cap.supports_audio is True + + cap2 = ProviderCapabilities( + supports=ProviderSupports(audio_input=False, audio_output=False), + ) + assert cap2.supports_audio is False + + def test_strengths_accessible(self): + cap = ProviderCapabilities( + strengths=ProviderStrengths(speed=0.9, quality=0.8), + ) + assert cap.strengths.speed == 0.9 + assert cap.strengths.quality == 0.8 + + def test_frozen(self): + cap = ProviderCapabilities() + with pytest.raises(Exception): + cap.supports = ProviderSupports() # type: ignore[misc] + + +class TestPrebuiltCapabilityProfiles: + """Pre-built profiles for common model families.""" + + def test_openai_gpt4o(self): + assert CAPABILITIES_OPENAI_GPT4O.supports.vision is True + assert CAPABILITIES_OPENAI_GPT4O.supports.streaming is True + assert CAPABILITIES_OPENAI_GPT4O.supports.function_calling is True + assert CAPABILITIES_OPENAI_GPT4O.supports.structured_output is True + assert CAPABILITIES_OPENAI_GPT4O.strengths.quality > 0.8 + + def test_anthropic_claude(self): + assert CAPABILITIES_ANTHROPIC_CLAUDE.supports.prompt_caching is True + assert CAPABILITIES_ANTHROPIC_CLAUDE.supports.document_input is True + assert CAPABILITIES_ANTHROPIC_CLAUDE.strengths.context_fidelity > 0.8 + + def test_gemini_flash(self): + assert CAPABILITIES_GEMINI_FLASH.supports.audio_input is True + assert CAPABILITIES_GEMINI_FLASH.strengths.speed > 0.8 + assert CAPABILITIES_GEMINI_FLASH.strengths.cost_efficiency > 0.8 + + def test_all_prebuilts_are_frozen(self): + for cap in ( + CAPABILITIES_OPENAI_GPT4O, + CAPABILITIES_ANTHROPIC_CLAUDE, + CAPABILITIES_GEMINI_FLASH, + ): + with pytest.raises(Exception): + cap.supports.vision = False # type: ignore[misc] + + +class TestCapabilitiesFromModelInfo: + """Bridge: migrate from existing model_cost_map flags.""" + + def test_all_false_by_default(self): + cap = capabilities_from_model_info() + assert cap.supports_vision is False + assert cap.supports_function_calling is False + + def test_maps_individual_flags(self): + cap = capabilities_from_model_info( + supports_vision=True, + supports_function_calling=True, + supports_streaming=True, + ) + assert cap.supports_vision is True + assert cap.supports_function_calling is True + assert cap.supports_streaming is True + + def test_ignores_unknown_kwargs(self): + """Should not crash on extra keys from model_cost_map.""" + cap = capabilities_from_model_info( + max_tokens=128000, # unknown key + input_cost_per_token=0.0001, # unknown key + supports_vision=True, + ) + assert cap.supports_vision is True + + def test_maps_response_schema_to_structured_output(self): + """supports_response_schema should map to structured_output.""" + cap = capabilities_from_model_info( + supports_response_schema=True, + ) + assert cap.supports_structured_output is True + + def test_maps_reasoning_flag(self): + """supports_reasoning should be preserved.""" + cap = capabilities_from_model_info( + supports_reasoning=True, + ) + assert cap.supports.reasoning is True