Refactor src/models/registry.py to reduce code duplication

Author: ishaanxgupta

src/models/registry.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import importlib
 import logging
 from typing import Optional
 
@@ -20,43 +21,36 @@
 
 logger = logging.getLogger("xmem.models")
 
+
+def _build_from_module(module_name: str, func_name: str, **kwargs) -> BaseChatModel:
+    module = importlib.import_module(f"src.models.{module_name}")
+    factory_fn = getattr(module, func_name)
+    return factory_fn(**kwargs)
+
+
 _BUILDERS = {
-    "gemini": lambda **kw: _build_gemini(**kw),
-    "claude": lambda **kw: _build_claude(**kw),
-    "openai": lambda **kw: _build_openai(**kw),
+    "gemini": lambda **kw: _build_from_module("gemini", "build_gemini_model", **kw),
+    "claude": lambda **kw: _build_from_module("claude", "build_claude_model", **kw),
+    "openai": lambda **kw: _build_from_module("openai", "build_openai_model", **kw),
 }
 
+
 _KEY_MAP = {
     "gemini": lambda: settings.gemini_api_key,
     "claude": lambda: settings.claude_api_key,
     "openai": lambda: settings.openai_api_key,
 }
 
 
-def _build_gemini(**kw) -> BaseChatModel:
-    from src.models.gemini import build_gemini_model
-    return build_gemini_model(**kw)
-
-
-def _build_claude(**kw) -> BaseChatModel:
-    from src.models.claude import build_claude_model
-    return build_claude_model(**kw)
-
-
-def _build_openai(**kw) -> BaseChatModel:
-    from src.models.openai import build_openai_model
-    return build_openai_model(**kw)
-
-
 def get_model(
     provider: Optional[Provider] = None,
     model_name: Optional[str] = None,
     temperature: Optional[float] = None,
 ) -> BaseChatModel:
     """Build and return a chat model.
 
-    If *provider* is None the first provider from ``settings.fallback_order``
-    whose API key is configured will be used.  Raises ``RuntimeError`` if no
+    If *provider* is None the first provider from settings.fallback_order
+    whose API key is configured will be used.  Raises RuntimeError if no
     provider can be initialised.
     """
     kw: dict = {}

tests/unit/models/test_registry.py

@@ -0,0 +1,90 @@
+import pytest
+from unittest.mock import MagicMock, patch
+import sys
+import importlib
+
+
+@pytest.fixture
+def mock_modules():
+    with patch.dict(
+        sys.modules,
+        {
+            "src.models.gemini": MagicMock(),
+            "src.models.claude": MagicMock(),
+            "src.models.openai": MagicMock(),
+        },
+    ):
+        yield
+
+
+def test_build_gemini(mock_modules):
+    from src.models import registry
+
+    importlib.reload(registry)  # Reload to ensure _BUILDERS is fresh if needed
+
+    mock_builder = MagicMock()
+    sys.modules["src.models.gemini"].build_gemini_model = mock_builder
+
+    # We are testing the _BUILDERS map effectively
+    registry._BUILDERS["gemini"](temperature=0.5)
+
+    mock_builder.assert_called_once_with(temperature=0.5)
+
+
+def test_build_claude(mock_modules):
+    from src.models import registry
+
+    importlib.reload(registry)
+
+    mock_builder = MagicMock()
+    sys.modules["src.models.claude"].build_claude_model = mock_builder
+
+    registry._BUILDERS["claude"](model_name="claude-3")
+
+    mock_builder.assert_called_once_with(model_name="claude-3")
+
+
+def test_build_openai(mock_modules):
+    from src.models import registry
+
+    importlib.reload(registry)
+
+    mock_builder = MagicMock()
+    sys.modules["src.models.openai"].build_openai_model = mock_builder
+
+    registry._BUILDERS["openai"]()
+
+    mock_builder.assert_called_once()
+
+
+def test_get_model_specific_provider(mock_modules):
+    from src.models import registry
+
+    importlib.reload(registry)
+
+    mock_builder = MagicMock()
+    sys.modules["src.models.gemini"].build_gemini_model = mock_builder
+
+    registry.get_model("gemini", temperature=0.7)
+
+    mock_builder.assert_called_once_with(temperature=0.7)
+
+
+def test_get_model_fallback(mock_modules):
+    from src.models import registry
+
+    importlib.reload(registry)
+
+    # Mock settings.fallback_order and API keys
+    # Note: registry imports settings, so we need to patch it in registry
+    with patch("src.models.registry.settings") as mock_settings:
+        mock_settings.fallback_order = ["openai", "gemini"]
+        mock_settings.openai_api_key = "sk-test"
+        mock_settings.gemini_api_key = None
+
+        mock_openai_builder = MagicMock()
+        sys.modules["src.models.openai"].build_openai_model = mock_openai_builder
+
+        registry.get_model()
+
+        mock_openai_builder.assert_called_once()