From baa0c11f9465f100772ba0c4d013b68c9c157ecd Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Mon, 18 Aug 2025 14:25:37 +0300 Subject: [PATCH 01/13] support_multiple_llms --- holmes/config.py | 27 ++++++++++++++++++++++++--- holmes/utils/robusta.py | 4 ++-- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/holmes/config.py b/holmes/config.py index 1b10347ada..8eaab019c7 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -6,6 +6,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, List, Optional, Union +import requests # type: ignore import yaml # type: ignore from pydantic import BaseModel, ConfigDict, FilePath, SecretStr @@ -24,7 +25,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.supabase_dal import SupabaseDal from holmes.core.tool_calling_llm import IssueInvestigator, ToolCallingLLM from holmes.plugins.destinations.slack import SlackDestination from holmes.plugins.sources.github import GitHubSource @@ -33,6 +33,7 @@ from holmes.plugins.sources.pagerduty import PagerDutySource from holmes.plugins.sources.prometheus.plugin import AlertManagerSource +from holmes.core.supabase_dal import SupabaseDal from holmes.core.config import config_path_dir from holmes.utils.definitions import RobustaConfig from holmes.utils.env import replace_env_vars_values @@ -136,8 +137,28 @@ def toolset_manager(self) -> ToolsetManager: def model_post_init(self, __context: Any) -> None: self._model_list = parse_models_file(MODEL_LIST_FILE_LOCATION) - if self._should_load_robusta_ai(): - logging.info("Loading Robusta AI model") + if not self._should_load_robusta_ai() or not self.cluster_name: + return + try: + dal = SupabaseDal(self.cluster_name) + account_id, token = dal.get_ai_credentials() + session_request = {"session_token": token, "account_id": account_id} + resp = requests.post( + f"{ROBUSTA_API_ENDPOINT}/api/llm/models", + json=session_request, + timeout=10, + ) + resp.raise_for_status() + response_json = resp.json() + models = response_json.get("models", [ROBUSTA_AI_MODEL_NAME]) + for model in models: + logging.info("Loading Robusta AI model") + self._model_list[model] = { + "base_url": ROBUSTA_API_ENDPOINT + f"/deployments/{model}", + } + except Exception: + logging.exception("Failed to get robusta models") + # defaulting to previous behavior self._model_list[ROBUSTA_AI_MODEL_NAME] = { "base_url": ROBUSTA_API_ENDPOINT, } diff --git a/holmes/utils/robusta.py b/holmes/utils/robusta.py index 57c320fa30..e86e1bf0b2 100644 --- a/holmes/utils/robusta.py +++ b/holmes/utils/robusta.py @@ -1,9 +1,9 @@ -from holmes.config import Config, ROBUSTA_AI_MODEL_NAME +from holmes.config import Config, ROBUSTA_AI from holmes.core.supabase_dal import SupabaseDal from pydantic import SecretStr def load_robusta_api_key(dal: SupabaseDal, config: Config): - if ROBUSTA_AI_MODEL_NAME in config._model_list: + if ROBUSTA_AI: account_id, token = dal.get_ai_credentials() config.api_key = SecretStr(f"{account_id} {token}") From 9e4a36371dfb3dafd4944dff4cf202609d1c5422 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Mon, 18 Aug 2025 14:30:20 +0300 Subject: [PATCH 02/13] refactoring --- holmes/config.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/holmes/config.py b/holmes/config.py index 8eaab019c7..8175cb3a0a 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -137,12 +137,21 @@ def toolset_manager(self) -> ToolsetManager: def model_post_init(self, __context: Any) -> None: self._model_list = parse_models_file(MODEL_LIST_FILE_LOCATION) - if not self._should_load_robusta_ai() or not self.cluster_name: + + if not self._should_load_robusta_ai(): return + + self._load_robusta_ai_models() + + def _load_robusta_ai_models(self) -> None: try: + if not self.cluster_name: + return + dal = SupabaseDal(self.cluster_name) account_id, token = dal.get_ai_credentials() session_request = {"session_token": token, "account_id": account_id} + resp = requests.post( f"{ROBUSTA_API_ENDPOINT}/api/llm/models", json=session_request, @@ -150,15 +159,18 @@ def model_post_init(self, __context: Any) -> None: ) resp.raise_for_status() response_json = resp.json() + models = response_json.get("models", [ROBUSTA_AI_MODEL_NAME]) for model in models: - logging.info("Loading Robusta AI model") + logging.info(f"Loading Robusta AI model: {model}") self._model_list[model] = { - "base_url": ROBUSTA_API_ENDPOINT + f"/deployments/{model}", + "base_url": f"{ROBUSTA_API_ENDPOINT}/deployments/{model}", } + except Exception: - logging.exception("Failed to get robusta models") - # defaulting to previous behavior + logging.exception("Failed to get all robusta models") + # fallback to default behavior + logging.info("Loading default Robusta AI model") self._model_list[ROBUSTA_AI_MODEL_NAME] = { "base_url": ROBUSTA_API_ENDPOINT, } From 0caf2d79ac8c5472ede6853e7298b241861c47a8 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 11:43:53 +0300 Subject: [PATCH 03/13] fixing tests --- .../test_config_load_robusta_ai.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 8e1748c36d..5e98c9b576 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -1,31 +1,33 @@ from unittest.mock import patch from holmes.config import Config +ROBUSTA_MODEL_NAME = "Robusta/gpt-4o" + @patch("holmes.config.ROBUSTA_AI", True) def test_cli_not_loading_robusta_ai(monkeypatch): config = Config.load_from_file(None) - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) def test_server_loads_robusta_ai_when_true(monkeypatch): config = Config.load_from_env() - assert "Robusta" in config._model_list + assert ROBUSTA_MODEL_NAME in config._model_list @patch("holmes.config.ROBUSTA_AI", None) def test_server_loads_robusta_ai_when_not_exists_and_not_other_models(monkeypatch): config = Config.load_from_env() assert len(config._model_list) == 1 - assert "Robusta" in config._model_list + assert ROBUSTA_MODEL_NAME in config._model_list @patch("holmes.config.ROBUSTA_AI", False) def test_server_not_loads_robusta_ai_when_false(monkeypatch): config = Config.load_from_env() assert len(config._model_list) == 0 - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) @@ -36,7 +38,7 @@ def test_server_not_loads_robusta_ai_when_false(monkeypatch): def test_server_loads_robusta_ai_when_true_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert "Robusta" in config._model_list + assert ROBUSTA_MODEL_NAME in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @@ -47,7 +49,7 @@ def test_server_loads_robusta_ai_when_true_and_model_list_exists(monkeypatch): def test_server_not_loads_robusta_ai_when_false_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list @patch("holmes.config.ROBUSTA_AI", None) @@ -58,21 +60,21 @@ def test_server_not_loads_robusta_ai_when_false_and_model_list_exists(monkeypatc def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) def test_server_loads_robusta_ai_when_model_var_exists(monkeypatch): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert "Robusta" in config._model_list + assert ROBUSTA_MODEL_NAME in config._model_list @patch("holmes.config.ROBUSTA_AI", None) def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var(monkeypatch): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @@ -81,4 +83,4 @@ def test_server_not_loads_robusta_ai_when_model_var_exists_and_false_env_var( ): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert "Robusta" not in config._model_list + assert ROBUSTA_MODEL_NAME not in config._model_list From 7851df921441970e70d47891e34de8d1017dc7b6 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 14:10:17 +0300 Subject: [PATCH 04/13] fix test --- holmes/config.py | 12 ++++++---- .../test_config_load_robusta_ai.py | 22 +++++++++---------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/holmes/config.py b/holmes/config.py index 8175cb3a0a..9b42086ef5 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -146,6 +146,7 @@ def model_post_init(self, __context: Any) -> None: def _load_robusta_ai_models(self) -> None: try: if not self.cluster_name: + self._load_default_robusta_config() return dal = SupabaseDal(self.cluster_name) @@ -170,10 +171,13 @@ def _load_robusta_ai_models(self) -> None: except Exception: logging.exception("Failed to get all robusta models") # fallback to default behavior - logging.info("Loading default Robusta AI model") - self._model_list[ROBUSTA_AI_MODEL_NAME] = { - "base_url": ROBUSTA_API_ENDPOINT, - } + self._load_default_robusta_config() + + def _load_default_robusta_config(self): + logging.info("Loading default Robusta AI model") + self._model_list[ROBUSTA_AI_MODEL_NAME] = { + "base_url": ROBUSTA_API_ENDPOINT, + } def _should_load_robusta_ai(self) -> bool: if not self.should_try_robusta_ai: diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 5e98c9b576..8e1748c36d 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -1,33 +1,31 @@ from unittest.mock import patch from holmes.config import Config -ROBUSTA_MODEL_NAME = "Robusta/gpt-4o" - @patch("holmes.config.ROBUSTA_AI", True) def test_cli_not_loading_robusta_ai(monkeypatch): config = Config.load_from_file(None) - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) def test_server_loads_robusta_ai_when_true(monkeypatch): config = Config.load_from_env() - assert ROBUSTA_MODEL_NAME in config._model_list + assert "Robusta" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) def test_server_loads_robusta_ai_when_not_exists_and_not_other_models(monkeypatch): config = Config.load_from_env() assert len(config._model_list) == 1 - assert ROBUSTA_MODEL_NAME in config._model_list + assert "Robusta" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) def test_server_not_loads_robusta_ai_when_false(monkeypatch): config = Config.load_from_env() assert len(config._model_list) == 0 - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) @@ -38,7 +36,7 @@ def test_server_not_loads_robusta_ai_when_false(monkeypatch): def test_server_loads_robusta_ai_when_true_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert ROBUSTA_MODEL_NAME in config._model_list + assert "Robusta" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @@ -49,7 +47,7 @@ def test_server_loads_robusta_ai_when_true_and_model_list_exists(monkeypatch): def test_server_not_loads_robusta_ai_when_false_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", None) @@ -60,21 +58,21 @@ def test_server_not_loads_robusta_ai_when_false_and_model_list_exists(monkeypatc def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists(monkeypatch): config = Config.load_from_env() assert "existing_model" in config._model_list - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) def test_server_loads_robusta_ai_when_model_var_exists(monkeypatch): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert ROBUSTA_MODEL_NAME in config._model_list + assert "Robusta" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var(monkeypatch): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @@ -83,4 +81,4 @@ def test_server_not_loads_robusta_ai_when_model_var_exists_and_false_env_var( ): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert ROBUSTA_MODEL_NAME not in config._model_list + assert "Robusta" not in config._model_list From 968bd8e1304371911de27554c0f28abe12aa1764 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 15:25:27 +0300 Subject: [PATCH 05/13] fix loading tools to db --- holmes/clients/robusta_client.py | 18 +++++++++++++- holmes/common/env_vars.py | 1 + holmes/config.py | 42 +++++++++++++++----------------- server.py | 1 + 4 files changed, 38 insertions(+), 24 deletions(-) diff --git a/holmes/clients/robusta_client.py b/holmes/clients/robusta_client.py index 94950af64c..56ba0b072c 100644 --- a/holmes/clients/robusta_client.py +++ b/holmes/clients/robusta_client.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import List, Optional import requests # type: ignore from functools import cache from pydantic import BaseModel, ConfigDict @@ -13,6 +13,22 @@ class HolmesInfo(BaseModel): latest_version: Optional[str] = None +@cache +def fetch_robusta_models(account_id, token) -> Optional[List[str]]: + try: + session_request = {"session_token": token, "account_id": account_id} + resp = requests.post( + f"{ROBUSTA_API_ENDPOINT}/api/llm/models", + json=session_request, + timeout=10, + ) + resp.raise_for_status() + response_json = resp.json() + return response_json.get("models") + except Exception: + return None + + @cache def fetch_holmes_info() -> Optional[HolmesInfo]: try: diff --git a/holmes/common/env_vars.py b/holmes/common/env_vars.py index 6077635ef7..55261f7671 100644 --- a/holmes/common/env_vars.py +++ b/holmes/common/env_vars.py @@ -27,6 +27,7 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]: STORE_PASSWORD = os.environ.get("STORE_PASSWORD", "") HOLMES_POST_PROCESSING_PROMPT = os.environ.get("HOLMES_POST_PROCESSING_PROMPT", "") ROBUSTA_AI = load_bool("ROBUSTA_AI", None) +ROBUSTA_ALL_AI_MODELS = load_bool("ROBUSTA_ALL_AI_MODELS", False) ROBUSTA_API_ENDPOINT = os.environ.get("ROBUSTA_API_ENDPOINT", "https://api.robusta.dev") LOG_PERFORMANCE = os.environ.get("LOG_PERFORMANCE", None) diff --git a/holmes/config.py b/holmes/config.py index 3c825e6a7e..c47db52c25 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -6,13 +6,18 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, List, Optional, Union -import requests # type: ignore import yaml # type: ignore from pydantic import BaseModel, ConfigDict, FilePath, SecretStr +from holmes.clients.robusta_client import fetch_robusta_models from holmes.core.llm import DefaultLLM -from holmes.common.env_vars import ROBUSTA_AI, ROBUSTA_API_ENDPOINT, ROBUSTA_CONFIG_PATH +from holmes.common.env_vars import ( + ROBUSTA_AI, + ROBUSTA_ALL_AI_MODELS, + ROBUSTA_API_ENDPOINT, + ROBUSTA_CONFIG_PATH, +) from holmes.core.tools_utils.tool_executor import ToolExecutor from holmes.core.toolset_manager import ToolsetManager from holmes.plugins.runbooks import ( @@ -141,31 +146,28 @@ def model_post_init(self, __context: Any) -> None: if not self._should_load_robusta_ai(): return - self._load_robusta_ai_models() + self.configure_robusta_ai_model() - def _load_robusta_ai_models(self) -> None: + def configure_robusta_ai_model(self) -> None: + if not self.api_key: + return try: - if not self.cluster_name: + if not self.cluster_name or not ROBUSTA_ALL_AI_MODELS: self._load_default_robusta_config() return dal = SupabaseDal(self.cluster_name) account_id, token = dal.get_ai_credentials() - session_request = {"session_token": token, "account_id": account_id} - - resp = requests.post( - f"{ROBUSTA_API_ENDPOINT}/api/llm/models", - json=session_request, - timeout=10, - ) - resp.raise_for_status() - response_json = resp.json() + models = fetch_robusta_models(account_id, token) + if not models: + self._load_default_robusta_config() + return - models = response_json.get("models", [ROBUSTA_AI_MODEL_NAME]) for model in models: logging.info(f"Loading Robusta AI model: {model}") self._model_list[model] = { - "base_url": f"{ROBUSTA_API_ENDPOINT}/deployments/{model}", + "base_url": f"{ROBUSTA_API_ENDPOINT}/llm/{model}", + "api_key": self.api_key.get_secret_value(), } except Exception: @@ -174,14 +176,8 @@ def _load_robusta_ai_models(self) -> None: self._load_default_robusta_config() def _load_default_robusta_config(self): - logging.info("Loading default Robusta AI model") - self._model_list[ROBUSTA_AI_MODEL_NAME] = { - "base_url": ROBUSTA_API_ENDPOINT, - } - - def configure_robusta_ai_model(self) -> None: if self._should_load_robusta_ai() and self.api_key: - logging.info("Loading Robusta AI model") + logging.info("Loading default Robusta AI model") self._model_list[ROBUSTA_AI_MODEL_NAME] = { "base_url": ROBUSTA_API_ENDPOINT, "api_key": self.api_key.get_secret_value(), diff --git a/server.py b/server.py index c160122acb..5b62698454 100644 --- a/server.py +++ b/server.py @@ -81,6 +81,7 @@ def init_logging(): def sync_before_server_start(): + load_robusta_api_key(dal=dal, config=config) try: update_holmes_status_in_db(dal, config) except Exception: From c24f4d8f4b1d936213e0dfbde1e02ecdf38fac74 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 15:46:30 +0300 Subject: [PATCH 06/13] change env var name --- holmes/common/env_vars.py | 2 +- holmes/config.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/holmes/common/env_vars.py b/holmes/common/env_vars.py index 55261f7671..aa4602cd17 100644 --- a/holmes/common/env_vars.py +++ b/holmes/common/env_vars.py @@ -27,7 +27,7 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]: STORE_PASSWORD = os.environ.get("STORE_PASSWORD", "") HOLMES_POST_PROCESSING_PROMPT = os.environ.get("HOLMES_POST_PROCESSING_PROMPT", "") ROBUSTA_AI = load_bool("ROBUSTA_AI", None) -ROBUSTA_ALL_AI_MODELS = load_bool("ROBUSTA_ALL_AI_MODELS", False) +LOAD_ALL_ROBUSTA_MODELS = load_bool("LOAD_ALL_ROBUSTA_MODELS", True) ROBUSTA_API_ENDPOINT = os.environ.get("ROBUSTA_API_ENDPOINT", "https://api.robusta.dev") LOG_PERFORMANCE = os.environ.get("LOG_PERFORMANCE", None) diff --git a/holmes/config.py b/holmes/config.py index c47db52c25..fb0faf58aa 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -14,7 +14,7 @@ from holmes.core.llm import DefaultLLM from holmes.common.env_vars import ( ROBUSTA_AI, - ROBUSTA_ALL_AI_MODELS, + LOAD_ALL_ROBUSTA_MODELS, ROBUSTA_API_ENDPOINT, ROBUSTA_CONFIG_PATH, ) @@ -152,7 +152,7 @@ def configure_robusta_ai_model(self) -> None: if not self.api_key: return try: - if not self.cluster_name or not ROBUSTA_ALL_AI_MODELS: + if not self.cluster_name or not LOAD_ALL_ROBUSTA_MODELS: self._load_default_robusta_config() return From 29b9b675e3cd5a639416a52728c09ae6053d623b Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 15:54:08 +0300 Subject: [PATCH 07/13] move credentials to api --- holmes/clients/robusta_client.py | 5 ++++- holmes/config.py | 4 +--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/holmes/clients/robusta_client.py b/holmes/clients/robusta_client.py index 56ba0b072c..897bf63ab6 100644 --- a/holmes/clients/robusta_client.py +++ b/holmes/clients/robusta_client.py @@ -3,6 +3,7 @@ from functools import cache from pydantic import BaseModel, ConfigDict from holmes.common.env_vars import ROBUSTA_API_ENDPOINT +from holmes.core.supabase_dal import SupabaseDal HOLMES_GET_INFO_URL = f"{ROBUSTA_API_ENDPOINT}/api/holmes/get_info" TIMEOUT = 0.5 @@ -14,8 +15,10 @@ class HolmesInfo(BaseModel): @cache -def fetch_robusta_models(account_id, token) -> Optional[List[str]]: +def fetch_robusta_models(cluster_name) -> Optional[List[str]]: try: + dal = SupabaseDal(cluster_name) + account_id, token = dal.get_ai_credentials() session_request = {"session_token": token, "account_id": account_id} resp = requests.post( f"{ROBUSTA_API_ENDPOINT}/api/llm/models", diff --git a/holmes/config.py b/holmes/config.py index 08de52d41e..09ac4d49b3 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -156,9 +156,7 @@ def configure_robusta_ai_model(self) -> None: self._load_default_robusta_config() return - dal = SupabaseDal(self.cluster_name) - account_id, token = dal.get_ai_credentials() - models = fetch_robusta_models(account_id, token) + models = fetch_robusta_models(self.cluster_name) if not models: self._load_default_robusta_config() return From 4314534271dbc01bc37ba56d37fa0b0a9336b210 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 16:05:15 +0300 Subject: [PATCH 08/13] fixing tests --- .../test_config_load_robusta_ai.py | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 8e1748c36d..2d60acf504 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -9,16 +9,28 @@ def test_cli_not_loading_robusta_ai(monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) +@patch("holmes.config.api_key", "test") +@patch("holmes.config.cluster_name", "test") +@patch( + "holmes.client.robusta_client.fetch_robusta_models", + return_value=["Robusta/test"], +) def test_server_loads_robusta_ai_when_true(monkeypatch): config = Config.load_from_env() - assert "Robusta" in config._model_list + assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) +@patch("holmes.config.api_key", "test") +@patch("holmes.config.cluster_name", "test") +@patch( + "holmes.client.robusta_client.fetch_robusta_models", + return_value=["Robusta/test"], +) def test_server_loads_robusta_ai_when_not_exists_and_not_other_models(monkeypatch): config = Config.load_from_env() assert len(config._model_list) == 1 - assert "Robusta" in config._model_list + assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @@ -29,6 +41,7 @@ def test_server_not_loads_robusta_ai_when_false(monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) +@patch("holmes.config.LOAD_ALL_ROBUSTA_MODELS", False) @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, @@ -62,10 +75,16 @@ def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists(monke @patch("holmes.config.ROBUSTA_AI", True) +@patch("holmes.config.api_key", "test") +@patch("holmes.config.cluster_name", "test") +@patch( + "holmes.client.robusta_client.fetch_robusta_models", + return_value=["Robusta/test"], +) def test_server_loads_robusta_ai_when_model_var_exists(monkeypatch): monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() - assert "Robusta" in config._model_list + assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) From d2151b575e2bcf367a3b52f8ba8e05ce54746bb1 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 16:06:28 +0300 Subject: [PATCH 09/13] fix test --- tests/config_class/test_config_load_robusta_ai.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 2d60acf504..89616d84a8 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -41,7 +41,12 @@ def test_server_not_loads_robusta_ai_when_false(monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.LOAD_ALL_ROBUSTA_MODELS", False) +@patch("holmes.config.api_key", "test") +@patch("holmes.config.cluster_name", "test") +@patch( + "holmes.client.robusta_client.fetch_robusta_models", + return_value=["Robusta/test"], +) @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, From b8e00db873f46335a5354772831eeea60998978d Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Tue, 19 Aug 2025 17:22:45 +0300 Subject: [PATCH 10/13] fix pytests --- .../test_config_load_robusta_ai.py | 80 +++++++++++-------- 1 file changed, 45 insertions(+), 35 deletions(-) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 89616d84a8..0d0d2c14c3 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -3,106 +3,116 @@ @patch("holmes.config.ROBUSTA_AI", True) -def test_cli_not_loading_robusta_ai(monkeypatch): +def test_cli_not_loading_robusta_ai(*, monkeypatch): config = Config.load_from_file(None) assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.api_key", "test") -@patch("holmes.config.cluster_name", "test") -@patch( - "holmes.client.robusta_client.fetch_robusta_models", - return_value=["Robusta/test"], -) -def test_server_loads_robusta_ai_when_true(monkeypatch): +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +def test_server_loads_robusta_ai_when_true(mock_fetch, mock_cluster, *, monkeypatch): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) -@patch("holmes.config.api_key", "test") -@patch("holmes.config.cluster_name", "test") -@patch( - "holmes.client.robusta_client.fetch_robusta_models", - return_value=["Robusta/test"], -) -def test_server_loads_robusta_ai_when_not_exists_and_not_other_models(monkeypatch): +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +def test_server_loads_robusta_ai_when_not_exists_and_not_other_models( + mock_fetch, mock_cluster, *, monkeypatch +): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert len(config._model_list) == 1 assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) -def test_server_not_loads_robusta_ai_when_false(monkeypatch): +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +def test_server_not_loads_robusta_ai_when_false(mock_cluster, *, monkeypatch): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert len(config._model_list) == 0 assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.api_key", "test") -@patch("holmes.config.cluster_name", "test") -@patch( - "holmes.client.robusta_client.fetch_robusta_models", - return_value=["Robusta/test"], -) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, ) -def test_server_loads_robusta_ai_when_true_and_model_list_exists(monkeypatch): +def test_server_loads_robusta_ai_when_true_and_model_list_exists( + mock_parse, mock_fetch, mock_cluster, *, monkeypatch +): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "existing_model" in config._model_list - assert "Robusta" in config._model_list + assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, ) -def test_server_not_loads_robusta_ai_when_false_and_model_list_exists(monkeypatch): +def test_server_not_loads_robusta_ai_when_false_and_model_list_exists( + mock_parse, mock_cluster, *, monkeypatch +): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "existing_model" in config._model_list assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", None) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, ) -def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists(monkeypatch): +def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists( + mock_parse, mock_cluster, *, monkeypatch +): + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "existing_model" in config._model_list assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.api_key", "test") -@patch("holmes.config.cluster_name", "test") -@patch( - "holmes.client.robusta_client.fetch_robusta_models", - return_value=["Robusta/test"], -) -def test_server_loads_robusta_ai_when_model_var_exists(monkeypatch): +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +@patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +def test_server_loads_robusta_ai_when_model_var_exists( + mock_fetch, mock_cluster, *, monkeypatch +): monkeypatch.setenv("MODEL", "some_model") + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) -def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var(monkeypatch): +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") +def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var( + mock_cluster, *, monkeypatch +): monkeypatch.setenv("MODEL", "some_model") + monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", False) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_not_loads_robusta_ai_when_model_var_exists_and_false_env_var( - monkeypatch, + mock_cluster, *, monkeypatch ): + monkeypatch.setenv("API_KEY", "test") monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() assert "Robusta" not in config._model_list From 4783ec505edbf176fe254d99899809960e349552 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Wed, 20 Aug 2025 10:05:19 +0300 Subject: [PATCH 11/13] refactoring and code review --- holmes/clients/robusta_client.py | 7 +++---- holmes/config.py | 26 ++++++++++++++++++++++---- holmes/core/investigation.py | 5 ++--- holmes/utils/robusta.py | 10 ---------- server.py | 11 +++++------ tests/test_server_endpoints.py | 8 ++++---- 6 files changed, 36 insertions(+), 31 deletions(-) delete mode 100644 holmes/utils/robusta.py diff --git a/holmes/clients/robusta_client.py b/holmes/clients/robusta_client.py index 897bf63ab6..e99722aee5 100644 --- a/holmes/clients/robusta_client.py +++ b/holmes/clients/robusta_client.py @@ -1,9 +1,9 @@ +import logging from typing import List, Optional import requests # type: ignore from functools import cache from pydantic import BaseModel, ConfigDict from holmes.common.env_vars import ROBUSTA_API_ENDPOINT -from holmes.core.supabase_dal import SupabaseDal HOLMES_GET_INFO_URL = f"{ROBUSTA_API_ENDPOINT}/api/holmes/get_info" TIMEOUT = 0.5 @@ -15,10 +15,8 @@ class HolmesInfo(BaseModel): @cache -def fetch_robusta_models(cluster_name) -> Optional[List[str]]: +def fetch_robusta_models(account_id, token) -> Optional[List[str]]: try: - dal = SupabaseDal(cluster_name) - account_id, token = dal.get_ai_credentials() session_request = {"session_token": token, "account_id": account_id} resp = requests.post( f"{ROBUSTA_API_ENDPOINT}/api/llm/models", @@ -29,6 +27,7 @@ def fetch_robusta_models(cluster_name) -> Optional[List[str]]: response_json = resp.json() return response_json.get("models") except Exception: + logging.exception("Failed to fetch robusta models") return None diff --git a/holmes/config.py b/holmes/config.py index 09ac4d49b3..2e42b5b10a 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -79,6 +79,9 @@ class Config(RobustaBaseConfig): api_key: Optional[SecretStr] = ( None # if None, read from OPENAI_API_KEY or AZURE_OPENAI_ENDPOINT env var ) + account_id: Optional[str] = None + session_token: Optional[SecretStr] = None + model: Optional[str] = "gpt-4o" max_steps: int = 40 cluster_name: Optional[str] = None @@ -149,14 +152,22 @@ def model_post_init(self, __context: Any) -> None: self.configure_robusta_ai_model() def configure_robusta_ai_model(self) -> None: - if not self.api_key: - return try: if not self.cluster_name or not LOAD_ALL_ROBUSTA_MODELS: self._load_default_robusta_config() return - models = fetch_robusta_models(self.cluster_name) + if not self.api_key: + dal = SupabaseDal(self.cluster_name) + self.load_robusta_api_key(dal) + + if not self.account_id or not self.session_token: + self._load_default_robusta_config() + return + + models = fetch_robusta_models( + self.account_id, self.session_token.get_secret_value() + ) if not models: self._load_default_robusta_config() return @@ -165,7 +176,7 @@ def configure_robusta_ai_model(self) -> None: logging.info(f"Loading Robusta AI model: {model}") self._model_list[model] = { "base_url": f"{ROBUSTA_API_ENDPOINT}/llm/{model}", - "api_key": self.api_key.get_secret_value(), + "api_key": self.api_key.get_secret_value(), # type: ignore } except Exception: @@ -527,6 +538,13 @@ def get_models_list(self) -> List[str]: return json.dumps([self.model]) # type: ignore + def load_robusta_api_key(self, dal: SupabaseDal): + if ROBUSTA_AI: + account_id, token = dal.get_ai_credentials() + self.api_key = SecretStr(f"{account_id} {token}") + self.account_id = account_id + self.session_token = SecretStr(token) + class TicketSource(BaseModel): config: Config diff --git a/holmes/core/investigation.py b/holmes/core/investigation.py index f44be456a5..1440b63891 100644 --- a/holmes/core/investigation.py +++ b/holmes/core/investigation.py @@ -9,7 +9,6 @@ from holmes.core.supabase_dal import SupabaseDal from holmes.core.tracing import DummySpan, SpanType from holmes.utils.global_instructions import add_global_instructions_to_user_prompt -from holmes.utils.robusta import load_robusta_api_key from holmes.core.todo_manager import get_todo_manager from holmes.core.investigation_structured_output import ( @@ -28,7 +27,7 @@ def investigate_issues( model: Optional[str] = None, trace_span=DummySpan(), ) -> InvestigationResult: - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) context = dal.get_issue_data(investigate_request.context.get("robusta_issue_id")) resource_instructions = dal.get_resource_instructions( @@ -82,7 +81,7 @@ def get_investigation_context( config: Config, request_structured_output_from_llm: Optional[bool] = None, ): - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) ai = config.create_issue_investigator(dal=dal, model=investigate_request.model) raw_data = investigate_request.model_dump() diff --git a/holmes/utils/robusta.py b/holmes/utils/robusta.py deleted file mode 100644 index efa94516ea..0000000000 --- a/holmes/utils/robusta.py +++ /dev/null @@ -1,10 +0,0 @@ -from holmes.config import Config, ROBUSTA_AI -from holmes.core.supabase_dal import SupabaseDal -from pydantic import SecretStr - - -def load_robusta_api_key(dal: SupabaseDal, config: Config): - if ROBUSTA_AI: - account_id, token = dal.get_ai_credentials() - config.api_key = SecretStr(f"{account_id} {token}") - config.configure_robusta_ai_model() diff --git a/server.py b/server.py index 979c4b975b..27ec89cf67 100644 --- a/server.py +++ b/server.py @@ -23,7 +23,6 @@ from litellm.exceptions import AuthenticationError from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse -from holmes.utils.robusta import load_robusta_api_key from holmes.utils.stream import stream_investigate_formatter, stream_chat_formatter from holmes.common.env_vars import ( HOLMES_HOST, @@ -82,7 +81,7 @@ def init_logging(): def sync_before_server_start(): - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) try: update_holmes_status_in_db(dal, config) except Exception: @@ -185,7 +184,7 @@ def stream_investigate_issues(req: InvestigateRequest): @app.post("/api/workload_health_check") def workload_health_check(request: WorkloadHealthRequest): - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) try: resource = request.resource workload_alerts: list[str] = [] @@ -252,7 +251,7 @@ def workload_health_conversation( request: WorkloadHealthChatRequest, ): try: - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) ai = config.create_toolcalling_llm(dal=dal, model=request.model) global_instructions = dal.get_global_instructions_for_account() @@ -281,7 +280,7 @@ def workload_health_conversation( @app.post("/api/issue_chat") def issue_conversation(issue_chat_request: IssueChatRequest): try: - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) ai = config.create_toolcalling_llm(dal=dal, model=issue_chat_request.model) global_instructions = dal.get_global_instructions_for_account() @@ -320,7 +319,7 @@ def already_answered(conversation_history: Optional[List[dict]]) -> bool: @app.post("/api/chat") def chat(chat_request: ChatRequest): try: - load_robusta_api_key(dal=dal, config=config) + config.load_robusta_api_key(dal=dal) ai = config.create_toolcalling_llm(dal=dal, model=chat_request.model) global_instructions = dal.get_global_instructions_for_account() diff --git a/tests/test_server_endpoints.py b/tests/test_server_endpoints.py index 673bb74283..3565a6b080 100644 --- a/tests/test_server_endpoints.py +++ b/tests/test_server_endpoints.py @@ -9,7 +9,7 @@ def client(): return TestClient(app) -@patch("server.load_robusta_api_key") +@patch("holmes.config.Config.load_robusta_api_key") @patch("holmes.config.Config.create_toolcalling_llm") @patch("holmes.core.supabase_dal.SupabaseDal.get_global_instructions_for_account") def test_api_chat_all_fields( @@ -78,7 +78,7 @@ def test_api_chat_all_fields( assert "pre_action_notification_text" in action -@patch("server.load_robusta_api_key") +@patch("holmes.config.Config.load_robusta_api_key") @patch("holmes.config.Config.create_toolcalling_llm") @patch("holmes.core.supabase_dal.SupabaseDal.get_global_instructions_for_account") def test_api_issue_chat_all_fields( @@ -140,7 +140,7 @@ def test_api_issue_chat_all_fields( assert "result" in tool_call -@patch("server.load_robusta_api_key") +@patch("holmes.config.Config.load_robusta_api_key") @patch("holmes.config.Config.create_toolcalling_llm") @patch("holmes.core.supabase_dal.SupabaseDal.get_global_instructions_for_account") def test_api_workload_health_chat( @@ -205,7 +205,7 @@ def test_api_workload_health_chat( assert "result" in tool_call -@patch("server.load_robusta_api_key") +@patch("holmes.config.Config.load_robusta_api_key") @patch("holmes.config.Config.create_toolcalling_llm") @patch("holmes.core.supabase_dal.SupabaseDal.get_global_instructions_for_account") @patch("holmes.core.supabase_dal.SupabaseDal.get_workload_issues") From e60f18b23e23962a133f2476e817088995dd8a87 Mon Sep 17 00:00:00 2001 From: "avi@robusta.dev" Date: Wed, 20 Aug 2025 10:38:48 +0300 Subject: [PATCH 12/13] tests fixes --- .../test_config_load_robusta_ai.py | 59 +++++++++++++------ 1 file changed, 40 insertions(+), 19 deletions(-) diff --git a/tests/config_class/test_config_load_robusta_ai.py b/tests/config_class/test_config_load_robusta_ai.py index 0d0d2c14c3..ab1797f93e 100644 --- a/tests/config_class/test_config_load_robusta_ai.py +++ b/tests/config_class/test_config_load_robusta_ai.py @@ -1,7 +1,14 @@ from unittest.mock import patch +from pydantic import SecretStr from holmes.config import Config +def fake_load_robusta_api_key(config, _): + config.account_id = "mock-account" + config.session_token = SecretStr("mock-token") + config.api_key = SecretStr("mock-token") + + @patch("holmes.config.ROBUSTA_AI", True) def test_cli_not_loading_robusta_ai(*, monkeypatch): config = Config.load_from_file(None) @@ -9,46 +16,58 @@ def test_cli_not_loading_robusta_ai(*, monkeypatch): @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) -def test_server_loads_robusta_ai_when_true(mock_fetch, mock_cluster, *, monkeypatch): - monkeypatch.setenv("API_KEY", "test") +@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): + self.account_id = "mock-account" + self.session_token = SecretStr("mock-token") + self.api_key = SecretStr("mock-token") + + monkeypatch.setattr(Config, "load_robusta_api_key", fake_loader) config = Config.load_from_env() assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", None) -@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@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_fetch, mock_cluster, *, monkeypatch + mock_cluster, mock_fetch, *, monkeypatch ): - monkeypatch.setenv("API_KEY", "test") + def fake_loader(self, dal): + self.account_id = "mock-account" + self.session_token = SecretStr("mock-token") + self.api_key = SecretStr("mock-token") + + monkeypatch.setattr(Config, "load_robusta_api_key", fake_loader) config = Config.load_from_env() - assert len(config._model_list) == 1 assert "Robusta/test" in config._model_list @patch("holmes.config.ROBUSTA_AI", False) @patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_not_loads_robusta_ai_when_false(mock_cluster, *, monkeypatch): - monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() - assert len(config._model_list) == 0 assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch( "holmes.config.parse_models_file", return_value={"existing_model": {"base_url": "http://foo"}}, ) def test_server_loads_robusta_ai_when_true_and_model_list_exists( - mock_parse, mock_fetch, mock_cluster, *, monkeypatch + mock_parse, mock_cluster, mock_fetch, *, monkeypatch ): - monkeypatch.setenv("API_KEY", "test") + def fake_loader(self, dal): + self.account_id = "mock-account" + self.session_token = SecretStr("mock-token") + self.api_key = SecretStr("mock-token") + + monkeypatch.setattr(Config, "load_robusta_api_key", fake_loader) config = Config.load_from_env() assert "existing_model" in config._model_list assert "Robusta/test" in config._model_list @@ -63,7 +82,6 @@ def test_server_loads_robusta_ai_when_true_and_model_list_exists( def test_server_not_loads_robusta_ai_when_false_and_model_list_exists( mock_parse, mock_cluster, *, monkeypatch ): - monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "existing_model" in config._model_list assert "Robusta" not in config._model_list @@ -78,20 +96,25 @@ def test_server_not_loads_robusta_ai_when_false_and_model_list_exists( def test_server_not_loads_robusta_ai_when_no_env_var_and_model_list_exists( mock_parse, mock_cluster, *, monkeypatch ): - monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "existing_model" in config._model_list assert "Robusta" not in config._model_list @patch("holmes.config.ROBUSTA_AI", True) -@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") @patch("holmes.config.fetch_robusta_models", return_value=["Robusta/test"]) +@patch("holmes.config.Config._Config__get_cluster_name", return_value="test") def test_server_loads_robusta_ai_when_model_var_exists( - mock_fetch, mock_cluster, *, monkeypatch + mock_cluster, mock_fetch, *, monkeypatch ): monkeypatch.setenv("MODEL", "some_model") - monkeypatch.setenv("API_KEY", "test") + + def fake_loader(self, dal): + self.account_id = "mock-account" + self.session_token = SecretStr("mock-token") + self.api_key = SecretStr("mock-token") + + monkeypatch.setattr(Config, "load_robusta_api_key", fake_loader) config = Config.load_from_env() assert "Robusta/test" in config._model_list @@ -102,7 +125,6 @@ def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var( mock_cluster, *, monkeypatch ): monkeypatch.setenv("MODEL", "some_model") - monkeypatch.setenv("API_KEY", "test") config = Config.load_from_env() assert "Robusta" not in config._model_list @@ -112,7 +134,6 @@ def test_server_not_loads_robusta_ai_when_model_var_exists_and_no_env_var( def test_server_not_loads_robusta_ai_when_model_var_exists_and_false_env_var( mock_cluster, *, monkeypatch ): - monkeypatch.setenv("API_KEY", "test") monkeypatch.setenv("MODEL", "some_model") config = Config.load_from_env() assert "Robusta" not in config._model_list From 83615df2bb2ef25907dc5cb08e52386336ea0acf Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Wed, 20 Aug 2025 12:57:50 +0300 Subject: [PATCH 13/13] Make sure use updated key --- holmes/config.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/holmes/config.py b/holmes/config.py index 2e42b5b10a..0613831f10 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -176,7 +176,7 @@ def configure_robusta_ai_model(self) -> None: logging.info(f"Loading Robusta AI model: {model}") self._model_list[model] = { "base_url": f"{ROBUSTA_API_ENDPOINT}/llm/{model}", - "api_key": self.api_key.get_secret_value(), # type: ignore + "is_robusta_model": True, } except Exception: @@ -189,7 +189,7 @@ def _load_default_robusta_config(self): logging.info("Loading default Robusta AI model") self._model_list[ROBUSTA_AI_MODEL_NAME] = { "base_url": ROBUSTA_API_ENDPOINT, - "api_key": self.api_key.get_secret_value(), + "is_robusta_model": True, } def _should_load_robusta_ai(self) -> bool: @@ -527,7 +527,10 @@ def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "LLM": if model_key else next(iter(self._model_list.values())).copy() ) - api_key = model_params.pop("api_key", api_key) + if model_params.get("is_robusta_model") and self.api_key: + api_key = self.api_key.get_secret_value() + else: + api_key = model_params.pop("api_key", api_key) model = model_params.pop("model", model) return DefaultLLM(model, api_key, model_params, tracer) # type: ignore