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

holmes/core/tool_calling_llm.py

@@ -4,6 +4,13 @@ import logging
 import textwrap
 from typing import Dict, List, Optional, Type, Union, Callable, Any
 
+from holmes.core.models import (
+    ToolApprovalDecision,
+    ToolCallResult,
+    TruncationResult,
+    TruncationMetadata,
+    PendingToolApproval,
+)
 
 import sentry_sdk
 from openai import BadRequestError
@@ -32,8 +39,13 @@ from holmes.core.performance_timing import PerformanceTiming
 from holmes.core.resource_instruction import ResourceInstructions
 from holmes.core.runbooks import RunbookManager
 from holmes.core.safeguards import prevent_overly_repeated_tool_call
-from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    ToolInvokeContext,
+)
 from holmes.core.tools_utils.tool_context_window_limiter import (
+    get_max_token_count_for_single_tool,
     prevent_overly_big_tool_response,
 )
 from holmes.plugins.prompts import load_and_render_prompt
@@ -44,11 +56,6 @@ 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.tools_utils.data_types import (
-    TruncationResult,
-    ToolCallResult,
-    TruncationMetadata,
-)
 from holmes.core.tracing import DummySpan
 from holmes.utils.colors import AI_COLOR
 from holmes.utils.stream import StreamEvents, StreamMessage
@@ -264,6 +271,99 @@ class ToolCallingLLM:
             Callable[[StructuredToolResult], tuple[bool, Optional[str]]]
         ] = None
 
+    def process_tool_decisions(
+        self, messages: List[Dict[str, Any]], tool_decisions: List[ToolApprovalDecision]
+    ) -> List[Dict[str, Any]]:
+        """
+        Process tool approval decisions and execute approved tools.
+
+        Args:
+            messages: Current conversation messages
+            tool_decisions: List of ToolApprovalDecision objects
+
+        Returns:
+            Updated messages list with tool execution results
+        """
+        # Import here to avoid circular imports
+
+        # Find the last message with pending approvals
+        pending_message_idx = None
+        pending_tool_calls = None
+
+        for i in reversed(range(len(messages))):
+            msg = messages[i]
+            if msg.get("role") == "assistant" and msg.get("pending_approval"):
+                pending_message_idx = i
+                pending_tool_calls = msg.get("tool_calls", [])
+                break
+
+        if pending_message_idx is None or not pending_tool_calls:
+            # No pending approvals found
+            if tool_decisions:
+                logging.warning(
+                    f"Received {len(tool_decisions)} tool decisions but no pending approvals found"
+                )
+            return messages
+
+        # Create decision lookup
+        decisions_by_id = {
+            decision.tool_call_id: decision for decision in tool_decisions
+        }
+
+        # Validate that all decisions have corresponding pending tool calls
+        pending_tool_ids = {tool_call["id"] for tool_call in pending_tool_calls}
+        invalid_decisions = [
+            decision.tool_call_id
+            for decision in tool_decisions
+            if decision.tool_call_id not in pending_tool_ids
+        ]
+
+        if invalid_decisions:
+            logging.warning(
+                f"Received decisions for non-pending tool calls: {invalid_decisions}"
+            )
+
+        # Process each tool call
+        for tool_call in pending_tool_calls:
+            tool_call_id = tool_call["id"]
+            decision = decisions_by_id.get(tool_call_id)
+
+            if decision and decision.approved:
+                try:
+                    tool_call_obj = ChatCompletionMessageToolCall(**tool_call)
+                    llm_tool_result = self._invoke_llm_tool_call(
+                        tool_to_call=tool_call_obj,
+                        previous_tool_calls=[],
+                        trace_span=DummySpan(),
+                        tool_number=None,
+                    )
+                    messages.append(llm_tool_result.as_tool_call_message())
+
+                except Exception as e:
+                    logging.error(
+                        f"Failed to execute approved tool {tool_call_id}: {e}"
+                    )
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call_id,
+                            "role": "tool",
+                            "name": tool_call["function"]["name"],
+                            "content": f"Tool execution failed: {str(e)}",
+                        }
+                    )
+            else:
+                # Tool was rejected or no decision found, add rejection message
+                messages.append(
+                    {
+                        "tool_call_id": tool_call_id,
+                        "role": "tool",
+                        "name": tool_call["function"]["name"],
+                        "content": "Tool execution was denied by the user.",
+                    }
+                )
+
+        return messages
+
     def prompt_call(
         self,
         system_prompt: str,
@@ -527,9 +627,13 @@ class ToolCallingLLM:
             )
 
         try:
-            tool_response = tool.invoke(
-                tool_params, tool_number=tool_number, user_approved=user_approved
+            invoke_context = ToolInvokeContext(
+                tool_number=tool_number,
+                user_approved=user_approved,
+                llm=self.llm,
+                max_token_count=get_max_token_count_for_single_tool(self.llm),
             )
+            tool_response = tool.invoke(tool_params, context=invoke_context)
         except Exception as e:
             logging.error(
                 f"Tool call to {tool_name} failed with an Exception", exc_info=True
@@ -587,7 +691,9 @@ class ToolCallingLLM:
         return ToolCallResult(
             tool_call_id=tool_call_id,
             tool_name=tool_name,
-            description=tool.get_parameterized_one_liner(tool_params) if tool else "",
+            description=str(tool.get_parameterized_one_liner(tool_params))
+            if tool
+            else "",
             result=tool_response,
         )
 
@@ -765,12 +871,13 @@ class ToolCallingLLM:
         response_format: Optional[Union[dict, Type[BaseModel]]] = None,
         sections: Optional[InputSectionsDataType] = None,
         msgs: Optional[list[dict]] = None,
+        enable_tool_approval: bool = False,
     ):
         """
         This function DOES NOT call llm.completion(stream=true).
         This function streams holmes one iteration at a time instead of waiting for all iterations to complete.
         """
-        messages = []
+        messages: list[dict] = []
         if system_prompt:
             messages.append({"role": "system", "content": system_prompt})
         if user_prompt:
@@ -890,6 +997,11 @@ class ToolCallingLLM:
                 )
 
             perf_timing.measure("pre-tool-calls")
+
+            # Check if any tools require approval first
+            pending_approvals = []
+            approval_required_tools = []
+
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 futures = []
                 for tool_index, t in enumerate(tools_to_call, 1):  # type: ignore
@@ -909,15 +1021,85 @@ class ToolCallingLLM:
 
                 for future in concurrent.futures.as_completed(futures):
                     tool_call_result: ToolCallResult = future.result()
-                    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}")
+                    if (
+                        tool_call_result.result.status
+                        == StructuredToolResultStatus.APPROVAL_REQUIRED
+                    ):
+                        if enable_tool_approval:
+                            pending_approvals.append(
+                                PendingToolApproval(
+                                    tool_call_id=tool_call_result.tool_call_id,
+                                    tool_name=tool_call_result.tool_name,
+                                    description=tool_call_result.description,
+                                    params=tool_call_result.result.params or {},
+                                )
+                            )
+                            approval_required_tools.append(tool_call_result)
+
+                            yield StreamMessage(
+                                event=StreamEvents.TOOL_RESULT,
+                                data=tool_call_result.as_streaming_tool_result_response(),
+                            )
+                        else:
+                            tool_call_result.result.status = (
+                                StructuredToolResultStatus.ERROR
+                            )
+                            tool_call_result.result.error = f"Tool call rejected for security reasons: {tool_call_result.result.error}"
+
+                            tool_calls.append(
+                                tool_call_result.as_tool_result_response()
+                            )
+                            messages.append(tool_call_result.as_tool_call_message())
 
+                            yield StreamMessage(
+                                event=StreamEvents.TOOL_RESULT,
+                                data=tool_call_result.as_streaming_tool_result_response(),
+                            )
+
+                    else:
+                        tool_calls.append(tool_call_result.as_tool_result_response())
+                        messages.append(tool_call_result.as_tool_call_message())
+
+                        yield StreamMessage(
+                            event=StreamEvents.TOOL_RESULT,
+                            data=tool_call_result.as_streaming_tool_result_response(),
+                        )
+
+                # If we have approval required tools, end the stream with pending approvals
+                if pending_approvals:
+                    # Add assistant message with pending tool calls
+                    assistant_msg = {
+                        "role": "assistant",
+                        "content": response_message.content,
+                        "tool_calls": [
+                            {
+                                "id": result.tool_call_id,
+                                "type": "function",
+                                "function": {
+                                    "name": result.tool_name,
+                                    "arguments": json.dumps(result.result.params or {}),
+                                },
+                            }
+                            for result in approval_required_tools
+                        ],
+                        "pending_approval": True,
+                    }
+                    messages.append(assistant_msg)
+
+                    # End stream with approvals required
                     yield StreamMessage(
-                        event=StreamEvents.TOOL_RESULT,
-                        data=tool_call_result.as_streaming_tool_result_response(),
+                        event=StreamEvents.APPROVAL_REQUIRED,
+                        data={
+                            "content": None,
+                            "messages": messages,
+                            "pending_approvals": [
+                                approval.model_dump() for approval in pending_approvals
+                            ],
+                            "requires_approval": True,
+                        },
                     )
+                    return
 
                 # Update the tool number offset for the next iteration
                 tool_number_offset += len(tools_to_call)

holmes/core/tools.py

@@ -31,6 +31,7 @@ from pydantic import (
 )
 from rich.console import Console
 
+from holmes.core.llm import LLM
 from holmes.core.openai_formatting import format_tool_to_open_ai_standard
 from holmes.plugins.prompts import load_and_render_prompt
 from holmes.core.transformers import (
@@ -159,6 +160,15 @@ class ToolParameter(BaseModel):
     items: Optional["ToolParameter"] = None  # For array item schemas
 
 
+class ToolInvokeContext(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    tool_number: Optional[int] = None
+    user_approved: bool = False
+    llm: LLM
+    max_token_count: int
+
+
 class Tool(ABC, BaseModel):
     name: str
     description: str
@@ -225,15 +235,14 @@ class Tool(ABC, BaseModel):
     def invoke(
         self,
         params: Dict,
-        tool_number: Optional[int] = None,
-        user_approved: bool = False,
+        context: ToolInvokeContext,
     ) -> StructuredToolResult:
-        tool_number_str = f"#{tool_number} " if tool_number else ""
+        tool_number_str = f"#{context.tool_number} " if context.tool_number else ""
         logger.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=params, user_approved=user_approved)
+        result = self._invoke(params=params, context=context)
         result.icon_url = self.icon_url
 
         # Apply transformers to the result
@@ -244,7 +253,7 @@ class Tool(ABC, BaseModel):
             if hasattr(transformed_result, "get_stringified_data")
             else str(transformed_result)
         )
