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

holmes/__init__.py

@@ -1,5 +1,5 @@
 # This is patched by github actions during release
-__version__ = "0.14.4-alpha"
+__version__ = "0.15.0"
 
 # Re-export version functions from version module for backward compatibility
 from .version import (

holmes/clients/robusta_client.py

@@ -1,8 +1,8 @@
 import logging
-from typing import List, Optional, Dict, Any
+from typing import Optional, Dict, Any
 import requests  # type: ignore
 from functools import cache
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict
 from holmes.common.env_vars import ROBUSTA_API_ENDPOINT
 
 HOLMES_GET_INFO_URL = f"{ROBUSTA_API_ENDPOINT}/api/holmes/get_info"
@@ -14,13 +14,15 @@ class HolmesInfo(BaseModel):
     latest_version: Optional[str] = None
 
 
-class RobustaModelsResponse(BaseModel):
+class RobustaModel(BaseModel):
     model_config = ConfigDict(extra="ignore")
-    models: List[str]
-    models_args: Dict[str, Any] = Field(
-        default_factory=dict, alias="models_holmes_args"
-    )
-    default_model: Optional[str] = None
+    model: str
+    holmes_args: Optional[dict[str, Any]] = None
+    is_default: bool = False
+
+
+class RobustaModelsResponse(BaseModel):
+    models: Dict[str, RobustaModel]
 
 
 @cache
@@ -30,13 +32,13 @@ def fetch_robusta_models(
     try:
         session_request = {"session_token": token, "account_id": account_id}
         resp = requests.post(
-            f"{ROBUSTA_API_ENDPOINT}/api/llm/models",
+            f"{ROBUSTA_API_ENDPOINT}/api/llm/models/v2",
             json=session_request,
             timeout=10,
         )
         resp.raise_for_status()
         response_json = resp.json()
-        return RobustaModelsResponse(**response_json)
+        return RobustaModelsResponse(**{"models": response_json})
     except Exception:
         logging.exception("Failed to fetch robusta models")
         return None

holmes/common/env_vars.py

@@ -2,6 +2,16 @@ import os
 import json
 from typing import Optional
 
+# Recommended models for different providers
+RECOMMENDED_OPENAI_MODEL = "gpt-4.1"
+RECOMMENDED_ANTHROPIC_MODEL = "anthropic/claude-opus-4-1-20250805"
+
+# Default model for HolmesGPT
+DEFAULT_MODEL = RECOMMENDED_OPENAI_MODEL
+FALLBACK_CONTEXT_WINDOW_SIZE = (
+    200000  # Fallback context window size if it can't be determined from the model
+)
+
 
 def load_bool(env_var, default: Optional[bool]) -> Optional[bool]:
     env_value = os.environ.get(env_var)
@@ -38,6 +48,7 @@ DEVELOPMENT_MODE = load_bool("DEVELOPMENT_MODE", False)
 SENTRY_DSN = os.environ.get("SENTRY_DSN", "")
 SENTRY_TRACES_SAMPLE_RATE = float(os.environ.get("SENTRY_TRACES_SAMPLE_RATE", "0.0"))
 
+EXTRA_HEADERS = os.environ.get("EXTRA_HEADERS", "")
 THINKING = os.environ.get("THINKING", "")
 REASONING_EFFORT = os.environ.get("REASONING_EFFORT", "").strip().lower()
 TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.00000001"))
@@ -82,6 +93,9 @@ TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT = float(
     os.environ.get("TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", 15)
 )
 
+# Absolute max tokens to allocate for a single tool response
+TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_TOKENS = 25000
+
 MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION = int(
     os.environ.get("MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", 3000)
 )

holmes/config.py

@@ -45,6 +45,9 @@ class SupportedTicketSources(str, Enum):
 
 class Config(RobustaBaseConfig):
     model: Optional[str] = None
+    api_key: Optional[SecretStr] = (
+        None  # if None, read from OPENAI_API_KEY or AZURE_OPENAI_ENDPOINT env var
+    )
     api_base: Optional[str] = None
     api_version: Optional[str] = None
     fast_model: Optional[str] = None
@@ -95,6 +98,7 @@ class Config(RobustaBaseConfig):
     mcp_servers: Optional[dict[str, dict[str, Any]]] = None
 
     _server_tool_executor: Optional[ToolExecutor] = None
+    _agui_tool_executor: Optional[ToolExecutor] = None
 
     # TODO: Separate those fields to facade class, this shouldn't be part of the config.
     _toolset_manager: Optional[ToolsetManager] = PrivateAttr(None)
@@ -242,6 +246,23 @@ class Config(RobustaBaseConfig):
         )
         return ToolExecutor(cli_toolsets)
 
+    def create_agui_tool_executor(self, dal: Optional["SupabaseDal"]) -> ToolExecutor:
+        """
+        Creates ToolExecutor for the AG-UI server endpoints
+        """
+
+        if self._agui_tool_executor:
+            return self._agui_tool_executor
+
+        # Use same toolset as CLI for AG-UI front-end.
+        agui_toolsets = self.toolset_manager.list_console_toolsets(
+            dal=dal, refresh_status=True
+        )
+
+        self._agui_tool_executor = ToolExecutor(agui_toolsets)
+
+        return self._agui_tool_executor
+
     def create_tool_executor(self, dal: Optional["SupabaseDal"]) -> ToolExecutor:
         """
         Creates ToolExecutor for the server endpoints
@@ -273,6 +294,19 @@ class Config(RobustaBaseConfig):
             tool_executor, self.max_steps, self._get_llm(tracer=tracer)
         )
 
+    def create_agui_toolcalling_llm(
+        self,
+        dal: Optional["SupabaseDal"] = None,
+        model: Optional[str] = None,
+        tracer=None,
+    ) -> "ToolCallingLLM":
+        tool_executor = self.create_agui_tool_executor(dal)
+        from holmes.core.tool_calling_llm import ToolCallingLLM
+
+        return ToolCallingLLM(
+            tool_executor, self.max_steps, self._get_llm(model, tracer)
+        )
+
     def create_toolcalling_llm(
         self,
         dal: Optional["SupabaseDal"] = None,
@@ -441,7 +475,8 @@ class Config(RobustaBaseConfig):
     # TODO: move this to the llm model registry
     def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> "DefaultLLM":
         sentry_sdk.set_tag("requested_model", model_key)
-        model_params = self.llm_model_registry.get_model_params(model_key)
+        model_entry = self.llm_model_registry.get_model_params(model_key)
+        model_params = model_entry.model_dump(exclude_none=True)
         api_base = self.api_base
         api_version = self.api_version
 
@@ -453,6 +488,8 @@ class Config(RobustaBaseConfig):
             api_key = f"{account_id} {token}"
         else:
             api_key = model_params.pop("api_key", None)
+            if api_key is not None:
+                api_key = api_key.get_secret_value()
 
         model = model_params.pop("model")
         # It's ok if the model does not have api base and api version, which are defaults to None.
@@ -463,10 +500,20 @@ class Config(RobustaBaseConfig):
         api_version = model_params.pop("api_version", api_version)
         model_name = model_params.pop("name", None) or model_key or model
         sentry_sdk.set_tag("model_name", model_name)
-        logging.info(f"Creating LLM with model: {model_name}")
-        return DefaultLLM(
-            model, api_key, api_base, api_version, model_params, tracer, model_name
+        llm = DefaultLLM(
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            api_version=api_version,
+            args=model_params,
+            tracer=tracer,
+            name=model_name,
+            is_robusta_model=is_robusta_model,
         )  # type: ignore
+        logging.info(
+            f"Using model: {model_name} ({llm.get_context_window_size():,} total tokens, {llm.get_maximum_output_token():,} output tokens)"
+        )
+        return llm
 
     def get_models_list(self) -> List[str]:
         if self.llm_model_registry and self.llm_model_registry.models:

holmes/core/conversations.py

@@ -26,7 +26,8 @@ def calculate_tool_size(
         return DEFAULT_TOOL_SIZE
 
     context_window = ai.llm.get_context_window_size()
-    message_size_without_tools = ai.llm.count_tokens_for_message(messages_without_tools)
+    tokens = ai.llm.count_tokens(messages_without_tools)
+    message_size_without_tools = tokens.total_tokens
     maximum_output_token = ai.llm.get_maximum_output_token()
 
     tool_size = min(
@@ -372,13 +373,13 @@ def build_chat_messages(
     )
 
     ask = add_global_instructions_to_user_prompt(ask, global_instructions)
-
     conversation_history.append(  # type: ignore
         {
             "role": "user",
             "content": ask,
         },
     )
+
     number_of_tools = len(
         [message for message in conversation_history if message.get("role") == "tool"]  # type: ignore
     )

holmes/core/llm.py

@@ -1,23 +1,31 @@
 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,
 )
 from holmes.core.supabase_dal import SupabaseDal
 from holmes.utils.env import environ_get_safe_int, replace_env_vars_values
@@ -36,6 +44,39 @@ OVERRIDE_MAX_CONTENT_SIZE = environ_get_safe_int("OVERRIDE_MAX_CONTENT_SIZE")
 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
+    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):
@@ -50,7 +91,9 @@ class LLM:
         pass
 
     @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 +116,7 @@ class DefaultLLM(LLM):
     api_base: Optional[str]
     api_version: Optional[str]
     args: Dict
+    is_robusta_model: bool
 
     def __init__(
         self,
@@ -83,6 +127,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 +136,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 +152,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 +195,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 +264,78 @@ 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
+        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
             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,
+        )
+
+    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 +357,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 +375,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 +434,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 +494,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 +516,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 +528,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
@@ -437,16 +551,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 +572,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 +599,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 +631,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 +663,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",