From 51e062b9d2e74c3d3a3cae5e254162d8f3115aea Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Thu, 4 Sep 2025 10:47:51 +0300 Subject: [PATCH 1/5] WIP --- conftest.py | 31 +++++++ holmes/clients/robusta_client.py | 10 ++- holmes/config.py | 80 +++++++++++++------ .../test_config_load_model_list.py | 46 +++++++++++ 4 files changed, 141 insertions(+), 26 deletions(-) diff --git a/conftest.py b/conftest.py index a67ea79121..7eb95d8e76 100644 --- a/conftest.py +++ b/conftest.py @@ -3,6 +3,8 @@ from tests.llm.conftest import show_llm_summary_report from holmes.core.tracing import readable_timestamp, get_active_branch_name from tests.llm.utils.braintrust import get_braintrust_url +from unittest.mock import MagicMock, patch +import pytest def pytest_addoption(parser): @@ -126,3 +128,32 @@ def pytest_report_header(config): # due to pytest quirks, we need to define this in the main conftest.py - when defined in the llm conftest.py it # is SOMETIMES picked up and sometimes not, depending on how the test was invokedr pytest_terminal_summary = show_llm_summary_report + + +@pytest.fixture(autouse=True) +def patch_supabase(monkeypatch): + monkeypatch.setattr("holmes.core.supabase_dal.ROBUSTA_ACCOUNT_ID", "test-cluster") + monkeypatch.setattr( + "holmes.core.supabase_dal.STORE_URL", "https://fakesupabaseref.supabase.co" + ) + monkeypatch.setattr( + "holmes.core.supabase_dal.STORE_API_KEY", + "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyb2xlIjoiYW5vbiIsImlhdCI6MTYzNTAwODQ4NywiZXhwIjoxOTUwNTg0NDg3fQ.l8IgkO7TQokGSc9OJoobXIVXsOXkilXl4Ak6SCX5qI8", + ) + monkeypatch.setattr("holmes.core.supabase_dal.STORE_EMAIL", "mock_store_user") + monkeypatch.setattr( + "holmes.core.supabase_dal.STORE_PASSWORD", "mock_store_password" + ) + + +@pytest.fixture(autouse=True, scope="session") +def storage_dal_mock(): + with patch("holmes.config.SupabaseDal") as MockSupabaseDal: + mock_supabase_dal_instance = MagicMock() + MockSupabaseDal.return_value = mock_supabase_dal_instance + mock_supabase_dal_instance.sign_in.return_value = "mock_supabase_user_id" + mock_supabase_dal_instance.get_ai_credentials.return_value = ( + "mock_account_id", + "mock_session_token", + ) + yield mock_supabase_dal_instance diff --git a/holmes/clients/robusta_client.py b/holmes/clients/robusta_client.py index e99722aee5..e38960a0a4 100644 --- a/holmes/clients/robusta_client.py +++ b/holmes/clients/robusta_client.py @@ -14,8 +14,14 @@ class HolmesInfo(BaseModel): latest_version: Optional[str] = None +class RobustaModelsResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + models: List[str] + default_model: Optional[str] = None + + @cache -def fetch_robusta_models(account_id, token) -> Optional[List[str]]: +def fetch_robusta_models(account_id, token) -> Optional[RobustaModelsResponse]: try: session_request = {"session_token": token, "account_id": account_id} resp = requests.post( @@ -25,7 +31,7 @@ def fetch_robusta_models(account_id, token) -> Optional[List[str]]: ) resp.raise_for_status() response_json = resp.json() - return response_json.get("models") + return RobustaModelsResponse(**response_json) except Exception: logging.exception("Failed to fetch robusta models") return None diff --git a/holmes/config.py b/holmes/config.py index 8d60e60101..a90bb50936 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -135,6 +135,7 @@ class Config(RobustaBaseConfig): _server_tool_executor: Optional[ToolExecutor] = None _toolset_manager: Optional[ToolsetManager] = None + _default_robusta_model: Optional[str] = None @property def toolset_manager(self) -> ToolsetManager: @@ -170,20 +171,27 @@ def configure_robusta_ai_model(self) -> None: self._load_default_robusta_config() return - models = fetch_robusta_models( + robusta_models = fetch_robusta_models( self.account_id, self.session_token.get_secret_value() ) - if not models: + if not robusta_models: self._load_default_robusta_config() return - for model in models: + for model in robusta_models.models: logging.info(f"Loading Robusta AI model: {model}") self._model_list[model] = { + "name": model, "base_url": f"{ROBUSTA_API_ENDPOINT}/llm/{model}", "is_robusta_model": True, } + if robusta_models.default_model: + logging.info( + f"Setting default Robusta AI model to: {robusta_models.default_model}" + ) + self._default_robusta_model = robusta_models.default_model + except Exception: logging.exception("Failed to get all robusta models") # fallback to default behavior @@ -193,6 +201,7 @@ def _load_default_robusta_config(self): if self._should_load_robusta_ai() and self.api_key: logging.info("Loading default Robusta AI model") self._model_list[ROBUSTA_AI_MODEL_NAME] = { + "name": ROBUSTA_AI_MODEL_NAME, "base_url": ROBUSTA_API_ENDPOINT, "is_robusta_model": True, } @@ -525,32 +534,55 @@ def create_slack_destination(self) -> "SlackDestination": raise ValueError("--slack-channel must be specified") return SlackDestination(self.slack_token.get_secret_value(), self.slack_channel) + def _get_model_params(self, model_key: Optional[str] = None) -> dict: + if not self._model_list: + logging.info("No model list setup, using config model") + return {} + + if model_key: + model_params = self._model_list.get(model_key) + if model_params is not None: + logging.info(f"Using model: {model_key}") + return model_params.copy() + + logging.error(f"Couldn't find model: {model_key} in model list") + + if self._default_robusta_model: + model_params = self._model_list.get(self._default_robusta_model) + if model_params is not None: + logging.info( + f"Using default Robusta AI model: {self._default_robusta_model}" + ) + return model_params.copy() + + logging.error( + f"Couldn't find default Robusta AI model: {self._default_robusta_model} in model list" + ) + + first_model_params = next(iter(self._model_list.values())).copy() + logging.info("Using first model") + return first_model_params + def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "LLM": api_key = self.api_key + model_params = self._get_model_params(model_key) model = self.model api_base = self.api_base api_version = self.api_version - model_params = {} - if self._model_list: - # get requested model or the first credentials if no model requested. - model_params = ( - self._model_list.get(model_key, {}).copy() - if model_key - else next(iter(self._model_list.values())).copy() - ) - is_robusta_model = model_params.pop("is_robusta_model", False) - if is_robusta_model and self.api_key: - # we set here the api_key since it is being refresh when exprided and not as part of the model loading. - api_key = self.api_key.get_secret_value() # type: ignore - else: - api_key = model_params.pop("api_key", api_key) - model = model_params.pop("model", model) - # It's ok if the model does not have api base and api version, which are defaults to None. - # Handle both api_base and base_url - api_base takes precedence - model_api_base = model_params.pop("api_base", None) - model_base_url = model_params.pop("base_url", None) - api_base = model_api_base or model_base_url or api_base - api_version = model_params.pop("api_version", api_version) + + is_robusta_model = model_params.pop("is_robusta_model", False) + if is_robusta_model and self.api_key: + # we set here the api_key since it is being refresh when exprided and not as part of the model loading. + api_key = self.api_key.get_secret_value() # type: ignore + else: + api_key = model_params.pop("api_key", api_key) + model = model_params.pop("model", model) + # It's ok if the model does not have api base and api version, which are defaults to None. + # Handle both api_base and base_url - api_base takes precedence + model_api_base = model_params.pop("api_base", None) + model_base_url = model_params.pop("base_url", None) + api_base = model_api_base or model_base_url or api_base + api_version = model_params.pop("api_version", api_version) return DefaultLLM(model, api_key, api_base, api_version, model_params, tracer) # type: ignore diff --git a/tests/config_class/test_config_load_model_list.py b/tests/config_class/test_config_load_model_list.py index d78373a778..4385170ebe 100644 --- a/tests/config_class/test_config_load_model_list.py +++ b/tests/config_class/test_config_load_model_list.py @@ -1,5 +1,7 @@ +import responses import yaml +from holmes.common.env_vars import ROBUSTA_API_ENDPOINT from holmes.config import Config @@ -19,3 +21,47 @@ def test_load_custom_toolsets_config_valid(monkeypatch, tmp_path): assert len(list(config._model_list.keys())) == 3 az = config._model_list.get("azure") assert az.get("model") == "test-value" + + +@responses.activate +def test_load_custom_toolsets_config_valid_with_robusta_ai(monkeypatch, tmp_path): + responses.post( + "https://api.robusta.dev/api/llm/models", + json={ + "models": [ + "Robusta/sonnet-4 preview", + "Robusta/gpt-5-mini preview (minimal reasoning)", + "Robusta/gpt-5 preview (minimal reasoning)", + "Robusta/gpt-4o", + ], + "default_model": "Robusta/gpt-5-mini preview (minimal reasoning)", + }, + ) + temp_config_file = tmp_path / "custom_toolset.yaml" + data = { + "bedrock": {"model": "bbbb", "api-key": "asfffd"}, + } + + temp_config_file.write_text(yaml.dump(data)) + monkeypatch.setattr("holmes.config.MODEL_LIST_FILE_LOCATION", str(temp_config_file)) + monkeypatch.setattr("holmes.config.ROBUSTA_AI", True) + monkeypatch.setenv("CLUSTER_NAME", "test-cluster") + + config = Config.load_from_env() + assert isinstance(config._model_list, dict) + assert len(list(config._model_list.keys())) == 5 + assert config._model_list.get("bedrock").get("model") == "bbbb" + + sonnet_model = config._model_list["Robusta/sonnet-4 preview"] + + assert ( + sonnet_model.get("base_url") + == f"{ROBUSTA_API_ENDPOINT}/llm/Robusta/sonnet-4 preview" + ) + assert sonnet_model.get("name") == "Robusta/sonnet-4 preview" + assert sonnet_model.get("is_robusta_model") + + assert ( + config._default_robusta_model + == "Robusta/gpt-5-mini preview (minimal reasoning)" + ) From 1a6fe17c150e5f7121d4aafa48cb9a57707b4422 Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Sun, 7 Sep 2025 15:30:14 +0300 Subject: [PATCH 2/5] WIP --- holmes/config.py | 16 +++--- holmes/core/llm.py | 4 +- poetry.lock | 19 ++++++- pyproject.toml | 1 + .../test_config_api_base_version.py | 15 ++++-- tests/config_class/test_config_get_llm.py | 43 ++++++++++++++++ .../test_config_load_model_list.py | 38 +++----------- .../test_config_load_robusta_ai.py | 12 +++-- tests/conftest.py | 49 +++++++++++++++++++ tests/plugins/toolsets/test_internet.py | 7 ++- 10 files changed, 156 insertions(+), 48 deletions(-) create mode 100644 tests/config_class/test_config_get_llm.py create mode 100644 tests/conftest.py diff --git a/holmes/config.py b/holmes/config.py index a90bb50936..cace5bcfce 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -29,7 +29,6 @@ # Source plugin imports moved to their respective create methods to speed up startup if TYPE_CHECKING: - from holmes.core.llm import LLM from holmes.core.tool_calling_llm import IssueInvestigator, ToolCallingLLM from holmes.plugins.destinations.slack import SlackDestination from holmes.plugins.sources.github import GitHubSource @@ -184,6 +183,7 @@ def configure_robusta_ai_model(self) -> None: "name": model, "base_url": f"{ROBUSTA_API_ENDPOINT}/llm/{model}", "is_robusta_model": True, + "model": "gpt-4o", # Robusta AI model is using openai like API. } if robusta_models.default_model: @@ -563,10 +563,8 @@ def _get_model_params(self, model_key: Optional[str] = None) -> dict: logging.info("Using first model") return first_model_params - def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "LLM": - api_key = self.api_key + def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "DefaultLLM": model_params = self._get_model_params(model_key) - model = self.model api_base = self.api_base api_version = self.api_version @@ -575,16 +573,20 @@ def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "LLM": # we set here the api_key since it is being refresh when exprided and not as part of the model loading. api_key = self.api_key.get_secret_value() # type: ignore else: - api_key = model_params.pop("api_key", api_key) - model = model_params.pop("model", model) + api_key = model_params.pop("api_key", self.api_key) + + model = model_params.pop("model", self.model) # It's ok if the model does not have api base and api version, which are defaults to None. # Handle both api_base and base_url - api_base takes precedence model_api_base = model_params.pop("api_base", None) model_base_url = model_params.pop("base_url", None) api_base = model_api_base or model_base_url or api_base api_version = model_params.pop("api_version", api_version) + model_name = model_params.pop("name", None) or model_key or model - return DefaultLLM(model, api_key, api_base, api_version, model_params, tracer) # type: ignore + return DefaultLLM( + model, api_key, api_base, api_version, model_params, tracer, model_name + ) # type: ignore def get_models_list(self) -> List[str]: if self._model_list: diff --git a/holmes/core/llm.py b/holmes/core/llm.py index 5bd4bdb30c..9cdf7e22bc 100644 --- a/holmes/core/llm.py +++ b/holmes/core/llm.py @@ -72,7 +72,8 @@ def __init__( api_base: Optional[str] = None, api_version: Optional[str] = None, args: Optional[Dict] = None, - tracer=None, + tracer: Optional[Any] = None, + name: Optional[str] = None, ): self.model = model self.api_key = api_key @@ -80,6 +81,7 @@ def __init__( self.api_version = api_version self.args = args or {} self.tracer = tracer + self.name = name self.check_llm(self.model, self.api_key, self.api_base, self.api_version) diff --git a/poetry.lock b/poetry.lock index 6a2915fc4a..6de322188c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3799,6 +3799,23 @@ files = [ pytest = ">=5.0.0" python-dotenv = ">=0.9.1" +[[package]] +name = "pytest-responses" +version = "0.5.1" +description = "py.test integration for responses" +optional = false +python-versions = "*" +files = [ + {file = "pytest_responses-0.5.1-py2.py3-none-any.whl", hash = "sha256:4172e565b94ac1ea3b10aba6e40855ad60cd7f141476b2d8a47e4b5f250be734"}, +] + +[package.dependencies] +pytest = ">=2.5" +responses = "*" + +[package.extras] +tests = ["flake8"] + [[package]] name = "pytest-shared-session-scope" version = "0.4.0" @@ -5319,4 +5336,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "c8ac23cbe9c68f1cd71d7a409c31b8e276f847bdb67d875d8928ef6b4b07e5dc" +content-hash = "34e951fa027f712d739155b3d6bbafe1c47183d4dc05645f367ba7d1bf051ffe" diff --git a/pyproject.toml b/pyproject.toml index 2174c7dcaa..d4aecc9152 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ pytest-cov = "^6.2.1" types-python-dateutil = "^2.9.0.20250708" pytest-dotenv = "^0.5.2" pytest-shared-session-scope = "^0.4.0" +pytest-responses = "^0.5.1" [build-system] requires = ["poetry-core"] diff --git a/tests/config_class/test_config_api_base_version.py b/tests/config_class/test_config_api_base_version.py index a10aaf0c3f..5e76098407 100644 --- a/tests/config_class/test_config_api_base_version.py +++ b/tests/config_class/test_config_api_base_version.py @@ -35,7 +35,6 @@ def test_config_get_llm_with_api_base_version(): """Test that Config._get_llm passes api_base and api_version to DefaultLLM.""" config = Config( model="test-model", - api_key="test-key", api_base="https://test.api.base", api_version="2023-12-01", ) @@ -49,7 +48,7 @@ def test_config_get_llm_with_api_base_version(): # Check that DefaultLLM was called with the right positional arguments call_args = mock_default_llm.call_args[0] assert call_args[0] == "test-model" - assert call_args[1].get_secret_value() == "test-key" # api_key is SecretStr + assert call_args[1] is None assert call_args[2] == "https://test.api.base" assert call_args[3] == "2023-12-01" assert call_args[4] == {} @@ -85,7 +84,8 @@ def test_config_get_llm_with_model_list_api_base_version(monkeypatch, tmp_path): "https://model.api.base", "2024-02-01", {}, - None, # tracer + None, # tracer, + "test-model", ) assert result == mock_llm_instance @@ -122,6 +122,7 @@ def test_config_get_llm_model_list_overrides_config_values(monkeypatch, tmp_path "2024-03-01", # from model list {}, None, # tracer + "test-model", ) @@ -156,6 +157,7 @@ def test_config_get_llm_model_list_defaults_to_config_values(monkeypatch, tmp_pa "2023-01-01", # from config {}, None, # tracer + "test-model", ) @@ -202,6 +204,7 @@ def test_config_get_llm_with_non_none_model_list_first_model_fallback( "2024-01-01", # from first model {}, None, # tracer + "gpt-4", ) @@ -255,7 +258,8 @@ def test_config_get_llm_with_specific_model_from_model_list(monkeypatch, tmp_pat "https://openai.api.base", # from openai-gpt35 model "2024-04-01", # from openai-gpt35 model {}, - None, # tracer + None, # tracer, + "openai-gpt35", ) @@ -289,6 +293,7 @@ def test_config_get_llm_with_base_url_only(monkeypatch, tmp_path): "2024-01-01", {}, None, # tracer + "test-model", ) @@ -323,6 +328,7 @@ def test_config_get_llm_api_base_overrides_base_url(monkeypatch, tmp_path): "2024-01-01", {}, None, # tracer + "test-model", ) @@ -358,6 +364,7 @@ def test_config_get_llm_neither_api_base_nor_base_url_uses_config( "2024-01-01", {}, None, # tracer + "test-model", ) diff --git a/tests/config_class/test_config_get_llm.py b/tests/config_class/test_config_get_llm.py new file mode 100644 index 0000000000..e428181bb5 --- /dev/null +++ b/tests/config_class/test_config_get_llm.py @@ -0,0 +1,43 @@ +from holmes.config import Config +from holmes.core.llm import DefaultLLM + +from tests.conftest import DEFAULT_ROBUSTA_MODEL + + +def test_config_get_llm_no_model_key_returns_default_model(server_config: Config): + llm: DefaultLLM = server_config._get_llm() + assert llm.name == DEFAULT_ROBUSTA_MODEL + assert llm.model == "gpt-4o" + assert llm.api_base == f"https://api.robusta.dev/llm/{DEFAULT_ROBUSTA_MODEL}" + + +def test_confgi_get_llm_with_model_key_returns_model_from_config( + server_config: Config, monkeypatch +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "access_key_id") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret_access_key") + llm: DefaultLLM = server_config._get_llm(model_key="our_local_model") + assert llm.name == "our_local_model" + assert llm.model == "bedrock/custom_ai_model" + assert llm.api_key == "existing_api_key" + + +def test_config_get_llm_unexisting_model_key_returns_default_model( + server_config: Config, +): + llm: DefaultLLM = server_config._get_llm(model_key="unexisting_model") + assert llm.name == DEFAULT_ROBUSTA_MODEL + assert llm.model == "gpt-4o" + assert llm.api_base == f"https://api.robusta.dev/llm/{DEFAULT_ROBUSTA_MODEL}" + + +def test_config_get_llm_no_default_model_fallback_to_first_available_model( + server_config: Config, monkeypatch +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "access_key_id") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret_access_key") + server_config._default_robusta_model = None + llm: DefaultLLM = server_config._get_llm() + assert llm.name == "bedrock/custom_ai_model" + assert llm.model == "bedrock/custom_ai_model" + assert llm.api_key == "existing_api_key" diff --git a/tests/config_class/test_config_load_model_list.py b/tests/config_class/test_config_load_model_list.py index 4385170ebe..bf028b9d09 100644 --- a/tests/config_class/test_config_load_model_list.py +++ b/tests/config_class/test_config_load_model_list.py @@ -1,4 +1,3 @@ -import responses import yaml from holmes.common.env_vars import ROBUSTA_API_ENDPOINT @@ -23,36 +22,15 @@ def test_load_custom_toolsets_config_valid(monkeypatch, tmp_path): assert az.get("model") == "test-value" -@responses.activate -def test_load_custom_toolsets_config_valid_with_robusta_ai(monkeypatch, tmp_path): - responses.post( - "https://api.robusta.dev/api/llm/models", - json={ - "models": [ - "Robusta/sonnet-4 preview", - "Robusta/gpt-5-mini preview (minimal reasoning)", - "Robusta/gpt-5 preview (minimal reasoning)", - "Robusta/gpt-4o", - ], - "default_model": "Robusta/gpt-5-mini preview (minimal reasoning)", - }, +def test_config_load_model_list_valid_with_robusta_ai(server_config): + assert isinstance(server_config._model_list, dict) + assert len(list(server_config._model_list.keys())) == 5 + assert ( + server_config._model_list.get("our_local_model").get("model") + == "bedrock/custom_ai_model" ) - temp_config_file = tmp_path / "custom_toolset.yaml" - data = { - "bedrock": {"model": "bbbb", "api-key": "asfffd"}, - } - - temp_config_file.write_text(yaml.dump(data)) - monkeypatch.setattr("holmes.config.MODEL_LIST_FILE_LOCATION", str(temp_config_file)) - monkeypatch.setattr("holmes.config.ROBUSTA_AI", True) - monkeypatch.setenv("CLUSTER_NAME", "test-cluster") - - config = Config.load_from_env() - assert isinstance(config._model_list, dict) - assert len(list(config._model_list.keys())) == 5 - assert config._model_list.get("bedrock").get("model") == "bbbb" - sonnet_model = config._model_list["Robusta/sonnet-4 preview"] + sonnet_model = server_config._model_list["Robusta/sonnet-4 preview"] assert ( sonnet_model.get("base_url") @@ -62,6 +40,6 @@ def test_load_custom_toolsets_config_valid_with_robusta_ai(monkeypatch, tmp_path assert sonnet_model.get("is_robusta_model") assert ( - config._default_robusta_model + server_config._default_robusta_model == "Robusta/gpt-5-mini preview (minimal reasoning)" ) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index ab1797f93e..11410bd60e 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -1,6 +1,10 @@ from unittest.mock import patch from pydantic import SecretStr from holmes.config import Config +from holmes.clients.robusta_client import RobustaModelsResponse + + +ROBUSTA_TEST_MODELS = RobustaModelsResponse(models=["Robusta/test"], default_model=None) def fake_load_robusta_api_key(config, _): @@ -16,7 +20,7 @@ def test_cli_not_loading_robusta_ai(*, monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.fetch_robusta_models", return_value=ROBUSTA_TEST_MODELS) @patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_loads_robusta_ai_when_true(mock_cluster, mock_fetch, *, monkeypatch): def fake_loader(self, dal): @@ -30,7 +34,7 @@ def fake_loader(self, dal): @patch("holmes.config.ROBUSTA_AI", None) -@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.fetch_robusta_models", return_value=ROBUSTA_TEST_MODELS) @patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_loads_robusta_ai_when_not_exists_and_not_other_models( mock_cluster, mock_fetch, *, monkeypatch @@ -53,7 +57,7 @@ def test_server_not_loads_robusta_ai_when_false(mock_cluster, *, monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.fetch_robusta_models", return_value=ROBUSTA_TEST_MODELS) @patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch( "holmes.config.parse_models_file", @@ -102,7 +106,7 @@ def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists( @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.fetch_robusta_models", return_value=ROBUSTA_TEST_MODELS) @patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_loads_robusta_ai_when_model_var_exists( mock_cluster, mock_fetch, *, monkeypatch diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000..a9249b4a0f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,49 @@ +import pytest +import yaml + +from holmes.config import Config + +DEFAULT_ROBUSTA_MODEL = "Robusta/gpt-5-mini preview (minimal reasoning)" +ROBUSTA_MODELS = [ + "Robusta/sonnet-4 preview", + "Robusta/gpt-5-mini preview (minimal reasoning)", + "Robusta/gpt-5 preview (minimal reasoning)", + "Robusta/gpt-4o", +] + + +@pytest.fixture(autouse=True) +def clear_all_caches(): + """Clear all function caches that may affect test isolation.""" + try: + import holmes.clients.robusta_client as rc + + rc.fetch_robusta_models.cache_clear() + rc.fetch_holmes_info.cache_clear() + except Exception: + pass + + +@pytest.fixture(autouse=False) +def server_config(tmp_path, monkeypatch, responses): + responses.post( + "https://api.robusta.dev/api/llm/models", + json={ + "models": ROBUSTA_MODELS, + "default_model": DEFAULT_ROBUSTA_MODEL, + }, + ) + temp_config_file = tmp_path / "custom_toolset.yaml" + data = { + "our_local_model": { + "model": "bedrock/custom_ai_model", + "api_key": "existing_api_key", + }, + } + + temp_config_file.write_text(yaml.dump(data)) + monkeypatch.setattr("holmes.config.MODEL_LIST_FILE_LOCATION", str(temp_config_file)) + monkeypatch.setattr("holmes.config.ROBUSTA_AI", True) + monkeypatch.setenv("CLUSTER_NAME", "test-cluster") + + return Config.load_from_env() diff --git a/tests/plugins/toolsets/test_internet.py b/tests/plugins/toolsets/test_internet.py index 570dfbe243..1277b2f5d7 100644 --- a/tests/plugins/toolsets/test_internet.py +++ b/tests/plugins/toolsets/test_internet.py @@ -112,7 +112,12 @@ def test_html_to_markdown(fixture: Fixture): assert match, f"Values mismatch. Run the following command to compare expected with actual: `diff {fixture.expected_output_file_path} {actual_file_path_for_debugging}`" -def test_fetch_webpage(): +def test_fetch_webpage(responses): + responses.get( + TEST_URL, + status=200, + body=EXPECTED_TEST_RESULT, + ) toolset = InternetToolset() toolset.status = ToolsetStatusEnum.ENABLED tool_executor = ToolExecutor(toolsets=[toolset]) From 24fcc7a1ab253d6e16e32c8b823812e9a14c2fc1 Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Sun, 7 Sep 2025 13:55:06 +0300 Subject: [PATCH 3/5] Fix api_key --- holmes/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/holmes/config.py b/holmes/config.py index cace5bcfce..4712b3d5af 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -573,7 +573,7 @@ def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "DefaultLLM" # we set here the api_key since it is being refresh when exprided and not as part of the model loading. api_key = self.api_key.get_secret_value() # type: ignore else: - api_key = model_params.pop("api_key", self.api_key) + api_key = model_params.pop("api_key", None) model = model_params.pop("model", self.model) # It's ok if the model does not have api base and api version, which are defaults to None. From cf0817e2da13bbf3a50da740f838107cb0e62af5 Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Sun, 7 Sep 2025 15:03:07 +0300 Subject: [PATCH 4/5] PR Comments --- holmes/clients/robusta_client.py | 4 +++- holmes/config.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/holmes/clients/robusta_client.py b/holmes/clients/robusta_client.py index e38960a0a4..5065167ec1 100644 --- a/holmes/clients/robusta_client.py +++ b/holmes/clients/robusta_client.py @@ -21,7 +21,9 @@ class RobustaModelsResponse(BaseModel): @cache -def fetch_robusta_models(account_id, token) -> Optional[RobustaModelsResponse]: +def fetch_robusta_models( + account_id: str, token: str +) -> Optional[RobustaModelsResponse]: try: session_request = {"session_token": token, "account_id": account_id} resp = requests.post( diff --git a/holmes/config.py b/holmes/config.py index 4712b3d5af..bbded99743 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -173,7 +173,7 @@ def configure_robusta_ai_model(self) -> None: robusta_models = fetch_robusta_models( self.account_id, self.session_token.get_secret_value() ) - if not robusta_models: + if not robusta_models or not robusta_models.models: self._load_default_robusta_config() return @@ -204,6 +204,7 @@ def _load_default_robusta_config(self): "name": ROBUSTA_AI_MODEL_NAME, "base_url": ROBUSTA_API_ENDPOINT, "is_robusta_model": True, + "model": "gpt-4o", } def _should_load_robusta_ai(self) -> bool: From eb7a4f77d5e2d61d3ab0f06f043f88ae58adacea Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Sun, 7 Sep 2025 15:22:30 +0300 Subject: [PATCH 5/5] PR Comments --- holmes/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/holmes/config.py b/holmes/config.py index bbded99743..58fccdd997 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -206,6 +206,7 @@ def _load_default_robusta_config(self): "is_robusta_model": True, "model": "gpt-4o", } + self._default_robusta_model = ROBUSTA_AI_MODEL_NAME def _should_load_robusta_ai(self) -> bool: if not self.should_try_robusta_ai: