# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for the provider registry module. Note: This file tests the deprecated registry module which is now an alias for router. The no-name-in-module warning for providers.registry is expected. Test helper classes also intentionally have few public methods. """ # pylint: disable=no-name-in-module import re from absl.testing import absltest from langextract import exceptions from langextract import providers as providers_module from langextract.core import base_model from langextract.core import types from langextract.providers import builtin_registry from langextract.providers import router class FakeProvider(base_model.BaseLanguageModel): """Fake provider for testing.""" def infer(self, batch_prompts, **kwargs): return [[types.ScoredOutput(score=1.0, output="test")]] def infer_batch(self, prompts, batch_size=32): return self.infer(prompts) class AnotherFakeProvider(base_model.BaseLanguageModel): """Another fake provider for testing.""" def infer(self, batch_prompts, **kwargs): return [[types.ScoredOutput(score=1.0, output="another")]] def infer_batch(self, prompts, batch_size=32): return self.infer(prompts) class RegistryTest(absltest.TestCase): def setUp(self): super().setUp() router.clear() def tearDown(self): super().tearDown() router.clear() def test_register_decorator(self): """Test registering a provider using the decorator.""" @router.register(r"^test-model") class TestProvider(FakeProvider): pass resolved = router.resolve("test-model-v1") self.assertEqual(resolved, TestProvider) def test_register_lazy(self): """Test lazy registration with string target.""" # Use direct registration for test provider to avoid module path issues router.register(r"^fake-model")(FakeProvider) resolved = router.resolve("fake-model-v2") self.assertEqual(resolved, FakeProvider) def test_multiple_patterns(self): """Test registering multiple patterns for one provider.""" # Use direct registration to avoid module path issues in Bazel router.register(r"^gemini", r"^palm")(FakeProvider) self.assertEqual(router.resolve("gemini-pro"), FakeProvider) self.assertEqual(router.resolve("palm-2"), FakeProvider) def test_priority_resolution(self): """Test that higher priority wins on conflicts.""" # Use direct registration to avoid module path issues in Bazel router.register(r"^model", priority=0)(FakeProvider) router.register(r"^model", priority=10)(AnotherFakeProvider) resolved = router.resolve("model-v1") self.assertEqual(resolved, AnotherFakeProvider) def test_no_provider_registered(self): """Test error when no provider matches.""" with self.assertRaisesRegex( exceptions.InferenceConfigError, "No provider registered for model_id='unknown-model'", ): router.resolve("unknown-model") def test_caching(self): """Test that resolve results are cached.""" # Use direct registration for test provider to avoid module path issues router.register(r"^cached")(FakeProvider) # First call result1 = router.resolve("cached-model") # Second call should return cached result result2 = router.resolve("cached-model") self.assertIs(result1, result2) def test_clear_registry(self): """Test clearing the router.""" # Use direct registration for test provider to avoid module path issues router.register(r"^temp")(FakeProvider) # Should resolve before clear resolved = router.resolve("temp-model") self.assertEqual(resolved, FakeProvider) # Clear registry router.clear() # Should fail after clear with self.assertRaises(exceptions.InferenceConfigError): router.resolve("temp-model") def test_load_builtins_once_deduplicates_and_recovers_after_clear(self): """Built-in lazy registration can be repeated or restored after clear.""" providers_module._builtins_loaded = False # pylint: disable=protected-access providers_module.load_builtins_once() first_providers = router.list_providers() providers_module.load_builtins_once() self.assertLen(first_providers, len(builtin_registry.BUILTIN_PROVIDERS)) self.assertEqual(router.list_providers(), first_providers) router.clear() providers_module.load_builtins_once() self.assertEqual(router.list_providers(), first_providers) def test_list_entries(self): """Test listing registered entries.""" router.register_lazy(r"^test1", target="fake:Target1", priority=5) router.register_lazy( r"^test2", r"^test3", target="fake:Target2", priority=10 ) entries = router.list_entries() self.assertEqual(len(entries), 2) patterns1, priority1 = entries[0] self.assertEqual(patterns1, ["^test1"]) self.assertEqual(priority1, 5) patterns2, priority2 = entries[1] self.assertEqual(set(patterns2), {"^test2", "^test3"}) self.assertEqual(priority2, 10) def test_lazy_loading_defers_import(self): """Test that lazy registration doesn't import until resolve.""" # Register with a module that would fail if imported router.register_lazy(r"^lazy", target="non.existent.module:Provider") # Registration should succeed without importing entries = router.list_entries() self.assertTrue(any("^lazy" in patterns for patterns, _ in entries)) # Only on resolve should it try to import and fail with self.assertRaises(ModuleNotFoundError): router.resolve("lazy-model") def test_regex_pattern_objects(self): """Test using pre-compiled regex patterns.""" pattern = re.compile(r"^custom-\d+") @router.register(pattern) class CustomProvider(FakeProvider): pass self.assertEqual(router.resolve("custom-123"), CustomProvider) # Should not match without digits with self.assertRaises(exceptions.InferenceConfigError): router.resolve("custom-abc") def test_resolve_provider_by_name(self): """Test resolving provider by exact name.""" @router.register(r"^test-model", r"^TestProvider$") class TestProvider(FakeProvider): pass # Resolve by exact class name pattern provider = router.resolve_provider("TestProvider") self.assertEqual(provider, TestProvider) # Resolve by partial name match provider = router.resolve_provider("test") self.assertEqual(provider, TestProvider) def test_resolve_provider_not_found(self): """Test resolve_provider raises for unknown provider.""" with self.assertRaises(exceptions.InferenceConfigError) as cm: router.resolve_provider("UnknownProvider") self.assertIn("No provider found matching", str(cm.exception)) def test_hf_style_model_id_patterns(self): """Test that Hugging Face style model ID patterns work. This addresses issue #129 where HF-style model IDs like 'meta-llama/Llama-3.2-1B-Instruct' weren't being recognized. """ @router.register( r"^meta-llama/[Ll]lama", r"^google/gemma", r"^mistralai/[Mm]istral", r"^microsoft/phi", r"^Qwen/", r"^TinyLlama/", priority=100, ) class TestHFProvider(base_model.BaseLanguageModel): # pylint: disable=too-few-public-methods def infer(self, batch_prompts, **kwargs): return [] hf_model_ids = [ "meta-llama/Llama-3.2-1B-Instruct", "meta-llama/llama-2-7b", "google/gemma-2b", "mistralai/Mistral-7B-v0.1", "microsoft/phi-3-mini", "Qwen/Qwen2.5-7B", "TinyLlama/TinyLlama-1.1B-Chat-v1.0", ] for model_id in hf_model_ids: with self.subTest(model_id=model_id): provider_class = router.resolve(model_id) self.assertEqual(provider_class, TestHFProvider) if __name__ == "__main__": absltest.main()