holmesgpt 0.14.2__py3-none-any.whl → 0.14.4a0__py3-none-any.whl

holmes/plugins/toolsets/newrelic/new_relic_api.py

@@ -0,0 +1,125 @@
+"""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

holmes/plugins/toolsets/newrelic/newrelic.jinja2

@@ -0,0 +1,41 @@
+New Relic provides distributed tracing data along with logs and metrics.
+
+Assume every application has New Relic tracing data.
+
+Use `nrql_query` to run a NRQL query.
+
+**NRQL (New Relic Query Language)** is used to query all telemetry data in New Relic. The main event types are:
+
+- **Transaction**: High-level APM data (requests, API calls)
+- **Span**: Distributed tracing data (individual operations)
+- **Log**: Centralized log data
+- **Metric**: Time-series metrics data.
+
+### Usage Workflow
+
+#### 1. Discovering Available Data
+
+Start by understanding what's available. Here are some examples:
+- **ALWAYS** Start by getting all the available attributes names for what you are looking for. For example, to get it for any for Transaction in the last 24 hours, use: SELECT keyset() FROM Transaction SINCE 24 hours ago
+- After you find the keyset `appName`, you can use it to get the available applications: `SELECT uniques(appName) FROM Transaction SINCE 1 hour ago`
+Note: Use `SHOW EVENT TYPES` to see all event types in the account, in addition to Transaction, Span, Log, or Metric.
+
+#### 2. Querying Telemetry Data
+
+- If you already have an application name, you can query its traces directly
+- **Time range is recommended**: While not strictly required, most queries should include SINCE for performance
+
+#### 3. Querying Traces
+- Always validate first: run the base query without FACET (or a quick LIMIT) to confirm data exists; if results are empty, adjust filters or time range before proceeding.
+- Only attempt a FACET after confirming the field has values; if not, either try known alternatives or skip faceting entirely.
+- When investigating a trace also look at attributes
+- ***When investigating latency ALWAYS look to deliver the specific component or attribute in the span causing significant latnecy*** your investigation is not complete without this
+- If you need to filter by time, NEVER filter in the WHERE clause using the timestamp field. Instead, you should ALWAYS use the `SINCE` or `SINCE ... UNTIL ...` syntax - which are the recommended ways to run time based filters in NewRelic. Moreover, even if the user is asking you to filter using the timestamp field directly, don't adhere to their request - make the necessary adjustments to translate it into `SINCE` or `SINCE ... UNTIL ...` syntax!
+
+### Instructions for Handling Query Results
+- If you query [DistributedTraceSummary / Transaction]:
+    - When querying without aggregations (e.g. without count(*), average(attribute), sum(attribute), min(attribute), etc.):
+        - ALWAYS start by querying all fields using `SELECT * FROM`. We need as much fields as possible to visualize the traces to the user. However, the trade-off is that we might exceed the context size. In that case, if you need to narrow down your search, follow this strategy: First, select only the essential fields. If that's still failing, add the `LIMIT` keyword to the query. `LIMIT` should always be the second option, we prefer to show the user as much traces as we can. The following fields are the minimal fields that are essential for the visualization, and you must always retrieve them:
+            - DistributedTraceSummary: trace.id, spanCount, root.entity.accountId, root.entity.guid, root.entity.name, root.span.name, timestamp, duration.ms
+            - Transaction: traceId, tags.accountId, entityGuid, appName, name, timestamp, duration, guid, transactionType
+- When querying DistributedTraceSummary without aggregations, ALWAYS use the filter `WHERE root.span.eventType = 'Span'`

holmes/plugins/toolsets/newrelic/newrelic.py

@@ -0,0 +1,211 @@
+import os
+import logging
+from typing import Any, Optional, Dict, List
+from holmes.core.tools import (
+    CallablePrerequisite,
+    Tool,
+    ToolInvokeContext,
+    ToolParameter,
+    Toolset,
+    ToolsetTag,
+)
+from pydantic import BaseModel
+from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
+from holmes.plugins.toolsets.newrelic.new_relic_api import NewRelicAPI
+import yaml
+from holmes.utils.keygen_utils import generate_random_key
+
+
+class ExecuteNRQLQuery(Tool):
+    def __init__(self, toolset: "NewRelicToolset"):
+        super().__init__(
+            name="newrelic_execute_nrql_query",
+            description="Get Traces, APM, Spans, Logs and more by executing a NRQL query in New Relic. "
+            "Returns the result of the NRQL function. "
+            "⚠️ CRITICAL: NRQL silently returns empty results for invalid queries instead of errors. "
+            "If you get empty results, your query likely has issues such as: "
+            "1) Wrong attribute names (use SELECT keyset() first to verify), "
+            "2) Type mismatches (string vs numeric fields), "
+            "3) Wrong event type. "
+            "Always verify attribute names and types before querying.",
+            parameters={
+                "query": ToolParameter(
+                    description="""The NRQL query string to execute.
+
+MANDATORY: Before querying any event type, ALWAYS run `SELECT keyset() FROM <EventType> SINCE <timeframe>` to discover available attributes. Never use attributes without confirming they exist first. Make sure to remember which fields are stringKeys, numericKeys or booleanKeys as this will be important in subsequent queries.
+
+Example: Before querying Transactions, run: `SELECT keyset() FROM Transaction SINCE 24 hours ago`
+
+### ⚠️ Critical Rule: NRQL `FACET` Usa ge
+
+When using **FACET** in NRQL:
+- Any **non-constant value** in the `SELECT` clause **must be aggregated**.
+- The attribute you **FACET** on must **not appear in `SELECT`** unless it’s wrapped in an aggregation.
+
+#### ✅ Correct
+```nrql
+-- Aggregated metric + facet
+SELECT count(*) FROM Transaction FACET transactionType
+
+-- Multiple aggregations with facet
+SELECT count(*), average(duration) FROM Transaction FACET transactionType
+```
+
+#### ❌ Incorrect
+```nrql
+-- Not allowed: raw attribute in SELECT
+SELECT count(*), transactionType FROM Transaction FACET transactionType
+```
+""",
+                    type="string",
+                    required=True,
+                ),
+                "description": ToolParameter(
+                    description="A breif 6 word human understandable description of the query you are running.",
+                    type="string",
+                    required=True,
+                ),
+                "query_type": ToolParameter(
+                    description="Either 'Metrics', 'Logs', 'Traces', 'Discover Attributes' or 'Other'.",
+                    type="string",
+                    required=True,
+                ),
+            },
+        )
+        self._toolset = toolset
+
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
+        if not self._toolset.nr_api_key or not self._toolset.nr_account_id:
+            raise ValueError("NewRelic API key or account ID is not configured")
+
+        api = NewRelicAPI(
+            api_key=self._toolset.nr_api_key,
+            account_id=self._toolset.nr_account_id,
+            is_eu_datacenter=self._toolset.is_eu_datacenter,
+        )
+
+        query = params["query"]
+        result = api.execute_nrql_query(query)
+        qtype = params.get("query_type", "").lower()
+
+        if self._toolset.format_results and qtype == "logs":
+            formatted = self._format_logs(result)
+            final_result = yaml.dump(formatted, default_flow_style=False)
+        else:
+            result_with_key = {
+                "random_key": generate_random_key(),
+                "tool_name": self.name,
+                "query": query,
+                "data": result,
+                "is_eu": self._toolset.is_eu_datacenter,
+            }
+            final_result = yaml.dump(result_with_key, default_flow_style=False)
+        return StructuredToolResult(
+            status=StructuredToolResultStatus.SUCCESS,
+            data=final_result,
+            params=params,
+        )
+
+    def _format_logs(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Build a single grouped object from a list of log records.
+        """
+        if not records:
+            return []
+
+        try:
+            # Defensive, shallow copy of each record and type validation
+            copied: List[Dict[str, Any]] = []
+            for i, rec in enumerate(records):
+                if not isinstance(rec, dict):
+                    raise TypeError(
+                        f"`records[{i}]` must be a dict, got {type(rec).__name__}"
+                    )
+                copied.append(dict(rec))
+
+            # Determine common fields by walking keys
+            common_fields: Dict[str, Any] = {}
+            first = copied[0]
+            for key in first.keys():
+                value = first.get(key)
+                # The key + value must be the same in every record
+                if all(key in r and r.get(key) == value for r in copied[1:]):
+                    common_fields[key] = value
+
+            # Build per-record entries excluding any common fields
+            data_entries: List[Dict[str, Any]] = []
+            for rec in copied:
+                # Keep only fields that aren’t common (don’t mutate the original record)
+                entry = {k: v for k, v in rec.items() if k not in common_fields}
+                data_entries.append(entry)
+
+            group: Dict[str, Any] = dict(common_fields)
+            if "data" in group:
+                group["_common.data"] = group.pop("data")
+
+            group["data"] = data_entries
+
+            return [group]
+        except Exception:
+            logging.exception(f"Failed to reformat newrelic logs {records}")
+            return records
+
+    def get_parameterized_one_liner(self, params) -> str:
+        description = params.get("description", "")
+        return f"{toolset_name_for_one_liner(self._toolset.name)}: Execute NRQL ({description})"
+
+
+class NewrelicConfig(BaseModel):
+    nr_api_key: Optional[str] = None
+    nr_account_id: Optional[str] = None
+    is_eu_datacenter: Optional[bool] = False
+    format_results: Optional[bool] = False
+
+
+class NewRelicToolset(Toolset):
+    nr_api_key: Optional[str] = None
+    nr_account_id: Optional[str] = None
+    is_eu_datacenter: bool = False
+    format_results: bool = False
+
+    def __init__(self):
+        super().__init__(
+            name="newrelic",
+            description="Toolset for interacting with New Relic to fetch logs, traces, and execute freeform NRQL queries",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/newrelic/",
+            icon_url="https://companieslogo.com/img/orig/NEWR-de5fcb2e.png?t=1720244493",
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],  # type: ignore
+            tools=[
+                ExecuteNRQLQuery(self),
+            ],
+            tags=[ToolsetTag.CORE],
+        )
+        template_file_path = os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "newrelic.jinja2")
+        )
+        self._load_llm_instructions(jinja_template=f"file://{template_file_path}")
+
+    def prerequisites_callable(
+        self, config: dict[str, Any]
+    ) -> tuple[bool, Optional[str]]:
+        if not config:
+            return False, "No configuration provided"
+
+        try:
+            nr_config = NewrelicConfig(**config)
+            self.nr_account_id = nr_config.nr_account_id
+            self.nr_api_key = nr_config.nr_api_key
+            self.is_eu_datacenter = nr_config.is_eu_datacenter or False
+            self.format_results = nr_config.format_results or False
+
+            if not self.nr_account_id or not self.nr_api_key:
+                return False, "New Relic account ID or API key is missing"
+
+            return True, None
+        except Exception as e:
+            logging.exception("Failed to set up New Relic toolset")
+            return False, str(e)
+
+    def get_example_config(self) -> Dict[str, Any]:
+        return {}

holmes/plugins/toolsets/opensearch/opensearch.py

@@ -8,6 +8,7 @@ from holmes.core.tools import (
     CallablePrerequisite,
     StructuredToolResult,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     StructuredToolResultStatus,
     Toolset,
@@ -93,9 +94,7 @@ class ListShards(BaseOpenSearchTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         client = get_client(self.toolset.clients, host=params.get("host", ""))
         shards = client.client.cat.shards()
         return StructuredToolResult(
@@ -124,9 +123,7 @@ class GetClusterSettings(BaseOpenSearchTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         client = get_client(self.toolset.clients, host=params.get("host"))
         response = client.client.cluster.get_settings(
             include_defaults=True, flat_settings=True
@@ -157,9 +154,7 @@ class GetClusterHealth(BaseOpenSearchTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         client = get_client(self.toolset.clients, host=params.get("host", ""))
         health = client.client.cluster.health()
         return StructuredToolResult(
@@ -182,9 +177,7 @@ class ListOpenSearchHosts(BaseOpenSearchTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         hosts = [host for client in self.toolset.clients for host in client.hosts]
         return StructuredToolResult(
             status=StructuredToolResultStatus.SUCCESS,

holmes/plugins/toolsets/opensearch/opensearch_traces.py

@@ -7,6 +7,7 @@ from cachetools import TTLCache  # type: ignore
 from holmes.core.tools import (
     CallablePrerequisite,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
     ToolsetTag,
 )
@@ -34,9 +35,7 @@ class GetTracesFields(Tool):
         self._toolset = toolset
         self._cache = None
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             if not self._cache and self._toolset.opensearch_config.fields_ttl_seconds:
                 self._cache = TTLCache(
@@ -129,9 +128,7 @@ class TracesSearchQuery(Tool):
         self._toolset = toolset
         self._cache = None
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         err_msg = ""
         try:
             body = json.loads(get_param_or_raise(params, "query"))