holmesgpt 0.14.1a0__py3-none-any.whl → 0.14.3a0__py3-none-any.whl

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))

holmes/interactive.py

@@ -26,9 +26,16 @@ from prompt_toolkit.widgets import TextArea
 from pygments.lexers import guess_lexer
 from rich.console import Console
 from rich.markdown import Markdown, Panel
+from rich.markup import escape
 
 from holmes.common.env_vars import ENABLE_CLI_TOOL_APPROVAL
 from holmes.core.config import config_path_dir
+from holmes.core.feedback import (
+    PRIVACY_NOTICE_BANNER,
+    Feedback,
+    FeedbackCallback,
+    UserFeedback,
+)
 from holmes.core.prompt import build_initial_ask_messages
 from holmes.core.tool_calling_llm import ToolCallingLLM, ToolCallResult
 from holmes.core.tools import StructuredToolResult, pretty_print_toolset_status
@@ -62,19 +69,25 @@ class SlashCommands(Enum):
     )
     CONTEXT = ("/context", "Show conversation context size and token count")
     SHOW = ("/show", "Show specific tool output in scrollable view")
+    FEEDBACK = ("/feedback", "Provide feedback on the agent's response")
 
     def __init__(self, command, description):
         self.command = command
         self.description = description
 
 
-SLASH_COMMANDS_REFERENCE = {cmd.command: cmd.description for cmd in SlashCommands}
-ALL_SLASH_COMMANDS = [cmd.command for cmd in SlashCommands]
-
-
 class SlashCommandCompleter(Completer):
-    def __init__(self):
-        self.commands = SLASH_COMMANDS_REFERENCE
+    def __init__(self, unsupported_commands: Optional[List[str]] = None):
+        # Build commands dictionary, excluding unsupported commands
+        all_commands = {cmd.command: cmd.description for cmd in SlashCommands}
+        if unsupported_commands:
+            self.commands = {
+                cmd: desc
+                for cmd, desc in all_commands.items()
+                if cmd not in unsupported_commands
+            }
+        else:
+            self.commands = all_commands
 
     def get_completions(self, document, complete_event):
         text = document.text_before_cursor
@@ -811,6 +824,88 @@ def handle_last_command(
         )
 
 
