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

holmes/plugins/toolsets/newrelic/new_relic_api.py

@@ -5,7 +5,6 @@ from typing import Any, Dict
 
 import requests  # type: ignore
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -123,3 +122,44 @@ class NewRelicAPI:
             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

holmes/plugins/toolsets/newrelic/newrelic.jinja2

@@ -1,5 +1,29 @@
 New Relic provides distributed tracing data along with logs and metrics.
 
+{% if config.enable_multi_account %}
+**MULTI-ACCOUNT MODE**: You have access to multiple New Relic accounts in this organization.
+
+### Important Multi-Account Workflow
+
+**Each NRQL query MUST include the account_id parameter.**
+1. A New Relic account ID is a numeric identifier, typically a 6–8 digit integer (e.g., 1234567).
+
+**Here's how to determine which account_id to use**
+
+1. **ALWAYS Check context first**: Look for common new relic labels or tags with the account id or name such as  `nrAccountId` `accountId` or `account` in the provided context 
+(e.g., from alerts, traces, or previous queries). If found, use that value.
+
+2. **ALWAYS CHECK if Account name provided**: If the user mentions a specific account name (e.g., "Production Account", "Staging"):
+   - YOU MUST First call `newrelic_list_organization_accounts` to get the list of all accounts
+   - Find the matching account by name and use its ID
+
+3. **No account specified**: If you can't find any account ID or name based on the context of the question.
+   - Use the function newrelic_execute_nrql_query default account id value as the account ID.
+   - Let the user know you have used the default account.
+
+**Important**: The context may contain account IDs in various places - check trace data, alert metadata, or previous query results for `nrAccountId`, `accountId`, `account.id` or similar fields.
+
+{% endif %}
 Assume every application has New Relic tracing data.
 
 Use `nrql_query` to run a NRQL query.

holmes/plugins/toolsets/newrelic/newrelic.py

@@ -1,37 +1,78 @@
-import os
+import base64
+import json
 import logging
-from typing import Any, Optional, Dict
+import os
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel
+
 from holmes.core.tools import (
     CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
     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
+from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
+
+
+def _build_newrelic_query_url(
+    base_url: str,
+    account_id: str,
+    nrql_query: str,
+) -> Optional[str]:
+    """Build a New Relic query URL for the NRQL query builder.
+
+    Note: URL links to queries are not officially supported by New Relic, so we are using
+    a workaround to open their overlay to the query builder with the query pre-filled.
+    This uses the dashboard launcher with an overlay parameter to open the query builder nerdlet.
+
+    """
+    try:
+        account_id_int = int(account_id) if isinstance(account_id, str) else account_id
+
+        overlay = {
+            "nerdletId": "data-exploration.query-builder",
+            "initialActiveInterface": "nrqlEditor",
+            "initialQueries": [
+                {
+                    "accountId": account_id_int,
+                    "nrql": nrql_query,
+                }
+            ],
+        }
+
+        overlay_json = json.dumps(overlay, separators=(",", ":"))
+        overlay_base64 = base64.b64encode(overlay_json.encode("utf-8")).decode("utf-8")
+
+        pane = {
+            "nerdletId": "dashboards.list",
+            "entityDomain": "VIZ",
+            "entityType": "DASHBOARD",
+        }
+        pane_json = json.dumps(pane, separators=(",", ":"))
+        pane_base64 = base64.b64encode(pane_json.encode("utf-8")).decode("utf-8")
+
+        url = (
+            f"{base_url}/launcher/dashboards.launcher"
+            f"?pane={pane_base64}"
+            f"&overlay={overlay_base64}"
+        )
+
+        return url
+    except Exception:
+        return None
 
 
 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.
+        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.
 
@@ -58,48 +99,78 @@ SELECT count(*), average(duration) FROM Transaction FACET transactionType
 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,
+                type="string",
+                required=True,
+            ),
+            "description": ToolParameter(
+                description="A brief 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,
+            ),
+        }
+
+        # Add account_id parameter only in multi-account mode
+        if toolset.enable_multi_account:
+            parameters["account_id"] = ToolParameter(
+                description=(
+                    f"A New Relic account ID is a numeric identifier, typically a 6-8 digit integer (e.g., 1234567). It contains only digits, has no prefixes or separators, and uniquely identifies a New Relic account. default: {toolset.nr_account_id}"
                 ),
-            },
+                type="integer",
+                required=True,
+            )
+
+        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=parameters,
         )
         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")
+        if self._toolset.enable_multi_account:
+            account_id = params.get("account_id") or self._toolset.nr_account_id
+            account_id = str(account_id)
+        else:
+            account_id = self._toolset.nr_account_id
 
-        api = NewRelicAPI(
-            api_key=self._toolset.nr_api_key,
-            account_id=self._toolset.nr_account_id,
-            is_eu_datacenter=self._toolset.is_eu_datacenter,
-        )
+        if not account_id:
+            raise ValueError("NewRelic account ID is not configured")
+
+        api = self._toolset.create_api_client(account_id)
 
         query = params["query"]
         result = api.execute_nrql_query(query)
 
         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)
+
+        # Build New Relic query URL
+        explore_url = _build_newrelic_query_url(
+            base_url=self._toolset.base_url,
+            account_id=account_id,
+            nrql_query=query,
+        )
+
         return StructuredToolResult(
             status=StructuredToolResultStatus.SUCCESS,
-            data=final_result,
+            data=result_with_key,
             params=params,
+            url=explore_url,
         )
 
     def get_parameterized_one_liner(self, params) -> str:
@@ -107,16 +178,100 @@ SELECT count(*), transactionType FROM Transaction FACET transactionType
         return f"{toolset_name_for_one_liner(self._toolset.name)}: Execute NRQL ({description})"
 
 
+class ListOrganizationAccounts(Tool):
+    def __init__(self, toolset: "NewRelicToolset"):
+        super().__init__(
+            name="newrelic_list_organization_accounts",
+            description=(
+                "List all account names and IDs accessible in the New Relic organization. "
+                "Use this tool to:\n"
+                "1. Find the account ID when given an account name\n"
+                "2. Map account names to IDs for running NRQL queries\n"
+                "Returns a list of accounts with 'id' and 'name' fields."
+            ),
+            parameters={},
+        )
+        self._toolset = toolset
+
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
+        api = self._toolset.create_api_client(
+            account_id="0"
+        )  # organization query does not need account_id
+
+        accounts = api.get_organization_accounts()
+
+        result_with_key = {
+            "accounts": accounts,
+            "total_count": len(accounts),
+            "is_eu": self._toolset.is_eu_datacenter,
+        }
+
+        # Build New Relic accounts URL
+        accounts_url = (
+            f"{self._toolset.base_url}/admin-portal/organizations/organization-detail"
+        )
+
+        return StructuredToolResult(
+            status=StructuredToolResultStatus.SUCCESS,
+            data=result_with_key,
+            params=params,
+            url=accounts_url,
+        )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return f"{toolset_name_for_one_liner(self._toolset.name)}: List organization accounts"
+
+
 class NewrelicConfig(BaseModel):
-    nr_api_key: Optional[str] = None
-    nr_account_id: Optional[str] = None
+    nr_api_key: str
+    nr_account_id: str
     is_eu_datacenter: Optional[bool] = False
+    enable_multi_account: Optional[bool] = False
 
 
 class NewRelicToolset(Toolset):
     nr_api_key: Optional[str] = None
     nr_account_id: Optional[str] = None
     is_eu_datacenter: bool = False
+    enable_multi_account: bool = False
+
+    @property
+    def base_url(self) -> str:
+        """Get the New Relic base URL based on datacenter region."""
+        return (
+            "https://one.eu.newrelic.com"
+            if self.is_eu_datacenter
+            else "https://one.newrelic.com"
+        )
+
+    def create_api_client(self, account_id: Optional[str] = None) -> NewRelicAPI:
+        """Create a NewRelicAPI client instance.
+
+        Args:
+            account_id: Account ID to use. If None, uses the default from config.
+                       Set to "0" for organization-level queries.
+
+        Returns:
+            Configured NewRelicAPI instance
+
+        Raises:
+            ValueError: If API key is not configured
+        """
+        if not self.nr_api_key:
+            raise ValueError("NewRelic API key is not configured")
+
+        effective_account_id = (
+            account_id if account_id is not None else self.nr_account_id
+        )
+
+        if not effective_account_id:
+            raise ValueError("NewRelic Account id is not configured")
+
+        return NewRelicAPI(
+            api_key=self.nr_api_key,
+            account_id=effective_account_id,
+            is_eu_datacenter=self.is_eu_datacenter,
+        )
 
     def __init__(self):
         super().__init__(
@@ -125,15 +280,9 @@ class NewRelicToolset(Toolset):
             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),
-            ],
+            tools=[],
             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]
@@ -146,9 +295,16 @@ class NewRelicToolset(Toolset):
             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.enable_multi_account = nr_config.enable_multi_account 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"
+            # Tool uses enable_multi_account flag.
+            self.tools = [ExecuteNRQLQuery(self)]
+            if self.enable_multi_account:
+                self.tools.append(ListOrganizationAccounts(self))
+            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}")
 
             return True, None
         except Exception as e:
@@ -160,4 +316,5 @@ class NewRelicToolset(Toolset):
             "nr_api_key": "NRAK-XXXXXXXXXXXXXXXXXXXXXXXXXX",
             "nr_account_id": "1234567",
             "is_eu_datacenter": False,
+            "enable_multi_account": False,
         }