diff --git a/src/model_capability_readers/__init__.py b/src/model_capability_readers/__init__.py index d25721d27..15b208163 100644 --- a/src/model_capability_readers/__init__.py +++ b/src/model_capability_readers/__init__.py @@ -128,7 +128,7 @@ def records_from_payload( ) shape = pcs.catalog_shape_for_id(resolution.shape_id) - if resolution.fallback or reader is generic_openai or shape is None: + if shape is None: records = generic_openai.records_from_payload( payload, vendor_id=record_vendor, @@ -140,16 +140,42 @@ def records_from_payload( shape_id=resolution.shape_id, fallback=resolution.fallback, ) + elif resolution.fallback: + normalized_records: list[ModelCapabilityRecord] = [] + for item in shape.items(payload): + fallback_record = generic_openai.record_from_model( + item, + vendor_id=record_vendor, + endpoint_id=endpoint_id, + base_url=base_url, + ) + if fallback_record: + normalized_records.extend( + annotate( + (fallback_record,), + shape_id=shape.shape_id, + fallback=True, + ) + ) + normalized = tuple(normalized_records) else: normalized_records: list[ModelCapabilityRecord] = [] for item in shape.items(payload): item_payload = shape.payload_for_item(payload, item) if shape.item_matches(item): - native_records = reader.records_from_payload( - item_payload, - endpoint_id=endpoint_id, - base_url=base_url, - ) + if reader is generic_openai: + native_records = reader.records_from_payload( + item_payload, + vendor_id=record_vendor, + endpoint_id=endpoint_id, + base_url=base_url, + ) + else: + native_records = reader.records_from_payload( + item_payload, + endpoint_id=endpoint_id, + base_url=base_url, + ) if native_records: normalized_records.extend( annotate(native_records, shape_id=shape.shape_id, fallback=False) diff --git a/src/provider_capability_schemas.py b/src/provider_capability_schemas.py index bafc06c05..0e7593fbc 100644 --- a/src/provider_capability_schemas.py +++ b/src/provider_capability_schemas.py @@ -122,7 +122,11 @@ class ProviderCatalogShape: if self.envelope == ENVELOPE_SINGLE: return item if isinstance(payload, Mapping): - narrowed = dict(payload) + narrowed = { + key: value + for key, value in payload.items() + if key not in {ENVELOPE_DATA, ENVELOPE_MODELS} + } narrowed[self.envelope] = [item] return narrowed return {self.envelope: [item]} @@ -613,8 +617,14 @@ def catalog_shape_for_id(shape_id: Any) -> ProviderCatalogShape | None: return next( ( shape - for schema in PROVIDER_SCHEMAS.values() - for shape in schema.catalog_shapes + for shape in ( + *FALLBACK_CATALOG_SHAPES, + *( + provider_shape + for schema in PROVIDER_SCHEMAS.values() + for provider_shape in schema.catalog_shapes + ), + ) if shape.shape_id == shape_id ), None, diff --git a/tests/test_provider_capability_schemas.py b/tests/test_provider_capability_schemas.py index 87d4244a7..4e61087b5 100644 --- a/tests/test_provider_capability_schemas.py +++ b/tests/test_provider_capability_schemas.py @@ -369,6 +369,45 @@ def test_provider_specific_reader_is_not_used_for_a_different_fallback_envelope( assert record.fallback is True +def test_selected_native_envelope_ignores_an_unrelated_alternate_envelope(): + record = records_from_payload( + { + "data": [{"id": "unrelated-openai-card"}], + "models": [ + { + "key": "native-lmstudio-card", + "type": "llm", + "capabilities": {"vision": True}, + } + ], + }, + vendor="lmstudio", + )[0] + + assert record.model_id == "native-lmstudio-card" + assert record.vendor == "lmstudio" + assert record.capability.family == mc.FAMILY_CHAT + assert record.capability.capabilities == (mc.CAP_VISION,) + assert record.catalog_shape_id == "lmstudio.models.native.v1" + assert record.fallback is False + + +def test_selected_fallback_envelope_is_not_shadowed_by_empty_data(): + record = records_from_payload( + { + "data": [], + "models": [{"id": "fallback-model", "capabilities": {"tools": True}}], + } + )[0] + + assert record.model_id == "fallback-model" + assert record.vendor == pcs.PROVIDER_UNKNOWN + assert record.capability.family == mc.FAMILY_UNKNOWN + assert record.capability.capabilities == () + assert record.catalog_shape_id == "fallback.models.envelope.v1" + assert record.fallback is True + + def test_fallback_reader_is_identity_only_even_for_dangerous_looking_fields(): payload = [ {