holmesgpt 0.14.0a0__py3-none-any.whl → 0.14.1__py3-none-any.whl

holmes/core/tools.py

@@ -8,40 +8,68 @@ import tempfile
 from abc import ABC, abstractmethod
 from datetime import datetime
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, OrderedDict, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    OrderedDict,
+    Tuple,
+    Union,
+)
 
 from jinja2 import Template
-from pydantic import BaseModel, ConfigDict, Field, FilePath, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    FilePath,
+    model_validator,
+    PrivateAttr,
+)
 from rich.console import Console
 
 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 (
+    registry,
+    TransformerError,
+    Transformer,
+)
+
+if TYPE_CHECKING:
+    from holmes.core.transformers import BaseTransformer
+from holmes.utils.config_utils import merge_transformers
 import time
 from rich.table import Table
 
+logger = logging.getLogger(__name__)
 
-class ToolResultStatus(str, Enum):
+
+class StructuredToolResultStatus(str, Enum):
     SUCCESS = "success"
     ERROR = "error"
     NO_DATA = "no_data"
     APPROVAL_REQUIRED = "approval_required"
 
     def to_color(self) -> str:
-        if self == ToolResultStatus.SUCCESS:
+        if self == StructuredToolResultStatus.SUCCESS:
             return "green"
-        elif self == ToolResultStatus.ERROR:
+        elif self == StructuredToolResultStatus.ERROR:
             return "red"
-        elif self == ToolResultStatus.APPROVAL_REQUIRED:
+        elif self == StructuredToolResultStatus.APPROVAL_REQUIRED:
             return "yellow"
         else:
             return "white"
 
     def to_emoji(self) -> str:
-        if self == ToolResultStatus.SUCCESS:
+        if self == StructuredToolResultStatus.SUCCESS:
             return "✔"
-        elif self == ToolResultStatus.ERROR:
+        elif self == StructuredToolResultStatus.ERROR:
             return "❌"
-        elif self == ToolResultStatus.APPROVAL_REQUIRED:
+        elif self == StructuredToolResultStatus.APPROVAL_REQUIRED:
             return "⚠️"
         else:
             return "⚪️"
@@ -49,7 +77,7 @@ class ToolResultStatus(str, Enum):
 
 class StructuredToolResult(BaseModel):
     schema_version: str = "robusta:v1.0.0"
-    status: ToolResultStatus
+    status: StructuredToolResultStatus
     error: Optional[str] = None
     return_code: Optional[int] = None
     data: Optional[Any] = None
@@ -143,6 +171,48 @@ class Tool(ABC, BaseModel):
         default=None,
         description="The URL of the icon for the tool, if None will get toolset icon",
     )
+    transformers: Optional[List[Transformer]] = None
+
+    # Private attribute to store initialized transformer instances for performance
+    _transformer_instances: Optional[List["BaseTransformer"]] = PrivateAttr(
+        default=None
+    )
+
+    def model_post_init(self, __context) -> None:
+        """Initialize transformer instances once during tool creation for better performance."""
+        logger.debug(
+            f"Tool '{self.name}' model_post_init: creating transformer instances"
+        )
+
+        if self.transformers:
+            logger.debug(
+                f"Tool '{self.name}' has {len(self.transformers)} transformers to initialize"
+            )
+            self._transformer_instances = []
+            for transformer in self.transformers:
+                if not transformer:
+                    continue
+                logger.debug(
+                    f"  Initializing transformer '{transformer.name}' with config: {transformer.config}"
+                )
+                try:
+                    # Create transformer instance once and cache it
+                    transformer_instance = registry.create_transformer(
+                        transformer.name, transformer.config
+                    )
+                    self._transformer_instances.append(transformer_instance)
+                    logger.debug(
+                        f"Initialized transformer '{transformer.name}' for tool '{self.name}'"
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to initialize transformer '{transformer.name}' for tool '{self.name}': {e}"
+                    )
+                    # Continue with other transformers, don't fail the entire initialization
+                    continue
+        else:
+            logger.debug(f"Tool '{self.name}' has no transformers")
+            self._transformer_instances = None
 
     def get_openai_format(self, target_model: str):
         return format_tool_to_open_ai_standard(
@@ -159,23 +229,113 @@ class Tool(ABC, BaseModel):
         user_approved: bool = False,
     ) -> StructuredToolResult:
         tool_number_str = f"#{tool_number} " if tool_number else ""
-        logging.info(
+        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.icon_url = self.icon_url
+
+        # Apply transformers to the result
+        transformed_result = self._apply_transformers(result)
         elapsed = time.time() - start_time
         output_str = (
-            result.get_stringified_data()
-            if hasattr(result, "get_stringified_data")
-            else str(result)
+            transformed_result.get_stringified_data()
+            if hasattr(transformed_result, "get_stringified_data")
+            else str(transformed_result)
         )
         show_hint = f"/show {tool_number}" if tool_number else "/show"
         line_count = output_str.count("\n") + 1 if output_str else 0
-        logging.info(
+        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]"
         )
+        return transformed_result
+
+    def _apply_transformers(self, result: StructuredToolResult) -> StructuredToolResult:
+        """
+        Apply configured transformers to the tool result.
+
+        Args:
+            result: The original tool result
+
+        Returns:
+            The tool result with transformed data, or original result if transformation fails
+        """
+        if (
+            not self._transformer_instances
+            or result.status != StructuredToolResultStatus.SUCCESS
+        ):
+            return result
+
+        # Get the output string to transform
+        original_data = result.get_stringified_data()
+        if not original_data:
+            return result
+
+        transformed_data = original_data
+        transformers_applied = []
+
+        # Use cached transformer instances instead of creating new ones
+        for transformer_instance in self._transformer_instances:
+            try:
+                # Check if transformer should be applied
+                if not transformer_instance.should_apply(transformed_data):
+                    logger.debug(
+                        f"Transformer '{transformer_instance.name}' skipped for tool '{self.name}' (conditions not met)"
+                    )
+                    continue
+
+                # Apply transformation
+                pre_transform_size = len(transformed_data)
+                transform_start_time = time.time()
+                original_data = transformed_data  # Keep a copy for potential reversion
+                transformed_data = transformer_instance.transform(transformed_data)
+                transform_elapsed = time.time() - transform_start_time
+
+                # Check if this is llm_summarize and revert if summary is not smaller
+                post_transform_size = len(transformed_data)
+                if (
+                    transformer_instance.name == "llm_summarize"
+                    and post_transform_size >= pre_transform_size
+                ):
+                    # Revert to original data if summary is not smaller
+                    transformed_data = original_data
+                    logger.debug(
+                        f"Transformer '{transformer_instance.name}' reverted for tool '{self.name}' "
+                        f"(output size {post_transform_size:,} >= input size {pre_transform_size:,})"
+                    )
+                    continue  # Don't mark as applied
+
+                transformers_applied.append(transformer_instance.name)
+
+                # Generic logging - transformers can override this with their own specific metrics
+                size_change = post_transform_size - pre_transform_size
+                logger.info(
+                    f"Applied transformer '{transformer_instance.name}' to tool '{self.name}' output "
+                    f"in {transform_elapsed:.2f}s (size: {pre_transform_size:,} → {post_transform_size:,} chars, "
+                    f"change: {size_change:+,})"
+                )
+
+            except TransformerError as e:
+                logger.warning(
+                    f"Transformer '{transformer_instance.name}' failed for tool '{self.name}': {e}"
+                )
+                # Continue with other transformers, don't fail the entire chain
+                continue
+            except Exception as e:
+                logger.error(
+                    f"Unexpected error applying transformer '{transformer_instance.name}' to tool '{self.name}': {e}"
+                )
+                # Continue with other transformers
+                continue
+
+        # If any transformers were applied, update the result
+        if transformers_applied:
+            # Create a copy of the result with transformed data
+            result_dict = result.model_dump(exclude={"data"})
+            result_dict["data"] = transformed_data
+            return StructuredToolResult(**result_dict)
+
         return result
 
     @abstractmethod
@@ -230,12 +390,14 @@ class YAMLTool(Tool, BaseModel):
         context = {**params}
         return context
 
-    def _get_status(self, return_code: int, raw_output: str) -> ToolResultStatus:
+    def _get_status(
+        self, return_code: int, raw_output: str
+    ) -> StructuredToolResultStatus:
         if return_code != 0:
-            return ToolResultStatus.ERROR
+            return StructuredToolResultStatus.ERROR
         if raw_output == "":
