holmesgpt 0.11.5__py3-none-any.whl → 0.12.0a0__py3-none-any.whl

holmes/__init__.py

@@ -9,7 +9,7 @@ this_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(this_path)
 
 # This is patched by github actions during release
-__version__ = "0.11.5"
+__version__ = "0.12.0-alpha"
 
 
 def is_official_release() -> bool:

holmes/common/env_vars.py

@@ -1,10 +1,14 @@
 import os
 import json
+from typing import Optional
 
 
-def load_bool(env_var, default: bool):
-    s = os.environ.get(env_var, str(default))
-    return json.loads(s.lower())
+def load_bool(env_var, default: Optional[bool]) -> Optional[bool]:
+    env_value = os.environ.get(env_var)
+    if env_value is None:
+        return default
+
+    return json.loads(env_value.lower())
 
 
 ENABLED_BY_DEFAULT_TOOLSETS = os.environ.get(
@@ -22,7 +26,7 @@ STORE_API_KEY = os.environ.get("STORE_API_KEY", "")
 STORE_EMAIL = os.environ.get("STORE_EMAIL", "")
 STORE_PASSWORD = os.environ.get("STORE_PASSWORD", "")
 HOLMES_POST_PROCESSING_PROMPT = os.environ.get("HOLMES_POST_PROCESSING_PROMPT", "")
-ROBUSTA_AI = load_bool("ROBUSTA_AI", False)
+ROBUSTA_AI = load_bool("ROBUSTA_AI", None)
 ROBUSTA_API_ENDPOINT = os.environ.get("ROBUSTA_API_ENDPOINT", "https://api.robusta.dev")
 
 LOG_PERFORMANCE = os.environ.get("LOG_PERFORMANCE", None)

holmes/config.py

@@ -39,6 +39,7 @@ DEFAULT_CONFIG_LOCATION = os.path.expanduser("~/.holmes/config.yaml")
 MODEL_LIST_FILE_LOCATION = os.environ.get(
     "MODEL_LIST_FILE_LOCATION", "/etc/holmes/config/model_list.yaml"
 )
+ROBUSTA_AI_MODEL_NAME = "Robusta"
 
 
 class SupportedTicketSources(str, Enum):
@@ -58,7 +59,7 @@ def is_old_toolset_config(
 def parse_models_file(path: str):
     models = load_yaml_file(path, raise_error=False, warn_not_found=False)
 
-    for model, params in models.items():
+    for _, params in models.items():
         params = replace_env_vars_values(params)
 
     return models
@@ -109,6 +110,7 @@ class Config(RobustaBaseConfig):
     # custom_toolsets_from_cli is passed from CLI option `--custom-toolsets` as 'experimental' custom toolsets.
     # The status of toolset here won't be cached, so the toolset from cli will always be loaded when specified in the CLI.
     custom_toolsets_from_cli: Optional[List[FilePath]] = None
+    should_try_robusta_ai: bool = False  # if True, we will try to load the Robusta AI model, in cli we aren't trying to load it.
 
     toolsets: Optional[dict[str, dict[str, Any]]] = None
 
@@ -148,11 +150,31 @@ class Config(RobustaBaseConfig):
         self._version = get_version()
         self._holmes_info = fetch_holmes_info()
         self._model_list = parse_models_file(MODEL_LIST_FILE_LOCATION)
-        if ROBUSTA_AI:
-            self._model_list["Robusta"] = {
+        if self._should_load_robusta_ai():
+            logging.info("Loading 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:
+            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._model_list:
+            return False
+
+        return True
+
     def log_useful_info(self):
         if self._model_list:
             logging.info(f"loaded models: {list(self._model_list.keys())}")
@@ -220,6 +242,7 @@ class Config(RobustaBaseConfig):
             if val is not None:
                 kwargs[field_name] = val
         kwargs["cluster_name"] = Config.__get_cluster_name()
+        kwargs["should_try_robusta_ai"] = True
         result = cls(**kwargs)
         result.log_useful_info()
         return result
@@ -249,7 +272,9 @@ class Config(RobustaBaseConfig):
         runbook_catalog = load_runbook_catalog()
         return runbook_catalog
 
-    def create_console_tool_executor(self, dal: Optional[SupabaseDal]) -> ToolExecutor:
+    def create_console_tool_executor(
+        self, dal: Optional[SupabaseDal], refresh_status: bool = False
+    ) -> ToolExecutor:
         """
         Creates a ToolExecutor instance configured for CLI usage. This executor manages the available tools
         and their execution in the command-line interface.
@@ -259,7 +284,9 @@ class Config(RobustaBaseConfig):
         2. toolsets from config file will override and be merged into built-in toolsets with the same name.
         3. Custom toolsets from config files which can not override built-in toolsets
         """
-        cli_toolsets = self.toolset_manager.list_console_toolsets(dal=dal)
+        cli_toolsets = self.toolset_manager.list_console_toolsets(
+            dal=dal, refresh_status=refresh_status
+        )
         return ToolExecutor(cli_toolsets)
 
     def create_tool_executor(self, dal: Optional[SupabaseDal]) -> ToolExecutor:
@@ -281,19 +308,32 @@ class Config(RobustaBaseConfig):
         return self._server_tool_executor
 
     def create_console_toolcalling_llm(
-        self, dal: Optional[SupabaseDal] = None
+        self,
+        dal: Optional[SupabaseDal] = None,
+        refresh_toolsets: bool = False,
+        tracer=None,
     ) -> ToolCallingLLM:
-        tool_executor = self.create_console_tool_executor(dal)
-        return ToolCallingLLM(tool_executor, self.max_steps, self._get_llm())
+        tool_executor = self.create_console_tool_executor(dal, refresh_toolsets)
+        return ToolCallingLLM(
+            tool_executor, self.max_steps, self._get_llm(tracer=tracer)
+        )
 
     def create_toolcalling_llm(
-        self, dal: Optional[SupabaseDal] = None, model: Optional[str] = None
+        self,
+        dal: Optional[SupabaseDal] = None,
+        model: Optional[str] = None,
+        tracer=None,
     ) -> ToolCallingLLM:
         tool_executor = self.create_tool_executor(dal)
-        return ToolCallingLLM(tool_executor, self.max_steps, self._get_llm(model))
+        return ToolCallingLLM(
+            tool_executor, self.max_steps, self._get_llm(model, tracer)
+        )
 
     def create_issue_investigator(
-        self, dal: Optional[SupabaseDal] = None, model: Optional[str] = None
+        self,
+        dal: Optional[SupabaseDal] = None,
+        model: Optional[str] = None,
+        tracer=None,
     ) -> IssueInvestigator:
         all_runbooks = load_builtin_runbooks()
         for runbook_path in self.custom_runbooks:
@@ -302,7 +342,7 @@ class Config(RobustaBaseConfig):
         runbook_manager = RunbookManager(all_runbooks)
         tool_executor = self.create_tool_executor(dal)
         return IssueInvestigator(
-            tool_executor, runbook_manager, self.max_steps, self._get_llm(model)
+            tool_executor, runbook_manager, self.max_steps, self._get_llm(model, tracer)
         )
 
     def create_console_issue_investigator(
@@ -411,7 +451,7 @@ class Config(RobustaBaseConfig):
             raise ValueError("--slack-channel must be specified")
         return SlackDestination(self.slack_token.get_secret_value(), self.slack_channel)
 
-    def _get_llm(self, model_key: Optional[str] = None) -> LLM:
+    def _get_llm(self, model_key: Optional[str] = None, tracer=None) -> LLM:
         api_key = self.api_key.get_secret_value() if self.api_key else None
         model = self.model
         model_params = {}
@@ -425,7 +465,7 @@ class Config(RobustaBaseConfig):
             api_key = model_params.pop("api_key", api_key)
             model = model_params.pop("model", model)
 
-        return DefaultLLM(model, api_key, model_params)  # type: ignore
+        return DefaultLLM(model, api_key, model_params, tracer)  # type: ignore
 
     def get_models_list(self) -> List[str]:
         if self._model_list:

holmes/core/investigation_structured_output.py

@@ -262,3 +262,10 @@ def is_response_an_incorrect_tool_call(
                     return False
             return True
     return False
+
+
+def clear_json_markdown(text: str):
+    if text and text.startswith("```json") and text.endswith("```"):
+        return text[8:-3]
+
+    return text

holmes/core/llm.py

@@ -64,12 +64,19 @@ class DefaultLLM(LLM):
     base_url: Optional[str]
     args: Dict
 
-    def __init__(self, model: str, api_key: Optional[str] = None, args: Dict = {}):
+    def __init__(
+        self,
+        model: str,
+        api_key: Optional[str] = None,
+        args: Optional[Dict] = None,
+        tracer=None,
+    ):
         self.model = model
         self.api_key = api_key
-        self.args = args
+        self.args = args or {}
+        self.tracer = tracer
 
-        if not args:
+        if not self.args:
             self.check_llm(self.model, self.api_key)
 
     def check_llm(self, model: str, api_key: Optional[str]):
@@ -214,7 +221,10 @@ class DefaultLLM(LLM):
         if self.args.get("thinking", None):
             litellm.modify_params = True
 
-        result = litellm.completion(
+        # Get the litellm module to use (wrapped or unwrapped)
+        litellm_to_use = self.tracer.wrap_llm(litellm) if self.tracer else litellm
+
+        result = litellm_to_use.completion(
             model=self.model,
             api_key=self.api_key,
             messages=messages,

holmes/core/models.py

@@ -155,3 +155,27 @@ class WorkloadHealthChatRequest(ChatRequestBaseModel):
     ask: str
     workload_health_result: WorkloadHealthInvestigationResult
     resource: dict
+
+
+workload_health_structured_output = {
+    "type": "json_schema",
+    "json_schema": {
+        "name": "WorkloadHealthResult",
+        "strict": False,
+        "schema": {
+            "type": "object",
+            "properties": {
+                "workload_healthy": {
+                    "type": "boolean",
+                    "description": "is the workload in healthy state or in error state",
+                },
+                "root_cause_summary": {
+                    "type": "string",
+                    "description": "concise short explaination leading to the workload_healthy result, pinpoint reason and root cause for the workload issues if any.",
+                },
+            },
+            "required": ["root_cause_summary", "workload_healthy"],
+            "additionalProperties": False,
+        },
+    },
+}

holmes/core/tool_calling_llm.py

@@ -39,6 +39,7 @@ from holmes.utils.global_instructions import (
 )
 from holmes.utils.tags import format_tags_in_string, parse_messages_tags
 from holmes.core.tools_utils.tool_executor import ToolExecutor
+from holmes.core.tracing import DummySpan, SpanType
 
 
 def format_tool_result_data(tool_result: StructuredToolResult) -> str:
@@ -200,9 +201,12 @@ class LLMResult(BaseModel):
 class ToolCallingLLM:
     llm: LLM
 
-    def __init__(self, tool_executor: ToolExecutor, max_steps: int, llm: LLM):
+    def __init__(
+        self, tool_executor: ToolExecutor, max_steps: int, llm: LLM, tracer=None
+    ):
         self.tool_executor = tool_executor
         self.max_steps = max_steps
+        self.tracer = tracer
         self.llm = llm
 
     def prompt_call(
@@ -230,8 +234,11 @@ class ToolCallingLLM:
         messages: List[Dict[str, str]],
         post_process_prompt: Optional[str] = None,
         response_format: Optional[Union[dict, Type[BaseModel]]] = None,
+        trace_span=DummySpan(),
     ) -> LLMResult:
-        return self.call(messages, post_process_prompt, response_format)
+        return self.call(
+            messages, post_process_prompt, response_format, trace_span=trace_span
+        )
 
     @sentry_sdk.trace
     def call(  # type: ignore
@@ -241,6 +248,7 @@ class ToolCallingLLM:
         response_format: Optional[Union[dict, Type[BaseModel]]] = None,
         user_prompt: Optional[str] = None,
         sections: Optional[InputSectionsDataType] = None,
+        trace_span=DummySpan(),
     ) -> LLMResult:
         perf_timing = PerformanceTiming("tool_calling_llm.call")
         tool_calls = []  # type: ignore
@@ -270,6 +278,7 @@ class ToolCallingLLM:
                 perf_timing.measure("truncate_messages_to_fit_context")
 
             logging.debug(f"sending messages={messages}\n\ntools={tools}")
+
             try:
                 full_response = self.llm.completion(
                     messages=parse_messages_tags(messages),
@@ -291,6 +300,7 @@ class ToolCallingLLM:
                     )
                 else:
                     raise
+
             response = full_response.choices[0]  # type: ignore
 
             response_message = response.message  # type: ignore
@@ -347,15 +357,17 @@ class ToolCallingLLM:
                 )
 
             perf_timing.measure("pre-tool-calls")
-            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 futures = []
-                for t in tools_to_call:
+                for tool_index, t in enumerate(tools_to_call, 1):
                     logging.debug(f"Tool to call: {t}")
                     futures.append(
                         executor.submit(
                             self._invoke_tool,
                             tool_to_call=t,
                             previous_tool_calls=tool_calls,
+                            trace_span=trace_span,
+                            tool_number=tool_index,
                         )
                     )
 
@@ -367,10 +379,16 @@ class ToolCallingLLM:
 
                     perf_timing.measure(f"tool completed {tool_call_result.tool_name}")
 
+                # Add a blank line after all tools in this batch complete
+                if tools_to_call:
+                    logging.info("")
+
     def _invoke_tool(
         self,
         tool_to_call: ChatCompletionMessageToolCall,
         previous_tool_calls: list[dict],
+        trace_span=DummySpan(),
+        tool_number=None,
     ) -> ToolCallResult:
         tool_name = tool_to_call.function.name
         tool_params = None
@@ -399,6 +417,10 @@ class ToolCallingLLM:
             )
 
         tool_response = None
+
+        # Create tool span if tracing is enabled
+        tool_span = trace_span.start_span(name=tool_name, type=SpanType.TOOL)
+
         try:
             tool_response = prevent_overly_repeated_tool_call(
                 tool_name=tool.name,
@@ -406,7 +428,7 @@ class ToolCallingLLM:
                 tool_calls=previous_tool_calls,
             )
             if not tool_response:
-                tool_response = tool.invoke(tool_params)
+                tool_response = tool.invoke(tool_params, tool_number=tool_number)
 
             if not isinstance(tool_response, StructuredToolResult):
                 # Should never be needed but ensure Holmes does not crash if one of the tools does not return the right type
@@ -419,6 +441,16 @@ class ToolCallingLLM:
                     params=tool_params,
                 )
 
+            # Log tool execution to trace span
+            tool_span.log(
+                input=tool_params,
+                output=tool_response.data,
+                metadata={
+                    "status": tool_response.status.value,
+                    "error": tool_response.error,
+                },
+            )
+
         except Exception as e:
             logging.error(
                 f"Tool call to {tool_name} failed with an Exception", exc_info=True
@@ -428,6 +460,14 @@ class ToolCallingLLM:
                 error=f"Tool call failed: {e}",
                 params=tool_params,
             )
+
+            # Log error to trace span
+            tool_span.log(
+                input=tool_params, output=str(e), metadata={"status": "ERROR"}
+            )
+        finally:
+            # End tool span
+            tool_span.end()
         return ToolCallResult(
             tool_call_id=tool_call_id,
             tool_name=tool_name,
@@ -650,12 +690,14 @@ class ToolCallingLLM:
             perf_timing.measure("pre-tool-calls")
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 futures = []
-                for t in tools_to_call:  # type: ignore
+                for tool_index, t in enumerate(tools_to_call, 1):  # type: ignore
                     futures.append(
                         executor.submit(
                             self._invoke_tool,
                             tool_to_call=t,  # type: ignore
                             previous_tool_calls=tool_calls,
+                            trace_span=DummySpan(),  # Streaming mode doesn't support tracing yet
+                            tool_number=tool_index,
                         )
                     )
                     yield create_sse_message(

holmes/core/tools.py

@@ -139,9 +139,12 @@ class Tool(ABC, BaseModel):
             tool_parameters=self.parameters,
         )
 
-    def invoke(self, params: Dict) -> StructuredToolResult:
+    def invoke(
+        self, params: Dict, tool_number: Optional[int] = None
+    ) -> StructuredToolResult:
+        tool_number_str = f"#{tool_number} " if tool_number else ""
         logging.info(
-            f"Running tool [bold]{self.name}[/bold]: {self.get_parameterized_one_liner(params)}"
+            f"Running tool {tool_number_str}[bold]{self.name}[/bold]: {self.get_parameterized_one_liner(params)}"
         )
         start_time = time.time()
         result = self._invoke(params)
@@ -152,7 +155,7 @@ class Tool(ABC, BaseModel):
             else str(result)
         )
         logging.info(
-            f"  [dim]Finished in {elapsed:.2f}s, output length: {len(output_str):,} characters[/dim]\n"
+            f"  [dim]Finished {tool_number_str}in {elapsed:.2f}s, output length: {len(output_str):,} characters - /show to view contents[/dim]"
         )
         return result
 
@@ -370,7 +373,7 @@ class Toolset(BaseModel):
             exclude_unset=True,
             exclude=("name"),  # type: ignore
         ).items():
-            if field in self.model_fields and value not in (None, [], {}, ""):
+            if field in self.__class__.model_fields and value not in (None, [], {}, ""):
                 setattr(self, field, value)
 
     @model_validator(mode="before")

holmes/core/toolset_manager.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import json
 import logging
 import os
@@ -113,14 +114,27 @@ class ToolsetManager:
         # check_prerequisites against each enabled toolset
         if not check_prerequisites:
             return list(toolsets_by_name.values())
+
+        enabled_toolsets: List[Toolset] = []
         for _, toolset in toolsets_by_name.items():
             if toolset.enabled:
-                toolset.check_prerequisites()
+                enabled_toolsets.append(toolset)
             else:
                 toolset.status = ToolsetStatusEnum.DISABLED
+        self.check_toolset_prerequisites(enabled_toolsets)
 
         return list(toolsets_by_name.values())
 
+    @classmethod
+    def check_toolset_prerequisites(cls, toolsets: list[Toolset]):
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+            futures = []
+            for toolset in toolsets:
+                futures.append(executor.submit(toolset.check_prerequisites))
+
+            for _ in concurrent.futures.as_completed(futures):
+                pass
+
     def _load_toolsets_from_config(
         self,
         toolsets: dict[str, dict[str, Any]],
@@ -231,6 +245,7 @@ class ToolsetManager:
             dal=dal, check_prerequisites=False, toolset_tags=toolset_tags
         )
 
+        enabled_toolsets_from_cache: List[Toolset] = []
         for toolset in all_toolsets_with_status:
             if toolset.name in toolsets_status_by_name:
                 # Update the status and error from the cached status
@@ -242,13 +257,15 @@ class ToolsetManager:
                     cached_status.get("type", ToolsetType.BUILTIN)
                 )
                 toolset.path = cached_status.get("path", None)
-            # check prerequisites for only enabled toolset when the toolset is loaded from cache
+            # check prerequisites for only enabled toolset when the toolset is loaded from cache. When the toolset is
+            # not loaded from cache, the prerequisites are checked in the refresh_toolset_status method.
             if (
                 toolset.enabled
                 and toolset.status == ToolsetStatusEnum.ENABLED
                 and using_cached
             ):
-                toolset.check_prerequisites()  # type: ignore
+                enabled_toolsets_from_cache.append(toolset)
+        self.check_toolset_prerequisites(enabled_toolsets_from_cache)
 
         # CLI custom toolsets status are not cached, and their prerequisites are always checked whenever the CLI runs.
         custom_toolsets_from_cli = self._load_toolsets_from_paths(
@@ -257,13 +274,15 @@ class ToolsetManager:
             check_conflict_default=True,
         )
         # custom toolsets from cli as experimental toolset should not override custom toolsets from config
+        enabled_toolsets_from_cli: List[Toolset] = []
         for custom_toolset_from_cli in custom_toolsets_from_cli:
             if custom_toolset_from_cli.name in toolsets_status_by_name:
                 raise ValueError(
                     f"Toolset {custom_toolset_from_cli.name} from cli is already defined in existing toolset"
                 )
-            # status of custom toolsets from cli is not cached, and we need to check prerequisites every time the cli runs.
-            custom_toolset_from_cli.check_prerequisites()
+            enabled_toolsets_from_cli.append(custom_toolset_from_cli)
+        # status of custom toolsets from cli is not cached, and we need to check prerequisites every time the cli runs.
+        self.check_toolset_prerequisites(enabled_toolsets_from_cli)
 
         all_toolsets_with_status.extend(custom_toolsets_from_cli)
         if using_cached: