holmesgpt 0.16.2a0__py3-none-any.whl → 0.18.4__py3-none-any.whl

holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py

@@ -1,15 +1,17 @@
+import json
 import os
-from typing import Any, Dict, Tuple, cast, List
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Tuple, cast
+from urllib.parse import quote
 
-import yaml  # type: ignore
-
-from holmes.common.env_vars import load_bool, MAX_GRAPH_POINTS
+from holmes.common.env_vars import MAX_GRAPH_POINTS, load_bool
 from holmes.core.tools import (
     StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
     ToolInvokeContext,
     ToolParameter,
-    StructuredToolResultStatus,
 )
 from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION
 from holmes.plugins.toolsets.grafana.base_grafana_toolset import BaseGrafanaToolset
@@ -21,17 +23,116 @@ from holmes.plugins.toolsets.logging_utils.logging_api import (
     DEFAULT_GRAPH_TIME_SPAN_SECONDS,
 )
 from holmes.plugins.toolsets.utils import (
-    toolset_name_for_one_liner,
-    process_timestamps_to_int,
-    standard_start_datetime_tool_param_description,
     adjust_step_for_max_points,
-    seconds_to_duration_string,
     duration_string_to_seconds,
+    process_timestamps_to_int,
+    seconds_to_duration_string,
+    standard_start_datetime_tool_param_description,
+    toolset_name_for_one_liner,
 )
 
 TEMPO_LABELS_ADD_PREFIX = load_bool("TEMPO_LABELS_ADD_PREFIX", True)
 
 
+def _build_grafana_explore_tempo_url(
+    config: GrafanaTempoConfig,
+    query: Optional[str] = None,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    limit: int = 20,
+    trace_id: Optional[str] = None,
+    filters: Optional[List[Dict[str, Any]]] = None,
+    tags: Optional[str] = None,
+) -> Optional[str]:
+    if not config.grafana_datasource_uid:
+        return None
+    try:
+        base_url = config.external_url or config.url
+        datasource_uid = config.grafana_datasource_uid
+        now_s = int(time.time())
+        start_ts = start if start else now_s - 3600
+        end_ts = end if end else now_s
+        start_delta = max(0, now_s - start_ts)
+        end_delta = max(0, now_s - end_ts)
+        from_str = f"now-{start_delta}s" if start_delta > 0 else "now-1h"
+        to_str = "now" if end_delta == 0 else f"now-{end_delta}s"
+        pane_id = "tmp"
+
+        if trace_id:
+            # Direct trace ID lookup - query is just the traceID string
+            query_obj = {
+                "refId": "A",
+                "datasource": {"type": "tempo", "uid": datasource_uid},
+                "queryType": "traceql",
+                "limit": limit,
+                "tableType": "traces",
+                "metricsQueryType": "range",
+                "query": trace_id,
+            }
+        elif tags:
+            # Build filters from tag name
+            scope = "resource" if tags.startswith("resource.") else "span"
+            filter_id = str(uuid.uuid4())[:8]
+            filters = [
+                {
+                    "id": filter_id,
+                    "operator": "=",
+                    "scope": scope,
+                    "tag": tags,
+                    "value": [],
+                }
+            ]
+            query_obj = {
+                "refId": "A",
+                "datasource": {"type": "tempo", "uid": datasource_uid},
+                "queryType": "traceqlSearch",
+                "limit": limit,
+                "tableType": "traces",
+                "metricsQueryType": "range",
+                "query": "",
+                "filters": filters,
+            }
+        elif filters:
+            # Tag filters - use traceqlSearch with filters array
+            query_obj = {
+                "refId": "A",
+                "datasource": {"type": "tempo", "uid": datasource_uid},
+                "queryType": "traceqlSearch",
+                "limit": limit,
+                "tableType": "traces",
+                "metricsQueryType": "range",
+                "query": "",
+                "filters": filters,
+            }
+        else:
+            # Regular TraceQL query
+            safe_query = query if query else "{}"
+            query_obj = {
+                "refId": "A",
+                "datasource": {"type": "tempo", "uid": datasource_uid},
+                "queryType": "traceql",
+                "limit": limit,
+                "tableType": "traces",
+                "metricsQueryType": "range",
+                "query": safe_query,
+            }
+
+        panes = {
+            pane_id: {
+                "datasource": datasource_uid,
+                "queries": [query_obj],
+                "range": {"from": from_str, "to": to_str},
+            }
+        }
+
+        panes_encoded = quote(
+            json.dumps(panes, separators=(",", ":"), ensure_ascii=False), safe=""
+        )
+        return f"{base_url}/explore?schemaVersion=1&panes={panes_encoded}&orgId=1"
+    except Exception:
+        return None
+
+
 class BaseGrafanaTempoToolset(BaseGrafanaToolset):
     config_class = GrafanaTempoConfig
 
@@ -47,22 +148,17 @@ class BaseGrafanaTempoToolset(BaseGrafanaToolset):
     def grafana_config(self) -> GrafanaTempoConfig:
         return cast(GrafanaTempoConfig, self._grafana_config)
 
-    def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
-        """Check Tempo connectivity using the echo endpoint."""
-        # First call parent to validate config
-        success, msg = super().prerequisites_callable(config)
-        if not success:
-            return success, msg
-
-        # Then check Tempo-specific echo endpoint
+    def health_check(self) -> Tuple[bool, str]:
+        """Test a dummy query to check if service available."""
         try:
-            api = GrafanaTempoAPI(self.grafana_config)
-            if api.query_echo_endpoint():
-                return True, "Successfully connected to Tempo"
-            else:
-                return False, "Failed to connect to Tempo echo endpoint"
+            _ = GrafanaTempoAPI(self.grafana_config).search_traces_by_query(
+                q='{ .service.name = "test-endpoint" }',
+                limit=1,
+            )
         except Exception as e:
-            return False, f"Failed to connect to Tempo: {str(e)}"
+            return False, f"Unable to connect to Tempo.\n{str(e)}"
+
+        return True, ""
 
     def build_k8s_filters(
         self, params: Dict[str, Any], use_exact_match: bool
@@ -335,11 +431,18 @@ Examples:
                 ],
             }
 
-            # Return as YAML for readability
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=f"{{{base_query}}}",
+                start=start,
+                end=end,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False, sort_keys=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
 
         except Exception as e:
@@ -428,10 +531,20 @@ class SearchTracesByQuery(Tool):
                 end=end,
                 spss=params.get("spss"),
             )
+
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=params["q"],
+                start=start,
+                end=end,
+                limit=params.get("limit") or 20,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -510,10 +623,21 @@ class SearchTracesByTags(Tool):
                 end=end,
                 spss=params.get("spss"),
             )
+
+            tag_filters = params["tags"].replace(" ", " && ")
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=f"{{{tag_filters}}}",
+                start=start,
+                end=end,
+                limit=params.get("limit") or 20,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -569,11 +693,18 @@ class QueryTraceById(Tool):
                 end=end,
             )
 
-            # Return raw trace data as YAML for readability
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                trace_id=params["trace_id"],
+                start=start,
+                end=end,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(trace_data, default_flow_style=False),
+                data=trace_data,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -646,10 +777,20 @@ class SearchTagNames(Tool):
                 limit=params.get("limit"),
                 max_stale_values=params.get("max_stale_values"),
             )
+
+            query_filter = params.get("q") or "{}"
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=query_filter,
+                start=start,
+                end=end,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -722,10 +863,19 @@ class SearchTagValues(Tool):
                 limit=params.get("limit"),
                 max_stale_values=params.get("max_stale_values"),
             )
+
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                start=start,
+                end=end,
+                tags=params["tag"],
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -797,10 +947,19 @@ class QueryMetricsInstant(Tool):
                 start=start,
                 end=end,
             )
+
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=params["q"],
+                start=start,
+                end=end,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(
@@ -893,10 +1052,19 @@ class QueryMetricsRange(Tool):
                 end=end,
                 exemplars=params.get("exemplars"),
             )
+
+            explore_url = _build_grafana_explore_tempo_url(
+                self._toolset.grafana_config,
+                query=params["q"],
+                start=start,
+                end=end,
+            )
+
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result, default_flow_style=False),
+                data=result,
                 params=params,
+                url=explore_url,
             )
         except Exception as e:
             return StructuredToolResult(

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
 
 

holmes/plugins/toolsets/internet/internet.py

@@ -1,25 +1,25 @@
-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, StructuredToolResultStatus
 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",

holmes/plugins/toolsets/internet/notion.py

@@ -1,8 +1,11 @@
-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,
@@ -12,10 +15,6 @@ from holmes.plugins.toolsets.internet.internet import (
     InternetBaseToolset,
     scrape,
 )
-from holmes.core.tools import (
-    StructuredToolResult,
-    StructuredToolResultStatus,
-)
 from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
 
 

holmes/plugins/toolsets/investigator/core_investigation.py

@@ -1,7 +1,6 @@
 import logging
 import os
 from typing import Any, Dict
-
 from uuid import uuid4
 
 from holmes.core.todo_tasks_formatter import format_tasks
@@ -39,7 +38,7 @@ class TodoWriteTool(Tool):
     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(
@@ -50,7 +49,7 @@ class TodoWriteTool(Tool):
                     "status": ToolParameter(
                         type="string",
                         required=True,
-                        enum=["pending", "in_progress", "completed"],
+                        enum=["pending", "in_progress", "completed", "failed"],
                     ),
                 },
             ),
@@ -67,6 +66,7 @@ class TodoWriteTool(Tool):
             "pending": "[ ]",
             "in_progress": "[~]",
             "completed": "[✓]",
+            "failed": "[✗]",
         }
 
         max_id_width = max(len(str(task.id)) for task in tasks)

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

holmes/plugins/toolsets/kafka.py

@@ -1,7 +1,10 @@
 import logging
+from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import yaml  # type: ignore
+from confluent_kafka import Consumer
+from confluent_kafka._model import Node
 from confluent_kafka.admin import (
     AdminClient,
     BrokerMetadata,
@@ -17,19 +20,16 @@ from confluent_kafka.admin import (
     PartitionMetadata,
     TopicMetadata,
 )
-from confluent_kafka import Consumer
-from confluent_kafka._model import Node
-from enum import Enum
 from confluent_kafka.admin import _TopicPartition as TopicPartition
 from pydantic import BaseModel, ConfigDict
 
 from holmes.core.tools import (
     CallablePrerequisite,
     StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
     ToolInvokeContext,
     ToolParameter,
-    StructuredToolResultStatus,
     Toolset,
     ToolsetTag,
 )
@@ -245,7 +245,7 @@ class DescribeConsumerGroup(BaseKafkaTool):
                 group_metadata = futures.get(group_id).result()
                 return StructuredToolResult(
                     status=StructuredToolResultStatus.SUCCESS,
-                    data=yaml.dump(convert_to_dict(group_metadata)),
+                    data=convert_to_dict(group_metadata),
                     params=params,
                 )
             else:
@@ -297,7 +297,7 @@ class ListTopics(BaseKafkaTool):
             topics = client.list_topics()
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(convert_to_dict(topics)),
+                data=convert_to_dict(topics),
                 params=params,
             )
         except Exception as e:
@@ -367,7 +367,7 @@ class DescribeTopic(BaseKafkaTool):
 
             return StructuredToolResult(
                 status=StructuredToolResultStatus.SUCCESS,
-                data=yaml.dump(result),
+                data=result,
                 params=params,
             )
         except Exception as e:
@@ -599,6 +599,8 @@ class KafkaToolset(Toolset):
                     admin_config = {
                         "bootstrap.servers": cluster.kafka_broker,
                         "client.id": cluster.kafka_client_id,
+                        "socket.timeout.ms": 15000,  # 15 second timeout
+                        "api.version.request.timeout.ms": 15000,  # 15 second API version timeout
                     }
 
                     if cluster.kafka_security_protocol:
@@ -612,6 +614,9 @@ class KafkaToolset(Toolset):
                         admin_config["sasl.password"] = cluster.kafka_password
 
                     client = AdminClient(admin_config)
+                    # Test the connection by trying to list topics with a timeout
+                    # This will fail fast if the broker is not reachable
+                    _ = client.list_topics(timeout=10)  # 10 second timeout
                     self.clients[cluster.name] = client  # Store in dictionary
                 except Exception as e:
                     message = (