holmesgpt 0.14.4a0__py3-none-any.whl → 0.16.0__py3-none-any.whl

holmes/core/llm.py

@@ -1,23 +1,33 @@
 import json
 import logging
+import os
 from abc import abstractmethod
 from math import floor
-from typing import Any, Dict, List, Optional, Type, Union, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
+import litellm
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse, TextCompletionResponse
 import sentry_sdk
+from pydantic import BaseModel, ConfigDict, SecretStr
+from typing_extensions import Self
+
+from holmes.clients.robusta_client import (
+    RobustaModel,
+    RobustaModelsResponse,
+    fetch_robusta_models,
+)
 
-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 (
+    FALLBACK_CONTEXT_WINDOW_SIZE,
     LOAD_ALL_ROBUSTA_MODELS,
     REASONING_EFFORT,
     ROBUSTA_AI,
     ROBUSTA_API_ENDPOINT,
     THINKING,
+    EXTRA_HEADERS,
+    TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT,
+    TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_TOKENS,
 )
 from holmes.core.supabase_dal import SupabaseDal
 from holmes.utils.env import environ_get_safe_int, replace_env_vars_values
@@ -33,9 +43,50 @@ MODEL_LIST_FILE_LOCATION = os.environ.get(
 
 OVERRIDE_MAX_OUTPUT_TOKEN = environ_get_safe_int("OVERRIDE_MAX_OUTPUT_TOKEN")
 OVERRIDE_MAX_CONTENT_SIZE = environ_get_safe_int("OVERRIDE_MAX_CONTENT_SIZE")
+
+
+def get_context_window_compaction_threshold_pct() -> int:
+    """Get the compaction threshold percentage at runtime to support test overrides."""
+    return environ_get_safe_int("CONTEXT_WINDOW_COMPACTION_THRESHOLD_PCT", default="95")
+
+
 ROBUSTA_AI_MODEL_NAME = "Robusta"
 
 
+class TokenCountMetadata(BaseModel):
+    total_tokens: int
+    tools_tokens: int
+    system_tokens: int
+    user_tokens: int
+    tools_to_call_tokens: int
+    assistant_tokens: int
+    other_tokens: int
+
+
+class ModelEntry(BaseModel):
+    """ModelEntry represents a single LLM model configuration."""
+
+    model: str
+    # TODO: the name field seems to be redundant, can we remove it?
+    name: Optional[str] = None
+    api_key: Optional[SecretStr] = None
+    base_url: Optional[str] = None
+    is_robusta_model: Optional[bool] = None
+    custom_args: Optional[Dict[str, Any]] = None
+
+    # LLM configurations used services like Azure OpenAI Service
+    api_base: Optional[str] = None
+    api_version: Optional[str] = None
+
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    @classmethod
+    def load_from_dict(cls, data: dict) -> Self:
+        return cls.model_validate(data)
+
+
 class LLM:
     @abstractmethod
     def __init__(self):
@@ -49,8 +100,23 @@ class LLM:
     def get_maximum_output_token(self) -> int:
         pass
 
+    def get_max_token_count_for_single_tool(self) -> int:
+        if (
+            0 < TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT
+            and TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT <= 100
+        ):
+            context_window_size = self.get_context_window_size()
+            calculated_max_tokens = int(
+                context_window_size * TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT // 100
+            )
+            return min(calculated_max_tokens, TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_TOKENS)
+        else:
+            return TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_TOKENS
+
     @abstractmethod
-    def count_tokens_for_message(self, messages: list[dict]) -> int:
+    def count_tokens(
+        self, messages: list[dict], tools: Optional[list[dict[str, Any]]] = None
+    ) -> TokenCountMetadata:
         pass
 
     @abstractmethod
@@ -73,6 +139,7 @@ class DefaultLLM(LLM):
     api_base: Optional[str]
     api_version: Optional[str]
     args: Dict
+    is_robusta_model: bool
 
     def __init__(
         self,
@@ -83,6 +150,7 @@ class DefaultLLM(LLM):
         args: Optional[Dict] = None,
         tracer: Optional[Any] = None,
         name: Optional[str] = None,
+        is_robusta_model: bool = False,
     ):
         self.model = model
         self.api_key = api_key
@@ -91,8 +159,11 @@ class DefaultLLM(LLM):
         self.args = args or {}
         self.tracer = tracer
         self.name = name
+        self.is_robusta_model = is_robusta_model
         self.update_custom_args()
-        self.check_llm(self.model, self.api_key, self.api_base, self.api_version)
+        self.check_llm(
+            self.model, self.api_key, self.api_base, self.api_version, self.args
+        )
 
     def update_custom_args(self):
         self.max_context_size = self.args.get("custom_args", {}).get("max_context_size")
@@ -104,7 +175,14 @@ class DefaultLLM(LLM):
         api_key: Optional[str],
         api_base: Optional[str],
         api_version: Optional[str],
+        args: Optional[dict] = None,
     ):
