diff --git a/src/vendor_capabilities.py b/src/vendor_capabilities.py index 716a479b..0ea18b78 100644 --- a/src/vendor_capabilities.py +++ b/src/vendor_capabilities.py @@ -15,6 +15,19 @@ class VendorCapabilities: cost_input_per_mtok: float = 0.0 cost_output_per_mtok: float = 0.0 notes: str = '' + # v2 fields (added 2026-06-11) + local: bool = False + reasoning: bool = False + structured_output: bool = False + code_execution: bool = False + web_search: bool = False + x_search: bool = False + file_search: bool = False + mcp_support: bool = False + audio: bool = False + video: bool = False + grounding: bool = False + computer_use: bool = False _REGISTRY: dict[tuple[str, str], VendorCapabilities] = {} diff --git a/tests/test_vendor_capabilities.py b/tests/test_vendor_capabilities.py index e8b8ac9a..7edf8438 100644 --- a/tests/test_vendor_capabilities.py +++ b/tests/test_vendor_capabilities.py @@ -38,3 +38,35 @@ def test_fallback_to_vendor_default(): def test_unknown_vendor_raises(): with pytest.raises(KeyError, match='No capabilities registered'): get_capabilities('nonexistent_vendor', 'anymodel') + +V2_FIELDS: list[str] = [ + 'local', 'reasoning', 'structured_output', 'code_execution', + 'web_search', 'x_search', 'file_search', 'mcp_support', + 'audio', 'video', 'grounding', 'computer_use', +] + +@pytest.mark.parametrize('field_name', V2_FIELDS) +def test_v2_field_default_is_false(field_name: str) -> None: + caps = VendorCapabilities(vendor='test', model='m') + assert getattr(caps, field_name) is False, f'{field_name} should default to False' + +@pytest.mark.parametrize('field_name', V2_FIELDS) +def test_v2_field_round_trip(field_name: str) -> None: + caps = VendorCapabilities(vendor='test', model='m', **{field_name: True}) + assert getattr(caps, field_name) is True, f'{field_name} should round-trip to True' + +def test_v2_local_flag_works_for_local_vendor() -> None: + register(VendorCapabilities(vendor='llama', model='llama-local-test-3.1', local=True)) + caps = get_capabilities('llama', 'llama-local-test-3.1') + assert caps.local is True + +def test_v2_local_flag_falls_back_to_wildcard() -> None: + register(VendorCapabilities(vendor='llama', model='*', local=True)) + caps = get_capabilities('llama', 'some-unregistered-model-3.1-future') + assert caps.local is True + +def test_v2_local_flag_does_not_affect_other_vendors() -> None: + register(VendorCapabilities(vendor='llama', model='*', local=True)) + register(VendorCapabilities(vendor='qwen', model='*')) + caps = get_capabilities('qwen', 'qwen-turbo') + assert caps.local is False