-        show_hint = f"/show {tool_number}" if tool_number else "/show"
+        show_hint = f"/show {context.tool_number}" if context.tool_number else "/show"
         line_count = output_str.count("\n") + 1 if output_str else 0
         logger.info(
             f"  [dim]Finished {tool_number_str}in {elapsed:.2f}s, output length: {len(output_str):,} characters ({line_count:,} lines) - {show_hint} to view contents[/dim]"
@@ -340,7 +349,9 @@ class Tool(ABC, BaseModel):
 
     @abstractmethod
     def _invoke(
-        self, params: dict, user_approved: bool = False
+        self,
+        params: dict,
+        context: ToolInvokeContext,
     ) -> StructuredToolResult:
         """
         params: the tool params
@@ -400,7 +411,9 @@ class YAMLTool(Tool, BaseModel):
         return StructuredToolResultStatus.SUCCESS
 
     def _invoke(
-        self, params: dict, user_approved: bool = False
+        self,
+        params: dict,
+        context: ToolInvokeContext,
     ) -> StructuredToolResult:
         if self.command is not None:
             raw_output, return_code, invocation = self.__invoke_command(params)

holmes/core/tools_utils/token_counting.py

@@ -0,0 +1,13 @@
+from holmes.core.llm import LLM
+from holmes.core.models import format_tool_result_data
+from holmes.core.tools import StructuredToolResult
+
+
+def count_tool_response_tokens(
+    llm: LLM, structured_tool_result: StructuredToolResult
+) -> int:
+    message = {
+        "role": "tool",
+        "content": format_tool_result_data(structured_tool_result),
+    }
+    return llm.count_tokens_for_message([message])

holmes/core/tools_utils/tool_context_window_limiter.py

@@ -1,33 +1,55 @@
+from typing import Optional
 from holmes.common.env_vars import TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT
 from holmes.core.llm import LLM
 from holmes.core.tools import StructuredToolResultStatus
-from holmes.core.tools_utils.data_types import ToolCallResult
+from holmes.core.models import ToolCallResult
 from holmes.utils import sentry_helper
 
 
+def get_pct_token_count(percent_of_total_context_window: float, llm: LLM) -> int:
+    context_window_size = llm.get_context_window_size()
+
+    if 0 < percent_of_total_context_window and percent_of_total_context_window <= 100:
+        return int(context_window_size * percent_of_total_context_window // 100)
+    else:
+        return context_window_size
+
+
+def get_max_token_count_for_single_tool(llm: LLM) -> int:
+    return get_pct_token_count(
+        percent_of_total_context_window=TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT, llm=llm
+    )
+
+
 def prevent_overly_big_tool_response(tool_call_result: ToolCallResult, llm: LLM):
-    if (
-        tool_call_result.result.status == StructuredToolResultStatus.SUCCESS
-        and 0 < TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT
-        and TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT <= 100
-    ):
-        message = tool_call_result.as_tool_call_message()
-
-        messages_token = llm.count_tokens_for_message(messages=[message])
-        context_window_size = llm.get_context_window_size()
-        max_tokens_allowed: int = int(
-            context_window_size * TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT // 100
+    max_tokens_allowed = get_max_token_count_for_single_tool(llm)
+
+    message = tool_call_result.as_tool_call_message()
+    messages_token = llm.count_tokens_for_message(messages=[message])
+
+    if messages_token > max_tokens_allowed:
+        relative_pct = ((messages_token - max_tokens_allowed) / messages_token) * 100
+
+        error_message: Optional[str] = (
+            f"The tool call result is too large to return: {messages_token} tokens.\nThe maximum allowed tokens is {max_tokens_allowed} which is {format(relative_pct, '.1f')}% smaller.\nInstructions for the LLM: try to repeat the query but proactively narrow down the result so that the tool answer fits within the allowed number of tokens."
         )
 
-        if messages_token > max_tokens_allowed:
-            relative_pct = (
-                (messages_token - max_tokens_allowed) / messages_token
-            ) * 100
-            error_message = f"The tool call result is too large to return: {messages_token} tokens.\nThe maximum allowed tokens is {max_tokens_allowed} which is {format(relative_pct, '.1f')}% smaller.\nInstructions for the LLM: try to repeat the query but proactively narrow down the result so that the tool answer fits within the allowed number of tokens."
-            tool_call_result.result.status = StructuredToolResultStatus.ERROR
-            tool_call_result.result.data = None
-            tool_call_result.result.error = error_message
-
-            sentry_helper.capture_toolcall_contains_too_many_tokens(
-                tool_call_result, messages_token, max_tokens_allowed
+        if tool_call_result.result.status == StructuredToolResultStatus.NO_DATA:
+            error_message = None
+            # tool_call_result.result.data is set to None below which is expected to fix the issue
+        elif tool_call_result.result.status == StructuredToolResultStatus.ERROR:
+            original_error = (
+                tool_call_result.result.error
+                or tool_call_result.result.data
+                or "Unknown error"
             )
+            truncated_error = str(original_error)[:100]
+            error_message = f"The tool call returned an error it is too large to return\nThe following original error is truncated:\n{truncated_error}"
+
+        tool_call_result.result.status = StructuredToolResultStatus.ERROR
+        tool_call_result.result.data = None
+        tool_call_result.result.error = error_message
+
+        sentry_helper.capture_toolcall_contains_too_many_tokens(
+            tool_call_result, messages_token, max_tokens_allowed
+        )

holmes/core/tools_utils/tool_executor.py

@@ -9,6 +9,7 @@ from holmes.core.tools import (
     StructuredToolResultStatus,
     Toolset,
     ToolsetStatusEnum,
+    ToolInvokeContext,
 )
 from holmes.core.tools_utils.toolset_utils import filter_out_default_logging_toolset
 
@@ -46,16 +47,20 @@ class ToolExecutor:
                     )
                 self.tools_by_name[tool.name] = tool
 
-    def invoke(self, tool_name: str, params: dict) -> StructuredToolResult:
+    def invoke(
+        self, tool_name: str, params: dict, context: ToolInvokeContext
+    ) -> StructuredToolResult:
+        """TODO: remove this function as it seems unused.
+        We call tool_executor.get_tool_by_name() and then tool.invoke() directly instead of this invoke function
+        """
         tool = self.get_tool_by_name(tool_name)
-        return (
-            tool.invoke(params)
-            if tool
-            else StructuredToolResult(
+        if not tool:
+            return StructuredToolResult(
                 status=StructuredToolResultStatus.ERROR,
                 error=f"Could not find tool named {tool_name}",
             )
-        )
+
+        return tool.invoke(params, context)
 
     def get_tool_by_name(self, name: str) -> Optional[Tool]:
         if name in self.tools_by_name:

holmes/core/toolset_manager.py

@@ -275,7 +275,11 @@ 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:
+            if toolset.enabled and (
+                toolset.status == ToolsetStatusEnum.ENABLED
+                or toolset.type == ToolsetType.MCP
+            ):
+                # MCP servers need to reload their tools even if previously failed, so rerun prerequisites
                 enabled_toolsets_from_cache.append(toolset)
         self.check_toolset_prerequisites(enabled_toolsets_from_cache)
 

holmes/core/truncation/dal_truncation_utils.py

@@ -0,0 +1,23 @@
+from holmes.common.env_vars import MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION
+
+
+def truncate_string(data_str: str) -> str:
+    if data_str and len(data_str) > MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION:
+        return (
+            data_str[:MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION]
+            + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS"
+        )
+    return data_str
+
+
+def truncate_evidences_entities_if_necessary(evidence_list: list[dict]):
+    if (
+        not MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION
+        or MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION <= 0
+    ):
+        return
+
+    for evidence in evidence_list:
+        data = evidence.get("data")
+        if data:
+            evidence["data"] = truncate_string(str(data))