holmesgpt 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

holmes/core/tool_calling_llm.py

@@ -2,18 +2,22 @@ import concurrent.futures
 import json
 import logging
 import textwrap
-import uuid
-from typing import Dict, List, Optional, Type, Union
+from typing import Dict, List, Optional, Type, Union, Callable
+
 
 import sentry_sdk
 from openai import BadRequestError
 from openai.types.chat.chat_completion_message_tool_call import (
     ChatCompletionMessageToolCall,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from rich.console import Console
 
-from holmes.common.env_vars import TEMPERATURE, MAX_OUTPUT_TOKEN_RESERVATION
+from holmes.common.env_vars import (
+    TEMPERATURE,
+    MAX_OUTPUT_TOKEN_RESERVATION,
+    LOG_LLM_USAGE_RESPONSE,
+)
 
 from holmes.core.investigation_structured_output import (
     DEFAULT_SECTIONS,
@@ -39,9 +43,80 @@ from holmes.core.tools_utils.tool_executor import ToolExecutor
 from holmes.core.tracing import DummySpan
 from holmes.utils.colors import AI_COLOR
 from holmes.utils.stream import StreamEvents, StreamMessage
-from holmes.core.todo_manager import (
-    get_todo_manager,
-)
+
+# Create a named logger for cost tracking
+cost_logger = logging.getLogger("holmes.costs")
+
+
+class LLMCosts(BaseModel):
+    """Tracks cost and token usage for LLM calls."""
+
+    total_cost: float = 0.0
+    total_tokens: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+def _extract_cost_from_response(full_response) -> float:
+    """Extract cost value from LLM response.
+
+    Args:
+        full_response: The raw LLM response object
+
+    Returns:
+        The cost as a float, or 0.0 if not available
+    """
+    try:
+        cost_value = (
+            full_response._hidden_params.get("response_cost", 0)
+            if hasattr(full_response, "_hidden_params")
+            else 0
+        )
+        # Ensure cost is a float
+        return float(cost_value) if cost_value is not None else 0.0
+    except Exception:
+        return 0.0
+
+
+def _process_cost_info(
+    full_response, costs: Optional[LLMCosts] = None, log_prefix: str = "LLM call"
+) -> None:
+    """Process cost and token information from LLM response.
+
+    Logs the cost information and optionally accumulates it into a costs object.
+
+    Args:
+        full_response: The raw LLM response object
+        costs: Optional LLMCosts object to accumulate costs into
+        log_prefix: Prefix for logging messages (e.g., "LLM call", "Post-processing")
+    """
+    try:
+        cost = _extract_cost_from_response(full_response)
+        usage = getattr(full_response, "usage", {})
+
+        if usage:
+            if LOG_LLM_USAGE_RESPONSE:  # shows stats on token cache usage
+                logging.info(f"LLM usage response:\n{usage}\n")
+            prompt_toks = usage.get("prompt_tokens", 0)
+            completion_toks = usage.get("completion_tokens", 0)
+            total_toks = usage.get("total_tokens", 0)
+            cost_logger.debug(
+                f"{log_prefix} cost: ${cost:.6f} | Tokens: {prompt_toks} prompt + {completion_toks} completion = {total_toks} total"
+            )
+            # Accumulate costs and tokens if costs object provided
+            if costs:
+                costs.total_cost += cost
+                costs.prompt_tokens += prompt_toks
+                costs.completion_tokens += completion_toks
+                costs.total_tokens += total_toks
+        elif cost > 0:
+            cost_logger.debug(
+                f"{log_prefix} cost: ${cost:.6f} | Token usage not available"
+            )
+            if costs:
+                costs.total_cost += cost
+    except Exception as e:
+        logging.debug(f"Could not extract cost information: {e}")
 
 
 def format_tool_result_data(tool_result: StructuredToolResult) -> str:
@@ -186,11 +261,11 @@ class ToolCallResult(BaseModel):
         }
 
 
-class LLMResult(BaseModel):
+class LLMResult(LLMCosts):
     tool_calls: Optional[List[ToolCallResult]] = None
     result: Optional[str] = None
     unprocessed_result: Optional[str] = None
-    instructions: List[str] = []
+    instructions: List[str] = Field(default_factory=list)
     # TODO: clean up these two
     prompt: Optional[str] = None
     messages: Optional[List[dict]] = None
@@ -211,7 +286,9 @@ class ToolCallingLLM:
         self.max_steps = max_steps
         self.tracer = tracer
         self.llm = llm
-        self.investigation_id = str(uuid.uuid4())
+        self.approval_callback: Optional[
+            Callable[[StructuredToolResult], tuple[bool, Optional[str]]]
+        ] = None
 
     def prompt_call(
         self,
@@ -259,6 +336,8 @@ class ToolCallingLLM:
     ) -> LLMResult:
         perf_timing = PerformanceTiming("tool_calling_llm.call")
         tool_calls = []  # type: ignore
+        costs = LLMCosts()
+
         tools = self.tool_executor.get_all_tools_openai_format(
             target_model=self.llm.model
         )
@@ -299,6 +378,9 @@ class ToolCallingLLM:
                 )
                 logging.debug(f"got response {full_response.to_json()}")  # type: ignore
 
+                # Extract and accumulate cost information
+                _process_cost_info(full_response, costs, "LLM call")
+
                 perf_timing.measure("llm.completion")
             # catch a known error that occurs with Azure and replace the error message with something more obvious to the user
             except BadRequestError as e:
@@ -352,11 +434,14 @@ class ToolCallingLLM:
                 if post_process_prompt and user_prompt:
                     logging.info("Running post processing on investigation.")
                     raw_response = text_response
-                    post_processed_response = self._post_processing_call(
-                        prompt=user_prompt,
-                        investigation=raw_response,
-                        user_prompt=post_process_prompt,
+                    post_processed_response, post_processing_cost = (
+                        self._post_processing_call(
+                            prompt=user_prompt,
+                            investigation=raw_response,
+                            user_prompt=post_process_prompt,
+                        )
                     )
+                    costs.total_cost += post_processing_cost
 
                     perf_timing.end(f"- completed in {i} iterations -")
                     return LLMResult(
@@ -365,6 +450,7 @@ class ToolCallingLLM:
                         tool_calls=tool_calls,
                         prompt=json.dumps(messages, indent=2),
                         messages=messages,
+                        **costs.model_dump(),  # Include all cost fields
                     )
 
                 perf_timing.end(f"- completed in {i} iterations -")
@@ -373,6 +459,7 @@ class ToolCallingLLM:
                     tool_calls=tool_calls,
                     prompt=json.dumps(messages, indent=2),
                     messages=messages,
+                    **costs.model_dump(),  # Include all cost fields
                 )
 
             if text_response and text_response.strip():
@@ -383,33 +470,106 @@ class ToolCallingLLM:
             perf_timing.measure("pre-tool-calls")
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 futures = []
+                futures_tool_numbers: dict[
+                    concurrent.futures.Future, Optional[int]
+                ] = {}
+                tool_number: Optional[int]
                 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_number_offset + tool_index,
-                        )
+                    tool_number = tool_number_offset + tool_index
+                    future = executor.submit(
+                        self._invoke_llm_tool_call,
+                        tool_to_call=t,
+                        previous_tool_calls=tool_calls,
+                        trace_span=trace_span,
+                        tool_number=tool_number,
                     )
+                    futures_tool_numbers[future] = tool_number
+                    futures.append(future)
 
                 for future in concurrent.futures.as_completed(futures):
                     tool_call_result: ToolCallResult = future.result()
 
+                    tool_number = (
+                        futures_tool_numbers[future]
+                        if future in futures_tool_numbers
+                        else None
+                    )
+                    tool_call_result = self.handle_tool_call_approval(
+                        tool_call_result=tool_call_result, tool_number=tool_number
+                    )
+
                     tool_calls.append(tool_call_result.as_tool_result_response())
                     messages.append(tool_call_result.as_tool_call_message())
 
                     perf_timing.measure(f"tool completed {tool_call_result.tool_name}")
 
+                # Update the tool number offset for the next iteration
+                tool_number_offset += len(tools_to_call)
+
                 # Add a blank line after all tools in this batch complete
                 if tools_to_call:
                     logging.info("")
 
         raise Exception(f"Too many LLM calls - exceeded max_steps: {i}/{max_steps}")
 