+        if self.is_robusta_model:
+            # The model is assumed correctly configured if it is a robusta model
+            # For robusta models, this code would fail because Holmes has no knowledge of the API keys
+            # to azure or bedrock as all completion API calls go through robusta's LLM proxy
+            return
+        args = args or {}
         logging.debug(f"Checking LiteLLM model {model}")
         lookup = litellm.get_llm_provider(model)
         if not lookup:
@@ -140,10 +218,17 @@ class DefaultLLM(LLM):
                     "environment variable for proper functionality. For more information, refer to the documentation: "
                     "https://docs.litellm.ai/docs/providers/watsonx#usage---models-in-deployment-spaces"
                 )
-        elif provider == "bedrock" and (
-            os.environ.get("AWS_PROFILE") or os.environ.get("AWS_BEARER_TOKEN_BEDROCK")
-        ):
-            model_requirements = {"keys_in_environment": True, "missing_keys": []}
+        elif provider == "bedrock":
+            if os.environ.get("AWS_PROFILE") or os.environ.get(
+                "AWS_BEARER_TOKEN_BEDROCK"
+            ):
+                model_requirements = {"keys_in_environment": True, "missing_keys": []}
+            elif args.get("aws_access_key_id") and args.get("aws_secret_access_key"):
+                return  # break fast.
+            else:
+                model_requirements = litellm.validate_environment(
+                    model=model, api_key=api_key, api_base=api_base
+                )
         else:
             model_requirements = litellm.validate_environment(
                 model=model, api_key=api_key, api_base=api_base
@@ -202,39 +287,82 @@ class DefaultLLM(LLM):
         # 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"using default {FALLBACK_CONTEXT_WINDOW_SIZE} tokens for max_input_tokens. "
             f"To override, set OVERRIDE_MAX_CONTENT_SIZE environment variable to the correct value for your model."
         )
-        return 128000
+        return FALLBACK_CONTEXT_WINDOW_SIZE
 
     @sentry_sdk.trace
-    def count_tokens_for_message(self, messages: list[dict]) -> int:
-        total_token_count = 0
+    def count_tokens(
+        self, messages: list[dict], tools: Optional[list[dict[str, Any]]] = None
+    ) -> TokenCountMetadata:
+        # TODO: Add a recount:bool flag to save time. When the flag is false, reuse 'message["token_count"]' for individual messages.
+        # It's only necessary to recount message tokens at the beginning of a session because the LLM model may have changed.
+        # Changing the model requires recounting tokens because the tokenizer may be different
+        total_tokens = 0
+        tools_tokens = 0
+        system_tokens = 0
+        assistant_tokens = 0
+        user_tokens = 0
+        other_tokens = 0
+        tools_to_call_tokens = 0
         for message in messages:
-            if "token_count" in message and message["token_count"]:
-                total_token_count += message["token_count"]
+            # count message tokens individually because it gives us fine grain information about each tool call/message etc.
+            # However be aware that the sum of individual message tokens is not equal to the overall messages token
+            token_count = litellm.token_counter(  # type: ignore
+                model=self.model, messages=[message]
+            )
+            message["token_count"] = token_count
+            role = message.get("role")
+            if role == "system":
+                system_tokens += token_count
+            elif role == "user":
+                user_tokens += token_count
+            elif role == "tool":
+                tools_tokens += token_count
+            elif role == "assistant":
+                assistant_tokens += token_count
             else:
