chore: allow cli to pick model from model list

holmes/core/llm.py

@@ -6,9 +6,9 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 import litellm
+import sentry_sdk
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse, TextCompletionResponse
-import sentry_sdk
 from pydantic import BaseModel, ConfigDict, SecretStr
 from typing_extensions import Self
 
@@ -17,15 +17,14 @@
     RobustaModelsResponse,
     fetch_robusta_models,
 )
-
 from holmes.common.env_vars import (
+    EXTRA_HEADERS,
     FALLBACK_CONTEXT_WINDOW_SIZE,
     LOAD_ALL_ROBUSTA_MODELS,
     REASONING_EFFORT,
     ROBUSTA_AI,
     ROBUSTA_API_ENDPOINT,
     THINKING,
-    EXTRA_HEADERS,
     TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT,
     TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_TOKENS,
 )
@@ -538,14 +537,25 @@ def _init_models(self):
             self.configure_robusta_ai_model()
 
         if self._should_load_config_model():
-            self._llms[self.config.model] = self._create_model_entry(
-                model=self.config.model,
-                model_name=self.config.model,
-                base_url=self.config.api_base,
-                is_robusta_model=False,
-                api_key=self.config.api_key,
-                api_version=self.config.api_version,
-            )
+            # Chances are that for the cli, the user configres model_list and specify which model to use in the flag,
+            # so we don't override existing model entry from the model list file.
+            existing_model = self._llms.get(self.config.model)
+            # When self.config.api_key and self.config.model are both set, we assume the user wants to override the
+            # model entry from the model list file.
+            if not existing_model or self.config.api_key:
+                logging.info(
+                    f"Loaded model from config: {self.config.model} (is_robusta_model=False)"
+                )
+                # it's ok to assign api_base to base_url for azure openai service, since base_url will be assigned back
+                # to api_base in Config._get_llm() when api_base is not set.
+                self._llms[self.config.model] = self._create_model_entry(
+                    model=self.config.model,
+                    model_name=self.config.model,
+                    base_url=self.config.api_base,
+                    is_robusta_model=False,
+                    api_key=self.config.api_key,
+                    api_version=self.config.api_version,
+                )
 
     def _should_load_config_model(self) -> bool:
         if self.config.model is not None:

tests/core/test_llm_model_registry.py

@@ -0,0 +1,168 @@
+import tempfile
+from unittest.mock import MagicMock, patch
+
+import pytest
+import yaml
+from pydantic import SecretStr
+
+from holmes.config import Config
+from holmes.core.llm import LLMModelRegistry, ModelEntry
+from holmes.core.supabase_dal import SupabaseDal
+
+
+class TestLLMModelRegistry:
+    """Test cases for LLMModelRegistry class focusing on _init_models method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.mock_dal = MagicMock(spec=SupabaseDal)
+        self.mock_dal.enabled = True
+        self.mock_dal.account_id = "test-account"
+
+    def create_temp_model_file(self, model_data: dict) -> str:
+        """Create a temporary model list file for testing."""
+        temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False)
+        yaml.dump(model_data, temp_file)
+        temp_file.flush()
+        return temp_file.name
+
+    def test_init_models_existing_model_with_api_key_overwrites(self):
+        """Test that existing model entry is overwritten when config.api_key is set."""
+        # Create a model file with existing model
+        model_data = {
+            "test-model": {
+                "model": "gpt-4",
+                "api_key": "original-key",
+                "api_base": "https://original.api.com",
+                "api_version": "2023-01-01"
+            }
+        }
+        temp_file = self.create_temp_model_file(model_data)
+
+        with patch("holmes.core.llm.MODEL_LIST_FILE_LOCATION", temp_file), \
+                patch.object(LLMModelRegistry, "_should_load_robusta_ai", return_value=False):
+
+            # Create config with api_key set - this should trigger overwrite
+            config = Config(
+                model="test-model",
+                api_key="new-key",
+                api_base="https://new.api.com",
+                api_version="2024-01-01"
+            )
+
+            registry = LLMModelRegistry(config, self.mock_dal)
+
+            # Check that model entry was completely overwritten with new values
+            model_entry = registry.models["test-model"]
+            assert model_entry.model == "test-model"
+            assert model_entry.name == "test-model"
+            assert model_entry.api_key == SecretStr("new-key")
+            assert model_entry.base_url == "https://new.api.com"
+            assert model_entry.api_version == "2024-01-01"
+            assert model_entry.is_robusta_model is False
+
+    def test_init_models_existing_model_without_api_key_preserves(self):
+        """Test that existing model entry is preserved when config.api_key is not set."""
+        # Create a model file with existing model
+        model_data = {
+            "test-model": {
+                "model": "gpt-4",
+                "api_key": "original-key",
+                "api_base": "https://original.api.com",
+                "api_version": "2023-01-01",
+                "is_robusta_model": True
+            }
+        }
+        temp_file = self.create_temp_model_file(model_data)
+
+        with patch("holmes.core.llm.MODEL_LIST_FILE_LOCATION", temp_file), \
+                patch.object(LLMModelRegistry, "_should_load_robusta_ai", return_value=False):
+
+            # Create config without api_key - this should preserve existing model
+            config = Config(
+                model="test-model",
+                api_key=None,  # No API key set
+                api_base="https://new.api.com",  # These won't be used
+                api_version="2024-01-01"  # These won't be used
+            )
+
+            registry = LLMModelRegistry(config, self.mock_dal)
+
+            # Check that original model entry was preserved
+            model_entry = registry.models["test-model"]
+            assert model_entry.model == "gpt-4"  # Original value
+            assert model_entry.api_key == SecretStr("original-key")  # Original value
+            assert model_entry.api_base == "https://original.api.com"  # Original value
+            assert model_entry.api_version == "2023-01-01"  # Original value
+            assert model_entry.is_robusta_model is True  # Original value
+
+    def test_init_models_nonexistent_model_creates_new_entry(self):
+        """Test that new model entry is created when model doesn't exist in file."""
+        # Create a model file without the target model
+        model_data = {
+            "other-model": {
+                "model": "gpt-3.5",
+                "api_key": "other-key"
+            }
+        }
+        temp_file = self.create_temp_model_file(model_data)
+
+        with patch("holmes.core.llm.MODEL_LIST_FILE_LOCATION", temp_file), \
+                patch.object(LLMModelRegistry, "_should_load_robusta_ai", return_value=False):
+
+            # Create config for a model that doesn't exist in file
+            config = Config(
+                model="new-model",
+                api_base="https://new.api.com",
+                api_key="new-key",
+                api_version="2024-01-01"
+            )
+
+            registry = LLMModelRegistry(config, self.mock_dal)
+
+            # Check that new model entry was created
+            assert "new-model" in registry.models
+            assert "other-model" in registry.models  # Original model should still exist
+
+            new_model_entry = registry.models["new-model"]
+            assert new_model_entry.model == "new-model"
+            assert new_model_entry.name == "new-model"
+            assert new_model_entry.base_url == "https://new.api.com"
+            assert new_model_entry.api_key == SecretStr("new-key")
+            assert new_model_entry.api_version == "2024-01-01"
+            assert new_model_entry.is_robusta_model is False
+
+    def test_init_models_nonexistent_model_without_api_key_still_creates(self):
+        """Test that new model entry is created even without api_key when model doesn't exist."""
+        # Create a model file without the target model
+        model_data = {
+            "other-model": {
+                "model": "gpt-3.5",
+                "api_key": "other-key"
+            }
+        }
+        temp_file = self.create_temp_model_file(model_data)
+
+        with patch("holmes.core.llm.MODEL_LIST_FILE_LOCATION", temp_file), \
+                patch.object(LLMModelRegistry, "_should_load_robusta_ai", return_value=False):
+
+            # Create config for nonexistent model without api_key
+            config = Config(
+                model="new-model",
+                api_base="https://new.api.com",
+                api_key=None,  # No API key
+                api_version="2024-01-01"
+            )
+
+            registry = LLMModelRegistry(config, self.mock_dal)
+
+            # Check that new model entry was still created (because not existing_model)
+            assert "new-model" in registry.models
+
+            new_model_entry = registry.models["new-model"]
+            assert new_model_entry.model == "new-model"
+            assert new_model_entry.name == "new-model"
+            assert new_model_entry.base_url == "https://new.api.com"
+            assert new_model_entry.api_key is None
+            assert new_model_entry.api_version == "2024-01-01"
+            assert new_model_entry.is_robusta_model is False