-    def _invoke_tool(
+    def _directly_invoke_tool(
+        self,
+        tool_name: str,
+        tool_params: dict,
+        user_approved: bool,
+        trace_span=DummySpan(),
+        tool_number: Optional[int] = None,
+    ) -> StructuredToolResult:
+        tool_span = trace_span.start_span(name=tool_name, type="tool")
+        tool = self.tool_executor.get_tool_by_name(tool_name)
+        tool_response = None
+        try:
+            if (not tool) or (tool_params is None):
+                logging.warning(
+                    f"Skipping tool execution for {tool_name}: args: {tool_params}"
+                )
+                tool_response = StructuredToolResult(
+                    status=ToolResultStatus.ERROR,
+                    error=f"Failed to find tool {tool_name}",
+                    params=tool_params,
+                )
+            else:
+                tool_response = tool.invoke(
+                    tool_params, tool_number=tool_number, user_approved=user_approved
+                )
+        except Exception as e:
+            logging.error(
+                f"Tool call to {tool_name} failed with an Exception", exc_info=True
+            )
+            tool_response = StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                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"}
+            )
+
+        tool_span.log(
+            input=tool_params,
+            output=tool_response.data,
+            metadata={
+                "status": tool_response.status.value,
+                "error": tool_response.error,
+                "description": tool.get_parameterized_one_liner(tool_params)
+                if tool
+                else "",
+                "structured_tool_result": tool_response,
+            },
+        )
+        tool_span.end()
+
+        return tool_response
+
+    def _invoke_llm_tool_call(
         self,
         tool_to_call: ChatCompletionMessageToolCall,
         previous_tool_calls: list[dict],
@@ -438,92 +598,97 @@ class ToolCallingLLM:
                 ),
             )
 
-        tool_params = None
+        tool_params = {}
         try:
             tool_params = json.loads(tool_arguments)
         except Exception:
             logging.warning(
                 f"Failed to parse arguments for tool: {tool_name}. args: {tool_arguments}"
             )
-        tool_call_id = tool_to_call.id
-        tool = self.tool_executor.get_tool_by_name(tool_name)
-
-        if (not tool) or (tool_params is None):
-            logging.warning(
-                f"Skipping tool execution for {tool_name}: args: {tool_arguments}"
-            )
-            return ToolCallResult(
-                tool_call_id=tool_call_id,
-                tool_name=tool_name,
-                description="NA",
-                result=StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
-                    error=f"Failed to find tool {tool_name}",
-                    params=tool_params,
-                ),
-            )
 
-        tool_response = None
+        tool_call_id = tool_to_call.id
 
-        # Create tool span if tracing is enabled
-        tool_span = trace_span.start_span(name=tool_name, type="tool")
+        tool_response = prevent_overly_repeated_tool_call(
+            tool_name=tool_name,
+            tool_params=tool_params,
+            tool_calls=previous_tool_calls,
+        )
 
-        try:
-            tool_response = prevent_overly_repeated_tool_call(
-                tool_name=tool.name,
+        if not tool_response:
+            tool_response = self._directly_invoke_tool(
+                tool_name=tool_name,
                 tool_params=tool_params,
-                tool_calls=previous_tool_calls,
-            )
-            if not tool_response:
-                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
-                logging.error(
-                    f"Tool {tool.name} return type is not StructuredToolResult. Nesting the tool result into StructuredToolResult..."
-                )
-                tool_response = StructuredToolResult(
-                    status=ToolResultStatus.SUCCESS,
-                    data=tool_response,
-                    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,
-                    "description": tool.get_parameterized_one_liner(tool_params),
-                    "structured_tool_result": tool_response,
-                },
+                user_approved=False,
+                trace_span=trace_span,
+                tool_number=tool_number,
             )
 
