holmesgpt 0.14.0a0__py3-none-any.whl → 0.14.1__py3-none-any.whl

holmes/core/llm.py

@@ -1,30 +1,38 @@
 import json
 import logging
 from abc import abstractmethod
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union, TYPE_CHECKING
 
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import ModelResponse, TextCompletionResponse
 import sentry_sdk
 
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from pydantic import BaseModel
 import litellm
 import os
+from holmes.clients.robusta_client import RobustaModelsResponse, fetch_robusta_models
 from holmes.common.env_vars import (
+    LOAD_ALL_ROBUSTA_MODELS,
     REASONING_EFFORT,
+    ROBUSTA_AI,
+    ROBUSTA_API_ENDPOINT,
     THINKING,
 )
+from holmes.core.supabase_dal import SupabaseDal
+from holmes.utils.env import environ_get_safe_int, replace_env_vars_values
+from holmes.utils.file_utils import load_yaml_file
 
+if TYPE_CHECKING:
+    from holmes.config import Config
 
-def environ_get_safe_int(env_var, default="0"):
-    try:
-        return max(int(os.environ.get(env_var, default)), 0)
-    except ValueError:
-        return int(default)
+MODEL_LIST_FILE_LOCATION = os.environ.get(
+    "MODEL_LIST_FILE_LOCATION", "/etc/holmes/config/model_list.yaml"
+)
 
 
 OVERRIDE_MAX_OUTPUT_TOKEN = environ_get_safe_int("OVERRIDE_MAX_OUTPUT_TOKEN")
 OVERRIDE_MAX_CONTENT_SIZE = environ_get_safe_int("OVERRIDE_MAX_CONTENT_SIZE")
+ROBUSTA_AI_MODEL_NAME = "Robusta"
 
 
 class LLM:
@@ -61,31 +69,43 @@ class LLM:
 class DefaultLLM(LLM):
     model: str
     api_key: Optional[str]
-    base_url: Optional[str]
+    api_base: Optional[str]
+    api_version: Optional[str]
     args: Dict
 
     def __init__(
         self,
         model: str,
         api_key: Optional[str] = None,
+        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
+        self.api_base = api_base
+        self.api_version = api_version
         self.args = args or {}
         self.tracer = tracer
+        self.name = name
+        self.update_custom_args()
+        self.check_llm(self.model, self.api_key, self.api_base, self.api_version)
 
-        if not self.args:
-            self.check_llm(self.model, self.api_key)
+    def update_custom_args(self):
+        self.max_context_size = self.args.get("custom_args", {}).get("max_context_size")
+        self.args.pop("custom_args", None)
 
-    def check_llm(self, model: str, api_key: Optional[str]):
+    def check_llm(
+        self,
+        model: str,
+        api_key: Optional[str],
+        api_base: Optional[str],
+        api_version: Optional[str],
+    ):
         logging.debug(f"Checking LiteLLM model {model}")
-        # TODO: this WAS a hack to get around the fact that we can't pass in an api key to litellm.validate_environment
-        # so without this hack it always complains that the environment variable for the api key is missing
-        # to fix that, we always set an api key in the standard format that litellm expects (which is ${PROVIDER}_API_KEY)
-        # TODO: we can now handle this better - see https://github.com/BerriAI/litellm/issues/4375#issuecomment-2223684750
-        lookup = litellm.get_llm_provider(self.model)
+        lookup = litellm.get_llm_provider(model)
         if not lookup:
             raise Exception(f"Unknown provider for model {model}")
         provider = lookup[1]
@@ -124,51 +144,67 @@ class DefaultLLM(LLM):
         ):
             model_requirements = {"keys_in_environment": True, "missing_keys": []}
         else:
-            #
-            api_key_env_var = f"{provider.upper()}_API_KEY"
-            if api_key:
-                os.environ[api_key_env_var] = api_key
-            model_requirements = litellm.validate_environment(model=model)
+            model_requirements = litellm.validate_environment(
+                model=model, api_key=api_key, api_base=api_base
+            )
+            # validate_environment does not accept api_version, and as a special case for Azure OpenAI Service,
+            # when all the other AZURE environments are set expect AZURE_API_VERSION, validate_environment complains
+            # the missing of it even after the api_version is set.
+            # TODO: There's an open PR in litellm to accept api_version in validate_environment, we can leverage this
+            # change if accepted to ignore the following check.
+            # https://github.com/BerriAI/litellm/pull/13808
+            if (
+                provider == "azure"
+                and ["AZURE_API_VERSION"] == model_requirements["missing_keys"]
+                and api_version is not None
+            ):
+                model_requirements["missing_keys"] = []
+                model_requirements["keys_in_environment"] = True
 
         if not model_requirements["keys_in_environment"]:
             raise Exception(
                 f"model {model} requires the following environment variables: {model_requirements['missing_keys']}"
             )
 
-    def _strip_model_prefix(self) -> str:
+    def _get_model_name_variants_for_lookup(self) -> list[str]:
         """
-        Helper function to strip 'openai/' prefix from model name if it exists.
-        model cost is taken from here which does not have the openai prefix
-        https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json
+        Generate model name variants to try when looking up in litellm.model_cost.
+        Returns a list of names to try in order: exact, lowercase, without prefix, etc.
         """
-        model_name = self.model
-        prefixes = ["openai/", "bedrock/", "vertex_ai/", "anthropic/"]
+        names_to_try = [self.model, self.model.lower()]
 
-        for prefix in prefixes:
-            if model_name.startswith(prefix):
-                return model_name[len(prefix) :]
+        # If there's a prefix, also try without it
+        if "/" in self.model:
+            base_model = self.model.split("/", 1)[1]
+            names_to_try.extend([base_model, base_model.lower()])
 
-        return model_name
-
-        # this unfortunately does not seem to work for azure if the deployment name is not a well-known model name
-        # if not litellm.supports_function_calling(model=model):
-        #    raise Exception(f"model {model} does not support function calling. You must use HolmesGPT with a model that supports function calling.")
+        # Remove duplicates while preserving order (dict.fromkeys maintains insertion order in Python 3.7+)
+        return list(dict.fromkeys(names_to_try))
 
     def get_context_window_size(self) -> int:
+        if self.max_context_size:
+            return self.max_context_size
+
         if OVERRIDE_MAX_CONTENT_SIZE:
             logging.debug(
                 f"Using override OVERRIDE_MAX_CONTENT_SIZE {OVERRIDE_MAX_CONTENT_SIZE}"
             )
             return OVERRIDE_MAX_CONTENT_SIZE
 
-        model_name = os.environ.get("MODEL_TYPE", self._strip_model_prefix())
-        try:
-            return litellm.model_cost[model_name]["max_input_tokens"]
-        except Exception:
-            logging.warning(
-                f"Couldn't find model's name {model_name} in litellm's model list, fallback to 128k tokens for max_input_tokens"
-            )
-            return 128000
+        # Try each name variant
+        for name in self._get_model_name_variants_for_lookup():
+            try:
+                return litellm.model_cost[name]["max_input_tokens"]
+            except Exception:
+                continue
+
+        # Log which lookups we tried
+        logging.warning(
+            f"Couldn't find model {self.model} in litellm's model list (tried: {', '.join(self._get_model_name_variants_for_lookup())}), "
+            f"using default 128k tokens for max_input_tokens. "
+            f"To override, set OVERRIDE_MAX_CONTENT_SIZE environment variable to the correct value for your model."
+        )
+        return 128000
 
     @sentry_sdk.trace
     def count_tokens_for_message(self, messages: list[dict]) -> int:
@@ -237,6 +273,8 @@ class DefaultLLM(LLM):
         result = litellm_to_use.completion(
             model=self.model,
             api_key=self.api_key,
+            base_url=self.api_base,
+            api_version=self.api_version,
             messages=messages,
             response_format=response_format,
             drop_params=drop_params,
@@ -260,14 +298,20 @@ class DefaultLLM(LLM):
             )
             return OVERRIDE_MAX_OUTPUT_TOKEN
 
-        model_name = os.environ.get("MODEL_TYPE", self._strip_model_prefix())
-        try:
-            return litellm.model_cost[model_name]["max_output_tokens"]
-        except Exception:
-            logging.warning(
-                f"Couldn't find model's name {model_name} in litellm's model list, fallback to 4096 tokens for max_output_tokens"
-            )
-            return 4096
+        # Try each name variant
+        for name in self._get_model_name_variants_for_lookup():
+            try:
+                return litellm.model_cost[name]["max_output_tokens"]
+            except Exception:
+                continue
+
+        # Log which lookups we tried
+        logging.warning(
+            f"Couldn't find model {self.model} in litellm's model list (tried: {', '.join(self._get_model_name_variants_for_lookup())}), "
+            f"using default 4096 tokens for max_output_tokens. "
+            f"To override, set OVERRIDE_MAX_OUTPUT_TOKEN environment variable to the correct value for your model."
+        )
+        return 4096
 
     def _add_cache_control_to_last_message(
         self, messages: List[Dict[str, Any]]
@@ -325,3 +369,202 @@ class DefaultLLM(LLM):
                 logging.debug(
                     f"Added cache_control to {target_msg.get('role')} message (structured content)"
                 )
+
+
+class LLMModelRegistry:
+    def __init__(self, config: "Config", dal: SupabaseDal) -> None:
+        self.config = config
+        self._llms: dict[str, dict[str, Any]] = {}
+        self._default_robusta_model = None
+        self.dal = dal
+
+        self._init_models()
+
+    @property
+    def default_robusta_model(self) -> Optional[str]:
+        return self._default_robusta_model
+
+    def _init_models(self):
+        self._llms = self._parse_models_file(MODEL_LIST_FILE_LOCATION)
+
+        if self._should_load_robusta_ai():
+            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,
+            )
+
+    def _should_load_config_model(self) -> bool:
+        if self.config.model is not None:
+            return True
+
+        # backward compatibility - in the past config.model was set by default to gpt-4o.
+        # so we need to check if the user has set an OPENAI_API_KEY to load the config model.
+        has_openai_key = os.environ.get("OPENAI_API_KEY")
+        if has_openai_key:
+            self.config.model = "gpt-4o"
+            return True
+
+        return False
+
+    def configure_robusta_ai_model(self) -> None:
+        try:
+            if not self.config.cluster_name or not LOAD_ALL_ROBUSTA_MODELS:
+                self._load_default_robusta_config()
+                return
+
+            if not self.dal.account_id or not self.dal.enabled:
+                self._load_default_robusta_config()
+                return
+
+            account_id, token = self.dal.get_ai_credentials()
+            robusta_models: RobustaModelsResponse | None = fetch_robusta_models(
+                account_id, token
+            )
+            if not robusta_models or not robusta_models.models:
+                self._load_default_robusta_config()
+                return
+
+            for model in robusta_models.models:
+                logging.info(f"Loading Robusta AI model: {model}")
+                args = robusta_models.models_args.get(model)
+                self._llms[model] = self._create_robusta_model_entry(model, args)
+
+            if robusta_models.default_model:
+                logging.info(
+                    f"Setting default Robusta AI model to: {robusta_models.default_model}"
+                )
+                self._default_robusta_model: str = robusta_models.default_model  # type: ignore
+
+        except Exception:
+            logging.exception("Failed to get all robusta models")
+            # fallback to default behavior
+            self._load_default_robusta_config()
+
+    def _load_default_robusta_config(self):
+        if self._should_load_robusta_ai():
+            logging.info("Loading default Robusta AI model")
+            self._llms[ROBUSTA_AI_MODEL_NAME] = {
+                "name": ROBUSTA_AI_MODEL_NAME,
+                "base_url": ROBUSTA_API_ENDPOINT,
+                "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.config.should_try_robusta_ai:
+            return False
+
+        # ROBUSTA_AI were set in the env vars, so we can use it directly
+        if ROBUSTA_AI is not None:
+            return ROBUSTA_AI
+
+        # MODEL is set in the env vars, e.g. the user is using a custom model
+        # so we don't need to load the robusta AI model and keep the behavior backward compatible
+        if "MODEL" in os.environ:
+            return False
+
+        # if the user has provided a model list, we don't need to load the robusta AI model
+        if self._llms:
+            return False
+
+        return True
+
+    def get_model_params(self, model_key: Optional[str] = None) -> dict:
+        if not self._llms:
+            raise Exception("No llm models were loaded")
+
+        if model_key:
+            model_params = self._llms.get(model_key)
+            if model_params is not None:
+                logging.info(f"Using selected 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._llms.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"
+            )
+
+        model_key, first_model_params = next(iter(self._llms.items()))
+        logging.debug(f"Using first available model: {model_key}")
+        return first_model_params.copy()
+
+    def get_llm(self, name: str) -> LLM:  # TODO: fix logic
+        return self._llms[name]  # type: ignore
+
+    @property
+    def models(self) -> dict[str, dict[str, Any]]:
+        return self._llms
+
+    def _parse_models_file(self, path: str):
+        models = load_yaml_file(path, raise_error=False, warn_not_found=False)
+        for _, params in models.items():
+            params = replace_env_vars_values(params)
+
+        return models
+
+    def _create_robusta_model_entry(
+        self, model_name: str, args: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        return self._create_model_entry(
+            model="gpt-4o",  # Robusta AI model is using openai like API.
+            model_name=model_name,
+            base_url=f"{ROBUSTA_API_ENDPOINT}/llm/{model_name}",
+            is_robusta_model=True,
+            args=args or {},
+        )
+
+    def _create_model_entry(
+        self,
+        model: str,
+        model_name: str,
+        base_url: Optional[str] = None,
+        is_robusta_model: Optional[bool] = None,
+        args: Optional[dict[str, Any]] = None,
+    ) -> dict[str, Any]:
+        entry = {
+            "name": model_name,
+            "base_url": base_url,
+            "is_robusta_model": is_robusta_model,
+            "model": model,
+        }
+        if args:
+            entry["custom_args"] = args  # type: ignore[assignment]
+
+        return entry
+
+
+def get_llm_usage(
+    llm_response: Union[ModelResponse, CustomStreamWrapper, TextCompletionResponse],
+) -> dict:
+    usage: dict = {}
+    if (
+        (
+            isinstance(llm_response, ModelResponse)
+            or isinstance(llm_response, TextCompletionResponse)
+        )
+        and hasattr(llm_response, "usage")
+        and llm_response.usage
+    ):  # type: ignore
+        usage["prompt_tokens"] = llm_response.usage.prompt_tokens  # type: ignore
+        usage["completion_tokens"] = llm_response.usage.completion_tokens  # type: ignore
+        usage["total_tokens"] = llm_response.usage.total_tokens  # type: ignore
+    elif isinstance(llm_response, CustomStreamWrapper):
+        complete_response = litellm.stream_chunk_builder(chunks=llm_response)  # type: ignore
+        if complete_response:
+            return get_llm_usage(complete_response)
+    return usage

holmes/core/models.py

@@ -10,6 +10,7 @@ class InvestigationResult(BaseModel):
     sections: Optional[Dict[str, Union[str, None]]] = None
     tool_calls: List[ToolCallResult] = []
     instructions: List[str] = []
+    metadata: Optional[Dict[Any, Any]] = None
 
 
 class InvestigateRequest(BaseModel):
@@ -145,6 +146,7 @@ class ChatResponse(BaseModel):
     conversation_history: list[dict]
     tool_calls: Optional[List[ToolCallResult]] = []
     follow_up_actions: Optional[List[FollowUpAction]] = []
+    metadata: Optional[Dict[Any, Any]] = None
 
 
 class WorkloadHealthInvestigationResult(BaseModel):

holmes/core/safeguards.py

@@ -5,7 +5,7 @@ from pydantic import ValidationError
 
 from holmes.common.env_vars import TOOL_CALL_SAFEGUARDS_ENABLED
 from holmes.plugins.toolsets.logging_utils.logging_api import POD_LOGGING_TOOL_NAME
-from holmes.core.tools import StructuredToolResult, ToolResultStatus
+from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
 from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams
 
 
@@ -39,7 +39,7 @@ def _has_previous_unfiltered_pod_logs_call(
             result = tool_call.get("result", {})
             if (
                 tool_call.get("tool_name") == POD_LOGGING_TOOL_NAME
-                and result.get("status") == ToolResultStatus.NO_DATA
+                and result.get("status") == StructuredToolResultStatus.NO_DATA
                 and result.get("params")
             ):
                 params = FetchPodLogsParams(**result.get("params"))
@@ -94,7 +94,7 @@ def prevent_overly_repeated_tool_call(
                 For example if Holmes checks if a resource is deployed, runs a command to deploy it and then checks again if it has deployed properly.
             """
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=(
                     "Refusing to run this tool call because it has already been called during this session with the exact same parameters.\n"
                     "Move on with your investigation to a different tool or change the parameter values."
@@ -106,7 +106,7 @@ def prevent_overly_repeated_tool_call(
             tool_name=tool_name, tool_params=tool_params, tool_calls=tool_calls
         ):
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=(
                     f"Refusing to run this tool call because the exact same {POD_LOGGING_TOOL_NAME} tool call without filter has already run and returned no data.\n"
                     "This tool call would also have returned no data.\n"

holmes/core/supabase_dal.py

@@ -37,6 +37,7 @@ from holmes.utils.global_instructions import Instructions
 SUPABASE_TIMEOUT_SECONDS = int(os.getenv("SUPABASE_TIMEOUT_SECONDS", 3600))
 
 ISSUES_TABLE = "Issues"
+GROUPED_ISSUES_TABLE = "GroupedIssues"
 EVIDENCE_TABLE = "Evidence"
 RUNBOOKS_TABLE = "HolmesRunbooks"
 SESSION_TOKENS_TABLE = "AuthTokens"
@@ -338,6 +339,14 @@ class SupabaseDal:
         data.extend(unzipped_files)
         return data
 
+    def get_issue_from_db(self, issue_id: str, table: str) -> Optional[Dict]:
+        issue_response = (
+            self.client.table(table).select("*").filter("id", "eq", issue_id).execute()
+        )
+        if len(issue_response.data):
+            return issue_response.data[0]
+        return None
+
     def get_issue_data(self, issue_id: Optional[str]) -> Optional[Dict]:
         # TODO this could be done in a single atomic SELECT, but there is no
         # foreign key relation between Issues and Evidence.
@@ -347,14 +356,11 @@ class SupabaseDal:
             return None
         issue_data = None
         try:
-            issue_response = (
-                self.client.table(ISSUES_TABLE)
-                .select("*")
-                .filter("id", "eq", issue_id)
-                .execute()
-            )
-            if len(issue_response.data):
-                issue_data = issue_response.data[0]
+            issue_data = self.get_issue_from_db(issue_id, ISSUES_TABLE)
+            if issue_data and issue_data["source"] == "prometheus":
+                logging.debug("Getting alert %s from GroupedIssuesTable", issue_id)
+                # This issue will have the complete alert duration information
+                issue_data = self.get_issue_from_db(issue_id, GROUPED_ISSUES_TABLE)
 
         except Exception:  # e.g. invalid id format
             logging.exception("Supabase error while retrieving issue data")