diff --git a/holmes/common/env_vars.py b/holmes/common/env_vars.py index b717978e68..43e564e6b8 100644 --- a/holmes/common/env_vars.py +++ b/holmes/common/env_vars.py @@ -74,3 +74,10 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]: ENABLE_CLI_TOOL_APPROVAL = load_bool("ENABLE_CLI_TOOL_APPROVAL", True) MAX_GRAPH_POINTS = float(os.environ.get("MAX_GRAPH_POINTS", 300)) + +# Limit each tool response to N% of the total context window. +# Number between 0 and 100 +# Setting to either 0 or any number above 100 disables the logic that limits tool response size +TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT = float( + os.environ.get("TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", 10) +) diff --git a/holmes/core/safeguards.py b/holmes/core/safeguards.py index 932b968902..775164b3a1 100644 --- a/holmes/core/safeguards.py +++ b/holmes/core/safeguards.py @@ -5,7 +5,7 @@ from holmes.common.env_vars import TOOL_CALL_SAFEGUARDS_ENABLED from holmes.plugins.toolsets.logging_utils.logging_api import POD_LOGGING_TOOL_NAME -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams @@ -39,7 +39,7 @@ def _has_previous_unfiltered_pod_logs_call( result = tool_call.get("result", {}) if ( tool_call.get("tool_name") == POD_LOGGING_TOOL_NAME - and result.get("status") == ToolResultStatus.NO_DATA + and result.get("status") == StructuredToolResultStatus.NO_DATA and result.get("params") ): params = FetchPodLogsParams(**result.get("params")) @@ -94,7 +94,7 @@ def prevent_overly_repeated_tool_call( For example if Holmes checks if a resource is deployed, runs a command to deploy it and then checks again if it has deployed properly. """ return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=( "Refusing to run this tool call because it has already been called during this session with the exact same parameters.\n" "Move on with your investigation to a different tool or change the parameter values." @@ -106,7 +106,7 @@ def prevent_overly_repeated_tool_call( tool_name=tool_name, tool_params=tool_params, tool_calls=tool_calls ): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=( f"Refusing to run this tool call because the exact same {POD_LOGGING_TOOL_NAME} tool call without filter has already run and returned no data.\n" "This tool call would also have returned no data.\n" diff --git a/holmes/core/tool_calling_llm.py b/holmes/core/tool_calling_llm.py index 3034947c4f..18425e45d9 100644 --- a/holmes/core/tool_calling_llm.py +++ b/holmes/core/tool_calling_llm.py @@ -32,7 +32,10 @@ from holmes.core.resource_instruction import ResourceInstructions from holmes.core.runbooks import RunbookManager from holmes.core.safeguards import prevent_overly_repeated_tool_call -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus +from holmes.core.tools_utils.tool_context_window_limiter import ( + prevent_overly_big_tool_response, +) from holmes.plugins.prompts import load_and_render_prompt from holmes.utils.global_instructions import ( Instructions, @@ -40,6 +43,11 @@ ) from holmes.utils.tags import format_tags_in_string, parse_messages_tags from holmes.core.tools_utils.tool_executor import ToolExecutor +from holmes.core.tools_utils.data_types import ( + TruncationResult, + ToolCallResult, + TruncationMetadata, +) from holmes.core.tracing import DummySpan from holmes.utils.colors import AI_COLOR from holmes.utils.stream import StreamEvents, StreamMessage @@ -119,34 +127,6 @@ def _process_cost_info( logging.debug(f"Could not extract cost information: {e}") -class TruncationMetadata(BaseModel): - tool_call_id: str - start_index: int - end_index: 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 == ToolResultStatus.ERROR: - tool_response = f"{tool_result.error or 'Tool execution failed'}:\n\n{tool_result.data or ''}".strip() - return tool_response - - # TODO: I think there's a bug here because we don't account for the 'role' or json structure like '{...}' when counting tokens # However, in practice it works because we reserve enough space for the output tokens that the minor inconsistency does not matter # We should fix this in the future @@ -249,52 +229,6 @@ def truncate_messages_to_fit_context( return TruncationResult(truncated_messages=messages, truncations=truncations) -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, - } - - class LLMResult(LLMCosts): tool_calls: Optional[List[ToolCallResult]] = None result: Optional[str] = None @@ -539,7 +473,7 @@ def call( # type: ignore if ( tool_call_result.result.status - == ToolResultStatus.APPROVAL_REQUIRED + == StructuredToolResultStatus.APPROVAL_REQUIRED ): with trace_span.start_span(type="tool") as tool_span: tool_call_result = self._handle_tool_call_approval( @@ -577,7 +511,7 @@ def _directly_invoke_tool_call( f"Skipping tool execution for {tool_name}: args: {tool_params}" ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed to find tool {tool_name}", params=tool_params, ) @@ -591,7 +525,7 @@ def _directly_invoke_tool_call( f"Tool call to {tool_name} failed with an Exception", exc_info=True ) tool_response = StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Tool call failed: {e}", params=tool_params, ) @@ -633,7 +567,7 @@ def _get_tool_call_result( f"Tool {tool_name} return type is not StructuredToolResult. Nesting the tool result into StructuredToolResult..." ) tool_response = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=tool_response, params=tool_params, ) @@ -683,7 +617,7 @@ def _invoke_llm_tool_call( tool_name=tool_name, description="NA", result=StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Custom tool calls are not supported", params=None, ), @@ -699,6 +633,11 @@ def _invoke_llm_tool_call( previous_tool_calls=previous_tool_calls, tool_number=tool_number, ) + + prevent_overly_big_tool_response( + tool_call_result=tool_call_result, llm=self.llm + ) + ToolCallingLLM._log_tool_call_result(tool_span, tool_call_result) return tool_call_result @@ -720,7 +659,7 @@ def _handle_tool_call_approval( # If no approval callback, convert to ERROR because it is assumed the client may not be able to handle approvals if not self.approval_callback: - tool_call_result.result.status = ToolResultStatus.ERROR + tool_call_result.result.status = StructuredToolResultStatus.ERROR return tool_call_result # Get approval from user @@ -740,7 +679,7 @@ def _handle_tool_call_approval( else: # User denied - update to error feedback_text = f" User feedback: {feedback}" if feedback else "" - tool_call_result.result.status = ToolResultStatus.ERROR + tool_call_result.result.status = StructuredToolResultStatus.ERROR tool_call_result.result.error = ( f"User denied command execution.{feedback_text}" ) @@ -952,7 +891,6 @@ def call_stream( for future in concurrent.futures.as_completed(futures): tool_call_result: ToolCallResult = future.result() - tool_calls.append(tool_call_result.as_tool_result_response()) messages.append(tool_call_result.as_tool_call_message()) diff --git a/holmes/core/tools.py b/holmes/core/tools.py index cdb5254bd6..2d13a7b7b4 100644 --- a/holmes/core/tools.py +++ b/holmes/core/tools.py @@ -48,28 +48,28 @@ 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 "⚪️" @@ -77,7 +77,7 @@ def to_emoji(self) -> str: 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 @@ -261,7 +261,10 @@ def _apply_transformers(self, result: StructuredToolResult) -> StructuredToolRes Returns: The tool result with transformed data, or original result if transformation fails """ - if not self._transformer_instances or result.status != ToolResultStatus.SUCCESS: + if ( + not self._transformer_instances + or result.status != StructuredToolResultStatus.SUCCESS + ): return result # Get the output string to transform @@ -387,12 +390,14 @@ def _build_context(self, params): 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 diff --git a/holmes/core/tools_utils/data_types.py b/holmes/core/tools_utils/data_types.py new file mode 100644 index 0000000000..04cb257ea6 --- /dev/null +++ b/holmes/core/tools_utils/data_types.py @@ -0,0 +1,79 @@ +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 + + +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, + } diff --git a/holmes/core/tools_utils/tool_context_window_limiter.py b/holmes/core/tools_utils/tool_context_window_limiter.py new file mode 100644 index 0000000000..ef487d901c --- /dev/null +++ b/holmes/core/tools_utils/tool_context_window_limiter.py @@ -0,0 +1,28 @@ +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 + + +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 = ( + 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 diff --git a/holmes/core/tools_utils/tool_executor.py b/holmes/core/tools_utils/tool_executor.py index c34c36cc8c..abfa88e4cf 100644 --- a/holmes/core/tools_utils/tool_executor.py +++ b/holmes/core/tools_utils/tool_executor.py @@ -6,7 +6,7 @@ from holmes.core.tools import ( StructuredToolResult, Tool, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetStatusEnum, ) @@ -52,7 +52,7 @@ def invoke(self, tool_name: str, params: dict) -> StructuredToolResult: tool.invoke(params) if tool else StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Could not find tool named {tool_name}", ) ) diff --git a/holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py b/holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py index 90f5d974cb..fe218cbe84 100644 --- a/holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py +++ b/holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py @@ -10,7 +10,7 @@ ) from pydantic import BaseModel, PrivateAttr -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from requests.auth import HTTPDigestAuth # type: ignore import gzip import io @@ -90,15 +90,15 @@ def return_result( if response.ok: res = response.json() return StructuredToolResult( - status=ToolResultStatus.SUCCESS + status=StructuredToolResultStatus.SUCCESS if res.get(field, []) - else ToolResultStatus.NO_DATA, + else StructuredToolResultStatus.NO_DATA, data=res, params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed {self.name}.\n{response.text}", return_code=response.status_code, params=params, @@ -130,7 +130,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=f"Exception {self.name}: {str(e)}", params=params, ) @@ -157,7 +157,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception {self.name}: {str(e)}", params=params, ) @@ -193,7 +193,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception {self.name}: {str(e)}", params=params, ) @@ -230,14 +230,14 @@ def _invoke( ) data = f"last 4 hours eventTypeName and # of occurrences list: {events_counter} \n to get more information about a given eventTypeName call atlas_return_events_type_from_project" status = ( - ToolResultStatus.SUCCESS + StructuredToolResultStatus.SUCCESS if events_counter - else ToolResultStatus.NO_DATA + else StructuredToolResultStatus.NO_DATA ) return StructuredToolResult(status=status, data=data, params=params) else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed {self.name}. \n{response.text}", return_code=response.status_code, params=params, @@ -245,7 +245,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception {self.name}: {str(e)}", params=params, ) @@ -287,11 +287,13 @@ def _invoke( with gzip.GzipFile(fileobj=io.BytesIO(response.content)) as gz: text_data = gz.read().decode("utf-8") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data=text_data, params=params + status=StructuredToolResultStatus.SUCCESS, + data=text_data, + params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed {self.name}. \n{response.text}", return_code=response.status_code, params=params, @@ -299,7 +301,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception {self.name}: {str(e)}", params=params, ) @@ -340,7 +342,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception {self.name}: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/analyze_connection_failures.py b/holmes/plugins/toolsets/azure_sql/tools/analyze_connection_failures.py index 3d6a1b1ea2..4ebf1047e9 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/analyze_connection_failures.py +++ b/holmes/plugins/toolsets/azure_sql/tools/analyze_connection_failures.py @@ -2,7 +2,11 @@ from typing import Dict, Tuple from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -242,7 +246,7 @@ def _invoke( # Check for errors if "error" in analysis_data: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=analysis_data["error"], params=params, ) @@ -253,7 +257,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -263,7 +267,7 @@ def _invoke( f"Error in analyze_connection_failures: {str(e)}", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed to analyze connection failures: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_connections.py b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_connections.py index 6fdba065ee..91c101d709 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_connections.py +++ b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_connections.py @@ -2,7 +2,11 @@ from typing import Any, Dict, Tuple from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -200,7 +204,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -208,7 +212,7 @@ def _invoke( error_msg = f"Failed to generate connection report: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_health_status.py b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_health_status.py index 573354bd5d..d6f36673af 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_health_status.py +++ b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_health_status.py @@ -2,7 +2,7 @@ from typing import Dict from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -145,7 +145,7 @@ def _invoke( report_text = self._build_health_report(health_data, db_config) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -153,7 +153,7 @@ def _invoke( error_msg = f"Failed to generate health report: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_performance.py b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_performance.py index a5e41a8b37..6f83a8b584 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_performance.py +++ b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_performance.py @@ -2,7 +2,7 @@ from typing import Any, Dict, List, Tuple, cast from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -206,7 +206,7 @@ def _invoke( report_text = self._build_performance_report(performance_data, db_config) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -214,7 +214,7 @@ def _invoke( error_msg = f"Failed to generate performance report: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_storage.py b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_storage.py index e9dd106220..c91fa9074d 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/analyze_database_storage.py +++ b/holmes/plugins/toolsets/azure_sql/tools/analyze_database_storage.py @@ -2,7 +2,11 @@ from typing import Any, Dict, Tuple from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -307,7 +311,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -315,7 +319,7 @@ def _invoke( error_msg = f"Failed to generate storage report: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/get_active_alerts.py b/holmes/plugins/toolsets/azure_sql/tools/get_active_alerts.py index 9bace4ebbe..e45ab7008a 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/get_active_alerts.py +++ b/holmes/plugins/toolsets/azure_sql/tools/get_active_alerts.py @@ -2,7 +2,7 @@ from typing import Dict from datetime import datetime, timezone -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -170,7 +170,7 @@ def _invoke( # Check for errors if "error" in alerts_data: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=alerts_data["error"], params=params, ) @@ -179,7 +179,7 @@ def _invoke( report_text = self._build_alerts_report(db_config, alerts_data, "active") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -187,7 +187,7 @@ def _invoke( error_msg = f"Failed to retrieve active alerts: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/get_slow_queries.py b/holmes/plugins/toolsets/azure_sql/tools/get_slow_queries.py index 77d8373d26..03908ab6fb 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/get_slow_queries.py +++ b/holmes/plugins/toolsets/azure_sql/tools/get_slow_queries.py @@ -1,7 +1,11 @@ import logging from typing import Dict, List, Tuple -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -125,7 +129,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -133,7 +137,7 @@ def _invoke( error_msg = f"Failed to get slow queries: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/get_top_cpu_queries.py b/holmes/plugins/toolsets/azure_sql/tools/get_top_cpu_queries.py index f624a751a8..dd3f9d9cc5 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/get_top_cpu_queries.py +++ b/holmes/plugins/toolsets/azure_sql/tools/get_top_cpu_queries.py @@ -1,7 +1,11 @@ import logging from typing import Dict, List, Tuple -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -123,7 +127,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -131,7 +135,7 @@ def _invoke( error_msg = f"Failed to get top CPU queries: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/get_top_data_io_queries.py b/holmes/plugins/toolsets/azure_sql/tools/get_top_data_io_queries.py index 3c35b8a8e8..f297a83822 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/get_top_data_io_queries.py +++ b/holmes/plugins/toolsets/azure_sql/tools/get_top_data_io_queries.py @@ -1,7 +1,11 @@ import logging from typing import Dict, List, Tuple -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -141,7 +145,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -149,7 +153,7 @@ def _invoke( error_msg = f"Failed to get top data I/O queries: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/azure_sql/tools/get_top_log_io_queries.py b/holmes/plugins/toolsets/azure_sql/tools/get_top_log_io_queries.py index 3690b5aef5..937f2430e1 100644 --- a/holmes/plugins/toolsets/azure_sql/tools/get_top_log_io_queries.py +++ b/holmes/plugins/toolsets/azure_sql/tools/get_top_log_io_queries.py @@ -1,7 +1,11 @@ import logging from typing import Dict, List, Tuple -from holmes.core.tools import StructuredToolResult, ToolParameter, ToolResultStatus +from holmes.core.tools import ( + StructuredToolResult, + ToolParameter, + StructuredToolResultStatus, +) from holmes.plugins.toolsets.azure_sql.azure_base_toolset import ( BaseAzureSQLTool, BaseAzureSQLToolset, @@ -133,7 +137,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report_text, params=params, ) @@ -141,7 +145,7 @@ def _invoke( error_msg = f"Failed to get top log I/O queries: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) diff --git a/holmes/plugins/toolsets/bash/bash_toolset.py b/holmes/plugins/toolsets/bash/bash_toolset.py index ef5e85b2ea..583e5831fe 100644 --- a/holmes/plugins/toolsets/bash/bash_toolset.py +++ b/holmes/plugins/toolsets/bash/bash_toolset.py @@ -17,7 +17,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -94,7 +94,7 @@ def _invoke( if namespace and not re.match(SAFE_NAMESPACE_PATTERN, namespace): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error: The namespace is invalid. Valid namespaces must match the following regexp: {SAFE_NAMESPACE_PATTERN}", params=params, ) @@ -118,7 +118,7 @@ def _invoke( } ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -172,14 +172,14 @@ def _invoke( if not command_str: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="The 'command' parameter is required and was not provided.", params=params, ) if not isinstance(command_str, str): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"The 'command' parameter must be a string, got {type(command_str).__name__}.", params=params, ) @@ -202,7 +202,7 @@ def _invoke( logging.info(f"Refusing LLM tool call {command_str}") return StructuredToolResult( - status=ToolResultStatus.APPROVAL_REQUIRED, + status=StructuredToolResultStatus.APPROVAL_REQUIRED, error=f"Refusing to execute bash command. {str(e)}", params=params, invocation=command_str, diff --git a/holmes/plugins/toolsets/bash/common/bash.py b/holmes/plugins/toolsets/bash/common/bash.py index b2645808c8..ae77dedcd0 100644 --- a/holmes/plugins/toolsets/bash/common/bash.py +++ b/holmes/plugins/toolsets/bash/common/bash.py @@ -1,5 +1,5 @@ import subprocess -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus def execute_bash_command(cmd: str, timeout: int, params: dict) -> StructuredToolResult: @@ -18,11 +18,11 @@ def execute_bash_command(cmd: str, timeout: int, params: dict) -> StructuredTool stdout = process.stdout.strip() if process.stdout else "" result_data = f"{cmd}\n" f"{stdout}" - status = ToolResultStatus.ERROR + status = StructuredToolResultStatus.ERROR if process.returncode == 0 and stdout: - status = ToolResultStatus.SUCCESS + status = StructuredToolResultStatus.SUCCESS elif not stdout: - status = ToolResultStatus.NO_DATA + status = StructuredToolResultStatus.NO_DATA return StructuredToolResult( status=status, @@ -33,20 +33,20 @@ def execute_bash_command(cmd: str, timeout: int, params: dict) -> StructuredTool ) except subprocess.TimeoutExpired: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error: Command '{cmd}' timed out after {timeout} seconds.", params=params, ) except FileNotFoundError: # This might occur if /bin/bash is not found, or if shell=False and command is not found return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Error: Bash executable or command not found. Ensure bash is installed and the command is valid.", params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error executing command '{cmd}': {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/coralogix/toolset_coralogix_logs.py b/holmes/plugins/toolsets/coralogix/toolset_coralogix_logs.py index 7f39086d2b..69ad6f7e11 100644 --- a/holmes/plugins/toolsets/coralogix/toolset_coralogix_logs.py +++ b/holmes/plugins/toolsets/coralogix/toolset_coralogix_logs.py @@ -3,7 +3,7 @@ from holmes.core.tools import ( CallablePrerequisite, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetTag, ) from holmes.plugins.toolsets.consts import ( @@ -74,7 +74,7 @@ def logger_name(self) -> str: def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if not self.coralogix_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"The {self.name} toolset is not configured", params=params.model_dump(), ) @@ -102,7 +102,9 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: return StructuredToolResult( status=( - ToolResultStatus.ERROR if logs_data.error else ToolResultStatus.SUCCESS + StructuredToolResultStatus.ERROR + if logs_data.error + else StructuredToolResultStatus.SUCCESS ), error=logs_data.error, data=data, diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_general.py b/holmes/plugins/toolsets/datadog/toolset_datadog_general.py index 602f1f14df..20f410e3ff 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_general.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_general.py @@ -13,7 +13,7 @@ ToolParameter, Toolset, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetTag, ) from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR @@ -333,7 +333,7 @@ def _invoke( if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -350,7 +350,7 @@ def _invoke( if not is_allowed: logging.error(f"Endpoint validation failed: {error_msg}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Endpoint validation failed: {error_msg}", params=params, ) @@ -381,13 +381,13 @@ def _invoke( > self.toolset.dd_config.max_response_size ): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Response too large (>{self.toolset.dd_config.max_response_size} bytes)", params=params, ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=response_str, params=params, ) @@ -407,7 +407,7 @@ def _invoke( error_msg = f"API error {e.status_code}: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=json.dumps({"url": url, "params": query_params}) @@ -418,7 +418,7 @@ def _invoke( except Exception as e: logging.exception(f"Failed to query Datadog API: {params}", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, ) @@ -469,7 +469,7 @@ def _invoke( if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -486,7 +486,7 @@ def _invoke( if not is_allowed: logging.error(f"Endpoint validation failed: {error_msg}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Endpoint validation failed: {error_msg}", params=params, ) @@ -517,13 +517,13 @@ def _invoke( > self.toolset.dd_config.max_response_size ): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Response too large (>{self.toolset.dd_config.max_response_size} bytes)", params=params, ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=response_str, params=params, ) @@ -543,7 +543,7 @@ def _invoke( error_msg = f"API error {e.status_code}: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=json.dumps({"url": url, "body": body}) if url else None, @@ -552,7 +552,7 @@ def _invoke( except Exception as e: logging.exception(f"Failed to query Datadog API: {params}", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, ) @@ -691,7 +691,7 @@ def _invoke( matching_categories = {k: v for k, v in resources.items() if category in k} if not matching_categories: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unknown category: {category}. Available: {', '.join(resources.keys())}", params=params, ) @@ -715,7 +715,7 @@ def _invoke( output.append("Example: datadog_api_get with endpoint='/api/v1/monitors'") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="\n".join(output), params=params, ) diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py index a46180f1e1..519d332c60 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py @@ -8,7 +8,7 @@ ToolsetTag, ) from pydantic import BaseModel, Field -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR from holmes.plugins.toolsets.datadog.datadog_api import ( DatadogBaseConfig, @@ -167,7 +167,7 @@ def logger_name(self) -> str: def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if not self.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=TOOLSET_CONFIG_MISSING_ERROR, params=params.model_dump(), ) @@ -182,13 +182,13 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if raw_logs: logs_str = format_logs(raw_logs) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=logs_str, params=params.model_dump(), ) return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params=params.model_dump(), ) @@ -202,7 +202,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params.model_dump(), invocation=json.dumps(e.payload), @@ -213,7 +213,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: f"Failed to query Datadog logs for params: {params}", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception while querying Datadog: {str(e)}", params=params.model_dump(), ) @@ -234,11 +234,11 @@ def _perform_healthcheck(self) -> Tuple[bool, str]: result = self.fetch_pod_logs(healthcheck_params) - if result.status == ToolResultStatus.ERROR: + if result.status == StructuredToolResultStatus.ERROR: error_msg = result.error or "Unknown error during healthcheck" logging.error(f"Datadog healthcheck failed: {error_msg}") return False, f"Datadog healthcheck failed: {error_msg}" - elif result.status == ToolResultStatus.NO_DATA: + elif result.status == StructuredToolResultStatus.NO_DATA: error_msg = "No logs were found in the last 48 hours using wildcards for pod and namespace. Is the configuration correct?" logging.error(f"Datadog healthcheck failed: {error_msg}") return False, f"Datadog healthcheck failed: {error_msg}" diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py index 47e942b1d5..f7aa192d89 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py @@ -7,7 +7,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -80,7 +80,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -121,7 +121,7 @@ def _invoke( metrics = data.get("metrics", []) if not metrics: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data="Your filter returned no metrics. Change your filter and try again", params=params, ) @@ -133,7 +133,7 @@ def _invoke( output.append(metric) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="\n".join(output), params=params, ) @@ -152,7 +152,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=json.dumps({"url": url, "params": query_params}) @@ -165,7 +165,7 @@ def _invoke( f"Failed to query Datadog metrics for params: {params}", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception while querying Datadog: {str(e)}", params=params, ) @@ -222,7 +222,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -262,7 +262,7 @@ def _invoke( if not series: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, error="The query returned no data. Please check your query syntax and time range.", params=params, ) @@ -317,7 +317,7 @@ def _invoke( data_str = json.dumps(response_data, indent=2) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=data_str, params=params, ) @@ -336,7 +336,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=json.dumps({"url": url, "params": query_params}) @@ -350,7 +350,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception while querying Datadog: {str(e)}", params=params, ) @@ -380,7 +380,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -396,7 +396,7 @@ def _invoke( if not metric_names: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="metric_names cannot be empty", params=params, ) @@ -442,14 +442,14 @@ def _invoke( if not results and errors: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Failed to retrieve metadata for all metrics", data=json.dumps(response_data, indent=2), params=params, ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=json.dumps(response_data, indent=2), params=params, ) @@ -461,7 +461,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception while querying Datadog: {str(e)}", params=params, ) @@ -496,7 +496,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -519,7 +519,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=data, params=params, ) @@ -540,7 +540,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=json.dumps({"url": url, "params": query_params}) @@ -554,7 +554,7 @@ def _invoke( exc_info=True, ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Exception while querying Datadog: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_rds.py b/holmes/plugins/toolsets/datadog/toolset_datadog_rds.py index 629bf02a1e..4677e4f5c1 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_rds.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_rds.py @@ -10,7 +10,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -97,7 +97,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -150,7 +150,7 @@ def _invoke( formatted_report = self._format_report(report) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=formatted_report, params=params, ) @@ -158,7 +158,7 @@ def _invoke( except Exception as e: logging.error(f"Error generating RDS performance report: {str(e)}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed to generate RDS performance report: {str(e)}", params=params, ) @@ -397,7 +397,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) @@ -416,7 +416,7 @@ def _invoke( if not instances: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, data="No RDS instances found with metrics in the specified time range", params=params, ) @@ -443,7 +443,7 @@ def _invoke( report += f"\n\nInstances:\n{json.dumps(worst_performers, indent=2)}" return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=report, params=params, ) @@ -451,7 +451,7 @@ def _invoke( except Exception as e: logging.error(f"Error getting top worst performing RDS instances: {str(e)}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Failed to get top worst performing RDS instances: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py index fac97bc5c0..a8104a3778 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py @@ -12,7 +12,7 @@ ToolParameter, Toolset, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetTag, ) from holmes.plugins.toolsets.datadog.datadog_api import ( @@ -216,7 +216,7 @@ def _invoke( """Execute the tool to fetch traces.""" if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Datadog configuration not initialized", params=params, ) @@ -305,13 +305,13 @@ def _invoke( formatted_output = format_traces_list(spans, limit=params.get("limit", 50)) if not formatted_output: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params=params, data="No matching traces found.", ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=formatted_output, params=params, ) @@ -330,7 +330,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=( @@ -343,7 +343,7 @@ def _invoke( except Exception as e: logging.exception(e, exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, invocation=( @@ -382,7 +382,7 @@ def _invoke( """Execute the tool to fetch trace details.""" if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Datadog configuration not initialized", params=params, ) @@ -390,7 +390,7 @@ def _invoke( trace_id = params.get("trace_id") if not trace_id: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="trace_id parameter is required", params=params, ) @@ -444,13 +444,13 @@ def _invoke( formatted_output = format_trace_hierarchy(trace_id, spans) if not formatted_output: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params=params, data=f"No trace found for trace_id: {trace_id}", ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=formatted_output, params=params, ) @@ -469,7 +469,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=( @@ -482,7 +482,7 @@ def _invoke( except Exception as e: logging.exception(e, exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, invocation=( @@ -565,7 +565,7 @@ def _invoke( """Execute the tool to search spans.""" if not self.toolset.dd_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Datadog configuration not initialized", params=params, ) @@ -653,13 +653,13 @@ def _invoke( formatted_output = format_spans_search(spans) if not formatted_output: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params=params, data="No matching spans found.", ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=formatted_output, params=params, ) @@ -677,7 +677,7 @@ def _invoke( error_msg = f"Exception while querying Datadog: {str(e)}" return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, invocation=( @@ -690,7 +690,7 @@ def _invoke( except Exception as e: logging.exception(e, exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, invocation=( diff --git a/holmes/plugins/toolsets/git.py b/holmes/plugins/toolsets/git.py index 0c52923cbd..31540159b2 100644 --- a/holmes/plugins/toolsets/git.py +++ b/holmes/plugins/toolsets/git.py @@ -4,7 +4,7 @@ import os from typing import Any, Optional, Dict, List, Tuple from pydantic import BaseModel -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.core.tools import ( Toolset, @@ -259,7 +259,7 @@ def _invoke( resp = requests.get(url, headers=headers) if resp.status_code != 200: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error( f"Error fetching file: {resp.text}" ), @@ -268,13 +268,13 @@ def _invoke( content = base64.b64decode(resp.json()["content"]).decode().splitlines() numbered = "\n".join(f"{i+1}: {line}" for i, line in enumerate(content)) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=numbered, params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error(str(e)), params=params, ) @@ -304,7 +304,7 @@ def _invoke( resp = requests.get(url, headers=headers) if resp.status_code != 200: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error( f"Error listing files: {resp.text}" ), @@ -312,13 +312,13 @@ def _invoke( ) paths = [entry["path"] for entry in resp.json()["tree"]] return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=paths, params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error(str(e)), params=params, ) @@ -353,13 +353,13 @@ def _invoke( for pr in prs ] return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=formatted, params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error(str(e)), params=params, ) @@ -413,14 +413,14 @@ def _invoke( ) -> StructuredToolResult: def error(msg: str) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=self.toolset._sanitize_error(msg), params=params, ) def success(msg: Any) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data=msg, params=params + status=StructuredToolResultStatus.SUCCESS, data=msg, params=params ) def modify_lines(lines: List[str]) -> List[str]: @@ -643,24 +643,24 @@ def _invoke( # Validate inputs if not commit_message.strip(): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Tool call failed to run: Commit message cannot be empty", ) if not filename.strip(): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Tool call failed to run: Filename cannot be empty", ) if line < 1: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Tool call failed to run: Line number must be positive", ) # Verify this is a PR created by our tool if not self.toolset.is_created_pr(pr_number): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Tool call failed to run: PR #{pr_number} was not created by this tool. Only PRs created using git_execute_changes can be updated.", ) @@ -714,7 +714,7 @@ def _invoke( del content_lines[line - 1] else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Tool call failed to run: Invalid command: {command}", ) @@ -722,7 +722,7 @@ def _invoke( if dry_run: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=f"DRY RUN: Updated content for PR #{pr_number}:\n\n{updated_content}", ) @@ -731,13 +731,13 @@ def _invoke( pr_number, filename, updated_content, commit_message ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=f"Added commit to PR #{pr_number} successfully", ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=self.toolset._sanitize_error( f"Tool call failed to run: Error updating PR: {str(e)}" ), @@ -745,14 +745,14 @@ def _invoke( except requests.exceptions.RequestException as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=self.toolset._sanitize_error( f"Tool call failed to run: Network error: {str(e)}" ), ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=self.toolset._sanitize_error( f"Tool call failed to run: Unexpected error: {str(e)}" ), diff --git a/holmes/plugins/toolsets/grafana/common.py b/holmes/plugins/toolsets/grafana/common.py index c39f590ea4..483712e415 100644 --- a/holmes/plugins/toolsets/grafana/common.py +++ b/holmes/plugins/toolsets/grafana/common.py @@ -3,7 +3,7 @@ from pydantic import BaseModel import datetime -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus class GrafanaConfig(BaseModel): @@ -61,7 +61,7 @@ def ensure_grafana_uid_or_return_error_result( ) -> Optional[StructuredToolResult]: if not config.grafana_datasource_uid: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="This tool only works when the toolset is configued ", ) else: diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana.py b/holmes/plugins/toolsets/grafana/toolset_grafana.py index b7c04dbc38..ccdb10a590 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana.py @@ -4,7 +4,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.plugins.toolsets.grafana.base_grafana_toolset import BaseGrafanaToolset import requests # type: ignore @@ -90,9 +90,9 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS + status=StructuredToolResultStatus.SUCCESS if formatted_dashboards - else ToolResultStatus.NO_DATA, + else StructuredToolResultStatus.NO_DATA, data="\n".join(formatted_dashboards) if formatted_dashboards else "No dashboards found.", @@ -102,7 +102,7 @@ def _invoke( except requests.RequestException as e: logging.error(f"Error fetching dashboards: {str(e)}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error fetching dashboards: {str(e)}", url=url, params=params, diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py index ca6340b9bc..85bf937927 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py @@ -22,7 +22,7 @@ from holmes.plugins.toolsets.grafana.loki_api import ( query_loki_logs_by_label, ) -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus class GrafanaLokiLabelsConfig(BaseModel): @@ -99,12 +99,12 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if logs: logs.sort(key=lambda x: x["timestamp"]) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="\n".join([format_log(log) for log in logs]), params=params.model_dump(), ) else: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params=params.model_dump(), ) diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py b/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py index 0a25517d8d..3778fa3852 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py @@ -8,7 +8,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION from holmes.plugins.toolsets.grafana.base_grafana_toolset import BaseGrafanaToolset @@ -219,7 +219,7 @@ def _invoke( ) if invalid_params_error: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=invalid_params_error, params=params, ) @@ -258,7 +258,7 @@ def _invoke( traces = all_traces_response.get("traces", []) if not traces: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="No traces found matching the query", params=params, ) @@ -335,14 +335,14 @@ def fetch_full_trace(trace_summary): # Return as YAML for readability return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False, sort_keys=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error fetching traces: {str(e)}", params=params, ) @@ -429,13 +429,13 @@ def _invoke( spss=params.get("spss"), ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -513,13 +513,13 @@ def _invoke( spss=params.get("spss"), ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -575,13 +575,13 @@ def _invoke( # Return raw trace data as YAML for readability return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(trace_data, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -653,13 +653,13 @@ def _invoke( max_stale_values=params.get("max_stale_values"), ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -731,13 +731,13 @@ def _invoke( max_stale_values=params.get("max_stale_values"), ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -808,13 +808,13 @@ def _invoke( end=end, ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) @@ -906,13 +906,13 @@ def _invoke( exemplars=params.get("exemplars"), ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result, default_flow_style=False), params=params, ) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e), params=params, ) diff --git a/holmes/plugins/toolsets/internet/internet.py b/holmes/plugins/toolsets/internet/internet.py index 320f39deba..17f8895f15 100644 --- a/holmes/plugins/toolsets/internet/internet.py +++ b/holmes/plugins/toolsets/internet/internet.py @@ -15,7 +15,7 @@ from bs4 import BeautifulSoup import requests # type: ignore -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.utils import toolset_name_for_one_liner @@ -199,7 +199,7 @@ def _invoke( 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 +211,7 @@ def _invoke( content = html_to_markdown(content) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=content, params=params, ) diff --git a/holmes/plugins/toolsets/internet/notion.py b/holmes/plugins/toolsets/internet/notion.py index 5953b0ac0b..03dc6d8743 100644 --- a/holmes/plugins/toolsets/internet/notion.py +++ b/holmes/plugins/toolsets/internet/notion.py @@ -13,7 +13,7 @@ ) from holmes.core.tools import ( StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.plugins.toolsets.utils import toolset_name_for_one_liner @@ -59,13 +59,13 @@ def _invoke( 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, ) diff --git a/holmes/plugins/toolsets/investigator/core_investigation.py b/holmes/plugins/toolsets/investigator/core_investigation.py index b4f353732f..ae99759138 100644 --- a/holmes/plugins/toolsets/investigator/core_investigation.py +++ b/holmes/plugins/toolsets/investigator/core_investigation.py @@ -10,7 +10,7 @@ ToolParameter, Tool, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.plugins.toolsets.investigator.model import Task, TaskStatus @@ -103,7 +103,7 @@ def _invoke( response_data += "No tasks currently in the investigation plan." return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=response_data, params=params, ) @@ -111,7 +111,7 @@ def _invoke( 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, ) diff --git a/holmes/plugins/toolsets/kafka.py b/holmes/plugins/toolsets/kafka.py index 205437d413..69180cdb6a 100644 --- a/holmes/plugins/toolsets/kafka.py +++ b/holmes/plugins/toolsets/kafka.py @@ -28,7 +28,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -161,7 +161,7 @@ def _invoke( client = self.get_kafka_client(kafka_cluster_name) if client is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="No admin_client on toolset. This toolset is misconfigured.", params=params, ) @@ -190,7 +190,7 @@ def _invoke( if errors_text: result_text = result_text + "\n\n" + errors_text return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=result_text, params=params, ) @@ -198,7 +198,7 @@ def _invoke( error_msg = f"Failed to list consumer groups: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) @@ -237,7 +237,7 @@ def _invoke( client = self.get_kafka_client(kafka_cluster_name) if client is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="No admin_client on toolset. This toolset is misconfigured.", params=params, ) @@ -247,13 +247,13 @@ def _invoke( if futures.get(group_id): group_metadata = futures.get(group_id).result() return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(convert_to_dict(group_metadata)), params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Group not found", params=params, ) @@ -261,7 +261,7 @@ def _invoke( error_msg = f"Failed to describe consumer group {group_id}: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) @@ -294,14 +294,14 @@ def _invoke( client = self.get_kafka_client(kafka_cluster_name) if client is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="No admin_client on toolset. This toolset is misconfigured.", params=params, ) topics = client.list_topics() return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(convert_to_dict(topics)), params=params, ) @@ -309,7 +309,7 @@ def _invoke( error_msg = f"Failed to list topics: {str(e)}" logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) @@ -353,7 +353,7 @@ def _invoke( client = self.get_kafka_client(kafka_cluster_name) if client is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="No admin_client on toolset. This toolset is misconfigured.", params=params, ) @@ -373,7 +373,7 @@ def _invoke( result["configuration"] = convert_to_dict(config) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=yaml.dump(result), params=params, ) @@ -381,7 +381,7 @@ def _invoke( error_msg = f"Failed to describe topic {topic_name}: {str(e)}" logging.error(error_msg, exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) @@ -478,7 +478,7 @@ def _invoke( client = self.get_kafka_client(kafka_cluster_name) if client is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="No admin_client on toolset. This toolset is misconfigured.", params=params, ) @@ -530,7 +530,7 @@ def _invoke( result_text = result_text + "\n\n" + errors_text return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=result_text, params=params, ) @@ -540,7 +540,7 @@ def _invoke( ) logging.error(error_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=error_msg, params=params, ) @@ -564,7 +564,7 @@ def _invoke( ) -> StructuredToolResult: cluster_names = list(self.toolset.clients.keys()) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Available Kafka Clusters:\n" + "\n".join(cluster_names), params=params, ) diff --git a/holmes/plugins/toolsets/kubernetes_logs.py b/holmes/plugins/toolsets/kubernetes_logs.py index 08c9adc678..2ee60f2dd1 100644 --- a/holmes/plugins/toolsets/kubernetes_logs.py +++ b/holmes/plugins/toolsets/kubernetes_logs.py @@ -10,7 +10,7 @@ from holmes.core.tools import ( StaticPrerequisite, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetTag, ) from holmes.plugins.toolsets.logging_utils.logging_api import ( @@ -140,7 +140,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: # Ensure both results are not None (they should always be set by the loop) if current_logs_result is None or previous_logs_result is None: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Internal error: Failed to fetch logs", params=params.model_dump(), ) @@ -162,7 +162,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: ): # Both commands failed - return error from current logs return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=current_logs_result.error, params=params.model_dump(), return_code=return_code, @@ -206,7 +206,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if len(filtered_logs) == 0: # Return NO_DATA status when there are no logs return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, data="\n".join( metadata_lines ), # Still include metadata for context @@ -218,7 +218,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: response_data = formatted_logs + "\n" + "\n".join(metadata_lines) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=response_data, params=params.model_dump(), return_code=return_code, @@ -226,7 +226,7 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: except Exception as e: logging.exception(f"Error fetching logs for pod {params.pod_name}") return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Error fetching logs: {str(e)}", params=params.model_dump(), ) diff --git a/holmes/plugins/toolsets/mcp/toolset_mcp.py b/holmes/plugins/toolsets/mcp/toolset_mcp.py index 37ccd5aa76..4493591587 100644 --- a/holmes/plugins/toolsets/mcp/toolset_mcp.py +++ b/holmes/plugins/toolsets/mcp/toolset_mcp.py @@ -3,7 +3,7 @@ Tool, ToolParameter, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, CallablePrerequisite, ) @@ -31,7 +31,7 @@ def _invoke( return asyncio.run(self._invoke_async(params)) except Exception as e: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=str(e.args), params=params, invocation=f"MCPtool {self.name} with params {params}", @@ -48,9 +48,9 @@ async def _invoke_async(self, params: Dict) -> StructuredToolResult: ) return StructuredToolResult( status=( - ToolResultStatus.ERROR + StructuredToolResultStatus.ERROR if tool_result.isError - else ToolResultStatus.SUCCESS + else StructuredToolResultStatus.SUCCESS ), data=merged_text, params=params, diff --git a/holmes/plugins/toolsets/newrelic.py b/holmes/plugins/toolsets/newrelic.py index d6d71562c0..4b658c71b2 100644 --- a/holmes/plugins/toolsets/newrelic.py +++ b/holmes/plugins/toolsets/newrelic.py @@ -9,7 +9,7 @@ ToolsetTag, ) from pydantic import BaseModel -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.utils import get_param_or_raise, toolset_name_for_one_liner @@ -42,14 +42,14 @@ def _invoke( ) -> StructuredToolResult: def success(msg: Any) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=msg, params=params, ) def error(msg: str) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=msg, params=params, ) @@ -122,14 +122,14 @@ def _invoke( ) -> StructuredToolResult: def success(msg: Any) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=msg, params=params, ) def error(msg: str) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=msg, params=params, ) diff --git a/holmes/plugins/toolsets/opensearch/opensearch.py b/holmes/plugins/toolsets/opensearch/opensearch.py index 41cee0519b..8e855ab50d 100644 --- a/holmes/plugins/toolsets/opensearch/opensearch.py +++ b/holmes/plugins/toolsets/opensearch/opensearch.py @@ -9,7 +9,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -99,7 +99,7 @@ def _invoke( client = get_client(self.toolset.clients, host=params.get("host", "")) shards = client.client.cat.shards() return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=str(shards), params=params, ) @@ -132,7 +132,7 @@ def _invoke( include_defaults=True, flat_settings=True ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=str(response), params=params, ) @@ -163,7 +163,7 @@ def _invoke( client = get_client(self.toolset.clients, host=params.get("host", "")) health = client.client.cluster.health() return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=str(health), params=params, ) @@ -187,7 +187,7 @@ def _invoke( ) -> StructuredToolResult: hosts = [host for client in self.toolset.clients for host in client.hosts] return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=str(hosts), params=params, ) diff --git a/holmes/plugins/toolsets/opensearch/opensearch_logs.py b/holmes/plugins/toolsets/opensearch/opensearch_logs.py index feab2860a6..c64f839089 100644 --- a/holmes/plugins/toolsets/opensearch/opensearch_logs.py +++ b/holmes/plugins/toolsets/opensearch/opensearch_logs.py @@ -8,7 +8,7 @@ from holmes.core.tools import ( CallablePrerequisite, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetTag, ) from holmes.plugins.toolsets.logging_utils.logging_api import ( @@ -79,7 +79,7 @@ def logger_name(self) -> str: def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: if not self.opensearch_config: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Missing OpenSearch configuration", params=params.model_dump(), ) @@ -126,13 +126,13 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: config=self.opensearch_config, ) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=logs, params=params.model_dump(), ) else: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, return_code=logs_response.status_code, error=logs_response.text, params=params.model_dump(), @@ -141,21 +141,21 @@ def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: except requests.Timeout: logging.warning("Timeout while fetching OpenSearch logs", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Request timed out while fetching OpenSearch logs", params=params.model_dump(), ) except RequestException as e: logging.warning("Failed to fetch OpenSearch logs", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Network error while fetching OpenSearch logs: {str(e)}", params=params.model_dump(), ) except Exception as e: logging.warning("Failed to process OpenSearch logs", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params.model_dump(), ) diff --git a/holmes/plugins/toolsets/opensearch/opensearch_traces.py b/holmes/plugins/toolsets/opensearch/opensearch_traces.py index 7c901d4c84..2e420bb9dc 100644 --- a/holmes/plugins/toolsets/opensearch/opensearch_traces.py +++ b/holmes/plugins/toolsets/opensearch/opensearch_traces.py @@ -18,7 +18,7 @@ add_auth_header, get_search_url, ) -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.utils import get_param_or_raise, toolset_name_for_one_liner TRACES_FIELDS_CACHE_KEY = "cached_traces_fields" @@ -48,7 +48,7 @@ def _invoke( if cached_response: logging.debug("traces fields returned from cache") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=cached_response, params=params, ) @@ -81,7 +81,7 @@ def _invoke( if self._cache: self._cache[TRACES_FIELDS_CACHE_KEY] = response return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=response, params=params, ) @@ -90,21 +90,21 @@ def _invoke( "Timeout while fetching opensearch traces fields", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Request timed out while fetching opensearch traces fields", params=params, ) except RequestException as e: logging.warning("Failed to fetch opensearch traces fields", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Network error while opensearch traces fields: {str(e)}", params=params, ) except Exception as e: logging.warning("Failed to process opensearch traces fields", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, ) @@ -157,7 +157,7 @@ def _invoke( logs_response.raise_for_status() return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=json.dumps(logs_response.json()), params=params, ) @@ -166,14 +166,14 @@ def _invoke( "Timeout while fetching opensearch traces search", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Request timed out while fetching opensearch traces search {err_msg}", params=params, ) except RequestException as e: logging.warning("Failed to fetch opensearch traces search", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Network error while opensearch traces search {err_msg} : {str(e)}", params=params, ) @@ -182,7 +182,7 @@ def _invoke( "Failed to process opensearch traces search ", exc_info=True ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error {err_msg}: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/prometheus/prometheus.py b/holmes/plugins/toolsets/prometheus/prometheus.py index 2ec3aab6a7..69d059825c 100644 --- a/holmes/plugins/toolsets/prometheus/prometheus.py +++ b/holmes/plugins/toolsets/prometheus/prometheus.py @@ -17,7 +17,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -496,13 +496,13 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.config or not self.toolset.config.prometheus_url: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Prometheus is not configured. Prometheus URL is missing", params=params, ) if self.toolset.config.is_amp(): return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Tool not supported in AMP", params=params, ) @@ -515,7 +515,7 @@ def _invoke( logging.debug("rules returned from cache") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=cached_rules, params=params, ) @@ -539,28 +539,28 @@ def _invoke( if self._cache: self._cache.set(PROMETHEUS_RULES_CACHE_KEY, data) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=data, params=params, ) except requests.Timeout: logging.warning("Timeout while fetching prometheus rules", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Request timed out while fetching rules", params=params, ) except RequestException as e: logging.warning("Failed to fetch prometheus rules", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Network error while fetching rules: {str(e)}", params=params, ) except Exception as e: logging.warning("Failed to process prometheus rules", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, ) @@ -595,7 +595,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.config or not self.toolset.config.prometheus_url: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Prometheus is not configured. Prometheus URL is missing", params=params, ) @@ -612,7 +612,7 @@ def _invoke( name_filter = params.get("name_filter") if not name_filter: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Error: cannot run tool 'list_available_metrics'. The param 'name_filter' is required but is missing.", params=params, ) @@ -646,7 +646,7 @@ def _invoke( table_output = "\n".join(output) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=table_output, params=params, ) @@ -654,21 +654,21 @@ def _invoke( except requests.Timeout: logging.warn("Timeout while fetching prometheus metrics", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Request timed out while fetching metrics", params=params, ) except RequestException as e: logging.warn("Failed to fetch prometheus metrics", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Network error while fetching metrics: {str(e)}", params=params, ) except Exception as e: logging.warn("Failed to process prometheus metrics", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error: {str(e)}", params=params, ) @@ -703,7 +703,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.config or not self.toolset.config.prometheus_url: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Prometheus is not configured. Prometheus URL is missing", params=params, ) @@ -748,7 +748,7 @@ def _invoke( data_str = json.dumps(response_data, indent=2) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=data_str, params=params, ) @@ -764,14 +764,14 @@ def _invoke( except json.JSONDecodeError: pass return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Query execution failed. HTTP {response.status_code}: {error_msg}", params=params, ) # For other status codes, just return the status code and content return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Query execution failed with unexpected status code: {response.status_code}. Response: {str(response.content)}", params=params, ) @@ -779,14 +779,14 @@ def _invoke( except RequestException as e: logging.info("Failed to connect to Prometheus", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Connection error to Prometheus: {str(e)}", params=params, ) except Exception as e: logging.info("Failed to connect to Prometheus", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error executing query: {str(e)}", params=params, ) @@ -843,7 +843,7 @@ def _invoke( ) -> StructuredToolResult: if not self.toolset.config or not self.toolset.config.prometheus_url: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Prometheus is not configured. Prometheus URL is missing", params=params, ) @@ -911,7 +911,7 @@ def _invoke( data_str = json.dumps(response_data, indent=2) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=data_str, params=params, ) @@ -926,13 +926,13 @@ def _invoke( except json.JSONDecodeError: pass return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Query execution failed. HTTP {response.status_code}: {error_msg}", params=params, ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Query execution failed with unexpected status code: {response.status_code}. Response: {str(response.content)}", params=params, ) @@ -940,14 +940,14 @@ def _invoke( except RequestException as e: logging.info("Failed to connect to Prometheus", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Connection error to Prometheus: {str(e)}", params=params, ) except Exception as e: logging.info("Failed to connect to Prometheus", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error executing query: {str(e)}", params=params, ) diff --git a/holmes/plugins/toolsets/rabbitmq/toolset_rabbitmq.py b/holmes/plugins/toolsets/rabbitmq/toolset_rabbitmq.py index 66258bdc36..fe39393167 100644 --- a/holmes/plugins/toolsets/rabbitmq/toolset_rabbitmq.py +++ b/holmes/plugins/toolsets/rabbitmq/toolset_rabbitmq.py @@ -8,7 +8,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -79,7 +79,7 @@ def _invoke( if c.connection_status == ClusterConnectionStatus.SUCCESS ] return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data=available_clusters + status=StructuredToolResultStatus.SUCCESS, data=available_clusters ) def get_parameterized_one_liner(self, params) -> str: @@ -112,12 +112,14 @@ def _invoke( cluster_id=params.get("cluster_id") ) result = get_cluster_status(cluster_config) - return StructuredToolResult(status=ToolResultStatus.SUCCESS, data=result) + return StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data=result + ) except Exception as e: logging.info("Failed to process RabbitMQ cluster status", exc_info=True) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=f"Unexpected error fetching RabbitMQ cluster status: {str(e)}", data=None, ) diff --git a/holmes/plugins/toolsets/robusta/robusta.py b/holmes/plugins/toolsets/robusta/robusta.py index a5308ecc05..0d38381dca 100644 --- a/holmes/plugins/toolsets/robusta/robusta.py +++ b/holmes/plugins/toolsets/robusta/robusta.py @@ -11,7 +11,7 @@ Toolset, ToolsetTag, ) -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus PARAM_FINDING_ID = "id" START_TIME = "start_datetime" @@ -53,13 +53,13 @@ def _invoke( finding = self._fetch_finding(finding_id) if finding: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=finding, params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, data=f"Could not find a finding with finding_id={finding_id}", params=params, ) @@ -70,7 +70,7 @@ def _invoke( ) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=f"There was an internal error while fetching finding {finding_id}", params=params, ) @@ -122,13 +122,13 @@ def _invoke( recommendations = self._resource_recommendation(params) if recommendations: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=recommendations, params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, data=f"Could not find recommendations for {params}", params=params, ) @@ -136,7 +136,7 @@ def _invoke( msg = f"There was an internal error while fetching recommendations for {params}. {str(e)}" logging.exception(msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=msg, params=params, ) @@ -182,13 +182,13 @@ def _invoke( changes = self._fetch_change_history(params) if changes: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=changes, params=params, ) else: return StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, data=f"Could not find changes for {params}", params=params, ) @@ -196,7 +196,7 @@ def _invoke( msg = f"There was an internal error while fetching changes for {params}. {str(e)}" logging.exception(msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=msg, params=params, ) diff --git a/holmes/plugins/toolsets/runbook/runbook_fetcher.py b/holmes/plugins/toolsets/runbook/runbook_fetcher.py index e567c8657c..a76d1dad1a 100644 --- a/holmes/plugins/toolsets/runbook/runbook_fetcher.py +++ b/holmes/plugins/toolsets/runbook/runbook_fetcher.py @@ -6,7 +6,7 @@ StructuredToolResult, Tool, ToolParameter, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetTag, ) @@ -52,7 +52,7 @@ def _invoke( ) logging.error(err_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=err_msg, params=params, ) @@ -96,7 +96,7 @@ def _invoke( """) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=wrapped_content, params=params, ) @@ -104,7 +104,7 @@ def _invoke( err_msg = f"Failed to read runbook {runbook_path}: {str(e)}" logging.error(err_msg) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error=err_msg, params=params, ) diff --git a/holmes/plugins/toolsets/servicenow/servicenow.py b/holmes/plugins/toolsets/servicenow/servicenow.py index c4f594f281..05bdb1741c 100644 --- a/holmes/plugins/toolsets/servicenow/servicenow.py +++ b/holmes/plugins/toolsets/servicenow/servicenow.py @@ -11,7 +11,7 @@ ) from pydantic import BaseModel, PrivateAttr -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.utils import ( process_timestamps_to_rfc3339, standard_start_datetime_tool_param_description, @@ -86,9 +86,9 @@ def return_result( response.raise_for_status() res = response.json() return StructuredToolResult( - status=ToolResultStatus.SUCCESS + status=StructuredToolResultStatus.SUCCESS if res.get(field, []) - else ToolResultStatus.NO_DATA, + else StructuredToolResultStatus.NO_DATA, data=res, params=params, ) @@ -139,7 +139,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=f"Exception {self.name}: {str(e)}", params=params, ) @@ -173,7 +173,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=f"Exception {self.name}: {str(e)}", params=params, ) @@ -213,7 +213,7 @@ def _invoke( except Exception as e: logging.exception(self.get_parameterized_one_liner(params)) return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, data=f"Exception {self.name}: {str(e)}", params=params, ) diff --git a/tests/core/test_safeguards.py b/tests/core/test_safeguards.py index ad45ce6836..b113b406b9 100644 --- a/tests/core/test_safeguards.py +++ b/tests/core/test_safeguards.py @@ -5,7 +5,7 @@ _has_previous_unfiltered_pod_logs_call, ) from holmes.core.tool_calling_llm import ToolCallResult -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus from holmes.plugins.toolsets.logging_utils.logging_api import POD_LOGGING_TOOL_NAME @@ -18,7 +18,7 @@ def test_is_redundant_fetch_pod_logs_with_filter_after_unfiltered_no_data(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "notification-consumer", "namespace": "services", @@ -55,7 +55,7 @@ def test_is_redundant_fetch_pod_logs_with_filter_after_unfiltered_no_data_no_dat tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "notification-consumer", "namespace": "services", @@ -86,7 +86,7 @@ def test_is_redundant_fetch_pod_logs_with_filter_after_unfiltered_has_data(self) tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="foobar", params={ "pod_name": "notification-consumer", @@ -122,7 +122,7 @@ def test_is_redundant_fetch_pod_logs_different_pod_name(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "other-pod", "namespace": "services", @@ -153,7 +153,7 @@ def test_is_redundant_fetch_pod_logs_no_filter_in_current_call(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "notification-consumer", "namespace": "services", @@ -196,7 +196,7 @@ def test_has_previous_unfiltered_pod_logs_call_found(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "my-pod", "namespace": "default", @@ -225,7 +225,7 @@ def test_has_previous_unfiltered_pod_logs_call_with_success_status(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params={ "pod_name": "my-pod", "namespace": "default", @@ -253,7 +253,7 @@ def test_has_previous_unfiltered_pod_logs_call_different_namespace(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "my-pod", "namespace": "other-namespace", @@ -288,7 +288,7 @@ def test_has_previous_exact_same_tool_call_found(self): tool_name="my_tool", description="Test tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params=params, ), ).as_tool_result_response() @@ -306,7 +306,7 @@ def test_has_previous_exact_same_tool_call_different_params(self): tool_name="my_tool", description="Test tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params={"different": "params"}, ), ).as_tool_result_response() @@ -334,7 +334,7 @@ def test_has_previous_exact_same_tool_call_different_tool_name(self): tool_name="different_tool", description="Test tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params=params, ), ).as_tool_result_response() @@ -356,7 +356,7 @@ def test_prevent_overly_repeated_tool_call_exact_duplicate(self): tool_name="my_tool", description="Test tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params=params, ), ).as_tool_result_response() @@ -367,7 +367,7 @@ def test_prevent_overly_repeated_tool_call_exact_duplicate(self): ) assert result is not None - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "already been called" in result.error def test_prevent_overly_repeated_tool_call_redundant_pod_logs(self): @@ -378,7 +378,7 @@ def test_prevent_overly_repeated_tool_call_redundant_pod_logs(self): tool_name=POD_LOGGING_TOOL_NAME, description="Fetch pod logs", result=StructuredToolResult( - status=ToolResultStatus.NO_DATA, + status=StructuredToolResultStatus.NO_DATA, params={ "pod_name": "my-pod", "namespace": "default", @@ -399,7 +399,7 @@ def test_prevent_overly_repeated_tool_call_redundant_pod_logs(self): ) assert result is not None - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert result.error assert "without filter has already run" in result.error @@ -422,7 +422,7 @@ def test_prevent_overly_repeated_tool_call_different_params_allowed(self): tool_name="my_tool", description="Test tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params={"different": "params"}, ), ).as_tool_result_response() @@ -453,7 +453,7 @@ def test_multiple_previous_calls(self): tool_name="other_tool", description="Other tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params={"different": "params"}, ), ).as_tool_result_response(), @@ -462,7 +462,7 @@ def test_multiple_previous_calls(self): tool_name="my_tool", description="My tool", result=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, params={"pod_name": "my-pod"}, ), ).as_tool_result_response(), diff --git a/tests/core/test_todo_write_tool.py b/tests/core/test_todo_write_tool.py index 79b754eff5..55f849f1e6 100644 --- a/tests/core/test_todo_write_tool.py +++ b/tests/core/test_todo_write_tool.py @@ -1,4 +1,4 @@ -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.investigator.core_investigation import TodoWriteTool from holmes.plugins.toolsets.investigator.model import TaskStatus @@ -16,7 +16,7 @@ def test_todo_write_tool_empty_params(self): tool = TodoWriteTool() result = tool._invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert isinstance(result.data, str) assert "0 tasks" in result.data assert "Investigation plan updated" in result.data @@ -43,7 +43,7 @@ def test_todo_write_tool_with_tasks(self): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert isinstance(result.data, str) assert "2 tasks" in result.data assert "Investigation plan updated" in result.data @@ -58,7 +58,7 @@ def test_todo_write_tool_default_values(self): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert isinstance(result.data, str) assert "1 tasks" in result.data assert "Investigation plan updated" in result.data @@ -81,7 +81,7 @@ def test_todo_write_tool_invalid_enum_values(self): result = tool._invoke(params) # Should handle gracefully and return error - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "Failed to process tasks" in result.error def test_get_parameterized_one_liner(self): diff --git a/tests/core/test_tool_transformers.py b/tests/core/test_tool_transformers.py index b82f82a17f..f8fc7466f4 100644 --- a/tests/core/test_tool_transformers.py +++ b/tests/core/test_tool_transformers.py @@ -11,7 +11,7 @@ YAMLToolset, ToolsetYamlFromConfig, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.core.transformers import ( registry, @@ -204,7 +204,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output", ) @@ -235,7 +235,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: # Tool should still work despite invalid transformer result = tool.invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS def test_yaml_tool_validation(self): """Test YAMLTool transformer validation.""" @@ -323,7 +323,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="This is a long test output that should be transformed", ) @@ -337,7 +337,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should have applied transformation - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "mock_transformed:" in result.data assert "This is a long test output that should be transformed" in result.data @@ -351,7 +351,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.ERROR, + status=StructuredToolResultStatus.ERROR, error="Tool execution failed", data="Some error output", ) @@ -366,7 +366,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should not have applied transformation due to error status - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "mock_transformed:" not in str(result.data) assert result.data == "Some error output" @@ -378,7 +378,9 @@ class ConcreteTestTool(Tool): def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: - return StructuredToolResult(status=ToolResultStatus.NO_DATA, data="") + return StructuredToolResult( + status=StructuredToolResultStatus.NO_DATA, data="" + ) def get_parameterized_one_liner(self, params: Dict) -> str: return "test command" @@ -390,7 +392,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should not have applied transformation due to empty data - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert result.data == "" def test_transformer_failure_handling(self): @@ -418,7 +420,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data="Test output" + status=StructuredToolResultStatus.SUCCESS, data="Test output" ) def get_parameterized_one_liner(self, params: Dict) -> str: @@ -434,7 +436,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should return original data when transformer fails - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "Test output" # Should log warning about transformer failure @@ -474,7 +476,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Original text that should be transformed twice", ) @@ -490,7 +492,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should have applied both transformations in sequence - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "second_transformed:" in result.data assert "mock_transformed:" in result.data @@ -525,7 +527,7 @@ def _invoke( ) -> StructuredToolResult: data = params.get("data", "short") return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data=data + status=StructuredToolResultStatus.SUCCESS, data=data ) def get_parameterized_one_liner(self, params: Dict) -> str: @@ -565,7 +567,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output to transform", url="http://example.com", invocation="test command", @@ -587,7 +589,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: assert "mock_transformed:" in result.data # Other fields should be preserved - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.url == "http://example.com" assert result.invocation == "test command" assert result.params == {"test": "param"} @@ -602,7 +604,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output that will be transformed for performance measurement", ) @@ -636,7 +638,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output without transformation", ) @@ -652,7 +654,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should return unchanged result - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "Test output without transformation" @@ -680,7 +682,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output for transformation", ) @@ -711,7 +713,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output for transformation", ) @@ -748,7 +750,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output without transformers", ) @@ -767,7 +769,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: # Verify tool still works correctly result = tool.invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "Test output without transformers" def test_transformer_initialization_failure_handling(self): @@ -784,7 +786,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output for transformation", ) @@ -813,7 +815,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: # Verify tool still works with the successful transformer result = tool.invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "mock_transformed:" in result.data def test_transformer_initialization_empty_transformers_list(self): @@ -825,7 +827,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output", ) @@ -841,7 +843,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: # Tool should still work correctly result = tool.invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "Test output" def test_performance_optimization_prevents_recreation(self): @@ -853,7 +855,7 @@ def _invoke( self, params: Dict, user_approved: bool = False ) -> StructuredToolResult: return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output for transformation", ) diff --git a/tests/core/tools_utils/test_tool_context_window_limiter.py b/tests/core/tools_utils/test_tool_context_window_limiter.py new file mode 100644 index 0000000000..d76b7f18b4 --- /dev/null +++ b/tests/core/tools_utils/test_tool_context_window_limiter.py @@ -0,0 +1,234 @@ +from unittest.mock import Mock, patch + +import pytest + +from holmes.core.llm import LLM +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus +from holmes.core.tools_utils.data_types import ToolCallResult +from holmes.core.tools_utils.tool_context_window_limiter import ( + prevent_overly_big_tool_response, +) + + +class TestPreventOverlyBigToolResponse: + @pytest.fixture + def mock_llm(self): + """Create a mock LLM instance.""" + llm = Mock(spec=LLM) + llm.get_context_window_size.return_value = 4096 + llm.count_tokens_for_message.return_value = 1000 + return llm + + @pytest.fixture + def success_tool_call_result(self): + """Create a successful tool call result.""" + result = StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, + data="Some successful output data", + ) + return ToolCallResult( + tool_call_id="test-id-1", + tool_name="test_tool", + description="Test tool description", + result=result, + ) + + def test_no_limit_configured(self, mock_llm, success_tool_call_result): + """Test that function does nothing when TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT is 0.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 0, + ): + original_status = success_tool_call_result.result.status + original_data = success_tool_call_result.result.data + original_error = success_tool_call_result.result.error + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should remain unchanged + assert success_tool_call_result.result.status == original_status + assert success_tool_call_result.result.data == original_data + assert success_tool_call_result.result.error == original_error + mock_llm.count_tokens_for_message.assert_not_called() + + def test_negative_limit_configured(self, mock_llm, success_tool_call_result): + """Test that function does nothing when TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT is negative.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + -10, + ): + original_status = success_tool_call_result.result.status + original_data = success_tool_call_result.result.data + original_error = success_tool_call_result.result.error + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should remain unchanged + assert success_tool_call_result.result.status == original_status + assert success_tool_call_result.result.data == original_data + assert success_tool_call_result.result.error == original_error + mock_llm.count_tokens_for_message.assert_not_called() + + def test_over_100_percent_limit(self, mock_llm, success_tool_call_result): + """Test that function does nothing when TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT is over 100.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 150, + ): + original_status = success_tool_call_result.result.status + original_data = success_tool_call_result.result.data + original_error = success_tool_call_result.result.error + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should remain unchanged + assert success_tool_call_result.result.status == original_status + assert success_tool_call_result.result.data == original_data + assert success_tool_call_result.result.error == original_error + mock_llm.count_tokens_for_message.assert_not_called() + + def test_within_token_limit(self, mock_llm, success_tool_call_result): + """Test that function does nothing when tool result is within token limit.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 50, + ): + # Context window: 4096, 50% = 2048 tokens allowed + # Token count: 1000 (within limit) + mock_llm.count_tokens_for_message.return_value = 1000 + + original_status = success_tool_call_result.result.status + original_data = success_tool_call_result.result.data + original_error = success_tool_call_result.result.error + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should remain unchanged + assert success_tool_call_result.result.status == original_status + assert success_tool_call_result.result.data == original_data + assert success_tool_call_result.result.error == original_error + mock_llm.count_tokens_for_message.assert_called_once() + + def test_exceeds_token_limit(self, mock_llm, success_tool_call_result): + """Test that function modifies result when tool result exceeds token limit.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 50, + ): + # Context window: 4096, 50% = 2048 tokens allowed + # Token count: 3000 (exceeds limit) + mock_llm.count_tokens_for_message.return_value = 3000 + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should be modified + assert ( + success_tool_call_result.result.status + == StructuredToolResultStatus.ERROR + ) + assert success_tool_call_result.result.data is None + assert "too large to return" in success_tool_call_result.result.error + assert "3000 tokens" in success_tool_call_result.result.error + assert "2048.0" in success_tool_call_result.result.error + assert ( + "31.7" in success_tool_call_result.result.error + ) # (3000-2048)/3000 * 100 + mock_llm.count_tokens_for_message.assert_called_once() + + def test_token_calculation_accuracy(self, mock_llm, success_tool_call_result): + """Test that token calculations are accurate.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 25, + ): + # Context window: 4096, 25% = 1024 tokens allowed + # Token count: 2000 (exceeds limit) + mock_llm.count_tokens_for_message.return_value = 2000 + mock_llm.get_context_window_size.return_value = 4096 + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Calculate expected percentage: (2000-1024)/2000 * 100 = 48.8% + assert "48.8" in success_tool_call_result.result.error + assert "2000 tokens" in success_tool_call_result.result.error + assert "1024.0" in success_tool_call_result.result.error + + def test_message_construction_calls_as_tool_call_message( + self, mock_llm, success_tool_call_result + ): + """Test that the function calls as_tool_call_message to get the message for token counting.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 50, + ): + mock_llm.count_tokens_for_message.return_value = 1000 # Within limit + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Verify that count_tokens_for_message was called with a list containing one message + mock_llm.count_tokens_for_message.assert_called_once() + call_args = mock_llm.count_tokens_for_message.call_args + assert ( + len(call_args[1]["messages"]) == 1 + ) # Should be called with messages kwarg containing 1 message + + def test_different_context_window_sizes(self, mock_llm, success_tool_call_result): + """Test with different context window sizes.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 40, + ): + # Test with smaller context window + mock_llm.get_context_window_size.return_value = 2048 + mock_llm.count_tokens_for_message.return_value = 1000 + # 40% of 2048 = 819.2 tokens allowed, 1000 exceeds this + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + assert ( + success_tool_call_result.result.status + == StructuredToolResultStatus.ERROR + ) + assert "1000 tokens" in success_tool_call_result.result.error + assert "819.2" in success_tool_call_result.result.error + + def test_edge_case_exactly_at_limit(self, mock_llm, success_tool_call_result): + """Test behavior when token count is exactly at the limit.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 50, + ): + mock_llm.get_context_window_size.return_value = 4096 + mock_llm.count_tokens_for_message.return_value = 2048 # Exactly 50% of 4096 + + original_status = success_tool_call_result.result.status + original_data = success_tool_call_result.result.data + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + # Should remain unchanged (not > max_tokens_allowed) + assert success_tool_call_result.result.status == original_status + assert success_tool_call_result.result.data == original_data + + def test_error_message_format(self, mock_llm, success_tool_call_result): + """Test that error message contains all expected components.""" + with patch( + "holmes.core.tools_utils.tool_context_window_limiter.TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", + 20, + ): + mock_llm.get_context_window_size.return_value = 5000 + mock_llm.count_tokens_for_message.return_value = 2000 + # 20% of 5000 = 1000 tokens allowed + + prevent_overly_big_tool_response(success_tool_call_result, mock_llm) + + error_msg = success_tool_call_result.result.error + assert "The tool call result is too large to return" in error_msg + assert "2000 tokens" in error_msg + assert "1000.0" in error_msg + assert "Instructions for the LLM" in error_msg + assert "try to repeat the query" in error_msg + assert "narrow down the result" in error_msg + + # Check percentage calculation: (2000-1000)/2000 * 100 = 50.0% + assert "50.0" in error_msg diff --git a/tests/core/transformers/test_llm_summarize.py b/tests/core/transformers/test_llm_summarize.py index e803bb32fa..9273a650a9 100644 --- a/tests/core/transformers/test_llm_summarize.py +++ b/tests/core/transformers/test_llm_summarize.py @@ -519,7 +519,11 @@ def test_non_expanding_behavior_integration(self, mock_default_llm): @patch("holmes.core.transformers.llm_summarize.DefaultLLM") def test_integration_with_tools_non_expanding_logic(self, mock_default_llm): """Test integration with the tools.py non-expanding logic.""" - from holmes.core.tools import Tool, StructuredToolResult, ToolResultStatus + from holmes.core.tools import ( + Tool, + StructuredToolResult, + StructuredToolResultStatus, + ) from holmes.core.transformers import Transformer # Create a mock expanding transformer response @@ -531,7 +535,8 @@ def test_integration_with_tools_non_expanding_logic(self, mock_default_llm): class TestTool(Tool): def _invoke(self, params, user_approved: bool = False): return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data="Original short data" + status=StructuredToolResultStatus.SUCCESS, + data="Original short data", ) def get_parameterized_one_liner(self, params): @@ -560,7 +565,11 @@ def get_parameterized_one_liner(self, params): @patch("holmes.core.transformers.llm_summarize.DefaultLLM") def test_non_expanding_logic_with_debug_logging(self, mock_default_llm): """Test that explicitly verifies the non-expanding logic is triggered.""" - from holmes.core.tools import Tool, StructuredToolResult, ToolResultStatus + from holmes.core.tools import ( + Tool, + StructuredToolResult, + StructuredToolResultStatus, + ) from holmes.core.transformers import Transformer # Create a mock expanding transformer response that's much larger than original @@ -574,7 +583,7 @@ def test_non_expanding_logic_with_debug_logging(self, mock_default_llm): class TestTool(Tool): def _invoke(self, params, user_approved: bool = False): return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=original_data, # This will be the input to the transformer ) @@ -624,7 +633,11 @@ def test_non_expanding_vs_successful_summarization_comparison( self, mock_default_llm ): """Test comparing expanding vs reducing scenarios to ensure logic works correctly.""" - from holmes.core.tools import Tool, StructuredToolResult, ToolResultStatus + from holmes.core.tools import ( + Tool, + StructuredToolResult, + StructuredToolResultStatus, + ) from holmes.core.transformers import Transformer original_data = "This is some longer original data that should be summarized by the transformer if it works properly and produces a smaller result" @@ -633,7 +646,7 @@ def test_non_expanding_vs_successful_summarization_comparison( class TestTool(Tool): def _invoke(self, params, user_approved: bool = False): return StructuredToolResult( - status=ToolResultStatus.SUCCESS, data=original_data + status=StructuredToolResultStatus.SUCCESS, data=original_data ) def get_parameterized_one_liner(self, params): diff --git a/tests/integration/test_kubernetes_transformer_execution.py b/tests/integration/test_kubernetes_transformer_execution.py index fb897f5657..f22892e7f4 100644 --- a/tests/integration/test_kubernetes_transformer_execution.py +++ b/tests/integration/test_kubernetes_transformer_execution.py @@ -7,7 +7,7 @@ from unittest.mock import patch from holmes.plugins.toolsets import load_toolsets_from_file -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.core.transformers import registry from holmes.core.transformers.base import BaseTransformer @@ -165,7 +165,7 @@ def test_kubectl_describe_with_large_output(self): ) # Should have applied transformation - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "SUMMARIZED:" in result.data assert f"Original length: {len(large_output)}" in result.data @@ -202,7 +202,7 @@ def test_kubectl_describe_with_small_output(self): ) # Should NOT have applied transformation - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == small_output assert "SUMMARIZED:" not in result.data @@ -259,7 +259,7 @@ def test_kubectl_logs_with_large_output(self): ) # Should have applied transformation - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "SUMMARIZED:" in result.data assert f"Original length: {len(large_log_output)}" in result.data @@ -309,7 +309,7 @@ def test_kubectl_get_cluster_with_transformer(self): result = kubectl_get_cluster.invoke({"kind": "pods"}) # Should have applied transformation - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "SUMMARIZED:" in result.data assert len(result.data) < len(large_get_output) @@ -372,7 +372,7 @@ def name(self) -> str: result = tool.invoke({}) # Should return original output when transformer fails - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == test_output # Should log error about transformer failure (generic Exception -> error log) @@ -416,7 +416,7 @@ def test_transformer_error_status_handling(self): ) # Should NOT have applied transformation due to error status - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert result.data == error_output assert "SUMMARIZED:" not in result.data @@ -485,7 +485,7 @@ def name(self) -> str: result = tool.invoke({}) # Should have applied both transformers in sequence - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None assert "SECOND_TRANSFORM:" in result.data assert "SUMMARIZED:" in result.data diff --git a/tests/integration/test_tool_execution_pipeline.py b/tests/integration/test_tool_execution_pipeline.py index 8eb9be8f51..956dc19d2c 100644 --- a/tests/integration/test_tool_execution_pipeline.py +++ b/tests/integration/test_tool_execution_pipeline.py @@ -13,7 +13,7 @@ Tool, YAMLTool, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ) from holmes.core.transformers import ( registry, @@ -88,7 +88,7 @@ def test_yaml_tool_with_transformer_integration(self): result = tool.invoke({}) # Should have executed successfully - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should have applied transformation (output should be summarized) assert result.data is not None @@ -121,7 +121,7 @@ def _invoke( large_log = "\n".join(log_entries) return StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data=large_log, invocation="tail -n 1000 /var/log/app.log", ) @@ -140,7 +140,7 @@ def get_parameterized_one_liner(self, params: Dict) -> str: result = tool.invoke({}) # Should have executed successfully - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should have applied transformation assert result.data is not None @@ -193,7 +193,7 @@ def name(self) -> str: result = tool.invoke({}) # Tool execution should still succeed - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should contain original output since first transformer failed # but second transformer should still be applied @@ -223,7 +223,7 @@ def test_conditional_transformer_application_integration(self): result = tool.invoke({}) # Should execute successfully - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should NOT have applied transformation (output too short) assert "SUMMARIZED:" not in result.data @@ -258,7 +258,7 @@ def name(self) -> str: result = tool.invoke({}) # Should execute successfully - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should have applied transformation assert "SLOW_TRANSFORMED:" in result.data @@ -312,7 +312,7 @@ def test_real_world_kubectl_scenario_integration(self): result = tool.invoke({}) # Should execute successfully - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Should have applied transformation (output should be summarized) assert "SUMMARIZED:" in result.data @@ -334,7 +334,7 @@ def test_error_handling_preserves_debugging_info_integration(self): # Even if transformer modifies the output, the result structure should be preserved result = tool.invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.invocation is not None # Debug info preserved assert "echo" in result.invocation # Original command preserved diff --git a/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/test_case.yaml b/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/test_case.yaml index 27b2f31548..48cc1c0ed0 100644 --- a/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/test_case.yaml +++ b/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/test_case.yaml @@ -15,7 +15,7 @@ tags: port_forwards: - namespace: app-115 service: tempo - local_port: 3200 + local_port: 3201 remote_port: 3200 before_test: | @@ -32,7 +32,7 @@ before_test: | echo "⏰ Waiting for Tempo to be fully ready (checking every 5s, timeout 60s)" TEMPO_READY=false for i in {1..12}; do - if kubectl exec -n app-115 deployment/tempo -- wget -q -O - http://localhost:3200/ready 2>/dev/null; then + if kubectl exec -n app-115 deployment/tempo -- wget -q -O - http://localhost:3201/ready 2>/dev/null; then echo "✅ Tempo is ready!" TEMPO_READY=true break diff --git a/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/toolsets.yaml b/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/toolsets.yaml index 7298e06876..671123486c 100644 --- a/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/toolsets.yaml +++ b/tests/llm/fixtures/test_ask_holmes/115_checkout_errors_tracing/toolsets.yaml @@ -6,5 +6,5 @@ toolsets: grafana/tempo: enabled: true config: - url: http://localhost:3200 + url: http://localhost:3201 healthcheck: "ready" diff --git a/tests/llm/utils/test_mock_toolset.py b/tests/llm/utils/test_mock_toolset.py index eba3570d8b..e8d2985d67 100644 --- a/tests/llm/utils/test_mock_toolset.py +++ b/tests/llm/utils/test_mock_toolset.py @@ -8,7 +8,7 @@ from holmes.core.tools import ( Tool, StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, ToolsetStatusEnum, ) from tests.llm.utils.mock_toolset import ( @@ -157,7 +157,7 @@ def test_write_and_read_mock(self): # Create test result result = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Test output data", metadata={"key": "value"}, ) @@ -178,7 +178,7 @@ def test_write_and_read_mock(self): assert mock.tool_name == "test_tool" assert mock.toolset_name == "test_toolset" assert mock.return_value.data == "Test output data" - assert mock.return_value.status == ToolResultStatus.SUCCESS + assert mock.return_value.status == StructuredToolResultStatus.SUCCESS def test_read_nonexistent_mock(self): """Test reading a mock that doesn't exist.""" @@ -224,7 +224,7 @@ def create_mock_tool(self): tool.user_description = None tool.invoke = Mock( return_value=StructuredToolResult( - status=ToolResultStatus.SUCCESS, data="Real tool output" + status=StructuredToolResultStatus.SUCCESS, data="Real tool output" ) ) tool.get_parameterized_one_liner = Mock(return_value="test_tool()") @@ -260,7 +260,7 @@ def test_mock_mode_with_existing_mock(self): # First, write a mock file mock_result = StructuredToolResult( - status=ToolResultStatus.SUCCESS, data="Mocked output" + status=StructuredToolResultStatus.SUCCESS, data="Mocked output" ) file_manager.write_mock( tool_name="test_tool", @@ -436,7 +436,7 @@ def test_mock_tools_exact_match(self): file_manager = MockFileManager(tmpdir, add_params_to_filename=True) params = {"field1": "1", "field2": "2"} mock_result = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="this tool is mocked", params=params, ) @@ -492,7 +492,7 @@ def test_mock_tools_without_params_in_filename(self, params): # Set up file-based mock WITHOUT params in filename file_manager = MockFileManager(tmpdir, add_params_to_filename=False) mock_result = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="this tool is mocked", params={}, # Will be ignored when matching ) @@ -551,7 +551,7 @@ def test_mock_tools_do_not_match(self, params): # Set up file-based mock with specific params file_manager = MockFileManager(tmpdir, add_params_to_filename=True) mock_result = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="this tool is mocked", params={"field1": "1", "field2": "2"}, ) @@ -623,7 +623,7 @@ def test_mock_tools_generate_mode_does_not_throw(self): original_invoke = kubectl_tool._tool._invoke kubectl_tool._tool._invoke = Mock( return_value=StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="Generated output from mock", ) ) @@ -635,7 +635,7 @@ def test_mock_tools_generate_mode_does_not_throw(self): ) # Should have called the mocked tool and saved the result - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "Generated output" in result.data # Verify mock was generated diff --git a/tests/mocks/toolset_mocks.py b/tests/mocks/toolset_mocks.py index 29faf7c38b..4d4f208852 100644 --- a/tests/mocks/toolset_mocks.py +++ b/tests/mocks/toolset_mocks.py @@ -1,6 +1,11 @@ from typing import Any, Dict, List -from holmes.core.tools import Tool, Toolset, StructuredToolResult, ToolResultStatus +from holmes.core.tools import ( + Tool, + Toolset, + StructuredToolResult, + StructuredToolResultStatus, +) class DummyTool(Tool): @@ -10,7 +15,7 @@ class DummyTool(Tool): def _invoke( self, params: dict, user_approved: bool = False ) -> StructuredToolResult: - return StructuredToolResult(status=ToolResultStatus.SUCCESS) + return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS) def get_parameterized_one_liner(self, params: Dict) -> str: return "" diff --git a/tests/plugins/prompt/test_toolsets_instructions.py b/tests/plugins/prompt/test_toolsets_instructions.py index 9c0d31e5f2..12b966a759 100644 --- a/tests/plugins/prompt/test_toolsets_instructions.py +++ b/tests/plugins/prompt/test_toolsets_instructions.py @@ -4,7 +4,7 @@ StaticPrerequisite, StructuredToolResult, Tool, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ) @@ -18,7 +18,7 @@ def __init__(self): def _invoke( self, params: dict, user_approved: bool = False ) -> StructuredToolResult: - return StructuredToolResult(status=ToolResultStatus.SUCCESS, data="") + return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data="") def get_parameterized_one_liner(self, params: Dict) -> str: return "" diff --git a/tests/plugins/toolsets/coralogix/test_coralogix_integration.py b/tests/plugins/toolsets/coralogix/test_coralogix_integration.py index 9a89fe749f..5dc932a6b3 100644 --- a/tests/plugins/toolsets/coralogix/test_coralogix_integration.py +++ b/tests/plugins/toolsets/coralogix/test_coralogix_integration.py @@ -6,7 +6,7 @@ import os import pytest -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.coralogix.api import health_check from holmes.plugins.toolsets.coralogix.toolset_coralogix_logs import ( CoralogixLogsToolset, @@ -93,7 +93,7 @@ def test_basic_query(coralogix_logs_toolset): FetchPodLogsParams(namespace=TEST_NAMESPACE, pod_name=TEST_POD_NAME) ) print(result.data) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error assert TEST_SEARCH_TERM in result.data @@ -106,7 +106,7 @@ def test_search_term(coralogix_logs_toolset): ) print(result.data) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error lines = result.data.split("\n")[2:] # skips headers lines for "link" and "query" # print(lines) @@ -125,7 +125,7 @@ def test_search_term_with_dates(coralogix_logs_toolset): ) ) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error print(result.data) lines = result.data.split("\n") diff --git a/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py b/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py index 350d1aebe5..5d37db33fc 100644 --- a/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py +++ b/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py @@ -1,5 +1,5 @@ from unittest.mock import Mock, patch -from holmes.core.tools import ToolResultStatus, ToolsetStatusEnum +from holmes.core.tools import StructuredToolResultStatus, ToolsetStatusEnum from holmes.plugins.toolsets.datadog.toolset_datadog_logs import ( DatadogLogsToolset, DataDogStorageTier, @@ -61,7 +61,7 @@ def test_check_prerequisites_successful_healthcheck(self, mock_fetch_pod_logs): """Test check_prerequisites with successful healthcheck""" # Mock successful healthcheck response mock_result = Mock() - mock_result.status = ToolResultStatus.SUCCESS + mock_result.status = StructuredToolResultStatus.SUCCESS mock_result.error = None mock_fetch_pod_logs.return_value = mock_result @@ -97,7 +97,7 @@ def test_check_prerequisites_healthcheck_error(self, mock_fetch_pod_logs): """Test check_prerequisites with healthcheck returning error""" # Mock healthcheck error response mock_result = Mock() - mock_result.status = ToolResultStatus.ERROR + mock_result.status = StructuredToolResultStatus.ERROR mock_result.error = "Authentication failed" mock_fetch_pod_logs.return_value = mock_result @@ -117,7 +117,7 @@ def test_check_prerequisites_healthcheck_no_data(self, mock_fetch_pod_logs): """Test check_prerequisites with healthcheck returning no data""" # Mock healthcheck no data response mock_result = Mock() - mock_result.status = ToolResultStatus.NO_DATA + mock_result.status = StructuredToolResultStatus.NO_DATA mock_result.error = None mock_fetch_pod_logs.return_value = mock_result @@ -157,7 +157,7 @@ def test_check_prerequisites_with_custom_config(self, mock_fetch_pod_logs): """Test check_prerequisites with custom configuration""" # Mock successful healthcheck response mock_result = Mock() - mock_result.status = ToolResultStatus.SUCCESS + mock_result.status = StructuredToolResultStatus.SUCCESS mock_result.error = None mock_fetch_pod_logs.return_value = mock_result diff --git a/tests/plugins/toolsets/datadog/logs/test_fetch_pod_logs.py b/tests/plugins/toolsets/datadog/logs/test_fetch_pod_logs.py index e76bbed7e7..695770b43b 100644 --- a/tests/plugins/toolsets/datadog/logs/test_fetch_pod_logs.py +++ b/tests/plugins/toolsets/datadog/logs/test_fetch_pod_logs.py @@ -1,5 +1,5 @@ from unittest.mock import Mock, patch -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.datadog.toolset_datadog_logs import ( DatadogLogsToolset, DatadogLogsConfig, @@ -75,7 +75,7 @@ def test_fetch_pod_logs_with_pagination(self, mock_post): result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None # Check that all 300 logs are present in reverse order (oldest first) @@ -126,7 +126,7 @@ def test_fetch_pod_logs_less_data_than_requested(self, mock_post): result = self.toolset.fetch_pod_logs(params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None # Check that only 80 logs are returned @@ -177,7 +177,7 @@ def test_fetch_pod_logs_storage_tier_fallback(self, mock_post): result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None # Check that logs from online archives are returned @@ -230,7 +230,7 @@ def test_fetch_pod_logs_rate_limiting(self, mock_sleep, mock_post): result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None # Check that logs are returned @@ -277,7 +277,7 @@ def test_fetch_pod_logs_rate_limiting_without_reset_header( result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None # Verify three API calls were made (two failed with 429, one succeeded) @@ -293,7 +293,7 @@ def test_fetch_pod_logs_no_config(self, mock_post): result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert result.data == "The toolset is missing its configuration" # Verify no API calls were made @@ -316,7 +316,7 @@ def test_fetch_pod_logs_rate_limit_exhausted(self, mock_sleep, mock_post): result = self.toolset.fetch_pod_logs(self.fetch_params) # Verify - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert ( "Datadog API rate limit exceeded. Failed after 5 retry attempts." in result.error @@ -353,7 +353,7 @@ def test_fetch_pod_logs_with_filter(self, mock_post): result = self.toolset.fetch_pod_logs(params) # Verify - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Check the API call included the filter call_args = mock_post.call_args diff --git a/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py b/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py index 925b26d34f..551dc907c6 100644 --- a/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py +++ b/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py @@ -1,6 +1,6 @@ import json from unittest.mock import Mock, patch -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.datadog.toolset_datadog_metrics import ( DatadogMetricsToolset, DatadogMetricsConfig, @@ -39,7 +39,7 @@ def test_list_active_metrics(self, mock_get): tool = self.toolset.tools[0] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "system.cpu.user" in result.data assert "system.mem.used" in result.data assert "Metric Name" in result.data @@ -69,7 +69,7 @@ def test_list_active_metrics_with_filters(self, mock_get): tool = self.toolset.tools[0] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS call_args = mock_get.call_args assert call_args[1]["params"]["host"] == "test-host" @@ -104,7 +104,7 @@ def test_query_metrics(self, mock_get): tool = self.toolset.tools[1] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "system.cpu.user" in result.data call_args = mock_get.call_args @@ -124,7 +124,7 @@ def test_query_metrics_no_data(self, mock_get): tool = self.toolset.tools[1] result = tool._invoke(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert "no data" in result.error.lower() @patch("holmes.plugins.toolsets.datadog.datadog_api.requests.get") @@ -145,7 +145,7 @@ def test_get_metric_metadata(self, mock_get): tool = self.toolset.tools[2] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS data = json.loads(result.data) assert "metrics_metadata" in data assert "system.cpu.user" in data["metrics_metadata"] @@ -166,7 +166,7 @@ def test_get_metric_metadata_not_found(self, mock_get): tool = self.toolset.tools[2] result = tool._invoke(params) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR data = json.loads(result.data) assert "errors" in data assert "nonexistent.metric" in data["errors"] @@ -196,7 +196,7 @@ def test_get_multiple_metrics_metadata(self, mock_get): tool = self.toolset.tools[2] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS data = json.loads(result.data) assert "metrics_metadata" in data assert "system.cpu.user" in data["metrics_metadata"] @@ -216,7 +216,7 @@ def test_get_multiple_metrics_metadata_partial_failure(self, mock_get): tool = self.toolset.tools[2] result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS data = json.loads(result.data) assert "system.cpu.user" in data["metrics_metadata"] assert "nonexistent.metric" in data["errors"] @@ -230,7 +230,7 @@ def test_no_config(self): tool = self.toolset.tools[0] result = tool._invoke(params) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert result.error == "The toolset is missing its configuration" @patch("holmes.plugins.toolsets.datadog.datadog_api.requests.get") @@ -246,7 +246,7 @@ def test_rate_limiting(self, mock_sleep, mock_get): tool = self.toolset.tools[0] result = tool._invoke(params) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "rate limit exceeded" in result.error.lower() assert "5 retry attempts" in result.error diff --git a/tests/plugins/toolsets/datadog/rds/test_datadog_rds_integration.py b/tests/plugins/toolsets/datadog/rds/test_datadog_rds_integration.py index 97c85c1feb..5a8d172d0d 100644 --- a/tests/plugins/toolsets/datadog/rds/test_datadog_rds_integration.py +++ b/tests/plugins/toolsets/datadog/rds/test_datadog_rds_integration.py @@ -15,7 +15,7 @@ DatadogRDSConfig, ) from holmes.plugins.toolsets.datadog.datadog_api import DataDogRequestError -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus @pytest.fixture @@ -142,7 +142,7 @@ def test_generate_performance_report_success(mock_execute, datadog_rds_toolset): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert isinstance(result.data, str) # Check report content @@ -226,7 +226,7 @@ def test_generate_performance_report_with_issues(mock_execute, datadog_rds_tools {"db_instance_identifier": "test-instance", "start_time": "-3600"} ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS report = result.data assert "ISSUES DETECTED" in report @@ -318,7 +318,7 @@ def mock_response(url, **kwargs): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert isinstance(result.data, str) # Check report format @@ -395,7 +395,7 @@ def mock_response(url, **kwargs): ) result = tool._invoke({"top_n": 10, "start_time": "-3600", "sort_by": "cpu"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS report = result.data # When sorted by CPU, instance-1 (95%) should appear before instance-2 (40%) @@ -430,7 +430,7 @@ def test_no_instances_found(mock_execute, datadog_rds_toolset): ) result = tool._invoke({"top_n": 5, "start_time": "-3600"}) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert "No RDS instances found" in result.data @@ -457,7 +457,7 @@ def test_api_error_handling(mock_execute, datadog_rds_toolset): ) # The tool should succeed but with no metrics collected due to API errors - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS report = result.data assert ( "Database is operating within normal parameters. No significant issues detected." @@ -481,7 +481,7 @@ def test_missing_required_parameter(datadog_rds_toolset): # Missing db_instance_identifier result = tool._invoke({"start_time": "-3600"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "db_instance_identifier" in result.error @@ -575,7 +575,7 @@ def test_performance_report_formatting(datadog_rds_toolset): {"db_instance_identifier": "test-db", "start_time": "-300"} ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS report = result.data # Check report structure @@ -621,7 +621,7 @@ def mock_response(url, **kwargs): ) result = tool._invoke({"top_n": 5, "start_time": "-3600"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS report = result.data # Should show 0 instances analyzed (no metrics found) diff --git a/tests/plugins/toolsets/datadog/rds/test_datadog_rds_live.py b/tests/plugins/toolsets/datadog/rds/test_datadog_rds_live.py index 3beb52c584..ea2db64307 100644 --- a/tests/plugins/toolsets/datadog/rds/test_datadog_rds_live.py +++ b/tests/plugins/toolsets/datadog/rds/test_datadog_rds_live.py @@ -13,7 +13,7 @@ import pytest from holmes.plugins.toolsets.datadog.toolset_datadog_rds import DatadogRDSToolset -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus @pytest.fixture @@ -70,7 +70,7 @@ def test_generate_performance_report(datadog_rds_toolset, test_rds_instance): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert test_rds_instance in result.data @@ -89,7 +89,7 @@ def test_get_top_worst_performing_instances(datadog_rds_toolset, test_rds_instan result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert test_rds_instance in result.data diff --git a/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py b/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py index 02a7814a15..33b6c20bcd 100644 --- a/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py +++ b/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py @@ -2,7 +2,7 @@ from unittest.mock import Mock, patch -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.datadog.toolset_datadog_general import ( DatadogGeneralToolset, is_endpoint_allowed, @@ -128,19 +128,19 @@ def test_list_api_resources_tool(self): # Test listing all resources result = list_tool._invoke({"category": "all"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "monitors" in result.data.lower() assert "dashboards" in result.data.lower() # Test filtering by category result = list_tool._invoke({"category": "monitors"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "monitor" in result.data.lower() assert "GET /api/v1/monitor" in result.data # Test invalid category result = list_tool._invoke({"category": "invalid_category"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "Unknown category" in result.error @patch( @@ -170,7 +170,7 @@ def test_api_get_tool(self, mock_headers, mock_execute): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "test_response" in result.data # Test blocked endpoint @@ -178,7 +178,7 @@ def test_api_get_tool(self, mock_headers, mock_execute): {"endpoint": "/api/v1/monitor/create", "description": "Create monitor"} ) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "blacklisted operation" in result.error def test_example_config(self): diff --git a/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py b/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py index 12c4424ab4..12d2f159df 100644 --- a/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py +++ b/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py @@ -6,7 +6,7 @@ FetchDatadogSpansByFilter, ) from holmes.plugins.toolsets.datadog.datadog_api import DataDogRequestError -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus class TestDatadogTracesToolset: @@ -154,7 +154,7 @@ def test_invoke_success_with_traces(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "Found 1 traces" in result.data assert "traceID=abc123" in result.data assert "durationMs=50.00" in result.data # 50ms total duration @@ -176,7 +176,7 @@ def test_invoke_no_traces_found(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert "No matching traces found" in result.data def test_invoke_no_config(self): @@ -185,7 +185,7 @@ def test_invoke_no_config(self): result = self.tool._invoke({}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "Datadog configuration not initialized" in result.error @patch( @@ -231,7 +231,7 @@ def test_invoke_rate_limit_error(self, mock_execute): result = self.tool._invoke({}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "rate limit exceeded" in result.error @@ -256,7 +256,7 @@ def test_invoke_missing_trace_id(self): """Test invocation without trace_id parameter.""" result = self.tool._invoke({}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "trace_id parameter is required" in result.error @patch( @@ -301,7 +301,7 @@ def test_invoke_success_with_spans(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "Trace ID: abc123" in result.data assert "GET /users (web-api)" in result.data assert "SELECT (database)" in result.data @@ -319,7 +319,7 @@ def test_invoke_no_trace_found(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert "No trace found for trace_id: nonexistent" in result.data @@ -373,7 +373,7 @@ def test_invoke_with_custom_query(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "Found 1 matching spans" in result.data assert "Trace ID: trace1" in result.data assert "GET /users (web-api)" in result.data @@ -411,5 +411,5 @@ def test_invoke_no_spans_found(self, mock_execute): result = self.tool._invoke(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA assert "No matching spans found" in result.data diff --git a/tests/plugins/toolsets/grafana/test_grafana_loki.py b/tests/plugins/toolsets/grafana/test_grafana_loki.py index 7fd23ae238..be6ca959e7 100644 --- a/tests/plugins/toolsets/grafana/test_grafana_loki.py +++ b/tests/plugins/toolsets/grafana/test_grafana_loki.py @@ -1,6 +1,6 @@ import os from typing import Any -from holmes.core.tools import ToolResultStatus, ToolsetStatusEnum +from holmes.core.tools import StructuredToolResultStatus, ToolsetStatusEnum from holmes.plugins.toolsets.grafana.grafana_api import grafana_health_check import pytest @@ -72,7 +72,7 @@ def test_basic_query(loki_toolset): FetchPodLogsParams(namespace=TEST_NAMESPACE, pod_name=TEST_POD_NAME) ) print(result.data) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error assert TEST_SEARCH_TERM in result.data @@ -85,7 +85,7 @@ def test_search_term(loki_toolset): ) print(result.data) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error lines = result.data.split("\n")[2:] # skips headers lines for "link" and "query" # print(lines) @@ -104,7 +104,7 @@ def test_search_term_with_dates(loki_toolset): ) ) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error print(result.data) assert TEST_SEARCH_TERM in result.data diff --git a/tests/plugins/toolsets/grafana/test_grafana_tempo_tools.py b/tests/plugins/toolsets/grafana/test_grafana_tempo_tools.py index f0b40c0235..ba406b8d20 100644 --- a/tests/plugins/toolsets/grafana/test_grafana_tempo_tools.py +++ b/tests/plugins/toolsets/grafana/test_grafana_tempo_tools.py @@ -2,7 +2,7 @@ import pytest -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.grafana.common import GrafanaTempoConfig from holmes.plugins.toolsets.grafana.grafana_tempo_api import TempoAPIError from holmes.plugins.toolsets.grafana.toolset_grafana_tempo import ( @@ -59,7 +59,7 @@ def test_search_traces_by_query_success(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "traces" in result.data mock_search.assert_called_once_with( q='{resource.service.name="api"}', @@ -80,7 +80,7 @@ def test_search_traces_by_query_error(self, tempo_toolset): result = tool._invoke({"q": "{invalid}"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "API Error" in result.error def test_search_traces_by_query_tempo_api_error(self, tempo_toolset): @@ -99,7 +99,7 @@ def test_search_traces_by_query_tempo_api_error(self, tempo_toolset): result = tool._invoke({"q": "{invalid syntax}"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "invalid TraceQL query: unexpected token" in result.error assert "400" in result.error @@ -124,7 +124,7 @@ def test_search_traces_by_tags_success(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_search.assert_called_once_with( tags='service.name="api" http.status_code="500"', min_duration="100ms", @@ -152,7 +152,7 @@ def test_query_trace_by_id_success(self, tempo_toolset): result = tool._invoke({"trace_id": "abc123"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_query.assert_called_once_with( trace_id="abc123", start=ANY, @@ -176,7 +176,7 @@ def test_query_trace_by_id_with_time_range(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_query.assert_called_once_with( trace_id="abc123", start=1234567890, @@ -204,7 +204,7 @@ def test_search_tag_names_success(self, tempo_toolset): result = tool._invoke({}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "scopes" in result.data def test_search_tag_names_with_filters(self, tempo_toolset): @@ -224,7 +224,7 @@ def test_search_tag_names_with_filters(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_search.assert_called_once_with( scope="resource", q='{resource.cluster="prod"}', @@ -254,7 +254,7 @@ def test_search_tag_values_success(self, tempo_toolset): result = tool._invoke({"tag": "service.name"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "tagValues" in result.data def test_search_tag_values_error(self, tempo_toolset): @@ -268,7 +268,7 @@ def test_search_tag_values_error(self, tempo_toolset): result = tool._invoke({"tag": "invalid.tag"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "Tag not found" in result.error @@ -296,7 +296,7 @@ def test_query_metrics_instant_success(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_query.assert_called_once_with( q="{ } | histogram_quantile(.95)", start=ANY, @@ -334,7 +334,7 @@ def test_query_metrics_range_success(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert "result" in result.data mock_query.assert_called_once_with( q='{ service.name="api" } | rate()', @@ -360,7 +360,7 @@ def test_query_metrics_range_with_exemplars(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS mock_query.assert_called_once_with( q="{ } | rate()", step=ANY, # step will be auto-calculated @@ -387,7 +387,7 @@ def test_search_traces_with_negative_start(self, tempo_toolset): {"q": '{resource.service.name="api"}', "start": "-3600"} ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Verify start was converted to a positive timestamp args, kwargs = mock_search.call_args assert kwargs["start"] > 0 @@ -413,7 +413,7 @@ def test_query_trace_with_negative_timestamps(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # Both should be positive timestamps assert kwargs["start"] > 0 @@ -442,7 +442,7 @@ def test_metrics_with_inverted_timestamps(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # Timestamps should be inverted assert kwargs["start"] == now @@ -464,7 +464,7 @@ def test_tag_search_with_rfc3339_and_negative(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_search.call_args # Start should be 3600 seconds before the RFC3339 end time assert kwargs["start"] > 0 @@ -515,7 +515,7 @@ def test_all_tools_handle_negative_start(self, tempo_toolset): mock_method.return_value = {"status": "success"} result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Verify the negative start was converted properly if mock_method.called: @@ -547,7 +547,7 @@ def test_metrics_range_with_no_step_auto_calculates(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # With 3600 seconds and MAX_GRAPH_POINTS=300, min step = ceil(3600/300) = 12 # The function should convert this to "12s" @@ -572,7 +572,7 @@ def test_metrics_range_with_small_step_gets_adjusted(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # With 86400 seconds and MAX_GRAPH_POINTS=300, min step = ceil(86400/300) = 288 # The function should adjust to "288s" = "4m48s" @@ -596,7 +596,7 @@ def test_metrics_range_with_large_step_unchanged(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # Step should remain "5m" since it's already large enough assert kwargs["step"] == "5m" @@ -619,7 +619,7 @@ def test_metrics_range_with_bare_number_step(self, tempo_toolset): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args # Step should be "30s" since 30 seconds is fine for 300 second range assert kwargs["step"] == "30s" @@ -649,7 +649,7 @@ def test_metrics_range_step_adjustment_various_ranges(self, tempo_toolset): result = tool._invoke(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS args, kwargs = mock_query.call_args assert ( kwargs["step"] == expected diff --git a/tests/plugins/toolsets/grafana/test_grafana_tempo_unit.py b/tests/plugins/toolsets/grafana/test_grafana_tempo_unit.py index 0f4e5f0bbe..0bdadbdf1f 100644 --- a/tests/plugins/toolsets/grafana/test_grafana_tempo_unit.py +++ b/tests/plugins/toolsets/grafana/test_grafana_tempo_unit.py @@ -2,7 +2,7 @@ import yaml -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.grafana.common import GrafanaTempoConfig from holmes.plugins.toolsets.grafana.toolset_grafana_tempo import ( FetchTracesSimpleComparison, @@ -42,7 +42,7 @@ def test_fetch_traces_simple_comparison_validation(): # Test with no parameters - should fail validation result = tool.invoke(params={}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "At least one of the following argument is expected" in result.error @@ -142,7 +142,7 @@ def test_fetch_traces_simple_comparison_with_mocked_data(): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None # Parse the YAML response @@ -205,7 +205,7 @@ def test_fetch_traces_simple_comparison_with_multiple_filters(): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "No traces found matching the query" # Verify all filters were included in the query @@ -251,7 +251,7 @@ def test_fetch_traces_simple_comparison_with_base_query(): custom_query = "span.http.status_code >= 400" result = tool.invoke(params={"base_query": custom_query}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Verify the custom query was used call_args = mock_api.search_traces_by_query.call_args[1] @@ -278,7 +278,7 @@ def test_fetch_traces_simple_comparison_error_handling(): result = tool.invoke(params={"service_name": "test-service"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert "Error fetching traces: API Error" in result.error @@ -310,7 +310,7 @@ def test_fetch_traces_simple_comparison_percentile_calculations(): mock_api.query_trace_by_id_v2.return_value = {"batches": []} result = tool.invoke(params={"service_name": "test"}) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None data = yaml.safe_load(result.data) @@ -402,7 +402,7 @@ def test_fetch_traces_simple_comparison_with_negative_start_time(): } ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Verify the search was called with positive timestamps call_args = mock_api.search_traces_by_query.call_args[1] diff --git a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_metadata.py b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_metadata.py index 2052c87b11..cc73c14dc5 100644 --- a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_metadata.py +++ b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_metadata.py @@ -6,7 +6,7 @@ from holmes.plugins.toolsets.kubernetes_logs import KubernetesLogsToolset from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus class TestKubernetesLogsMetadata: @@ -52,7 +52,7 @@ def test_logs_found_and_filtered(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS print("\n=== SCENARIO: Logs found and filtered ===") print(result.data) @@ -92,7 +92,7 @@ def test_hitting_display_limit(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS print("\n=== SCENARIO: Hitting display limit ===") print(result.data) @@ -126,7 +126,7 @@ def test_no_logs_match_filter(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA print("\n=== SCENARIO: No logs match filter ===") print(result.data) @@ -158,7 +158,7 @@ def test_no_logs_exist_with_time_range(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA print("\n=== SCENARIO: No logs exist (with time range) ===") print(result.data) @@ -185,7 +185,7 @@ def test_no_logs_exist_without_time_range(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA print("\n=== SCENARIO: No logs exist (no time range) ===") print(result.data) @@ -219,7 +219,7 @@ def test_regex_fallback_warnings(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA print("\n=== SCENARIO: Regex fallback warnings ===") print(result.data) @@ -257,7 +257,7 @@ def test_multi_container_logs(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS print("\n=== SCENARIO: Multi-container logs ===") print(result.data) @@ -288,7 +288,7 @@ def test_exclude_filter_effectiveness(self, mock_datetime_module, mock_run): result = self.toolset.fetch_pod_logs(params) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS print("\n=== SCENARIO: Exclude filter effectiveness ===") print(result.data) diff --git a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_mocked.py b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_mocked.py index d7bdf1b9c5..472680a5f6 100644 --- a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_mocked.py +++ b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_mocked.py @@ -1,7 +1,7 @@ import unittest from unittest.mock import patch, MagicMock -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.kubernetes_logs import KubernetesLogsToolset from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams @@ -106,7 +106,7 @@ def test_single_container(self): result = self.toolset.fetch_pod_logs(params=params) - self.assertEqual(result.status, ToolResultStatus.SUCCESS) + self.assertEqual(result.status, StructuredToolResultStatus.SUCCESS) self.assertEqual(result.return_code, 0) self.assertIsNone(result.error) assert result.data @@ -137,7 +137,7 @@ def test_multi_containers(self): result = self.toolset.fetch_pod_logs(params=params) - self.assertEqual(result.status, ToolResultStatus.SUCCESS) + self.assertEqual(result.status, StructuredToolResultStatus.SUCCESS) self.assertEqual(result.return_code, 0) self.assertIsNone(result.error) @@ -177,7 +177,7 @@ def test_pod_not_found(self): # With kubectl, we get an ERROR status when pod is not found self.assertEqual(result.return_code, 1) - self.assertEqual(result.status, ToolResultStatus.ERROR) + self.assertEqual(result.status, StructuredToolResultStatus.ERROR) self.assertIn("not found", result.error) def test_filter_logs(self): @@ -192,7 +192,7 @@ def test_filter_logs(self): result = self.toolset.fetch_pod_logs(params=params) - self.assertEqual(result.status, ToolResultStatus.SUCCESS) + self.assertEqual(result.status, StructuredToolResultStatus.SUCCESS) self.assertEqual(result.return_code, 0) self.assertIsNone(result.error) @@ -216,7 +216,7 @@ def test_limit_logs(self): result = self.toolset.fetch_pod_logs(params=params) - self.assertEqual(result.status, ToolResultStatus.SUCCESS) + self.assertEqual(result.status, StructuredToolResultStatus.SUCCESS) self.assertEqual(result.return_code, 0) self.assertIsNone(result.error) diff --git a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_no_data.py b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_no_data.py index 6929b2458e..8feaf0355d 100644 --- a/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_no_data.py +++ b/tests/plugins/toolsets/kubernetes/test_kubernetes_logs_no_data.py @@ -1,6 +1,6 @@ from holmes.plugins.toolsets.kubernetes_logs import KubernetesLogsToolset from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from unittest.mock import Mock, patch @@ -23,7 +23,7 @@ def test_no_logs_returns_no_data_status(self, mock_run): result = toolset.fetch_pod_logs(params) # Should return NO_DATA status when there are no logs - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA # May include metadata about the query even when no logs found @patch("subprocess.run") @@ -47,7 +47,7 @@ def test_logs_with_filter_no_matches_returns_no_data(self, mock_run): result = toolset.fetch_pod_logs(params) # Should return NO_DATA when filter matches no logs - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA @patch("subprocess.run") def test_logs_with_time_range_no_matches_returns_no_data(self, mock_run): @@ -70,7 +70,7 @@ def test_logs_with_time_range_no_matches_returns_no_data(self, mock_run): result = toolset.fetch_pod_logs(params) # Should return NO_DATA when no logs in time range - assert result.status == ToolResultStatus.NO_DATA + assert result.status == StructuredToolResultStatus.NO_DATA @patch("subprocess.run") def test_logs_exist_returns_success_with_metadata(self, mock_run): @@ -88,7 +88,7 @@ def test_logs_exist_returns_success_with_metadata(self, mock_run): result = toolset.fetch_pod_logs(params) # Should return SUCCESS when logs exist - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data is not None and len(result.data) > 0 # Should contain the log content assert "Something went wrong" in result.data diff --git a/tests/plugins/toolsets/opensearch/test_opensearch_logs_integration.py b/tests/plugins/toolsets/opensearch/test_opensearch_logs_integration.py index b449372e44..8abe912b71 100644 --- a/tests/plugins/toolsets/opensearch/test_opensearch_logs_integration.py +++ b/tests/plugins/toolsets/opensearch/test_opensearch_logs_integration.py @@ -6,7 +6,7 @@ import pytest import os -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.logging_utils.logging_api import FetchPodLogsParams from holmes.plugins.toolsets.opensearch.opensearch_logs import ( OpenSearchLogsToolset, @@ -71,7 +71,7 @@ def test_basic_query(opensearch_logs_toolset): FetchPodLogsParams(namespace=TEST_NAMESPACE, pod_name=TEST_POD_NAME) ) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error print(result.data) assert TEST_SEARCH_TERM in result.data @@ -84,7 +84,7 @@ def test_search_term(opensearch_logs_toolset): ) ) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error print(result.data) for line in result.data.split("\n"): @@ -102,7 +102,7 @@ def test_search_term_with_dates(opensearch_logs_toolset): ) ) - assert result.status == ToolResultStatus.SUCCESS, result.error + assert result.status == StructuredToolResultStatus.SUCCESS, result.error assert not result.error print(result.data) diff --git a/tests/plugins/toolsets/test_logging_api.py b/tests/plugins/toolsets/test_logging_api.py index fb4337118c..58a1d386e0 100644 --- a/tests/plugins/toolsets/test_logging_api.py +++ b/tests/plugins/toolsets/test_logging_api.py @@ -8,7 +8,7 @@ FetchPodLogsParams, LoggingCapability, ) -from holmes.core.tools import StructuredToolResult, ToolResultStatus +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus class TestPodLoggingTool: @@ -21,7 +21,7 @@ def test_tool_handles_integer_start_time(self): mock_toolset.name = "test-logging-backend" mock_toolset.supported_capabilities = set() mock_toolset.fetch_pod_logs.return_value = StructuredToolResult( - data="Sample logs", status=ToolResultStatus.SUCCESS + data="Sample logs", status=StructuredToolResultStatus.SUCCESS ) # Create the tool @@ -36,7 +36,7 @@ def test_tool_handles_integer_start_time(self): result = tool._invoke(params) # Verify the result - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.data == "Sample logs" # Verify toolset.fetch_pod_logs was called once @@ -60,7 +60,7 @@ def test_tool_handles_string_start_time(self): mock_toolset.name = "test-logging-backend" mock_toolset.supported_capabilities = set() mock_toolset.fetch_pod_logs.return_value = StructuredToolResult( - data="Sample logs", status=ToolResultStatus.SUCCESS + data="Sample logs", status=StructuredToolResultStatus.SUCCESS ) # Create the tool @@ -75,7 +75,7 @@ def test_tool_handles_string_start_time(self): result = tool._invoke(params) # Verify the result - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS # Get the actual params passed to fetch_pod_logs call_args = mock_toolset.fetch_pod_logs.call_args @@ -92,7 +92,7 @@ def test_tool_handles_rfc3339_start_time(self): mock_toolset.name = "test-logging-backend" mock_toolset.supported_capabilities = set() mock_toolset.fetch_pod_logs.return_value = StructuredToolResult( - data="Sample logs", status=ToolResultStatus.SUCCESS + data="Sample logs", status=StructuredToolResultStatus.SUCCESS ) # Create the tool @@ -121,7 +121,7 @@ def test_tool_handles_no_start_time(self): mock_toolset.name = "test-logging-backend" mock_toolset.supported_capabilities = set() mock_toolset.fetch_pod_logs.return_value = StructuredToolResult( - data="Sample logs", status=ToolResultStatus.SUCCESS + data="Sample logs", status=StructuredToolResultStatus.SUCCESS ) # Create the tool @@ -148,7 +148,7 @@ def test_tool_with_all_parameters(self): LoggingCapability.EXCLUDE_FILTER, } mock_toolset.fetch_pod_logs.return_value = StructuredToolResult( - data="Filtered logs", status=ToolResultStatus.SUCCESS + data="Filtered logs", status=StructuredToolResultStatus.SUCCESS ) # Create the tool diff --git a/tests/plugins/toolsets/test_runbook.py b/tests/plugins/toolsets/test_runbook.py index f191f90510..6772b5d9f9 100644 --- a/tests/plugins/toolsets/test_runbook.py +++ b/tests/plugins/toolsets/test_runbook.py @@ -1,4 +1,4 @@ -from holmes.core.tools import ToolResultStatus +from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.runbook.runbook_fetcher import ( RunbookFetcher, RunbookToolset, @@ -8,14 +8,14 @@ def test_RunbookFetcher(): runbook_fetch_tool = RunbookFetcher(RunbookToolset()) result = runbook_fetch_tool._invoke({"link": "wrong_runbook_path"}) - assert result.status == ToolResultStatus.ERROR + assert result.status == StructuredToolResultStatus.ERROR assert result.error is not None result = runbook_fetch_tool._invoke( {"link": "networking/dns_troubleshooting_instructions.md"} ) - assert result.status == ToolResultStatus.SUCCESS + assert result.status == StructuredToolResultStatus.SUCCESS assert result.error is None assert result.data is not None assert ( diff --git a/tests/plugins/toolsets/test_toolset_utils.py b/tests/plugins/toolsets/test_toolset_utils.py index c81bd4d21d..2842be060e 100644 --- a/tests/plugins/toolsets/test_toolset_utils.py +++ b/tests/plugins/toolsets/test_toolset_utils.py @@ -2,7 +2,7 @@ from dateutil import parser # type: ignore from holmes.core.tools import ( StructuredToolResult, - ToolResultStatus, + StructuredToolResultStatus, Toolset, ToolsetStatusEnum, ) @@ -155,7 +155,7 @@ def __init__(self): def _invoke( self, params: dict, user_approved: bool = False ) -> StructuredToolResult: - return StructuredToolResult(status=ToolResultStatus.SUCCESS) + return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS) def get_parameterized_one_liner(self, params: dict) -> str: """Generate a one-line description of this tool invocation""" @@ -178,7 +178,7 @@ def supported_capabilities(self) -> set: return set() # No advanced capabilities for dummy toolset def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult: - return StructuredToolResult(status=ToolResultStatus.SUCCESS) + return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS) def get_example_config(self): return {} diff --git a/tests/test_structured_toolcall_result.py b/tests/test_structured_toolcall_result.py index 4a4b2e1373..05a908f1b4 100644 --- a/tests/test_structured_toolcall_result.py +++ b/tests/test_structured_toolcall_result.py @@ -2,9 +2,8 @@ import pytest from pydantic import BaseModel -from holmes.core.tools import StructuredToolResult, ToolResultStatus - -from holmes.core.tool_calling_llm import format_tool_result_data, ToolCallResult +from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus +from holmes.core.tools_utils.data_types import ToolCallResult, format_tool_result_data class DummyResult(BaseModel): @@ -25,13 +24,13 @@ def __str__(self): ], ) def test_get_stringified_data_none_and_str(data, expected): - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data=data) + result = StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data=data) assert result.get_stringified_data() == expected def test_get_stringified_data_base_model(): dummy = DummyResult(x=10, y="hello") - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data=dummy) + result = StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data=dummy) expected = dummy.model_dump_json(indent=2) assert result.get_stringified_data() == expected @@ -44,22 +43,29 @@ def test_get_stringified_data_base_model(): ], ) def test_get_stringified_data_json_serializable(data): - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data=data) + result = StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data=data) expected = json.dumps(data, indent=2) assert result.get_stringified_data() == expected def test_get_stringified_data_unserializable_object(): obj = Unserializable() - result = StructuredToolResult(status=ToolResultStatus.ERROR, data=obj) + result = StructuredToolResult(status=StructuredToolResultStatus.ERROR, data=obj) assert result.get_stringified_data() == "unserializable_str" @pytest.mark.parametrize( "status,error,return_code,url,invocation,params", [ - (ToolResultStatus.SUCCESS, None, None, None, None, None), - (ToolResultStatus.ERROR, "oops", 1, "http://example.com", "invoke", {"a": 1}), + (StructuredToolResultStatus.SUCCESS, None, None, None, None, None), + ( + StructuredToolResultStatus.ERROR, + "oops", + 1, + "http://example.com", + "invoke", + {"a": 1}, + ), ], ) def test_default_and_custom_fields(status, error, return_code, url, invocation, params): @@ -85,15 +91,25 @@ def test_default_and_custom_fields(status, error, return_code, url, invocation, @pytest.mark.parametrize( "status,error,data,expected", [ - (ToolResultStatus.SUCCESS, None, "test", "test"), + (StructuredToolResultStatus.SUCCESS, None, "test", "test"), ( - ToolResultStatus.NO_DATA, + StructuredToolResultStatus.NO_DATA, None, DummyResult(x=2, y="test"), DummyResult(x=2, y="test").model_dump_json(indent=2), ), - (ToolResultStatus.SUCCESS, None, {"k": 1}, json.dumps({"k": 1}, indent=2)), - (ToolResultStatus.SUCCESS, None, Unserializable(), str(Unserializable())), + ( + StructuredToolResultStatus.SUCCESS, + None, + {"k": 1}, + json.dumps({"k": 1}, indent=2), + ), + ( + StructuredToolResultStatus.SUCCESS, + None, + Unserializable(), + str(Unserializable()), + ), ], ) def test_format_tool_result_data_non_error(status, error, data, expected): @@ -102,49 +118,57 @@ def test_format_tool_result_data_non_error(status, error, data, expected): def test_format_tool_result_data_str_non_error(): - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data="hello") + result = StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data="hello" + ) assert format_tool_result_data(result) == "hello" def test_format_tool_result_data_base_model_non_error(): dummy = DummyResult(x=2, y="b") - result = StructuredToolResult(status=ToolResultStatus.NO_DATA, data=dummy) + result = StructuredToolResult(status=StructuredToolResultStatus.NO_DATA, data=dummy) assert format_tool_result_data(result) == dummy.model_dump_json(indent=2) def test_format_tool_result_data_json_serializable_non_error(): data = {"k": 3} - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data=data) + result = StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data=data) assert format_tool_result_data(result) == json.dumps(data, indent=2) def test_format_tool_result_data_unserializable_non_error(): obj = Unserializable() - result = StructuredToolResult(status=ToolResultStatus.SUCCESS, data=obj) + result = StructuredToolResult(status=StructuredToolResultStatus.SUCCESS, data=obj) assert format_tool_result_data(result) == str(obj) def test_format_tool_result_data_error_with_message_and_data(): result = StructuredToolResult( - status=ToolResultStatus.ERROR, error="fail", data="oops" + status=StructuredToolResultStatus.ERROR, error="fail", data="oops" ) assert format_tool_result_data(result) == "fail:\n\noops" def test_format_tool_result_data_error_without_message_or_data(): - result = StructuredToolResult(status=ToolResultStatus.ERROR, error=None, data=None) + result = StructuredToolResult( + status=StructuredToolResultStatus.ERROR, error=None, data=None + ) assert format_tool_result_data(result) == "Tool execution failed:" def test_format_tool_result_data_error_without_message_with_unserializable(): obj = Unserializable() - result = StructuredToolResult(status=ToolResultStatus.ERROR, error=None, data=obj) + result = StructuredToolResult( + status=StructuredToolResultStatus.ERROR, error=None, data=obj + ) expected = f"Tool execution failed:\n\n{str(obj)}" assert format_tool_result_data(result) == expected def test_as_tool_call_message_without_params(): - structured = StructuredToolResult(status=ToolResultStatus.SUCCESS, data="hello") + structured = StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data="hello" + ) tcr = ToolCallResult( tool_call_id="call1", tool_name="toolX", @@ -162,7 +186,7 @@ def test_as_tool_call_message_without_params(): def test_as_tool_call_message_with_params(): structured = StructuredToolResult( - status=ToolResultStatus.SUCCESS, + status=StructuredToolResultStatus.SUCCESS, data="hello", params={"pod_name": "my-pod", "namespace": "my-namespace"}, ) @@ -185,7 +209,9 @@ def test_as_tool_call_message_with_params(): def test_as_tool_result_response(): - structured = StructuredToolResult(status=ToolResultStatus.SUCCESS, data="hello") + structured = StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data="hello" + ) tcr = ToolCallResult( tool_call_id="call1", tool_name="toolX", @@ -204,7 +230,9 @@ def test_as_tool_result_response(): def test_as_streaming_tool_result_response(): - structured = StructuredToolResult(status=ToolResultStatus.SUCCESS, data="hello") + structured = StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data="hello" + ) tcr = ToolCallResult( tool_call_id="call2", tool_name="toolY",