-            return ToolResultStatus.NO_DATA
-        return ToolResultStatus.SUCCESS
+            return StructuredToolResultStatus.NO_DATA
+        return StructuredToolResultStatus.SUCCESS
 
     def _invoke(
         self, params: dict, user_approved: bool = False
@@ -246,7 +408,7 @@ class YAMLTool(Tool, BaseModel):
             raw_output, return_code, invocation = self.__invoke_script(params)  # type: ignore
 
         if self.additional_instructions and return_code == 0:
-            logging.info(
+            logger.info(
                 f"Applying additional instructions: {self.additional_instructions}"
             )
             output_with_instructions = self.__apply_additional_instructions(raw_output)
@@ -281,7 +443,7 @@ class YAMLTool(Tool, BaseModel):
             )
             return result.stdout.strip()
         except subprocess.CalledProcessError as e:
-            logging.error(
+            logger.error(
                 f"Failed to apply additional instructions: {self.additional_instructions}. "
                 f"Error: {e.stderr}"
             )
@@ -316,7 +478,7 @@ class YAMLTool(Tool, BaseModel):
 
     def __execute_subprocess(self, cmd) -> Tuple[str, int]:
         try:
-            logging.debug(f"Running `{cmd}`")
+            logger.debug(f"Running `{cmd}`")
             result = subprocess.run(
                 cmd,
                 shell=True,
@@ -329,7 +491,7 @@ class YAMLTool(Tool, BaseModel):
 
             return result.stdout.strip(), result.returncode
         except Exception as e:
-            logging.error(
+            logger.error(
                 f"An unexpected error occurred while running '{cmd}': {e}",
                 exc_info=True,
             )
@@ -381,6 +543,7 @@ class Toolset(BaseModel):
     config: Optional[Any] = None
     is_default: bool = False
     llm_instructions: Optional[str] = None
+    transformers: Optional[List[Transformer]] = None
 
     # warning! private attributes are not copied, which can lead to subtle bugs.
     # e.g. l.extend([some_tool]) will reset these private attribute to None
@@ -406,13 +569,85 @@ class Toolset(BaseModel):
     @model_validator(mode="before")
     def preprocess_tools(cls, values):
         additional_instructions = values.get("additional_instructions", "")
+        transformers = values.get("transformers", None)
         tools_data = values.get("tools", [])
+
+        # Convert raw dict transformers to Transformer objects BEFORE merging
+        if transformers:
+            converted_transformers = []
+            for t in transformers:
+                if isinstance(t, dict):
+                    try:
+                        transformer_obj = Transformer(**t)
+                        # Check if transformer is registered
+                        from holmes.core.transformers import registry
+
+                        if not registry.is_registered(transformer_obj.name):
+                            logger.warning(
+                                f"Invalid toolset transformer configuration: Transformer '{transformer_obj.name}' is not registered"
+                            )
+                            continue  # Skip invalid transformer
+                        converted_transformers.append(transformer_obj)
+                    except Exception as e:
+                        # Log warning and skip invalid transformer
+                        logger.warning(
+                            f"Invalid toolset transformer configuration: {e}"
+                        )
+                        continue
+                else:
+                    # Already a Transformer object
+                    converted_transformers.append(t)
+            transformers = converted_transformers if converted_transformers else None
+
         tools = []
         for tool in tools_data:
             if isinstance(tool, dict):
                 tool["additional_instructions"] = additional_instructions
+
+                # Convert tool-level transformers to Transformer objects
+                tool_transformers = tool.get("transformers")
+                if tool_transformers:
+                    converted_tool_transformers = []
+                    for t in tool_transformers:
+                        if isinstance(t, dict):
+                            try:
+                                transformer_obj = Transformer(**t)
+                                # Check if transformer is registered
+                                from holmes.core.transformers import registry
+
+                                if not registry.is_registered(transformer_obj.name):
+                                    logger.warning(
+                                        f"Invalid tool transformer configuration: Transformer '{transformer_obj.name}' is not registered"
+                                    )
+                                    continue  # Skip invalid transformer
+                                converted_tool_transformers.append(transformer_obj)
+                            except Exception as e:
+                                # Log warning and skip invalid transformer
+                                logger.warning(
+                                    f"Invalid tool transformer configuration: {e}"
+                                )
+                                continue
+                        else:
+                            # Already a Transformer object
+                            converted_tool_transformers.append(t)
+                    tool_transformers = (
+                        converted_tool_transformers
+                        if converted_tool_transformers
+                        else None
+                    )
+
+                # Merge toolset-level transformers with tool-level configs
+                tool["transformers"] = merge_transformers(
+                    base_transformers=transformers,
+                    override_transformers=tool_transformers,
+                )
             if isinstance(tool, Tool):
                 tool.additional_instructions = additional_instructions
+                # Merge toolset-level transformers with tool-level configs
+                tool.transformers = merge_transformers(  # type: ignore
+                    base_transformers=transformers,
+                    override_transformers=tool.transformers,
+                )
             tools.append(tool)
         values["tools"] = tools
 
@@ -482,11 +717,11 @@ class Toolset(BaseModel):
                 self.status == ToolsetStatusEnum.DISABLED
                 or self.status == ToolsetStatusEnum.FAILED
             ):
-                logging.info(f"❌ Toolset {self.name}: {self.error}")
+                logger.info(f"❌ Toolset {self.name}: {self.error}")
                 # no point checking further prerequisites if one failed
                 return
 
-        logging.info(f"✅ Toolset {self.name}")
+        logger.info(f"✅ Toolset {self.name}")
 
     @abstractmethod
     def get_example_config(self) -> Dict[str, Any]:

holmes/core/tools_utils/data_types.py

@@ -0,0 +1,81 @@
+import json
+from typing import Optional
+from pydantic import BaseModel
+
+from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+
+
+class TruncationMetadata(BaseModel):
+    tool_call_id: str
+    start_index: int
+    end_index: int
+    tool_name: str
+    original_token_count: int
+
+
+class TruncationResult(BaseModel):
+    truncated_messages: list[dict]
+    truncations: list[TruncationMetadata]
+
+
+def format_tool_result_data(tool_result: StructuredToolResult) -> str:
+    tool_response = tool_result.data
+    if isinstance(tool_result.data, str):
+        tool_response = tool_result.data
+    else:
+        try:
+            if isinstance(tool_result.data, BaseModel):
+                tool_response = tool_result.data.model_dump_json(indent=2)
+            else:
+                tool_response = json.dumps(tool_result.data, indent=2)
+        except Exception:
+            tool_response = str(tool_result.data)
+    if tool_result.status == StructuredToolResultStatus.ERROR:
+        tool_response = f"{tool_result.error or 'Tool execution failed'}:\n\n{tool_result.data or ''}".strip()
+    return tool_response
+
+
+class ToolCallResult(BaseModel):
+    tool_call_id: str
+    tool_name: str
+    description: str
+    result: StructuredToolResult
+    size: Optional[int] = None
+
+    def as_tool_call_message(self):
+        content = format_tool_result_data(self.result)
+        if self.result.params:
+            content = (
+                f"Params used for the tool call: {json.dumps(self.result.params)}. The tool call output follows on the next line.\n"
+                + content
+            )
+        return {
+            "tool_call_id": self.tool_call_id,
+            "role": "tool",
+            "name": self.tool_name,
+            "content": content,
+        }
+
+    def as_tool_result_response(self):
+        result_dump = self.result.model_dump()
+        result_dump["data"] = self.result.get_stringified_data()
+
+        return {
+            "tool_call_id": self.tool_call_id,
+            "tool_name": self.tool_name,
+            "description": self.description,
+            "role": "tool",
+            "result": result_dump,
+        }
+
+    def as_streaming_tool_result_response(self):
+        result_dump = self.result.model_dump()
+        result_dump["data"] = self.result.get_stringified_data()
+
+        return {
+            "tool_call_id": self.tool_call_id,
+            "role": "tool",
+            "description": self.description,
+            "name": self.tool_name,
+            "result": result_dump,
+        }

holmes/core/tools_utils/tool_context_window_limiter.py

@@ -0,0 +1,33 @@
+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.utils import sentry_helper
+
+
+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
+        )
+
+        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
+            )

holmes/core/tools_utils/tool_executor.py

@@ -6,7 +6,7 @@ import sentry_sdk
 from holmes.core.tools import (
     StructuredToolResult,
     Tool,
-    ToolResultStatus,
+    StructuredToolResultStatus,
     Toolset,
     ToolsetStatusEnum,
 )
@@ -52,7 +52,7 @@ class ToolExecutor:
             tool.invoke(params)
             if tool
             else StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=f"Could not find tool named {tool_name}",
             )
         )