+def handle_feedback_command(
+    style: Style,
+    console: Console,
+    feedback: Feedback,
+    feedback_callback: FeedbackCallback,
+) -> None:
+    """Handle the /feedback command to collect user feedback."""
+    try:
+        # Create a temporary session without history for feedback prompts
+        temp_session = PromptSession(history=InMemoryHistory())  # type: ignore
+        # Prominent privacy notice to users
+        console.print(
+            f"[bold {HELP_COLOR}]Privacy Notice:[/bold {HELP_COLOR}] {PRIVACY_NOTICE_BANNER}"
+        )
+        # A "Cancel" button of equal discoverability to "Sent" or "Submit" buttons must be made available
+        console.print(
+            "[bold yellow]💡 Tip: Press Ctrl+C at any time to cancel feedback[/bold yellow]"
+        )
+
+        # Ask for thumbs up/down rating with validation
+        while True:
+            rating_prompt = temp_session.prompt(
+                [("class:prompt", "Was this response useful to you? 👍(y)/👎(n): ")],
+                style=style,
+            )
+
+            rating_lower = rating_prompt.lower().strip()
+            if rating_lower in ["y", "n"]:
+                break
+            else:
+                console.print(
+                    "[bold red]Please enter only 'y' for yes or 'n' for no.[/bold red]"
+                )
+
+        # Determine rating
+        is_positive = rating_lower == "y"
+
+        # Ask for additional comments
+        comment_prompt = temp_session.prompt(
+            [
+                (
+                    "class:prompt",
+                    "Do you want to provide any additional comments for feedback? (press Enter to skip):\n",
+                )
+            ],
+            style=style,
+        )
+
+        comment = comment_prompt.strip() if comment_prompt.strip() else None
+
+        # Create UserFeedback object
+        user_feedback = UserFeedback(is_positive, comment)
+
+        if comment:
+            console.print(
+                f'[bold green]✓ Feedback recorded (rating={user_feedback.rating_emoji}, "{escape(comment)}")[/bold green]'
+            )
+        else:
+            console.print(
+                f"[bold green]✓ Feedback recorded (rating={user_feedback.rating_emoji}, no comment)[/bold green]"
+            )
+
+        # Final confirmation before submitting
+        final_confirmation = temp_session.prompt(
+            [("class:prompt", "\nDo you want to submit this feedback? (Y/n): ")],
+            style=style,
+        )
+
+        # If user says no, cancel the feedback
+        if final_confirmation.lower().strip().startswith("n"):
+            console.print("[dim]Feedback cancelled.[/dim]")
+            return
+
+        feedback.user_feedback = user_feedback
+        feedback_callback(feedback)
+        console.print("[bold green]Thank you for your feedback! 🙏[/bold green]")
+
+    except KeyboardInterrupt:
+        console.print("[dim]Feedback cancelled.[/dim]")
+        return
+
+
 def display_recent_tool_outputs(
     tool_calls: List[ToolCallResult],
     console: Console,
@@ -823,7 +918,10 @@ def display_recent_tool_outputs(
     for tool_call in tool_calls:
         tool_index = find_tool_index_in_history(tool_call, all_tool_calls_history)
         preview_output = format_tool_call_output(tool_call, tool_index)
-        title = f"{tool_call.result.status.to_emoji()} {tool_call.description} -> returned {tool_call.result.return_code}"
+        title = (
+            f"{tool_call.result.status.to_emoji()} {tool_call.description} -> "
+            f"returned {tool_call.result.return_code}"
+        )
 
         console.print(
             Panel(
@@ -846,6 +944,7 @@ def run_interactive_loop(
     runbooks=None,
     system_prompt_additions: Optional[str] = None,
     check_version: bool = True,
+    feedback_callback: Optional[FeedbackCallback] = None,
 ) -> None:
     # Initialize tracer - use DummyTracer if no tracer provided
     if tracer is None:
@@ -874,7 +973,11 @@ def run_interactive_loop(
         ai.approval_callback = approval_handler
 
     # Create merged completer with slash commands, conditional executables, show command, and smart paths
-    slash_completer = SlashCommandCompleter()
+    # TODO: remove unsupported_commands support once we implement feedback callback
+    unsupported_commands = []
+    if feedback_callback is None:
+        unsupported_commands.append(SlashCommands.FEEDBACK.command)
+    slash_completer = SlashCommandCompleter(unsupported_commands)
     executable_completer = ConditionalExecutableCompleter()
     show_completer = ShowCommandCompleter()
     path_completer = SmartPathCompleter()
@@ -891,6 +994,9 @@ def run_interactive_loop(
     if initial_user_input:
         history.append_string(initial_user_input)
 
+    feedback = Feedback()
+    feedback.metadata.update_llm(ai.llm)
+
     # Create custom key bindings for Ctrl+C behavior
     bindings = KeyBindings()
     status_message = ""
@@ -963,7 +1069,15 @@ def run_interactive_loop(
 
     input_prompt = [("class:prompt", "User: ")]
 
-    console.print(WELCOME_BANNER)
+    # TODO: merge the /feedback command description to WELCOME_BANNER once we implement feedback callback
+    welcome_banner = WELCOME_BANNER
+    if feedback_callback:
+        welcome_banner = (
+            welcome_banner.rstrip(".")
+            + f", '{SlashCommands.FEEDBACK.command}' to share your thoughts."
+        )
+    console.print(welcome_banner)
+
     if initial_user_input:
         console.print(
             f"[bold {USER_COLOR}]User:[/bold {USER_COLOR}] {initial_user_input}"
@@ -985,14 +1099,18 @@ def run_interactive_loop(
             if user_input.startswith("/"):
                 original_input = user_input.strip()
                 command = original_input.lower()
-
                 # Handle prefix matching for slash commands
-                matches = [cmd for cmd in ALL_SLASH_COMMANDS if cmd.startswith(command)]
+                matches = [
+                    cmd
+                    for cmd in slash_completer.commands.keys()
+                    if cmd.startswith(command)
+                ]
                 if len(matches) == 1:
                     command = matches[0]
                 elif len(matches) > 1:
                     console.print(
-                        f"[bold {ERROR_COLOR}]Ambiguous command '{command}'. Matches: {', '.join(matches)}[/bold {ERROR_COLOR}]"
+                        f"[bold {ERROR_COLOR}]Ambiguous command '{command}'. "
+                        f"Matches: {', '.join(matches)}[/bold {ERROR_COLOR}]"
                     )
                     continue
 
@@ -1002,13 +1120,20 @@ def run_interactive_loop(
                     console.print(
                         f"[bold {HELP_COLOR}]Available commands:[/bold {HELP_COLOR}]"
                     )
-                    for cmd, description in SLASH_COMMANDS_REFERENCE.items():
+                    for cmd, description in slash_completer.commands.items():
+                        # Only show feedback command if callback is available
+                        if (
+                            cmd == SlashCommands.FEEDBACK.command
+                            and feedback_callback is None
+                        ):
+                            continue
                         console.print(f"  [bold]{cmd}[/bold] - {description}")
                     continue
                 elif command == SlashCommands.CLEAR.command:
                     console.clear()
                     console.print(
-                        f"[bold {STATUS_COLOR}]Screen cleared and context reset. You can now ask a new question.[/bold {STATUS_COLOR}]"
+                        f"[bold {STATUS_COLOR}]Screen cleared and context reset. "
+                        f"You can now ask a new question.[/bold {STATUS_COLOR}]"
                     )
                     messages = None
                     last_response = None
@@ -1052,6 +1177,12 @@ def run_interactive_loop(
                     if shared_input is None:
                         continue  # User chose not to share or no output, continue to next input
                     user_input = shared_input
+                elif (
+                    command == SlashCommands.FEEDBACK.command
+                    and feedback_callback is not None
+                ):
+                    handle_feedback_command(style, console, feedback, feedback_callback)
+                    continue
                 else:
                     console.print(f"Unknown command: {command}")
                     continue
@@ -1091,6 +1222,7 @@ def run_interactive_loop(
 
             messages = response.messages  # type: ignore
             last_response = response
+            feedback.metadata.add_llm_response(user_input, response.result)
 
             if response.tool_calls:
                 all_tool_calls_history.extend(response.tool_calls)

holmes/plugins/prompts/_fetch_logs.jinja2

@@ -11,6 +11,7 @@
 * IMPORTANT: ALWAYS inform the user about what logs you fetched. For example: "Here are pod logs for ..."
 * IMPORTANT: If logs commands have limits mention them. For example: "Showing last 100 lines of logs:"
 * IMPORTANT: If a filter was used, mention the filter. For example: "Logs filtered for 'error':"
+* IMPORTANT: If a date range was used (even if just the default one and you didn't specify the parameter, mention the date range. For example: "Logs from last 1 hour..."
 
 {% if loki_ts and loki_ts.status == "enabled" -%}
 * For any logs, including for investigating kubernetes problems, use Loki
@@ -34,8 +35,18 @@ Tools to search and fetch logs from Coralogix.
 ### datadog/logs
 #### Datadog Logs Toolset
 Tools to search and fetch logs from Datadog.
-{% include '_default_log_prompt.jinja2' %}
+* Use the tool `fetch_pod_logs` to access an application's logs.
+* Do fetch application logs yourself and DO not ask users to do so
+* If you have an alert/monitor try to figure out the time it fired
+** Then, use `start_time=-300` (5 minutes before `end_time`) and `end_time=<time monitor started firing>`  when calling `fetch_pod_logs`.
+** If there are too many logs, or not enough, narrow or widen the timestamps
+* If the user did not explicitly ask about a given timeframe, ignore the `start_time` and `end_time` so it will use the default.
+* IMPORTANT: ALWAYS inform the user about the actual time period fetched (e.g., "Looking at logs from the last <X> days")
+* IMPORTANT: If a limit was applied, ALWAYS tell the user how many logs were shown vs total (e.g., "Showing latest <Y> of <Z> logs")
+* IMPORTANT: If any filters were applied, ALWAYS mention them explicitly
 {%- elif k8s_yaml_ts and k8s_yaml_ts.status == "enabled" -%}
+### Logs from newrelic
+* you can fetch logs from newrelic if this is toolset is enabled
 ### kubernetes/logs
 #### Kubernetes Logs Toolset
 Tools to search and fetch logs from Kubernetes.
@@ -56,4 +67,5 @@ DO NOT use `--tail` or `| tail` when calling `kubectl logs` because you may miss
 ** 'opensearch/logs'
 ** 'coralogix/logs'
 ** 'datadog/logs'
+** 'newrelic'
 {%- endif -%}

holmes/plugins/runbooks/__init__.py

@@ -108,9 +108,14 @@ def get_runbook_by_path(
     Returns:
         Full path to the runbook if found, None otherwise
     """
+    # Validate runbook_relative_path is not empty
+    if not runbook_relative_path or not runbook_relative_path.strip():
+        return None
+
     for search_path in search_paths:
         runbook_path = os.path.join(search_path, runbook_relative_path)
-        if os.path.exists(runbook_path):
+        # Ensure it's a file, not a directory
+        if os.path.isfile(runbook_path):
             return runbook_path
 
     return None

holmes/plugins/toolsets/__init__.py

@@ -7,7 +7,10 @@ import yaml  # type: ignore
 from pydantic import ValidationError
 
 import holmes.utils.env as env_utils
-from holmes.common.env_vars import USE_LEGACY_KUBERNETES_LOGS
+from holmes.common.env_vars import (
+    USE_LEGACY_KUBERNETES_LOGS,
+    DISABLE_PROMETHEUS_TOOLSET,
+)
 from holmes.core.supabase_dal import SupabaseDal
 from holmes.core.tools import Toolset, ToolsetType, ToolsetYamlFromConfig, YAMLToolset
 from holmes.plugins.toolsets.atlas_mongodb.mongodb_atlas import MongoDBAtlasToolset
@@ -38,11 +41,10 @@ from holmes.plugins.toolsets.internet.notion import NotionToolset
 from holmes.plugins.toolsets.kafka import KafkaToolset
 from holmes.plugins.toolsets.kubernetes_logs import KubernetesLogsToolset
 from holmes.plugins.toolsets.mcp.toolset_mcp import RemoteMCPToolset
-from holmes.plugins.toolsets.newrelic import NewRelicToolset
+from holmes.plugins.toolsets.newrelic.newrelic import NewRelicToolset
 from holmes.plugins.toolsets.opensearch.opensearch import OpenSearchToolset
 from holmes.plugins.toolsets.opensearch.opensearch_logs import OpenSearchLogsToolset
 from holmes.plugins.toolsets.opensearch.opensearch_traces import OpenSearchTracesToolset
-from holmes.plugins.toolsets.prometheus.prometheus import PrometheusToolset
 from holmes.plugins.toolsets.rabbitmq.toolset_rabbitmq import RabbitMQToolset
 from holmes.plugins.toolsets.robusta.robusta import RobustaToolset
 from holmes.plugins.toolsets.runbook.runbook_fetcher import RunbookToolset
@@ -89,7 +91,6 @@ def load_python_toolsets(dal: Optional[SupabaseDal]) -> List[Toolset]:
         DatadogMetricsToolset(),
         DatadogTracesToolset(),
         DatadogRDSToolset(),
-        PrometheusToolset(),
         OpenSearchLogsToolset(),
         OpenSearchTracesToolset(),
         CoralogixLogsToolset(),
@@ -101,6 +102,12 @@ def load_python_toolsets(dal: Optional[SupabaseDal]) -> List[Toolset]:
         AzureSQLToolset(),
         ServiceNowToolset(),
     ]
+
+    if not DISABLE_PROMETHEUS_TOOLSET:
+        from holmes.plugins.toolsets.prometheus.prometheus import PrometheusToolset
+
+        toolsets.append(PrometheusToolset())
+
     if not USE_LEGACY_KUBERNETES_LOGS:
         toolsets.append(KubernetesLogsToolset())
 

holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, Tuple, List
 from holmes.core.tools import (
     CallablePrerequisite,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     Toolset,
     ToolsetTag,
@@ -66,8 +67,8 @@ class MongoDBAtlasToolset(Toolset):
                 {"Accept": "application/vnd.atlas.2025-03-12+json"}
             )
             self._session.auth = HTTPDigestAuth(
-                self.config.get("public_key"),
-                self.config.get("private_key"),
+                self.config.get("public_key"),  # type: ignore
+                self.config.get("private_key"),  # type: ignore
             )
             return True, ""
         except Exception:
@@ -118,9 +119,7 @@ class ReturnProjectAlerts(MongoDBAtlasBaseTool):
         project_id = self.toolset.config.get("project_id", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Project Alerts ({project_id})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             url = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/alerts".format(
                 project_id=self.toolset.config.get("project_id")
@@ -145,9 +144,7 @@ class ReturnProjectProcesses(MongoDBAtlasBaseTool):
         project_id = self.toolset.config.get("project_id", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Project Processes ({project_id})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             url = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/processes".format(
                 project_id=self.toolset.config.get("project_id")
@@ -180,9 +177,7 @@ class ReturnProjectSlowQueries(MongoDBAtlasBaseTool):
         process_id = params.get("process_id", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Slow Queries ({process_id})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             url = self.url.format(
                 project_id=self.toolset.config.get("project_id"),
@@ -209,9 +204,7 @@ class ReturnEventsFromProject(MongoDBAtlasBaseTool):
         project_id = self.toolset.config.get("project_id", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Project Events ({project_id})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         params.update({"itemsPerPage": 500})
         try:
             now_utc = datetime.now(timezone.utc)
@@ -268,9 +261,7 @@ class ReturnLogsForProcessInProject(MongoDBAtlasBaseTool):
         hostname = params.get("hostName", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Host Logs ({hostname})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         one_hour_ago = datetime.now(timezone.utc) - timedelta(hours=1)
         try:
             url = self.url.format(
@@ -324,9 +315,7 @@ class ReturnEventTypeFromProject(MongoDBAtlasBaseTool):
         event_type = params.get("eventType", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Event Details ({event_type})"
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             url = self.url.format(projectId=self.toolset.config.get("project_id"))
 

holmes/plugins/toolsets/azure_sql/tools/analyze_connection_failures.py

@@ -4,6 +4,7 @@ from datetime import datetime, timezone
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -217,9 +218,7 @@ class AnalyzeConnectionFailures(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             # Get configuration
             db_config = self.toolset.database_config()

holmes/plugins/toolsets/azure_sql/tools/analyze_database_connections.py

@@ -4,6 +4,7 @@ from datetime import datetime, timezone
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -155,9 +156,7 @@ class AnalyzeDatabaseConnections(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             hours_back = params.get("hours_back", 2)
 

holmes/plugins/toolsets/azure_sql/tools/analyze_database_health_status.py

@@ -2,7 +2,11 @@ import logging
 from typing import Dict
 from datetime import datetime, timezone
 
-from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    ToolInvokeContext,
+)
 from holmes.plugins.toolsets.azure_sql.azure_base_toolset import (
     BaseAzureSQLTool,
     BaseAzureSQLToolset,
@@ -131,9 +135,7 @@ class AnalyzeDatabaseHealthStatus(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             db_config = self.toolset.database_config()
             client = self.toolset.api_client()

holmes/plugins/toolsets/azure_sql/tools/analyze_database_performance.py

@@ -2,7 +2,11 @@ import logging
 from typing import Any, Dict, List, Tuple, cast
 from datetime import datetime, timezone
 
-from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    ToolInvokeContext,
+)
 from holmes.plugins.toolsets.azure_sql.azure_base_toolset import (
     BaseAzureSQLTool,
     BaseAzureSQLToolset,
@@ -192,9 +196,7 @@ class AnalyzeDatabasePerformance(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             db_config = self.toolset.database_config()
             client = self.toolset.api_client()

holmes/plugins/toolsets/azure_sql/tools/analyze_database_storage.py

@@ -4,6 +4,7 @@ from datetime import datetime, timezone
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -253,9 +254,7 @@ class AnalyzeDatabaseStorage(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             hours_back = params.get("hours_back", 24)
             top_tables = params.get("top_tables", 20)

holmes/plugins/toolsets/azure_sql/tools/get_active_alerts.py

@@ -2,7 +2,11 @@ import logging
 from typing import Dict
 from datetime import datetime, timezone
 
-from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+from holmes.core.tools import (
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    ToolInvokeContext,
+)
 from holmes.plugins.toolsets.azure_sql.azure_base_toolset import (
     BaseAzureSQLTool,
     BaseAzureSQLToolset,
@@ -147,9 +151,7 @@ class GetActiveAlerts(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             db_config = self.toolset.database_config()
             api_client = self.toolset.api_client()

holmes/plugins/toolsets/azure_sql/tools/get_slow_queries.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Tuple
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -103,9 +104,7 @@ class GetSlowQueries(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             top_count = params.get("top_count", 15)
             hours_back = params.get("hours_back", 2)

holmes/plugins/toolsets/azure_sql/tools/get_top_cpu_queries.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Tuple
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -101,9 +102,7 @@ class GetTopCPUQueries(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             top_count = params.get("top_count", 15)
             hours_back = params.get("hours_back", 2)

holmes/plugins/toolsets/azure_sql/tools/get_top_data_io_queries.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Tuple
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -119,9 +120,7 @@ class GetTopDataIOQueries(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             top_count = params.get("top_count", 15)
             hours_back = params.get("hours_back", 2)

holmes/plugins/toolsets/azure_sql/tools/get_top_log_io_queries.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Tuple
 
 from holmes.core.tools import (
     StructuredToolResult,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
 )
@@ -111,9 +112,7 @@ class GetTopLogIOQueries(BaseAzureSQLTool):
 
         return "\n".join(report_sections)
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             top_count = params.get("top_count", 15)
             hours_back = params.get("hours_back", 2)