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

holmes/plugins/toolsets/logging_utils/logging_api.py

@@ -1,27 +1,36 @@
-from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
 import logging
-from typing import Optional, Set
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta, timezone
 from enum import Enum
+from math import ceil
+from typing import Optional, Set
 
 from pydantic import BaseModel, field_validator
-from datetime import timezone
+
+from holmes.core.llm import LLM
 from holmes.core.tools import (
     StructuredToolResult,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     Toolset,
 )
+from holmes.core.tools_utils.token_counting import count_tool_response_tokens
 from holmes.plugins.toolsets.utils import get_param_or_raise
 
 # Default values for log fetching
 DEFAULT_LOG_LIMIT = 100
 SECONDS_PER_DAY = 24 * 60 * 60
 DEFAULT_TIME_SPAN_SECONDS = 7 * SECONDS_PER_DAY  # 1 week in seconds
-DEFAULT_GRAPH_TIME_SPAN_SECONDS = 1 * SECONDS_PER_DAY  # 1 day in seconds
+DEFAULT_GRAPH_TIME_SPAN_SECONDS = 1 * 60 * 60  # 1 hour in seconds
 
 POD_LOGGING_TOOL_NAME = "fetch_pod_logs"
 
+TRUNCATION_PROMPT_PREFIX = "[... PREVIOUS LOGS ABOVE THIS LINE HAVE BEEN TRUNCATED]"
+MIN_NUMBER_OF_CHARACTERS_TO_TRUNCATE: int = (
+    50 + len(TRUNCATION_PROMPT_PREFIX)
+)  # prevents the truncation algorithm from going too slow once the actual token count gets close to the expected limit
+
 
 class LoggingCapability(str, Enum):
     """Optional advanced logging capabilities"""
@@ -74,6 +83,76 @@ class BasePodLoggingToolset(Toolset, ABC):
         return ""
 
 
+def truncate_logs(
+    logging_structured_tool_result: StructuredToolResult,
+    llm: LLM,
+    token_limit: int,
+    structured_params: FetchPodLogsParams,
+    tool_call_id: str,
+    tool_name: str,
+):
+    original_token_count = count_tool_response_tokens(
+        llm=llm,
+        structured_tool_result=logging_structured_tool_result,
+        tool_call_id=tool_call_id,
+        tool_name=tool_name,
+    )
+    token_count = original_token_count
+    text = None
+    while token_count > token_limit:
+        # Loop because we are counting tokens but trimming characters. This means we try to trim a number of
+        # characters proportional to the number of tokens but we may still have too many tokens
+        if not text:
+            text = logging_structured_tool_result.get_stringified_data()
+        if not text:
+            # Weird scenario where the result exceeds the token allowance but there is not data.
+            # Exit and do nothing because I don't know how to handle such scenario.
+            logging.warning(
+                f"The calculated token count for logs is {token_count} but the limit is {token_limit}. However the data field is empty so there are no logs to truncate."
+            )
+            return
+        ratio = token_count / token_limit
+        character_count = len(text)
+        number_of_characters_to_truncate = character_count - ceil(
+            character_count / ratio
+        )
+        number_of_characters_to_truncate = max(
+            MIN_NUMBER_OF_CHARACTERS_TO_TRUNCATE, number_of_characters_to_truncate
+        )
+
+        if len(text) <= number_of_characters_to_truncate:
+            logging.warning(
+                f"The calculated token count for logs is {token_count} (max allowed tokens={token_limit}) but the logs are only {len(text)} characters which is below the intended truncation of {number_of_characters_to_truncate} characters. Logs will no longer be truncated"
+            )
+            return
+        else:
+            linefeed_truncation_offset = max(
+                text[number_of_characters_to_truncate:].find("\n"), 0
+            )  # keep log lines atomic
+
+            # Tentatively add the truncation prefix.
+            # When counting tokens, we want to include the TRUNCATION_PROMPT_PREFIX because it will be part of the tool response.
+            # Because we're truncating based on character counts but ultimately checking tokens count,
+            # it is possible that the character truncation is incorrect and more need to be truncated.
+            # This will be caught in the next iteration and the truncation prefix will be truncated
+            # because MIN_NUMBER_OF_CHARACTERS_TO_TRUNCATE cannot be smaller than TRUNCATION_PROMPT_PREFIX
+            text = (
+                TRUNCATION_PROMPT_PREFIX
+                + text[number_of_characters_to_truncate + linefeed_truncation_offset :]
+            )
+            logging_structured_tool_result.data = text
+            token_count = count_tool_response_tokens(
+                llm=llm,
+                structured_tool_result=logging_structured_tool_result,
+                tool_call_id=tool_call_id,
+                tool_name=tool_name,
+            )
+    if token_count < original_token_count:
+        logging.info(
+            f"Logs for pod {structured_params.pod_name}/{structured_params.namespace} have been truncated from {original_token_count} tokens down to {token_count} tokens."
+        )
+
+
 class PodLoggingTool(Tool):
     """Common tool for fetching pod logs across different logging backends"""
 
@@ -175,9 +254,7 @@ If you hit the log limit and see lots of repetitive INFO logs, use exclude_filte
 
         return params
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         structured_params = FetchPodLogsParams(
             namespace=get_param_or_raise(params, "namespace"),
             pod_name=get_param_or_raise(params, "pod_name"),
@@ -192,6 +269,15 @@ If you hit the log limit and see lots of repetitive INFO logs, use exclude_filte
             params=structured_params,
         )
 
+        truncate_logs(
+            logging_structured_tool_result=result,
+            llm=context.llm,
+            token_limit=context.max_token_count,
+            structured_params=structured_params,
+            tool_call_id=context.tool_call_id,
+            tool_name=context.tool_name,
+        )
+
         return result
 
     def get_parameterized_one_liner(self, params: dict) -> str:

holmes/plugins/toolsets/mcp/toolset_mcp.py

@@ -1,71 +1,168 @@
+import asyncio
+import json
+import logging
+import threading
+from contextlib import asynccontextmanager
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from mcp.client.session import ClientSession
+from mcp.client.sse import sse_client
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.types import Tool as MCP_Tool
+from pydantic import AnyUrl, BaseModel, Field, model_validator
+
+from holmes.common.env_vars import SSE_READ_TIMEOUT
 from holmes.core.tools import (
-    Toolset,
+    CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
-    StructuredToolResult,
-    ToolResultStatus,
-    CallablePrerequisite,
+    Toolset,
 )
 
-from typing import Dict, Any, List, Optional
-from mcp.client.session import ClientSession
-from mcp.client.sse import sse_client
+# Lock per MCP server URL to serialize calls to the same server
+_server_locks: Dict[str, threading.Lock] = {}
+_locks_lock = threading.Lock()
 
-from mcp.types import Tool as MCP_Tool
-from mcp.types import CallToolResult
 
-import asyncio
-from pydantic import Field, AnyUrl, field_validator
-from typing import Tuple
-import logging
+def get_server_lock(url: str) -> threading.Lock:
+    """Get or create a lock for a specific MCP server URL."""
+    with _locks_lock:
+        if url not in _server_locks:
+            _server_locks[url] = threading.Lock()
+        return _server_locks[url]
 
 
-class RemoteMCPTool(Tool):
-    url: str
+class MCPMode(str, Enum):
+    SSE = "sse"
+    STREAMABLE_HTTP = "streamable-http"
+    STDIO = "stdio"
+
+
+class MCPConfig(BaseModel):
+    url: AnyUrl
+    mode: MCPMode = MCPMode.SSE
     headers: Optional[Dict[str, str]] = None
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def get_lock_string(self) -> str:
+        return str(self.url)
+
+
+class StdioMCPConfig(BaseModel):
+    mode: MCPMode = MCPMode.STDIO
+    command: str
+    args: Optional[List[str]] = None
+    env: Optional[Dict[str, str]] = None
+
+    def get_lock_string(self) -> str:
+        return str(self.command)
+
+
+@asynccontextmanager
+async def get_initialized_mcp_session(toolset: "RemoteMCPToolset"):
+    if toolset._mcp_config is None:
+        raise ValueError("MCP config is not initialized")
+
+    if isinstance(toolset._mcp_config, StdioMCPConfig):
+        server_params = StdioServerParameters(
+            command=toolset._mcp_config.command,
+            args=toolset._mcp_config.args or [],
+            env=toolset._mcp_config.env,
+        )
+        async with stdio_client(server_params) as (
+            read_stream,
+            write_stream,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                _ = await session.initialize()
+                yield session
+    elif toolset._mcp_config.mode == MCPMode.SSE:
+        url = str(toolset._mcp_config.url)
+        async with sse_client(
+            url, toolset._mcp_config.headers, sse_read_timeout=SSE_READ_TIMEOUT
+        ) as (
+            read_stream,
+            write_stream,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                _ = await session.initialize()
+                yield session
+    else:
+        url = str(toolset._mcp_config.url)
+        async with streamablehttp_client(
+            url, headers=toolset._mcp_config.headers, sse_read_timeout=SSE_READ_TIMEOUT
+        ) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                _ = await session.initialize()
+                yield session
+
+
+class RemoteMCPTool(Tool):
+    toolset: "RemoteMCPToolset" = Field(exclude=True)
+
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
-            return asyncio.run(self._invoke_async(params))
+            # Serialize calls to the same MCP server to prevent SSE conflicts
+            # Different servers can still run in parallel
+            if not self.toolset._mcp_config:
+                raise ValueError("MCP config not initialized")
+
+            lock = get_server_lock(str(self.toolset._mcp_config.get_lock_string()))
+            with lock:
+                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}",
             )
 
+    @staticmethod
+    def _is_content_error(content: str) -> bool:
+        try:  # aws mcp sometimes returns an error in content - status code != 200
+            json_content: dict = json.loads(content)
+            status_code = json_content.get("response", {}).get("status_code", 200)
+            return status_code >= 300
+        except Exception:
+            return False
+
     async def _invoke_async(self, params: Dict) -> StructuredToolResult:
-        async with sse_client(self.url, self.headers) as (read_stream, write_stream):
-            async with ClientSession(read_stream, write_stream) as session:
-                _ = await session.initialize()
-                tool_result: CallToolResult = await session.call_tool(self.name, params)
+        async with get_initialized_mcp_session(self.toolset) as session:
+            tool_result = await session.call_tool(self.name, params)
 
-                merged_text = " ".join(
-                    c.text for c in tool_result.content if c.type == "text"
-                )
-                return StructuredToolResult(
-                    status=(
-                        ToolResultStatus.ERROR
-                        if tool_result.isError
-                        else ToolResultStatus.SUCCESS
-                    ),
-                    data=merged_text,
-                    params=params,
-                    invocation=f"MCPtool {self.name} with params {params}",
-                )
+        merged_text = " ".join(c.text for c in tool_result.content if c.type == "text")
+        return StructuredToolResult(
+            status=(
+                StructuredToolResultStatus.ERROR
+                if (tool_result.isError or self._is_content_error(merged_text))
+                else StructuredToolResultStatus.SUCCESS
+            ),
+            data=merged_text,
+            params=params,
+            invocation=f"MCPtool {self.name} with params {params}",
+        )
 
     @classmethod
-    def create(cls, url: str, tool: MCP_Tool, headers: Optional[Dict[str, str]] = None):
+    def create(
+        cls,
+        tool: MCP_Tool,
+        toolset: "RemoteMCPToolset",
+    ):
         parameters = cls.parse_input_schema(tool.inputSchema)
         return cls(
-            url=url,
             name=tool.name,
             description=tool.description or "",
             parameters=parameters,
-            headers=headers,
+            toolset=toolset,
         )
 
     @classmethod
@@ -85,53 +182,110 @@ class RemoteMCPTool(Tool):
         return parameters
 
     def get_parameterized_one_liner(self, params: Dict) -> str:
-        return f"Call MCP Server ({self.url} - {self.name})"
+        # AWS MCP cli_command
+        if params and params.get("cli_command"):
+            return f"{params.get('cli_command')}"
+
+        # gcloud MCP run_gcloud_command
+        if self.name == "run_gcloud_command" and params and "args" in params:
+            args = params.get("args", [])
+            if isinstance(args, list):
+                return f"gcloud {' '.join(str(arg) for arg in args)}"
+
+        return f"{self.toolset.name}: {self.name} {params}"
 
 
 class RemoteMCPToolset(Toolset):
-    url: AnyUrl
     tools: List[RemoteMCPTool] = Field(default_factory=list)  # type: ignore
     icon_url: str = "https://registry.npmmirror.com/@lobehub/icons-static-png/1.46.0/files/light/mcp.png"
+    _mcp_config: Optional[Union[MCPConfig, StdioMCPConfig]] = None
 
     def model_post_init(self, __context: Any) -> None:
-        self.prerequisites = [CallablePrerequisite(callable=self.init_server_tools)]
+        self.prerequisites = [
+            CallablePrerequisite(callable=self.prerequisites_callable)
+        ]
 
-    def get_headers(self) -> Optional[Dict[str, str]]:
-        return self.config and self.config.get("headers")
+    @model_validator(mode="before")
+    @classmethod
+    def migrate_url_to_config(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """
+        Migrates url from field parameter to config object.
+        If url is passed as a parameter, it's moved to config (or config is created if it doesn't exist).
+        """
+        if not isinstance(values, dict) or "url" not in values:
+            return values
+
+        url_value = values.pop("url")
+        if url_value is None:
+            return values
+
+        config = values.get("config")
+        if config is None:
+            config = {}
+            values["config"] = config
+
+        toolset_name = values.get("name", "unknown")
+        if "url" in config:
+            logging.warning(
+                f"Toolset {toolset_name}: has two urls defined, remove the 'url' field from the toolset configuration and keep the 'url' in the config section."
+            )
+            return values
 
-    @field_validator("url", mode="before")
-    def append_sse_if_missing(cls, v):
-        if isinstance(v, str) and not v.rstrip("/").endswith("/sse"):
-            v = v.rstrip("/") + "/sse"
-        return v
+        logging.warning(
+            f"Toolset {toolset_name}: 'url' field has been migrated to config. "
+            "Please move 'url' to the config section."
+        )
+        config["url"] = url_value
+        return values
 
-    # used as a CallablePrerequisite, config added for that case.
-    def init_server_tools(self, config: dict[str, Any]) -> Tuple[bool, str]:
+    def prerequisites_callable(self, config) -> Tuple[bool, str]:
         try:
+            if not config:
+                return (False, f"Config is required for {self.name}")
+
+            mode_value = config.get("mode", MCPMode.SSE.value)
+            allowed_modes = [e.value for e in MCPMode]
+            if mode_value not in allowed_modes:
+                return (
+                    False,
+                    f'Invalid mode "{mode_value}", allowed modes are {", ".join(allowed_modes)}',
+                )
+
+            if mode_value == MCPMode.STDIO.value:
+                self._mcp_config = StdioMCPConfig(**config)
+            else:
+                self._mcp_config = MCPConfig(**config)
+                clean_url_str = str(self._mcp_config.url).rstrip("/")
+
+                if self._mcp_config.mode == MCPMode.SSE and not clean_url_str.endswith(
+                    "/sse"
+                ):
+                    self._mcp_config.url = AnyUrl(clean_url_str + "/sse")
+
             tools_result = asyncio.run(self._get_server_tools())
+
             self.tools = [
-                RemoteMCPTool.create(str(self.url), tool, self.get_headers())
-                for tool in tools_result.tools
+                RemoteMCPTool.create(tool, self) for tool in tools_result.tools
             ]
 
             if not self.tools:
                 logging.warning(f"mcp server {self.name} loaded 0 tools.")
+
             return (True, "")
         except Exception as e:
-            # using e.args, the asyncio wrapper could stack another exception this helps printing them all.
             return (
                 False,
-                f"Failed to load mcp server {self.name} {self.url} {str(e.args)}",
+                f"Failed to load mcp server {self.name}: {str(e)}",
             )
 
     async def _get_server_tools(self):
-        async with sse_client(str(self.url), headers=self.get_headers()) as (
-            read_stream,
-            write_stream,
-        ):
-            async with ClientSession(read_stream, write_stream) as session:
-                _ = await session.initialize()
-                return await session.list_tools()
+        async with get_initialized_mcp_session(self) as session:
+            return await session.list_tools()
 
     def get_example_config(self) -> Dict[str, Any]:
-        return {}
+        example_config = MCPConfig(
+            url=AnyUrl("http://example.com:8000/mcp/messages"),
+            mode=MCPMode.STREAMABLE_HTTP,
+            headers={"Authorization": "Bearer YOUR_TOKEN"},
+        )
+        return example_config.model_dump()

holmes/plugins/toolsets/newrelic/new_relic_api.py

@@ -0,0 +1,165 @@
+"""NewRelic API wrapper for executing NRQL queries via GraphQL."""
+
+import logging
+from typing import Any, Dict
+
+import requests  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+class NewRelicAPI:
+    """Python wrapper for NewRelic GraphQL API.
+
+    This class provides a clean interface to execute NRQL queries via the NewRelic GraphQL API,
+    supporting both US and EU datacenters.
+    """
+
+    def __init__(self, api_key: str, account_id: str, is_eu_datacenter: bool = False):
+        """Initialize the NewRelic API wrapper.
+
+        Args:
+            api_key: NewRelic API key
+            account_id: NewRelic account ID
+            is_eu_datacenter: If True, use EU datacenter URL. Defaults to False (US).
+        """
+        self.api_key = api_key
+        # Validate account_id is numeric to prevent injection
+        try:
+            self.account_id = int(account_id)
+        except ValueError:
+            raise ValueError(f"Invalid account_id: must be numeric, got '{account_id}'")
+        self.is_eu_datacenter = is_eu_datacenter
+
+    def _get_api_url(self) -> str:
+        """Get the appropriate API URL based on datacenter location.
+
+        Returns:
+            str: The GraphQL API endpoint URL
+        """
+        if self.is_eu_datacenter:
+            return "https://api.eu.newrelic.com/graphql"
+        return "https://api.newrelic.com/graphql"
+
+    def _make_request(
+        self, graphql_query: Dict[str, Any], timeout: int = 30
+    ) -> Dict[str, Any]:
+        """Make HTTP POST request to NewRelic GraphQL API.
+
+        Args:
+            graphql_query: The GraphQL query as a dictionary
+            timeout: Request timeout in seconds
+
+        Returns:
+            JSON response from the API
+
+        Raises:
+            requests.exceptions.HTTPError: If the request fails
+            Exception: If GraphQL returns errors
+        """
+        url = self._get_api_url()
+        headers = {
+            "Content-Type": "application/json",
+            "Api-Key": self.api_key,
+        }
+
+        response = requests.post(
+            url,
+            headers=headers,
+            json=graphql_query,
+            timeout=timeout,
+        )
+        response.raise_for_status()
+
+        # Parse JSON response
+        data = response.json()
+
+        # Check for GraphQL errors even on 200 responses
+        if "errors" in data and data["errors"]:
+            error_msg = data["errors"][0].get("message", "Unknown GraphQL error")
+            raise Exception(f"NewRelic GraphQL error: {error_msg}")
+
+        return data
+
+    def execute_nrql_query(self, nrql_query: str) -> list:
+        """Execute an NRQL query via the NewRelic GraphQL API.
+
+        Args:
+            nrql_query: The NRQL query string to execute
+
+        Returns:
+            list: The query results from NewRelic (extracted from the nested response)
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails
+            Exception: If GraphQL returns errors
+        """
+        # Build the GraphQL query using variables to prevent injection
+        # Note: New Relic's GraphQL API requires the account ID to be inline, but we can use variables for the NRQL query
+        graphql_query = {
+            "query": f"""
+            query ExecuteNRQL($nrqlQuery: Nrql!) {{
+                actor {{
+                    account(id: {self.account_id}) {{
+                        nrql(query: $nrqlQuery) {{
+                            results
+                        }}
+                    }}
+                }}
+            }}
+            """,
+            "variables": {"nrqlQuery": nrql_query},
+        }
+
+        logger.info(f"Executing NRQL query: {nrql_query}")
+        response = self._make_request(graphql_query)
+
+        # Extract just the results array from the nested response
+        try:
+            results = response["data"]["actor"]["account"]["nrql"]["results"]
+            return results
+        except (KeyError, TypeError) as e:
+            raise Exception(
+                f"Failed to extract results from NewRelic response: {e}"
+            ) from e
+
+    def get_organization_accounts(self) -> list:
+        """Get all accounts accessible in the organization.
+
+        Returns:
+            list: List of account dictionaries with id and name
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails
+            Exception: If GraphQL returns errors
+        """
+        graphql_query = {
+            "query": """
+            query GetOrganizationAccounts {
+                actor {
+                    organization {
+                        accountManagement {
+                            managedAccounts {
+                                id
+                                name
+                            }
+                        }
+                    }
+                }
+            }
+            """
+        }
+
+        logger.info("Querying organization accounts")
+        response = self._make_request(graphql_query)
+
+        # Extract accounts from the nested response
+        try:
+            accounts = response["data"]["actor"]["organization"]["accountManagement"][
+                "managedAccounts"
+            ]
+            return accounts
+        except (KeyError, TypeError) as e:
+            raise Exception(
+                f"Failed to extract accounts from NewRelic response: {e}"
+            ) from e