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

holmes/plugins/toolsets/grafana/trace_parser.py

@@ -1,6 +1,7 @@
-from typing import Dict, List, Optional, Any
-from dataclasses import dataclass, field
 import base64
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
 from holmes.plugins.toolsets.utils import unix_nano_to_rfc3339
 
 
@@ -187,7 +188,7 @@ def format_traces_list(trace_data: Dict) -> str:
                 else "\n"
             )
             trace_str += f"\tstartTime={unix_nano_to_rfc3339(int(trace.get('startTimeUnixNano')))}"
-            trace_str += f" rootServiceName={trace.get('trootServiceName')}"
+            trace_str += f" rootServiceName={trace.get('rootServiceName')}"
             trace_str += f" rootTraceName={trace.get('rootTraceName')}"
             traces_str.append(trace_str)
         return "\n".join(traces_str)

holmes/plugins/toolsets/internet/internet.py

@@ -1,31 +1,32 @@
-import re
-import os
 import logging
-from typing import Any, Optional, Tuple, Dict, List
+import os
+import re
+from typing import Any, Dict, List, Optional, Tuple
 
+import requests  # type: ignore
+from bs4 import BeautifulSoup
+from markdownify import markdownify
 from requests import RequestException, Timeout  # type: ignore
+
 from holmes.core.tools import (
+    CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     Toolset,
     ToolsetTag,
-    CallablePrerequisite,
 )
-from markdownify import markdownify
-from bs4 import BeautifulSoup
-
-import requests  # type: ignore
-from holmes.core.tools import StructuredToolResult, ToolResultStatus
 from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
 
-
 # TODO: change and make it holmes
 INTERNET_TOOLSET_USER_AGENT = os.environ.get(
     "INTERNET_TOOLSET_USER_AGENT",
     "Mozilla/5.0 (X11; Linux x86_64; rv:128.0; holmesgpt;) Gecko/20100101 Firefox/128.0",
 )
 INTERNET_TOOLSET_TIMEOUT_SECONDS = int(
-    os.environ.get("INTERNET_TOOLSET_TIMEOUT_SECONDS", "60")
+    os.environ.get("INTERNET_TOOLSET_TIMEOUT_SECONDS", "5")
 )
 
 SELECTORS_TO_REMOVE = [
@@ -186,9 +187,7 @@ class FetchWebpage(Tool):
             toolset=toolset,  # type: ignore
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         url: str = params["url"]
 
         additional_headers = (
@@ -199,7 +198,7 @@ class FetchWebpage(Tool):
         if not content:
             logging.error(f"Failed to retrieve content from {url}")
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=f"Failed to retrieve content from {url}",
                 params=params,
             )
@@ -211,7 +210,7 @@ class FetchWebpage(Tool):
             content = html_to_markdown(content)
 
         return StructuredToolResult(
-            status=ToolResultStatus.SUCCESS,
+            status=StructuredToolResultStatus.SUCCESS,
             data=content,
             params=params,
         )

holmes/plugins/toolsets/internet/notion.py

@@ -1,9 +1,13 @@
-import re
-import logging
 import json
+import logging
+import re
 from typing import Any, Dict, Tuple
+
 from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     ToolsetTag,
 )
@@ -11,10 +15,6 @@ from holmes.plugins.toolsets.internet.internet import (
     InternetBaseToolset,
     scrape,
 )
-from holmes.core.tools import (
-    StructuredToolResult,
-    ToolResultStatus,
-)
 from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
 
 
@@ -44,9 +44,7 @@ class FetchNotion(Tool):
             return f"https://api.notion.com/v1/blocks/{notion_id}/children"
         return url  # Return original URL if no match is found
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         url: str = params["url"]
 
         # Get headers from the toolset configuration
@@ -59,13 +57,13 @@ class FetchNotion(Tool):
         if not content:
             logging.error(f"Failed to retrieve content from {url}")
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=f"Failed to retrieve content from {url}",
                 params=params,
             )
 
         return StructuredToolResult(
-            status=ToolResultStatus.SUCCESS,
+            status=StructuredToolResultStatus.SUCCESS,
             data=self.parse_notion_content(content),
             params=params,
         )

holmes/plugins/toolsets/investigator/core_investigation.py

@@ -1,26 +1,44 @@
 import logging
 import os
 from typing import Any, Dict
-
 from uuid import uuid4
+
 from holmes.core.todo_tasks_formatter import format_tasks
 from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    Tool,
+    ToolInvokeContext,
+    ToolParameter,
     Toolset,
     ToolsetTag,
-    ToolParameter,
-    Tool,
-    StructuredToolResult,
-    ToolResultStatus,
 )
 from holmes.plugins.toolsets.investigator.model import Task, TaskStatus
 
+TODO_WRITE_TOOL_NAME = "TodoWrite"
+
+
+def parse_tasks(todos_data: Any) -> list[Task]:
+    tasks = []
+
+    for todo_item in todos_data:
+        if isinstance(todo_item, dict):
+            task = Task(
+                id=todo_item.get("id", str(uuid4())),
+                content=todo_item.get("content", ""),
+                status=TaskStatus(todo_item.get("status", "pending")),
+            )
+            tasks.append(task)
+
+    return tasks
+
 
 class TodoWriteTool(Tool):
-    name: str = "TodoWrite"
+    name: str = TODO_WRITE_TOOL_NAME
     description: str = "Save investigation tasks to break down complex problems into manageable sub-tasks. ALWAYS provide the COMPLETE list of all tasks, not just the ones being updated."
     parameters: Dict[str, ToolParameter] = {
         "todos": ToolParameter(
-            description="COMPLETE list of ALL tasks on the task list. Each task should have: id (string), content (string), status (pending/in_progress/completed)",
+            description="COMPLETE list of ALL tasks on the task list. Each task should have: id (string), content (string), status (pending/in_progress/completed/failed)",
             type="array",
             required=True,
             items=ToolParameter(
@@ -28,7 +46,11 @@ class TodoWriteTool(Tool):
                 properties={
                     "id": ToolParameter(type="string", required=True),
                     "content": ToolParameter(type="string", required=True),
-                    "status": ToolParameter(type="string", required=True),
+                    "status": ToolParameter(
+                        type="string",
+                        required=True,
+                        enum=["pending", "in_progress", "completed", "failed"],
+                    ),
                 },
             ),
         ),
@@ -44,6 +66,7 @@ class TodoWriteTool(Tool):
             "pending": "[ ]",
             "in_progress": "[~]",
             "completed": "[✓]",
+            "failed": "[✗]",
         }
 
         max_id_width = max(len(str(task.id)) for task in tasks)
@@ -57,41 +80,28 @@ class TodoWriteTool(Tool):
         content_width = max(max_content_width, len("Content"))
         status_width = max(max_status_display_width, len("Status"))
 
-        # Build table
         separator = f"+{'-' * (id_width + 2)}+{'-' * (content_width + 2)}+{'-' * (status_width + 2)}+"
         header = f"| {'ID':<{id_width}} | {'Content':<{content_width}} | {'Status':<{status_width}} |"
-
-        # Log the table
-        logging.info("Updated Investigation Tasks:")
-        logging.info(separator)
-        logging.info(header)
-        logging.info(separator)
+        tasks_to_display = []
 
         for task in tasks:
             status_display = f"{status_icons[task.status.value]} {task.status.value}"
             row = f"| {task.id:<{id_width}} | {task.content:<{content_width}} | {status_display:<{status_width}} |"
-            logging.info(row)
+            tasks_to_display.append(row)
 
-        logging.info(separator)
+        logging.info(
+            f"Task List:\n{separator}\n{header}\n{separator}\n"
+            + "\n".join(tasks_to_display)
+            + f"\n{separator}"
+        )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             todos_data = params.get("todos", [])
 
-            tasks = []
-
-            for todo_item in todos_data:
-                if isinstance(todo_item, dict):
-                    task = Task(
-                        id=todo_item.get("id", str(uuid4())),
-                        content=todo_item.get("content", ""),
-                        status=TaskStatus(todo_item.get("status", "pending")),
-                    )
-                    tasks.append(task)
+            tasks = parse_tasks(todos_data=todos_data)
 
-            logging.info(f"Tasks: {len(tasks)}")
+            logging.debug(f"Tasks: {len(tasks)}")
 
             self.print_tasks_table(tasks)
             formatted_tasks = format_tasks(tasks)
@@ -103,7 +113,7 @@ class TodoWriteTool(Tool):
                 response_data += "No tasks currently in the investigation plan."
 
             return StructuredToolResult(
-                status=ToolResultStatus.SUCCESS,
+                status=StructuredToolResultStatus.SUCCESS,
                 data=response_data,
                 params=params,
             )
@@ -111,14 +121,13 @@ class TodoWriteTool(Tool):
         except Exception as e:
             logging.exception("error using todowrite tool")
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=f"Failed to process tasks: {str(e)}",
                 params=params,
             )
 
     def get_parameterized_one_liner(self, params: Dict) -> str:
-        todos = params.get("todos", [])
-        return f"Write {todos} investigation tasks"
+        return "Update investigation tasks"
 
 
 class CoreInvestigationToolset(Toolset):
@@ -133,7 +142,6 @@ class CoreInvestigationToolset(Toolset):
             tags=[ToolsetTag.CORE],
             is_default=True,
         )
-        logging.info("Core investigation toolset loaded")
 
     def get_example_config(self) -> Dict[str, Any]:
         return {}

holmes/plugins/toolsets/investigator/model.py

@@ -1,12 +1,14 @@
 from enum import Enum
-from pydantic import BaseModel, Field
 from uuid import uuid4
 
+from pydantic import BaseModel, Field
+
 
 class TaskStatus(str, Enum):
     PENDING = "pending"
     IN_PROGRESS = "in_progress"
     COMPLETED = "completed"
+    FAILED = "failed"
 
 
 class Task(BaseModel):

holmes/plugins/toolsets/json_filter_mixin.py

@@ -0,0 +1,134 @@
+import json
+import logging
+from typing import Any, Dict, Optional, Tuple
+
+import jq
+
+from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    ToolParameter,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _truncate_to_depth(value: Any, max_depth: Optional[int], current_depth: int = 0):
+    """Recursively truncate dictionaries/lists beyond the requested depth."""
+    if max_depth is None or max_depth < 0:
+        return value
+
+    if current_depth >= max_depth:
+        if isinstance(value, (dict, list)):
+            return f"...truncated at depth {max_depth}"
+        return value
+
+    if isinstance(value, dict):
+        return {
+            key: _truncate_to_depth(sub_value, max_depth, current_depth + 1)
+            for key, sub_value in value.items()
+        }
+    if isinstance(value, list):
+        return [
+            _truncate_to_depth(item, max_depth, current_depth + 1) for item in value
+        ]
+
+    return value
+
+
+def _apply_jq_filter(data: Any, expression: str) -> Tuple[Optional[Any], Optional[str]]:
+    try:
+        compiled = jq.compile(expression)
+        matches = compiled.input(data).all()
+        if len(matches) == 1:
+            return matches[0], None
+        return matches, None
+    except Exception as exc:  # pragma: no cover - defensive
+        logger.debug("Failed to apply jq filter", exc_info=exc)
+        return None, f"Invalid jq expression: {exc}"
+
+
+class JsonFilterMixin:
+    """Opt-in mixin for tools that return JSON and want filtering controls."""
+
+    filter_parameters: Dict[str, ToolParameter] = {
+        "max_depth": ToolParameter(
+            description="Maximum nesting depth to return from the JSON (0 returns only top-level keys). Leave empty for full response.",
+            type="integer",
+            required=False,
+        ),
+        "jq": ToolParameter(
+            description="Optional jq expression to extract specific parts of the JSON. Supports full jq syntax including filters, slicing, transformations, and more (e.g., '.items[] | select(.price > 10)', '.items[0:5]', '.items[].name').",
+            type="string",
+            required=False,
+        ),
+    }
+
+    @classmethod
+    def extend_parameters(
+        cls, existing: Dict[str, ToolParameter]
+    ) -> Dict[str, ToolParameter]:
+        merged = dict(cls.filter_parameters)
+        merged.update(existing)
+        return merged
+
+    def _filter_result_data(self, data: Any, params: Dict) -> Tuple[Any, Optional[str]]:
+        parsed_data = data
+        if isinstance(data, str):
+            try:
+                parsed_data = json.loads(data)
+            except Exception:
+                # Not JSON, leave as-is
+                return data, None
+
+        if params.get("jq"):
+            parsed_data, error = _apply_jq_filter(parsed_data, params["jq"])
+            if error:
+                return None, error
+
+        parsed_data = _truncate_to_depth(parsed_data, params.get("max_depth"))
+        return parsed_data, None
+
+    @staticmethod
+    def _safe_string(value: Any) -> Optional[str]:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            return value
+        try:
+            return str(value)
+        except Exception:
+            return None
+
+    def filter_result(
+        self, result: StructuredToolResult, params: Dict
+    ) -> StructuredToolResult:
+        base_result = result if isinstance(result, StructuredToolResult) else None
+        if base_result is None:
+            base_result = StructuredToolResult(
+                status=getattr(result, "status", StructuredToolResultStatus.SUCCESS),
+                data=getattr(result, "data", None),
+                params=getattr(result, "params", params),
+                url=self._safe_string(getattr(result, "url", None)),
+                invocation=self._safe_string(getattr(result, "invocation", None)),
+                icon_url=self._safe_string(getattr(result, "icon_url", None)),
+            )
+        else:
+            # Normalize string fields to avoid MagicMock validation failures
+            base_result.url = self._safe_string(base_result.url)
+            base_result.invocation = self._safe_string(base_result.invocation)
+            base_result.icon_url = self._safe_string(base_result.icon_url)
+
+        filtered_data, error = self._filter_result_data(base_result.data, params)
+        if error:
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error=error,
+                params=params,
+                url=base_result.url,
+                invocation=base_result.invocation,
+                icon_url=base_result.icon_url,
+            )
+
+        base_result.data = filtered_data
+        return base_result