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

holmes/plugins/toolsets/datadog/toolset_datadog_logs.py

@@ -1,268 +1,137 @@
-import os
-from enum import Enum
 import json
 import logging
-from typing import Any, Optional, Dict, Tuple, Set
+import os
+from typing import Any, Dict, Optional, Tuple
+
+from pydantic import AnyUrl
+
 from holmes.core.tools import (
     CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    Tool,
+    ToolInvokeContext,
+    ToolParameter,
     ToolsetTag,
 )
-from pydantic import BaseModel, Field
-from holmes.core.tools import StructuredToolResult, ToolResultStatus
-from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR
+from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION
 from holmes.plugins.toolsets.datadog.datadog_api import (
-    DatadogBaseConfig,
+    MAX_RETRY_COUNT_ON_RATE_LIMIT,
     DataDogRequestError,
-    execute_paginated_datadog_http_request,
+    execute_datadog_http_request,
     get_headers,
-    MAX_RETRY_COUNT_ON_RATE_LIMIT,
 )
+from holmes.plugins.toolsets.datadog.datadog_models import (
+    DatadogLogsConfig,
+)
+from holmes.plugins.toolsets.datadog.datadog_url_utils import generate_datadog_logs_url
 from holmes.plugins.toolsets.logging_utils.logging_api import (
-    DEFAULT_TIME_SPAN_SECONDS,
     DEFAULT_LOG_LIMIT,
-    BasePodLoggingToolset,
-    FetchPodLogsParams,
-    LoggingCapability,
-    PodLoggingTool,
+    DEFAULT_TIME_SPAN_SECONDS,
+    Toolset,
+)
+from holmes.plugins.toolsets.utils import (
+    process_timestamps_to_int,
+    standard_start_datetime_tool_param_description,
+    toolset_name_for_one_liner,
 )
-from holmes.plugins.toolsets.utils import process_timestamps_to_rfc3339
-
-
-class DataDogLabelsMapping(BaseModel):
-    pod: str = "pod_name"
-    namespace: str = "kube_namespace"
-
-
-class DataDogStorageTier(str, Enum):
-    INDEXES = "indexes"
-    ONLINE_ARCHIVES = "online-archives"
-    FLEX = "flex"
-
-
-DEFAULT_STORAGE_TIERS = [DataDogStorageTier.INDEXES]
-
-
-class DatadogLogsConfig(DatadogBaseConfig):
-    indexes: list[str] = ["*"]
-    # Ordered list of storage tiers. Works as fallback. Subsequent tiers are queried only if the previous tier yielded no result
-    storage_tiers: list[DataDogStorageTier] = Field(
-        default=DEFAULT_STORAGE_TIERS, min_length=1
-    )
-    labels: DataDogLabelsMapping = DataDogLabelsMapping()
-    page_size: int = 300
-    default_limit: int = DEFAULT_LOG_LIMIT
-
-
-def calculate_page_size(
-    params: FetchPodLogsParams, dd_config: DatadogLogsConfig, logs: list
-) -> int:
-    logs_count = len(logs)
-
-    max_logs_count = dd_config.default_limit
-    if params.limit:
-        max_logs_count = params.limit
-
-    return min(dd_config.page_size, max(0, max_logs_count - logs_count))
-
-
-def fetch_paginated_logs(
-    params: FetchPodLogsParams,
-    dd_config: DatadogLogsConfig,
-    storage_tier: DataDogStorageTier,
-) -> list[dict]:
-    limit = params.limit or dd_config.default_limit
-
-    (from_time, to_time) = process_timestamps_to_rfc3339(
-        start_timestamp=params.start_time,
-        end_timestamp=params.end_time,
-        default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS,
-    )
-
-    url = f"{dd_config.site_api_url}/api/v2/logs/events/search"
-    headers = get_headers(dd_config)
-
-    query = f"{dd_config.labels.namespace}:{params.namespace}"
-    query += f" {dd_config.labels.pod}:{params.pod_name}"
-    if params.filter:
-        filter = params.filter.replace('"', '\\"')
-        query += f' "{filter}"'
-
-    payload: Dict[str, Any] = {
-        "filter": {
-            "from": from_time,
-            "to": to_time,
-            "query": query,
-            "indexes": dd_config.indexes,
-            "storage_tier": storage_tier.value,
-        },
-        "sort": "-timestamp",
-        "page": {"limit": calculate_page_size(params, dd_config, [])},
-    }
-
-    logs, cursor = execute_paginated_datadog_http_request(
-        url=url,
-        headers=headers,
-        payload_or_params=payload,
-        timeout=dd_config.request_timeout,
-    )
-
-    while cursor and len(logs) < limit:
-        payload["page"]["cursor"] = cursor
-        new_logs, cursor = execute_paginated_datadog_http_request(
-            url=url,
-            headers=headers,
-            payload_or_params=payload,
-            timeout=dd_config.request_timeout,
-        )
-        logs += new_logs
-        payload["page"]["limit"] = calculate_page_size(params, dd_config, logs)
-
-    # logs are fetched descending order. Unified logging API follows the pattern of kubectl logs where oldest logs are first
-    logs.reverse()
-
-    if len(logs) > limit:
-        logs = logs[-limit:]
-    return logs
 
 
 def format_logs(raw_logs: list[dict]) -> str:
+    # Use similar structure to Datadog Log Explorer
     logs = []
 
     for raw_log_item in raw_logs:
-        message = raw_log_item.get("attributes", {}).get(
-            "message", json.dumps(raw_log_item)
-        )
-        logs.append(message)
+        attrs = raw_log_item.get("attributes", {})
+
+        timestamp = attrs.get("timestamp") or attrs.get("@timestamp", "")
+        host = attrs.get("host", "")
+        service = attrs.get("service", "")
+        status = attrs.get("attributes", {}).get("status") or attrs.get("status", "")
+        message = attrs.get("message", json.dumps(raw_log_item))
+        tags = attrs.get("tags", [])
+
+        pod_name_tag = next((t for t in tags if t.startswith("pod_")), "")
+
+        log_line = f"{timestamp} {host} {pod_name_tag} {service} {status} {message}"
+        logs.append(log_line)
 
     return "\n".join(logs)
 
 
-class DatadogLogsToolset(BasePodLoggingToolset):
-    dd_config: Optional[DatadogLogsConfig] = None
+class DatadogLogsToolset(Toolset):
+    """Toolset for working with Datadog logs data."""
 
-    @property
-    def supported_capabilities(self) -> Set[LoggingCapability]:
-        """Datadog logs API supports historical data and substring matching"""
-        return {
-            LoggingCapability.HISTORICAL_DATA
-        }  # No regex support, no exclude filter, but supports historical data
+    dd_config: Optional[DatadogLogsConfig] = None
 
     def __init__(self):
         super().__init__(
             name="datadog/logs",
             description="Toolset for fetching logs from Datadog, including historical data for pods no longer in the cluster",
-            docs_url="https://docs.datadoghq.com/api/latest/logs/",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/datadog/",
             icon_url="https://imgix.datadoghq.com//img/about/presskit/DDlogo.jpg",
             prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
             tools=[],  # Initialize with empty tools first
-            experimental=True,
             tags=[ToolsetTag.CORE],
         )
         # Now that parent is initialized and self.name exists, create the tool
-        self.tools = [PodLoggingTool(self)]
+        self.tools = [GetLogs(toolset=self)]
         self._reload_instructions()
 
-    def logger_name(self) -> str:
-        return "DataDog"
-
-    def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult:
-        if not self.dd_config:
-            return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
-                data=TOOLSET_CONFIG_MISSING_ERROR,
-                params=params.model_dump(),
-            )
-
-        try:
-            raw_logs = []
-            for storage_tier in self.dd_config.storage_tiers:
-                raw_logs = fetch_paginated_logs(
-                    params, self.dd_config, storage_tier=storage_tier
-                )
-
-                if raw_logs:
-                    logs_str = format_logs(raw_logs)
-                    return StructuredToolResult(
-                        status=ToolResultStatus.SUCCESS,
-                        data=logs_str,
-                        params=params.model_dump(),
-                    )
-
-            return StructuredToolResult(
-                status=ToolResultStatus.NO_DATA,
-                params=params.model_dump(),
-            )
-
-        except DataDogRequestError as e:
-            logging.exception(e, exc_info=True)
-
-            # Provide more specific error message for rate limiting failures
-            if e.status_code == 429:
-                error_msg = f"Datadog API rate limit exceeded. Failed after {MAX_RETRY_COUNT_ON_RATE_LIMIT} retry attempts."
-            else:
-                error_msg = f"Exception while querying Datadog: {str(e)}"
-
-            return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
-                error=error_msg,
-                params=params.model_dump(),
-                invocation=json.dumps(e.payload),
-            )
-
-        except Exception as e:
-            logging.exception(
-                f"Failed to query Datadog logs for params: {params}", exc_info=True
-            )
-            return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
-                error=f"Exception while querying Datadog: {str(e)}",
-                params=params.model_dump(),
-            )
-
     def _perform_healthcheck(self) -> Tuple[bool, str]:
-        """
-        Perform a healthcheck by fetching a single log from Datadog.
-        Returns (success, error_message).
-        """
+        """Perform health check on Datadog logs API."""
+        if not self.dd_config:
+            return False, "Datadog configuration not initialized"
         try:
-            logging.info("Performing Datadog configuration healthcheck...")
-            healthcheck_params = FetchPodLogsParams(
-                namespace="*",
-                pod_name="*",
-                limit=1,
-                start_time="-172800",  # 48 hours in seconds
+            logging.info("Performing Datadog logs configuration healthcheck...")
+            headers = get_headers(self.dd_config)
+            payload = {
+                "filter": {
+                    "from": "now-1m",
+                    "to": "now",
+                    "query": "*",
+                    "indexes": self.dd_config.indexes,
+                },
+                "page": {"limit": 1},
+            }
+
+            search_url = f"{self.dd_config.site_api_url}/api/v2/logs/events/search"
+            execute_datadog_http_request(
+                url=search_url,
+                headers=headers,
+                payload_or_params=payload,
+                timeout=self.dd_config.request_timeout,
+                method="POST",
             )
 
-            result = self.fetch_pod_logs(healthcheck_params)
-
-            if result.status == ToolResultStatus.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:
-                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}"
-
-            logging.info("Datadog healthcheck completed successfully")
             return True, ""
 
+        except DataDogRequestError as e:
+            logging.error(
+                f"Datadog API error during healthcheck: {e.status_code} - {e.response_text}"
+            )
+            if e.status_code == 403:
+                return (
+                    False,
+                    "API key lacks required permissions. Make sure your API key has 'apm_read' scope.",
+                )
+            else:
+                return False, f"Datadog API error: {e.status_code} - {e.response_text}"
         except Exception as e:
-            logging.exception("Failed during Datadog healthcheck")
+            logging.exception("Failed during Datadog traces healthcheck")
             return False, f"Healthcheck failed with exception: {str(e)}"
 
     def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
         if not config:
             return (
                 False,
-                TOOLSET_CONFIG_MISSING_ERROR,
+                "Missing config for dd_api_key, dd_app_key, or site_api_url. For details: https://holmesgpt.dev/data-sources/builtin-toolsets/datadog/",
             )
 
         try:
             dd_config = DatadogLogsConfig(**config)
             self.dd_config = dd_config
 
-            # Perform healthcheck
             success, error_msg = self._perform_healthcheck()
             return success, error_msg
 
@@ -271,11 +140,13 @@ class DatadogLogsToolset(BasePodLoggingToolset):
             return (False, f"Failed to parse Datadog configuration: {str(e)}")
 
     def get_example_config(self) -> Dict[str, Any]:
-        return {
-            "dd_api_key": "your-datadog-api-key",
-            "dd_app_key": "your-datadog-application-key",
-            "site_api_url": "https://api.datadoghq.com",
-        }
+        """Get example configuration for this toolset."""
+        example_config = DatadogLogsConfig(
+            dd_api_key="<your_datadog_api_key>",
+            dd_app_key="<your_datadog_app_key>",
+            site_api_url=AnyUrl("https://api.datadoghq.com"),
+        )
+        return example_config.model_dump(mode="json")
 
     def _reload_instructions(self):
         """Load Datadog logs specific troubleshooting instructions."""
@@ -283,3 +154,152 @@ class DatadogLogsToolset(BasePodLoggingToolset):
             os.path.join(os.path.dirname(__file__), "datadog_logs_instructions.jinja2")
         )
         self._load_llm_instructions(jinja_template=f"file://{template_file_path}")
+
+
+class GetLogs(Tool):
+    """Tool to search for logs with specific search query."""
+
+    toolset: "DatadogLogsToolset"
+    name: str = "fetch_datadog_logs"
+    description: str = "Search for logs in Datadog using search query syntax"
+    "Uses the DataDog api endpoint: POST /api/v2/logs/events/search with 'query' parameter. (e.g., 'service:web-app @http.status_code:500')"
+    parameters: Dict[str, ToolParameter] = {
+        "query": ToolParameter(
+            description="The search query - following the logs search syntax. default: *",
+            type="string",
+            required=False,
+        ),
+        "start_datetime": ToolParameter(
+            description=standard_start_datetime_tool_param_description(
+                DEFAULT_TIME_SPAN_SECONDS
+            ),
+            type="string",
+            required=False,
+        ),
+        "end_datetime": ToolParameter(
+            description=STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION,
+            type="string",
+            required=False,
+        ),
+        "cursor": ToolParameter(
+            description="The returned paging point to use to get the next results. IMPORTANT: Cursors are single-use and stateful - never reuse the same cursor value multiple times or parallelize cursor-based calls. Each response provides a new cursor for the subsequent request.",
+            type="string",
+            required=False,
+        ),
+        "limit": ToolParameter(
+            description=f"Maximum number of log records to return. Defaults to {DEFAULT_LOG_LIMIT}. This value is user-configured and represents the maximum allowed limit.",
+            type="integer",
+            required=False,
+        ),
+        "sort_desc": ToolParameter(
+            description="Get the results in descending order. default: true",
+            type="boolean",
+            required=False,
+        ),
+    }
+
+    def get_parameterized_one_liner(self, params: dict) -> str:
+        """Get a one-liner description of the tool invocation."""
+        return f"{toolset_name_for_one_liner(self.toolset.name)}: Search Logs ({params['query'] if 'query' in params else ''})"
+
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
+        """Execute the tool to search logs."""
+        if not self.toolset.dd_config:
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error="Datadog configuration not initialized",
+                params=params,
+            )
+        url = None
+        payload: Optional[Dict[str, Any]] = None
+        try:
+            # Process timestamps
+            from_time_int, to_time_int = process_timestamps_to_int(
+                start=params.get("start_datetime"),
+                end=params.get("end_datetime"),
+                default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS,
+            )
+
+            # Convert to milliseconds for Datadog API
+            from_time_ms = from_time_int * 1000
+            to_time_ms = to_time_int * 1000
+
+            config_limit = self.toolset.dd_config.default_limit
+            limit = min(params.get("limit", config_limit), config_limit)
+            params["limit"] = limit
+            sort = "timestamp" if params.get("sort_desc", False) else "-timestamp"
+
+            url = f"{self.toolset.dd_config.site_api_url}/api/v2/logs/events/search"
+            headers = get_headers(self.toolset.dd_config)
+
+            storage = self.toolset.dd_config.storage_tiers[-1]
+            payload = {
+                "filter": {
+                    "query": params.get("query", "*"),
+                    "from": str(from_time_ms),
+                    "to": str(to_time_ms),
+                    "storage_tier": storage,
+                    "indexes": self.toolset.dd_config.indexes,
+                },
+                "page": {
+                    "limit": limit,
+                },
+                "sort": sort,
+            }
+
+            if params.get("cursor"):
+                payload["page"]["cursor"] = params["cursor"]
+
+            response = execute_datadog_http_request(
+                url=url,
+                headers=headers,
+                payload_or_params=payload,
+                timeout=self.toolset.dd_config.request_timeout,
+                method="POST",
+            )
+
+            if self.toolset.dd_config.compact_logs and response.get("data"):
+                response["data"] = format_logs(response["data"])
+
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.SUCCESS,
+                data=response,
+                params=params,
+                url=generate_datadog_logs_url(self.toolset.dd_config, payload),
+            )
+
+        except DataDogRequestError as e:
+            logging.exception(e, exc_info=True)
+            if e.status_code == 429:
+                error_msg = f"Datadog API rate limit exceeded. Failed after {MAX_RETRY_COUNT_ON_RATE_LIMIT} retry attempts."
+            elif e.status_code == 403:
+                error_msg = (
+                    f"Permission denied. Ensure your Datadog Application Key has the 'apm_read' "
+                    f"permission. Error: {str(e)}"
+                )
+            else:
+                error_msg = f"Exception while querying Datadog: {str(e)}"
+
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error=error_msg,
+                params=params,
+                invocation=(
+                    json.dumps({"url": url, "payload": payload})
+                    if url and payload
+                    else None
+                ),
+            )
+
+        except Exception as e:
+            logging.exception(e, exc_info=True)
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error=f"Unexpected error: {str(e)}",
+                params=params,
+                invocation=(
+                    json.dumps({"url": url, "payload": payload})
+                    if url and payload
+                    else None
+                ),
+            )