diff --git a/llmrelic/registry.py b/llmrelic/registry.py index a0bb12b..1eb8d2a 100644 --- a/llmrelic/registry.py +++ b/llmrelic/registry.py @@ -54,15 +54,21 @@ def get_supported_models(self) -> List[str]: return sorted(self._supported_models) def get_supported_by_provider(self) -> Dict[str, List[str]]: - """Get supported models organized by provider.""" - result = {} + """ + Get supported models organized by provider. + A slightly optimized version is used to set intersections instead of calling ``is_supported`` for every model in the provider. This reduces Python overhead when the registry contains many models and providers have large model lists. + """ + result: Dict[str, List[str]] = {} + # local reference to avoid global lookups in loop + supported = self._supported_models for provider_name, provider in PROVIDERS.items(): - supported = [ - model for model in provider.list_models() - if self.is_supported(model) - ] - if supported: - result[provider_name] = supported + # intersect the two sets; provider.list_models() returns a list so + # convert to set once + provider_models = set(provider.list_models()) + matches = supported & provider_models + if matches: + # sort for deterministic output (tests rely on some ordering) + result[provider_name] = sorted(matches) return result def clear(self) -> "ModelRegistry": diff --git a/tests/test_registry.py b/tests/test_registry.py index 2675a21..510b64b 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -153,6 +153,21 @@ def test_get_supported_by_provider_excludes_empty_providers(self): assert "cohere" not in result assert "mistral" not in result + def test_get_supported_by_provider_output_is_sorted(self): + """The list returned for each provider should be sorted alphabetically. + + This test also indirectly verifies that our optimized intersection logic + doesn't break existing functionality when models are added in an + arbitrary order. + """ + registry = ModelRegistry() + # insert models in reverse order intentionally + for model in reversed(OpenAI.list_models()): + registry.add_model(model) + result = registry.get_supported_by_provider() + openai_models = result.get("openai", []) + assert openai_models == sorted(openai_models) + def test_clear(self): """clear() should remove all models.""" registry = ModelRegistry()