test_sam_component_base_model_init.py
1 """Unit tests for SamComponentBase model initialization (lazy loading).""" 2 3 import unittest 4 from typing import Any 5 from unittest.mock import patch, MagicMock 6 7 from sam_test_infrastructure.feature_flags import mock_flags 8 from solace_agent_mesh.common.features import core as feature_flags 9 from solace_agent_mesh.common.sac.sam_component_base import SamComponentBase 10 from solace_agent_mesh.agent.adk.models.lite_llm import LiteLlm 11 12 13 class ConcreteSamComponent(SamComponentBase): 14 """Concrete test implementation of SamComponentBase.""" 15 16 def __init__(self, info: dict[str, Any], **kwargs: Any): 17 super().__init__(info, **kwargs) 18 19 async def _handle_message_async(self, message, topic: str) -> None: 20 pass 21 22 def _get_component_id(self) -> str: 23 return "test_component_id" 24 25 def _get_component_type(self) -> str: 26 return "test_component" 27 28 def _pre_async_cleanup(self) -> None: 29 pass 30 31 async def _async_setup_and_run(self) -> None: 32 pass 33 34 def _on_model_status_change(self, old_status: str, new_status: str): 35 self.last_status_change = (old_status, new_status) 36 37 38 class TestSamComponentBaseModelInit(unittest.TestCase): 39 """Test _initialize_model and get_lite_llm_model.""" 40 41 def setUp(self): 42 feature_flags.initialize() 43 self.test_info = { 44 "component_name": "test_component", 45 "component_module": "test_module", 46 "component_config": { 47 "namespace": "test/namespace", 48 "max_message_size_bytes": 1024000, 49 }, 50 } 51 52 def _make_component(self, model_config=None, model_provider=None, lazy=False): 53 """Create a ConcreteSamComponent with mocked config.""" 54 config_map = { 55 "namespace": "test/namespace", 56 "max_message_size_bytes": 1024000, 57 "model": model_config, 58 "model_provider": model_provider, 59 } 60 with mock_flags(model_config_ui=lazy): 61 with patch.object( 62 SamComponentBase, 63 "get_config", 64 side_effect=lambda key, *args: config_map.get( 65 key, args[0] if args else None 66 ), 67 ): 68 component = ConcreteSamComponent(self.test_info) 69 # Re-patch get_config for subsequent calls 70 self._config_patcher = patch.object( 71 SamComponentBase, 72 "get_config", 73 side_effect=lambda key, *args: config_map.get( 74 key, args[0] if args else None 75 ), 76 ) 77 self._config_patcher.start() 78 return component 79 80 def tearDown(self): 81 if hasattr(self, "_config_patcher"): 82 self._config_patcher.stop() 83 feature_flags._reset_for_testing() 84 85 def test_initialize_model_with_string_config(self): 86 """String model config should create a LiteLlm instance.""" 87 component = self._make_component(model_config="gpt-4") 88 model = component._initialize_model() 89 assert isinstance(model, LiteLlm) 90 assert model.status == "ready" 91 assert component.adk_model_instance is model 92 93 def test_initialize_model_with_dict_config(self): 94 """Dict model config should create a LiteLlm instance with settings.""" 95 component = self._make_component( 96 model_config={"model": "gpt-4", "timeout": 300} 97 ) 98 model = component._initialize_model() 99 assert isinstance(model, LiteLlm) 100 assert model.status == "ready" 101 assert model._model_config["timeout"] == 300 102 103 def test_initialize_model_dict_applies_default_resilience(self): 104 """Dict config without 'type' should get default num_retries and timeout.""" 105 component = self._make_component(model_config={"model": "gpt-4"}) 106 model = component._initialize_model() 107 assert model._model_config["num_retries"] == 3 108 assert model._model_config["timeout"] == 120 109 110 def test_initialize_model_dict_no_defaults_when_type_set(self): 111 """Dict config with 'type' should not get default resilience settings.""" 112 component = self._make_component( 113 model_config={"model": "gpt-4", "type": "custom"} 114 ) 115 model = component._initialize_model() 116 assert "num_retries" not in model._model_config 117 assert "timeout" not in model._model_config 118 119 def test_initialize_model_invalid_type_raises(self): 120 """Non-string, non-dict model config should raise ValueError.""" 121 component = self._make_component(model_config=12345) 122 with self.assertRaises(ValueError): 123 component._initialize_model() 124 125 def test_initialize_model_lazy_mode_with_provider(self): 126 """In lazy mode with model_provider, LiteLlm starts in 'initializing'.""" 127 component = self._make_component( 128 model_provider=["dynamic-provider"], lazy=True 129 ) 130 model = component._initialize_model() 131 assert isinstance(model, LiteLlm) 132 assert model.status == "initializing" 133 134 def test_initialize_model_lazy_mode_with_string_config(self): 135 """In lazy mode with string config, LiteLlm should be created with model=string.""" 136 component = self._make_component( 137 model_config="gpt-4", model_provider=["dynamic-provider"], lazy=True 138 ) 139 model = component._initialize_model() 140 assert isinstance(model, LiteLlm) 141 # With a real model name, configure_model sets status to "ready" 142 assert model._model_config.get("model") == "gpt-4" 143 144 def test_initialize_model_lazy_mode_with_dict_config(self): 145 """In lazy mode with dict config, LiteLlm should unpack the dict.""" 146 component = self._make_component( 147 model_config={"model": "gpt-4", "timeout": 300}, 148 model_provider=["dynamic-provider"], 149 lazy=True, 150 ) 151 model = component._initialize_model() 152 assert isinstance(model, LiteLlm) 153 # Dict config should be unpacked, preserving extra settings 154 assert model._model_config["timeout"] == 300 155 assert model._model_config.get("model") == "gpt-4" 156 157 def test_initialize_model_lazy_mode_dict_preserves_model_name(self): 158 """In lazy mode with dict config, the model name should be preserved.""" 159 component = self._make_component( 160 model_config={"model": "claude-3-opus"}, 161 model_provider=["dynamic-provider"], 162 lazy=True, 163 ) 164 model = component._initialize_model() 165 assert isinstance(model, LiteLlm) 166 assert model._model_config.get("model") == "claude-3-opus" 167 168 def test_initialize_model_lazy_mode_no_config_starts_initializing(self): 169 """In lazy mode with no model config, LiteLlm starts in 'initializing' state.""" 170 component = self._make_component( 171 model_config=None, model_provider=["dynamic-provider"], lazy=True 172 ) 173 model = component._initialize_model() 174 assert isinstance(model, LiteLlm) 175 assert model.status == "initializing" 176 177 def test_get_lite_llm_model_returns_cached_instance(self): 178 """get_lite_llm_model should return the same instance on repeated calls.""" 179 component = self._make_component(model_config="gpt-4") 180 model1 = component.get_lite_llm_model() 181 model2 = component.get_lite_llm_model() 182 assert model1 is model2 183 184 def test_get_lite_llm_model_initializes_on_first_call(self): 185 """get_lite_llm_model should call _initialize_model if no instance exists.""" 186 component = self._make_component(model_config="gpt-4") 187 assert component.adk_model_instance is None 188 model = component.get_lite_llm_model() 189 assert model is not None 190 assert component.adk_model_instance is model 191 192 def test_get_lite_llm_model_returns_none_when_no_config(self): 193 """get_lite_llm_model returns None when no model or provider configured.""" 194 component = self._make_component(model_config=None, model_provider=None) 195 result = component.get_lite_llm_model() 196 assert result is None 197 198 def test_on_model_status_change_callback_wired(self): 199 """_on_model_status_change should be wired to LiteLlm instance.""" 200 component = self._make_component(model_config="gpt-4") 201 model = component._initialize_model() 202 # Trigger unconfigure -> should call component's callback 203 model.unconfigure_model() 204 assert hasattr(component, "last_status_change") 205 assert component.last_status_change == ("ready", "none") 206 207 208 class TestSamComponentBaseLazyModelMode(unittest.TestCase): 209 """Test the _lazy_model_mode flag.""" 210 211 def setUp(self): 212 feature_flags.initialize() 213 self.test_info = { 214 "component_name": "test_component", 215 "component_module": "test_module", 216 "component_config": { 217 "namespace": "test/namespace", 218 "max_message_size_bytes": 1024000, 219 }, 220 } 221 222 def tearDown(self): 223 feature_flags._reset_for_testing() 224 225 def _make_component_with_config(self, config_map, lazy=False): 226 with mock_flags(model_config_ui=lazy): 227 with patch.object( 228 SamComponentBase, 229 "get_config", 230 side_effect=lambda key, *args: config_map.get( 231 key, args[0] if args else None 232 ), 233 ): 234 return ConcreteSamComponent(self.test_info) 235 236 def test_lazy_mode_enabled_when_env_true(self): 237 """_lazy_model_mode should be True when MODEL_CONFIG_UI=true.""" 238 config_map = { 239 "namespace": "test/namespace", 240 "max_message_size_bytes": 1024000, 241 "model_provider": None, 242 } 243 component = self._make_component_with_config(config_map, lazy=True) 244 assert component._lazy_model_mode is True 245 246 def test_lazy_mode_disabled_by_default(self): 247 """_lazy_model_mode should be False when flag is not enabled.""" 248 config_map = { 249 "namespace": "test/namespace", 250 "max_message_size_bytes": 1024000, 251 "model_provider": None, 252 } 253 component = self._make_component_with_config(config_map, lazy=False) 254 assert component._lazy_model_mode is False 255 256 def test_model_provider_extracted_from_list(self): 257 """model_provider should be extracted as the first element of the list.""" 258 config_map = { 259 "namespace": "test/namespace", 260 "max_message_size_bytes": 1024000, 261 "model_provider": ["provider-a", "provider-b"], 262 } 263 component = self._make_component_with_config(config_map) 264 assert component.model_provider == "provider-a" 265 266 def test_model_provider_none_when_empty_list(self): 267 """model_provider should be None when config is an empty list.""" 268 config_map = { 269 "namespace": "test/namespace", 270 "max_message_size_bytes": 1024000, 271 "model_provider": [], 272 } 273 component = self._make_component_with_config(config_map) 274 assert component.model_provider is None 275 276 def test_model_provider_none_when_not_configured(self): 277 """model_provider should be None when not in config.""" 278 config_map = { 279 "namespace": "test/namespace", 280 "max_message_size_bytes": 1024000, 281 "model_provider": None, 282 } 283 component = self._make_component_with_config(config_map) 284 assert component.model_provider is None