-        except Exception as e:
+        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
             logging.error(
-                f"Tool call to {tool_name} failed with an Exception", exc_info=True
+                f"Tool {tool_name} return type is not StructuredToolResult. Nesting the tool result into StructuredToolResult..."
             )
             tool_response = StructuredToolResult(
-                status=ToolResultStatus.ERROR,
-                error=f"Tool call failed: {e}",
+                status=ToolResultStatus.SUCCESS,
+                data=tool_response,
                 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()
+        tool = self.tool_executor.get_tool_by_name(tool_name)
         return ToolCallResult(
             tool_call_id=tool_call_id,
             tool_name=tool_name,
-            description=tool.get_parameterized_one_liner(tool_params),
+            description=tool.get_parameterized_one_liner(tool_params) if tool else "",
             result=tool_response,
         )
 
+    def handle_tool_call_approval(
+        self, tool_call_result: ToolCallResult, tool_number: Optional[int]
+    ) -> ToolCallResult:
+        """
+        Handle approval for a single tool call if required.
+
+        Args:
+            tool_call_result: A single tool call result that may require approval
+
+        Returns:
+            Updated tool call result with approved/denied status
+        """
+
+        if tool_call_result.result.status != ToolResultStatus.APPROVAL_REQUIRED:
+            return tool_call_result
+
+        # If no approval callback, convert to ERROR because it is assumed the client may not be able to handle approvals
+        if not self.approval_callback:
+            tool_call_result.result.status = ToolResultStatus.ERROR
+            return tool_call_result
+
+        # Get approval from user
+        approved, feedback = self.approval_callback(tool_call_result.result)
+
+        if approved:
+            logging.debug(
+                f"User approved command: {tool_call_result.result.invocation}"
+            )
+
+            new_response = self._directly_invoke_tool(
+                tool_name=tool_call_result.tool_name,
+                tool_params=tool_call_result.result.params or {},
+                user_approved=True,
+                trace_span=DummySpan(),
+                tool_number=tool_number,
+            )
+            tool_call_result.result = new_response
+        else:
+            # User denied - update to error
+            feedback_text = f" User feedback: {feedback}" if feedback else ""
+            tool_call_result.result.status = ToolResultStatus.ERROR
+            tool_call_result.result.error = (
+                f"User denied command execution.{feedback_text}"
+            )
+
+        return tool_call_result
+
     @staticmethod
     def __load_post_processing_user_prompt(
         input_prompt, investigation, user_prompt: Optional[str] = None
@@ -540,7 +705,7 @@ class ToolCallingLLM:
         investigation,
         user_prompt: Optional[str] = None,
         system_prompt: str = "You are an AI assistant summarizing Kubernetes issues.",
-    ) -> Optional[str]:
+    ) -> tuple[Optional[str], float]:
         try:
             user_prompt = ToolCallingLLM.__load_post_processing_user_prompt(
                 prompt, investigation, user_prompt
@@ -559,10 +724,18 @@ class ToolCallingLLM:
             ]
             full_response = self.llm.completion(messages=messages, temperature=0)
             logging.debug(f"Post processing response {full_response}")
-            return full_response.choices[0].message.content  # type: ignore
+
+            # Extract and log cost information for post-processing
+            post_processing_cost = _extract_cost_from_response(full_response)
+            if post_processing_cost > 0:
+                cost_logger.debug(
+                    f"Post-processing LLM cost: ${post_processing_cost:.6f}"
+                )
+
+            return full_response.choices[0].message.content, post_processing_cost  # type: ignore
         except Exception:
             logging.exception("Failed to run post processing", exc_info=True)
-            return investigation
+            return investigation, 0.0
 
     @sentry_sdk.trace
     def truncate_messages_to_fit_context(
@@ -602,6 +775,7 @@ class ToolCallingLLM:
         perf_timing.measure("get_all_tools_openai_format")
         max_steps = self.max_steps
         i = 0
+        tool_number_offset = 0
 
         while i < max_steps:
             i += 1
@@ -634,6 +808,10 @@ class ToolCallingLLM:
                     stream=False,
                     drop_params=True,
                 )
+
+                # Log cost information for this iteration (no accumulation in streaming)
+                _process_cost_info(full_response, log_prefix="LLM iteration")
+
                 perf_timing.measure("llm.completion")
             # catch a known error that occurs with Azure and replace the error message with something more obvious to the user
             except BadRequestError as e:
@@ -689,15 +867,15 @@ class ToolCallingLLM:
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 futures = []
                 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,
-                        )
+                    tool_number = tool_number_offset + tool_index
+                    future = executor.submit(
+                        self._invoke_llm_tool_call,
+                        tool_to_call=t,  # type: ignore
+                        previous_tool_calls=tool_calls,
+                        trace_span=DummySpan(),  # Streaming mode doesn't support tracing yet
+                        tool_number=tool_number,
                     )
+                    futures.append(future)
                     yield StreamMessage(
                         event=StreamEvents.START_TOOL,
                         data={"tool_name": t.function.name, "id": t.id},
@@ -716,6 +894,9 @@ class ToolCallingLLM:
                         data=tool_call_result.as_streaming_tool_result_response(),
                     )
 
+                # Update the tool number offset for the next iteration
+                tool_number_offset += len(tools_to_call)
+
         raise Exception(
             f"Too many LLM calls - exceeded max_steps: {i}/{self.max_steps}"
         )
@@ -793,9 +974,6 @@ class IssueInvestigator(ToolCallingLLM):
                 "[bold]No runbooks found for this issue. Using default behaviour. (Add runbooks to guide the investigation.)[/bold]"
             )
 
-        todo_manager = get_todo_manager()
-        todo_context = todo_manager.format_tasks_for_prompt(self.investigation_id)
-
         system_prompt = load_and_render_prompt(
             prompt,
             {
@@ -804,8 +982,6 @@ class IssueInvestigator(ToolCallingLLM):
                 "structured_output": request_structured_output_from_llm,
                 "toolsets": self.tool_executor.toolsets,
                 "cluster_name": self.cluster_name,
-                "todo_list": todo_context,
-                "investigation_id": self.investigation_id,
             },
         )
 

holmes/core/tools.py

@@ -24,12 +24,15 @@ class ToolResultStatus(str, Enum):
     SUCCESS = "success"
     ERROR = "error"
     NO_DATA = "no_data"
+    APPROVAL_REQUIRED = "approval_required"
 
     def to_color(self) -> str:
         if self == ToolResultStatus.SUCCESS:
             return "green"
         elif self == ToolResultStatus.ERROR:
             return "red"
+        elif self == ToolResultStatus.APPROVAL_REQUIRED:
+            return "yellow"
         else:
             return "white"
 
@@ -38,6 +41,8 @@ class ToolResultStatus(str, Enum):
             return "✔"
         elif self == ToolResultStatus.ERROR:
             return "❌"
+        elif self == ToolResultStatus.APPROVAL_REQUIRED:
+            return "⚠️"
         else:
             return "⚪️"
 
@@ -148,14 +153,17 @@ class Tool(ABC, BaseModel):
         )
 
     def invoke(
-        self, params: Dict, tool_number: Optional[int] = None
+        self,
+        params: Dict,
+        tool_number: Optional[int] = None,
+        user_approved: bool = False,
     ) -> StructuredToolResult:
         tool_number_str = f"#{tool_number} " if tool_number else ""
         logging.info(
             f"Running tool {tool_number_str}[bold]{self.name}[/bold]: {self.get_parameterized_one_liner(params)}"
         )
         start_time = time.time()
-        result = self._invoke(params)
+        result = self._invoke(params=params, user_approved=user_approved)
         result.icon_url = self.icon_url
         elapsed = time.time() - start_time
         output_str = (
@@ -171,7 +179,13 @@ class Tool(ABC, BaseModel):
         return result
 
     @abstractmethod
-    def _invoke(self, params: Dict) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
+        """
+        params: the tool params
+        user_approved: whether the tool call is approved by the user. Can be used to confidently execute unsafe actions.
+        """
         pass
 
     @abstractmethod
@@ -223,7 +237,9 @@ class YAMLTool(Tool, BaseModel):
             return ToolResultStatus.NO_DATA
         return ToolResultStatus.SUCCESS
 
-    def _invoke(self, params) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if self.command is not None:
             raw_output, return_code, invocation = self.__invoke_command(params)
         else:

holmes/core/toolset_manager.py

@@ -266,11 +266,7 @@ class ToolsetManager:
                 toolset.path = cached_status.get("path", None)
             # 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
-            ):
+            if toolset.enabled and toolset.status == ToolsetStatusEnum.ENABLED:
                 enabled_toolsets_from_cache.append(toolset)
         self.check_toolset_prerequisites(enabled_toolsets_from_cache)
 

holmes/core/tracing.py

@@ -120,7 +120,7 @@ class DummySpan:
 class DummyTracer:
     """A no-op tracer implementation for when tracing is disabled."""
 
-    def start_experiment(self, experiment_name=None, metadata=None):
+    def start_experiment(self, experiment_name=None, additional_metadata=None):
         """No-op experiment creation."""
         return None