-                # message can be counted by this method only if message contains a "content" key
-                if "content" in message:
-                    if isinstance(message["content"], str):
-                        message_to_count = [
-                            {"type": "text", "text": message["content"]}
-                        ]
-                    elif isinstance(message["content"], list):
-                        message_to_count = [
-                            {"type": "text", "text": json.dumps(message["content"])}
-                        ]
-                    elif isinstance(message["content"], dict):
-                        if "type" not in message["content"]:
-                            message_to_count = [
-                                {"type": "text", "text": json.dumps(message["content"])}
-                            ]
-                    token_count = litellm.token_counter(
-                        model=self.model, messages=message_to_count
-                    )
-                    message["token_count"] = token_count
-                    total_token_count += token_count
-        return total_token_count
+                # although this should not be needed,
+                # it is defensive code so that all tokens are accounted for
+                # and can potentially make debugging easier
+                other_tokens += token_count
+
+        messages_token_count_without_tools = litellm.token_counter(  # type: ignore
+            model=self.model, messages=messages
+        )
+
+        total_tokens = litellm.token_counter(  # type: ignore
+            model=self.model,
+            messages=messages,
+            tools=tools,  # type: ignore
+        )
+        tools_to_call_tokens = max(0, total_tokens - messages_token_count_without_tools)
+
+        return TokenCountMetadata(
+            total_tokens=total_tokens,
+            system_tokens=system_tokens,
+            user_tokens=user_tokens,
+            tools_tokens=tools_tokens,
+            tools_to_call_tokens=tools_to_call_tokens,
+            other_tokens=other_tokens,
+            assistant_tokens=assistant_tokens,
+        )
+
+    def get_litellm_corrected_name_for_robusta_ai(self) -> str:
+        if self.is_robusta_model:
+            # For robusta models, self.model is the underlying provider/model used by Robusta AI
+            # To avoid litellm modifying the API URL according to the provider, the provider name
+            # is replaced with 'openai/' just before doing a completion() call
+            # Cf. https://docs.litellm.ai/docs/providers/openai_compatible
+            split_model_name = self.model.split("/")
+            return (
+                split_model_name[0]
+                if len(split_model_name) == 1
+                else f"openai/{split_model_name[1]}"
+            )
+        else:
+            return self.model
 
     def completion(
         self,
@@ -256,6 +384,9 @@ class DefaultLLM(LLM):
         if THINKING:
             self.args.setdefault("thinking", json.loads(THINKING))
 
+        if EXTRA_HEADERS:
+            self.args.setdefault("extra_headers", json.loads(EXTRA_HEADERS))
+
         if self.args.get("thinking", None):
             litellm.modify_params = True
 
@@ -271,8 +402,10 @@ class DefaultLLM(LLM):
 
         # Get the litellm module to use (wrapped or unwrapped)
         litellm_to_use = self.tracer.wrap_llm(litellm) if self.tracer else litellm
+
+        litellm_model_name = self.get_litellm_corrected_name_for_robusta_ai()
         result = litellm_to_use.completion(
-            model=self.model,
+            model=litellm_model_name,
             api_key=self.api_key,
             base_url=self.api_base,
             api_version=self.api_version,
@@ -328,6 +461,12 @@ class DefaultLLM(LLM):
         Add cache_control to the last non-user message for Anthropic prompt caching.
         Removes any existing cache_control from previous messages to avoid accumulation.
         """
+        # Skip cache_control for VertexAI/Gemini models as they don't support it with tools
+        if self.model and (
+            "vertex" in self.model.lower() or "gemini" in self.model.lower()
+        ):
+            return
+
         # First, remove any existing cache_control from all messages
         for msg in messages:
             content = msg.get("content")
@@ -382,7 +521,7 @@ class DefaultLLM(LLM):
 class LLMModelRegistry:
     def __init__(self, config: "Config", dal: SupabaseDal) -> None:
         self.config = config
-        self._llms: dict[str, dict[str, Any]] = {}
+        self._llms: dict[str, ModelEntry] = {}
         self._default_robusta_model = None
         self.dal = dal
 
@@ -404,6 +543,8 @@ class LLMModelRegistry:
                 model_name=self.config.model,
                 base_url=self.config.api_base,
                 is_robusta_model=False,
+                api_key=self.config.api_key,
+                api_version=self.config.api_version,
             )
 
     def _should_load_config_model(self) -> bool:
@@ -414,7 +555,7 @@ class LLMModelRegistry:
         # 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"
+            self.config.model = "gpt-4.1"
             return True
 
         return False
@@ -430,6 +571,7 @@ class LLMModelRegistry:
                 return
 
             account_id, token = self.dal.get_ai_credentials()
+
             robusta_models: RobustaModelsResponse | None = fetch_robusta_models(
                 account_id, token
             )
@@ -437,16 +579,18 @@ class LLMModelRegistry:
                 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}"
+            default_model = None
+            for model_name, model_data in robusta_models.models.items():
+                logging.info(f"Loading Robusta AI model: {model_name}")
+                self._llms[model_name] = self._create_robusta_model_entry(
+                    model_name=model_name, model_data=model_data
                 )
-                self._default_robusta_model: str = robusta_models.default_model  # type: ignore
+                if model_data.is_default:
+                    default_model = model_name
+
+            if default_model:
+                logging.info(f"Setting default Robusta AI model to: {default_model}")
+                self._default_robusta_model: str = default_model  # type: ignore
 
         except Exception:
             logging.exception("Failed to get all robusta models")
@@ -456,12 +600,12 @@ class LLMModelRegistry:
     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._llms[ROBUSTA_AI_MODEL_NAME] = ModelEntry(
+                name=ROBUSTA_AI_MODEL_NAME,
+                model="gpt-4o",  # TODO: tech debt, this isn't really
+                base_url=ROBUSTA_API_ENDPOINT,
+                is_robusta_model=True,
+            )
             self._default_robusta_model = ROBUSTA_AI_MODEL_NAME
 
     def _should_load_robusta_ai(self) -> bool:
@@ -483,7 +627,7 @@ class LLMModelRegistry:
 
         return True
 
-    def get_model_params(self, model_key: Optional[str] = None) -> dict:
+    def get_model_params(self, model_key: Optional[str] = None) -> ModelEntry:
         if not self._llms:
             raise Exception("No llm models were loaded")
 
@@ -515,26 +659,30 @@ class LLMModelRegistry:
         return self._llms[name]  # type: ignore
 
     @property
-    def models(self) -> dict[str, dict[str, Any]]:
+    def models(self) -> dict[str, ModelEntry]:
         return self._llms
 
-    def _parse_models_file(self, path: str):
+    def _parse_models_file(self, path: str) -> dict[str, ModelEntry]:
         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
+        llms = {}
+        for model_name, params in models.items():
+            llms[model_name] = ModelEntry.model_validate(params)
+
+        return llms
 
     def _create_robusta_model_entry(
-        self, model_name: str, args: Optional[dict[str, Any]] = None
-    ) -> dict[str, Any]:
+        self, model_name: str, model_data: RobustaModel
+    ) -> ModelEntry:
         entry = self._create_model_entry(
-            model="gpt-4o",  # Robusta AI model is using openai like API.
+            model=model_data.model,
             model_name=model_name,
             base_url=f"{ROBUSTA_API_ENDPOINT}/llm/{model_name}",
             is_robusta_model=True,
         )
-        entry["custom_args"] = args or {}  # type: ignore[assignment]
+        entry.custom_args = model_data.holmes_args or {}  # type: ignore[assignment]
         return entry
 
     def _create_model_entry(
@@ -543,13 +691,19 @@ class LLMModelRegistry:
         model_name: str,
         base_url: Optional[str] = None,
         is_robusta_model: Optional[bool] = None,
-    ) -> dict[str, Any]:
-        return {
-            "name": model_name,
-            "base_url": base_url,
-            "is_robusta_model": is_robusta_model,
-            "model": model,
-        }
+        api_key: Optional[SecretStr] = None,
+        api_base: Optional[str] = None,
+        api_version: Optional[str] = None,
+    ) -> ModelEntry:
+        return ModelEntry(
+            name=model_name,
+            model=model,
+            base_url=base_url,
+            is_robusta_model=is_robusta_model,
+            api_key=api_key,
+            api_base=api_base,
+            api_version=api_version,
+        )
 
 
 def get_llm_usage(

holmes/core/openai_formatting.py

@@ -80,6 +80,19 @@ def format_tool_to_open_ai_standard(
         )
         if param_attributes.description is not None:
             tool_properties[param_name]["description"] = param_attributes.description
+        # Add enum constraint if specified
+        if hasattr(param_attributes, "enum") and param_attributes.enum:
+            enum_values = list(
+                param_attributes.enum
+            )  # Create a copy to avoid modifying original
+            # In strict mode, optional parameters need None in their enum to match the type allowing null
+            if (
+                strict_mode
+                and not param_attributes.required
+                and None not in enum_values
+            ):
+                enum_values.append(None)
+            tool_properties[param_name]["enum"] = enum_values
 
     result: dict[str, Any] = {
         "type": "function",

holmes/core/supabase_dal.py

@@ -237,70 +237,61 @@ class SupabaseDal:
             logging.exception("Supabase error while retrieving efficiency data")
             return None
 
-    def get_configuration_changes(
-        self, start_datetime: str, end_datetime: str
+    def get_configuration_changes_metadata(
+        self,
+        start_datetime: str,
+        end_datetime: str,
+        limit: int = 100,
+        workload: Optional[str] = None,
+        ns: Optional[str] = None,
     ) -> Optional[List[Dict]]:
         if not self.enabled:
             return []
 
         try:
-            changes_response = (
+            query = (
                 self.client.table(ISSUES_TABLE)
-                .select("id", "subject_name", "subject_namespace", "description")
+                .select(
+                    "id",
+                    "title",
+                    "subject_name",
+                    "subject_namespace",
+                    "subject_type",
+                    "description",
+                    "starts_at",
+                    "ends_at",
+                )
                 .eq("account_id", self.account_id)
                 .eq("cluster", self.cluster)
                 .eq("finding_type", "configuration_change")
                 .gte("creation_date", start_datetime)
                 .lte("creation_date", end_datetime)
-                .execute()
+                .limit(limit)
             )
-            if not len(changes_response.data):
-                return None
 
-        except Exception:
-            logging.exception("Supabase error while retrieving change data")
-            return None
+            if workload:
+                query.eq("subject_name", workload)
+            if ns:
+                query.eq("subject_namespace", ns)
 
-        changes_ids = [change["id"] for change in changes_response.data]
-        try:
-            change_data_response = (
-                self.client.table(EVIDENCE_TABLE)
-                .select("*")
-                .eq("account_id", self.account_id)
-                .in_("issue_id", changes_ids)
-                .not_.in_("enrichment_type", ENRICHMENT_BLACKLIST)
-                .execute()
-            )
-            if not len(change_data_response.data):
+            res = query.execute()
+            if not res.data:
                 return None
 
-            truncate_evidences_entities_if_necessary(change_data_response.data)
-
         except Exception:
-            logging.exception("Supabase error while retrieving change content")
+            logging.exception("Supabase error while retrieving change data")
             return None
 
-        changes_data = []
-        change_data_map = {
-            change["issue_id"]: change for change in change_data_response.data
-        }
-
-        for change in changes_response.data:
-            change_content = change_data_map.get(change["id"])
-            if change_content:
-                changes_data.append(
-                    {
-                        "change": change_content["data"],
-                        "evidence_id": change_content["id"],
-                        **change,
-                    }
-                )
-
         logging.debug(
-            "Change history for %s-%s: %s", start_datetime, end_datetime, changes_data
+            "Change history metadata for %s-%s workload %s in ns %s: %s",
+            start_datetime,
+            end_datetime,
+            workload,
+            ns,
+            res.data,
         )
 
-        return changes_data
+        return res.data
 
     def unzip_evidence_file(self, data):
         try: