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

holmes/plugins/toolsets/coralogix/api.py

@@ -1,25 +1,12 @@
-from enum import Enum
+import json
 import logging
-from typing import Any, Tuple
+from enum import Enum
+from typing import Any, Optional, Tuple
 from urllib.parse import urljoin
 
 import requests  # type: ignore
 
-from holmes.plugins.toolsets.coralogix.utils import (
-    CoralogixConfig,
-    CoralogixQueryResult,
-    merge_log_results,
-    parse_logs,
-    CoralogixLogsMethodology,
-)
-from holmes.plugins.toolsets.logging_utils.logging_api import (
-    FetchPodLogsParams,
-    DEFAULT_TIME_SPAN_SECONDS,
-    DEFAULT_LOG_LIMIT,
-)
-from holmes.plugins.toolsets.utils import (
-    process_timestamps_to_rfc3339,
-)
+from holmes.plugins.toolsets.coralogix.utils import parse_json_lines
 
 
 class CoralogixTier(str, Enum):
@@ -31,130 +18,156 @@ def get_dataprime_base_url(domain: str) -> str:
     return f"https://ng-api-http.{domain}"
 
 
-def execute_http_query(domain: str, api_key: str, query: dict[str, Any]):
-    base_url = get_dataprime_base_url(domain)
-    url = urljoin(base_url, "api/v1/dataprime/query")
-    headers = {
+def _get_auth_headers(api_key: str) -> dict[str, str]:
+    return {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
     }
 
-    return requests.post(url, headers=headers, json=query)
 
+def execute_coralogix_query(
+    domain: str, api_key: str, query: dict[str, Any]
+) -> Tuple[requests.Response, str]:
+    base_url = get_dataprime_base_url(domain).rstrip("/") + "/"
+    url = urljoin(base_url, "api/v1/dataprime/query")
+    response = requests.post(
+        url,
+        headers=_get_auth_headers(api_key),
+        json=query,
+        timeout=(10, 120),
+    )
+    return response, url
 
-def health_check(domain: str, api_key: str) -> Tuple[bool, str]:
-    query = {"query": "source logs | limit 1"}
 
-    response = execute_http_query(domain=domain, api_key=api_key, query=query)
+def _parse_ndjson_response(response_text: str) -> Optional[Any]:
+    """Parse NDJSON response from Coralogix API."""
+    json_objects = parse_json_lines(response_text)
+    if not json_objects:
+        return None
 
-    if response.status_code == 200:
-        return True, ""
-    else:
-        return False, f"Failed with status_code={response.status_code}. {response.text}"
+    results: list[Any] = []
 
+    for obj in json_objects:
+        if not isinstance(obj, dict):
+            continue
 
-def build_query_string(config: CoralogixConfig, params: FetchPodLogsParams) -> str:
-    query_filters = []
-    query_filters.append(f'{config.labels.namespace}:"{params.namespace}"')
-    query_filters.append(f'{config.labels.pod}:"{params.pod_name}"')
+        if any(k in obj for k in ("result", "results", "batches", "records")):
+            results.append(obj)
 
-    if params.filter:
-        query_filters.append(f'{config.labels.log_message}:"{params.filter}"')
+    if not results:
+        return None
 
-    query_string = " AND ".join(query_filters)
-    query_string = f"source logs | lucene '{query_string}' | limit {params.limit or DEFAULT_LOG_LIMIT}"
-    return query_string
+    return results
 
 
-def get_start_end(params: FetchPodLogsParams):
-    (start, end) = process_timestamps_to_rfc3339(
-        start_timestamp=params.start_time,
-        end_timestamp=params.end_time,
-        default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS,
-    )
-    return (start, end)
+def _build_query_dict(
+    dataprime_query: str,
+    start_date: Optional[str] = None,
+    end_date: Optional[str] = None,
+    tier: Optional[CoralogixTier] = None,
+) -> dict[str, Any]:
+    metadata: dict[str, Any] = {"syntax": "QUERY_SYNTAX_DATAPRIME"}
+    if start_date:
+        metadata["startDate"] = start_date
+    if end_date:
+        metadata["endDate"] = end_date
+    if tier:
+        metadata["tier"] = tier.value
 
+    return {"query": dataprime_query, "metadata": metadata}
 
-def build_query(
-    config: CoralogixConfig, params: FetchPodLogsParams, tier: CoralogixTier
-):
-    (start, end) = get_start_end(params)
 
-    query_string = build_query_string(config, params)
-    return {
-        "query": query_string,
-        "metadata": {
-            "tier": tier.value,
-            "syntax": "QUERY_SYNTAX_DATAPRIME",
-            "startDate": start,
-            "endDate": end,
-        },
-    }
+def _get_error_body(response: requests.Response) -> str:
+    """Extract error body from response."""
+    try:
+        return (response.text or "").strip()
+    except Exception:
+        return ""
+
+
+def _cleanup_coralogix_results(parsed: list[Any]) -> Any:
+    """Clean up and normalize parsed Coralogix results structure."""
+    # Extract nested results if present
+    if len(parsed) == 1 and isinstance(parsed[0], dict) and "result" in parsed[0]:
+        nested_result = parsed[0]["result"]
+        if isinstance(nested_result, dict) and "results" in nested_result:
+            parsed = nested_result["results"]
+
+    # Replace items with userData JSON if present
+    # userData has additional data that is missing in the main result object along with the actual result
+    for i, item in enumerate(parsed):
+        if isinstance(item, dict) and "userData" in item:
+            try:
+                parsed[i] = json.loads(item["userData"])
+            except (json.JSONDecodeError, TypeError):
+                # If parsing fails, keep the original item
+                pass
+
+    return parsed
+
+
+def execute_dataprime_query(
+    domain: str,
+    api_key: str,
+    dataprime_query: str,
+    start_date: Optional[str] = None,
+    end_date: Optional[str] = None,
+    tier: Optional[CoralogixTier] = None,
+    max_poll_attempts: int = 60,
+    poll_interval_seconds: float = 1.0,
+) -> Tuple[Optional[Any], Optional[str]]:
+    try:
+        query_dict = _build_query_dict(dataprime_query, start_date, end_date, tier)
+        response, submit_url = execute_coralogix_query(domain, api_key, query_dict)
+
+        if response.status_code != 200:
+            body = _get_error_body(response)
+            if "Compiler error" in body or "Compilation errors" in body:
+                return (
+                    None,
+                    f"Compilation errors: {body}\nUse lucene instead of filter and verify that all labels are present before using them.",
+                )
+            return (
+                None,
+                f"Failed to submit query: status_code={response.status_code}, {body}\nURL: {submit_url}",
+            )
 
+        raw = response.text.strip()
+        if not raw:
+            return None, f"Empty 200 response from query submission\nURL: {submit_url}"
 
-def query_logs_for_tier(
-    config: CoralogixConfig, params: FetchPodLogsParams, tier: CoralogixTier
-) -> CoralogixQueryResult:
-    http_status = None
-    try:
-        query = build_query(config, params, tier)
+        parsed = _parse_ndjson_response(raw)
 
-        response = execute_http_query(
-            domain=config.domain,
-            api_key=config.api_key,
-            query=query,
-        )
-        http_status = response.status_code
-        if http_status == 200:
-            logs = parse_logs(
-                raw_logs=response.text.strip(), labels_config=config.labels
-            )
-            return CoralogixQueryResult(logs=logs, http_status=http_status, error=None)
-        else:
-            return CoralogixQueryResult(
-                logs=[], http_status=http_status, error=response.text
-            )
-    except Exception as e:
-        logging.error("Failed to fetch coralogix logs", exc_info=True)
-        return CoralogixQueryResult(logs=[], http_status=http_status, error=str(e))
-
-
-def query_logs_for_all_tiers(
-    config: CoralogixConfig, params: FetchPodLogsParams
-) -> CoralogixQueryResult:
-    methodology = config.logs_retrieval_methodology
-    result: CoralogixQueryResult
-
-    if methodology in [
-        CoralogixLogsMethodology.FREQUENT_SEARCH_ONLY,
-        CoralogixLogsMethodology.BOTH_FREQUENT_SEARCH_AND_ARCHIVE,
-        CoralogixLogsMethodology.ARCHIVE_FALLBACK,
-    ]:
-        result = query_logs_for_tier(
-            config=config, params=params, tier=CoralogixTier.FREQUENT_SEARCH
-        )
+        # Usually if someone ran query that returns no results
+        if not parsed and response.status_code in [200, 204]:
+            return [], None
 
-        if (
-            methodology == CoralogixLogsMethodology.ARCHIVE_FALLBACK and not result.logs
-        ) or methodology == CoralogixLogsMethodology.BOTH_FREQUENT_SEARCH_AND_ARCHIVE:
-            archive_search_results = query_logs_for_tier(
-                config=config, params=params, tier=CoralogixTier.ARCHIVE
+        if not parsed:
+            return None, (
+                f"Query submission:\n"
+                f"URL: {submit_url}\n"
+                f"Response status: {response.status_code}\n"
+                f"Response body (first 2000 chars): {raw[:2000]}\n\n"
             )
-            result = merge_log_results(result, archive_search_results)
 
-    else:
-        # methodology in [CoralogixLogsMethodology.ARCHIVE_ONLY, CoralogixLogsMethodology.FREQUENT_SEARCH_FALLBACK]:
-        result = query_logs_for_tier(
-            config=config, params=params, tier=CoralogixTier.ARCHIVE
-        )
+        cleaned_results = _cleanup_coralogix_results(parsed)
+        return cleaned_results, None
 
-        if (
-            methodology == CoralogixLogsMethodology.FREQUENT_SEARCH_FALLBACK
-            and not result.logs
-        ):
-            frequent_search_results = query_logs_for_tier(
-                config=config, params=params, tier=CoralogixTier.FREQUENT_SEARCH
-            )
-            result = merge_log_results(result, frequent_search_results)
+    except Exception as e:
+        logging.error("Failed to execute DataPrime query", exc_info=True)
+        return None, str(e)
+
+
+def health_check(domain: str, api_key: str) -> Tuple[bool, str]:
+    query_dict = _build_query_dict("source logs | limit 1")
+    response, submit_url = execute_coralogix_query(
+        domain=domain, api_key=api_key, query=query_dict
+    )
 
-    return result
+    if response.status_code != 200:
+        body = _get_error_body(response)
+        return (
+            False,
+            f"Failed with status_code={response.status_code}. {body}\nURL: {submit_url}",
+        )
+    return True, ""

holmes/plugins/toolsets/coralogix/coralogix.jinja2

@@ -0,0 +1,14 @@
+Coralogix DataPrime (tool: coralogix_execute_dataprime_query) queries logs/traces (sources: logs, spans).
+
+Rules:
+- Always set explicit time ranges: start_date/end_date
+- Include `limit` on every query (start with limit 100 and adjust accordingly)
+- Use `source logs | lucene 'text'` for log searches; prefer lucene over filter
+- Start broad, then narrow down
+- Never assume labels: start with plain text (e.g., `source logs | lucene 'app-173'`); only use label equality after seeing the exact field path in results
+- Traces: use `source spans`
+- Never query archive tier; only query frequent search data
+
+Example of a bad query: `source logs | filter $l.k8s.namespace_name == 'my_namespace'` (uses filter, unverified labels, no limit)
+
+To discover custom labels: run `source logs | limit 1` or `source spans | limit 1`

holmes/plugins/toolsets/coralogix/toolset_coralogix.py

@@ -0,0 +1,219 @@
+import json
+import os
+from typing import Any, Optional, Tuple
+from urllib.parse import quote
+
+from holmes.core.tools import (
+    CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    Tool,
+    ToolInvokeContext,
+    ToolParameter,
+    Toolset,
+    ToolsetTag,
+)
+from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR
+from holmes.plugins.toolsets.coralogix.api import (
+    CoralogixTier,
+    execute_dataprime_query,
+    health_check,
+)
+from holmes.plugins.toolsets.coralogix.utils import CoralogixConfig, normalize_datetime
+from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
+
+
+def _build_coralogix_query_url(
+    config: CoralogixConfig,
+    query: str,
+    start_date: str,
+    end_date: str,
+    tier: Optional[CoralogixTier] = None,
+) -> Optional[str]:
+    try:
+        if tier == CoralogixTier.ARCHIVE:
+            data_pipeline = "archive-logs"
+        else:
+            # due to a bug in Coralogix, we always use the logs pipeline
+            # since the tracing url does not support the query parameter
+            # https://coralogix.com/docs/user-guides/monitoring-and-insights/logs-screen/query-urls/
+            data_pipeline = "logs"
+
+        time_range = f"from:{start_date},to:{end_date}"
+
+        encoded_query = quote(query)
+        encoded_time = quote(time_range)
+        base_url = f"https://{config.team_hostname}.{config.domain}"
+
+        url = (
+            f"{base_url}/#/query-new/{data_pipeline}"
+            f"?querySyntax=dataprime"
+            f"&time={encoded_time}"
+            f"&query={encoded_query}"
+            f"&permalink=true"
+        )
+        return url
+
+    except Exception:
+        return None
+
+
+class ExecuteDataPrimeQuery(Tool):
+    def __init__(self, toolset: "CoralogixToolset"):
+        super().__init__(
+            name="coralogix_execute_dataprime_query",
+            description="Execute a DataPrime query against Coralogix to fetch logs, traces, metrics, and other telemetry data. "
+            "Returns the raw query results from Coralogix.",
+            parameters={
+                "query": ToolParameter(
+                    description="DataPrime query string. Examples: `source logs | lucene 'error' | limit 100`, `source spans | lucene 'my-service' | limit 100`. Always include a `limit` clause.",
+                    type="string",
+                    required=True,
+                ),
+                "description": ToolParameter(
+                    description="Brief 6-word description of the query.",
+                    type="string",
+                    required=True,
+                ),
+                "query_type": ToolParameter(
+                    description="'Logs', 'Traces', 'Metrics', 'Discover Data' or 'Other'.",
+                    type="string",
+                    required=True,
+                ),
+                "start_date": ToolParameter(
+                    description="Optional start date in RFC3339 format (e.g., '2024-01-01T00:00:00Z').",
+                    type="string",
+                    required=True,
+                ),
+                "end_date": ToolParameter(
+                    description="Optional end date in RFC3339 format (e.g., '2024-01-01T23:59:59Z').",
+                    type="string",
+                    required=True,
+                ),
+                "tier": ToolParameter(
+                    description="Optional tier: 'FREQUENT_SEARCH' or 'ARCHIVE'.",
+                    type="string",
+                    required=False,
+                ),
+            },
+        )
+        self._toolset = toolset
+
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
+        if not self._toolset.coralogix_config:
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error="Coralogix toolset is not configured",
+                params=params,
+            )
+
+        tier = None
+        if tier_str := params.get("tier"):
+            try:
+                tier = CoralogixTier[tier_str]
+            except KeyError:
+                return StructuredToolResult(
+                    status=StructuredToolResultStatus.ERROR,
+                    error=f"Invalid tier '{tier_str}'. Must be 'FREQUENT_SEARCH' or 'ARCHIVE'",
+                    params=params,
+                )
+
+        start_time = normalize_datetime(params.get("start_date"))
+        end_time = normalize_datetime(params.get("end_date"))
+        if start_time == "UNKNOWN_TIMESTAMP" or end_time == "UNKNOWN_TIMESTAMP":
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error=f"Invalid start or end date: {params.get('start_date')} or {params.get('end_date')}. Please provide valid dates in RFC3339 format (e.g., '2024-01-01T00:00:00Z').",
+                params=params,
+            )
+
+        if start_time > end_time:
+            start_time, end_time = end_time, start_time
+
+        result, error = execute_dataprime_query(
+            domain=self._toolset.coralogix_config.domain,
+            api_key=self._toolset.coralogix_config.api_key,
+            dataprime_query=params["query"],
+            start_date=start_time,
+            end_date=end_time,
+            tier=tier,
+        )
+
+        if error:
+            return StructuredToolResult(
+                status=StructuredToolResultStatus.ERROR,
+                error=error,
+                params=params,
+            )
+
+        result_dict = {
+            "tool_name": self.name,
+            "data": result,
+        }
+        status = StructuredToolResultStatus.SUCCESS
+
+        if not result:
+            results_msg = "No results found, it is possible that the query is not correct, using incorrect labels or filters."
+            result_dict["results_msg"] = results_msg
+            status = StructuredToolResultStatus.NO_DATA
+
+        # Build Coralogix query URL
+        explore_url = _build_coralogix_query_url(
+            config=self._toolset.coralogix_config,
+            query=params["query"],
+            start_date=start_time,
+            end_date=end_time,
+            tier=tier,
+        )
+
+        # Return a pretty-printed JSON string for readability by the model/user.
+        final_result = json.dumps(result_dict, indent=2, sort_keys=False)
+        return StructuredToolResult(
+            status=status,
+            data=final_result,
+            params=params,
+            url=explore_url,
+        )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        description = params.get("description", "")
+        return f"{toolset_name_for_one_liner(self._toolset.name)}: Execute DataPrime ({description})"
+
+
+class CoralogixToolset(Toolset):
+    def __init__(self):
+        super().__init__(
+            name="coralogix",
+            description="Toolset for interacting with Coralogix to fetch logs, traces, metrics, and execute DataPrime queries",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/coralogix-logs/",
+            icon_url="https://avatars.githubusercontent.com/u/35295744?s=200&v=4",
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
+            tools=[ExecuteDataPrimeQuery(self)],
+            tags=[ToolsetTag.CORE],
+        )
+        template_path = os.path.join(os.path.dirname(__file__), "coralogix.jinja2")
+        if os.path.exists(template_path):
+            self._load_llm_instructions(
+                jinja_template=f"file://{os.path.abspath(template_path)}"
+            )
+
+    def get_example_config(self):
+        example_config = CoralogixConfig(
+            api_key="<cxuw_...>", team_hostname="my-team", domain="eu2.coralogix.com"
+        )
+        return example_config.model_dump()
+
+    def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
+        if not config:
+            return False, TOOLSET_CONFIG_MISSING_ERROR
+
+        self.config = CoralogixConfig(**config)
+
+        if not self.config.api_key:
+            return False, "Missing configuration field 'api_key'"
+
+        return health_check(domain=self.config.domain, api_key=self.config.api_key)
+
+    @property
+    def coralogix_config(self) -> Optional[CoralogixConfig]:
+        return self.config

holmes/plugins/toolsets/coralogix/utils.py

@@ -1,9 +1,7 @@
-from enum import Enum
 import json
 import logging
-import urllib.parse
 from datetime import datetime
-from typing import Any, NamedTuple, Optional, Dict, List
+from typing import Any, Dict, List, NamedTuple, Optional
 
 from pydantic import BaseModel
 
@@ -26,45 +24,41 @@ class CoralogixLabelsConfig(BaseModel):
     timestamp: str = "logRecord.attributes.time"
 
 
-class CoralogixLogsMethodology(str, Enum):
-    FREQUENT_SEARCH_ONLY = "FREQUENT_SEARCH_ONLY"
-    ARCHIVE_ONLY = "ARCHIVE_ONLY"
-    ARCHIVE_FALLBACK = "ARCHIVE_FALLBACK"
-    FREQUENT_SEARCH_FALLBACK = "FREQUENT_SEARCH_FALLBACK"
-    BOTH_FREQUENT_SEARCH_AND_ARCHIVE = "BOTH_FREQUENT_SEARCH_AND_ARCHIVE"
-
-
 class CoralogixConfig(BaseModel):
     team_hostname: str
     domain: str
     api_key: str
     labels: CoralogixLabelsConfig = CoralogixLabelsConfig()
-    logs_retrieval_methodology: CoralogixLogsMethodology = (
-        CoralogixLogsMethodology.ARCHIVE_FALLBACK
-    )
 
 
 def parse_json_lines(raw_text) -> List[Dict[str, Any]]:
-    """Parses JSON objects from a raw text response."""
+    """Parses JSON objects from a raw text response and removes duplicate userData fields from child objects."""
     json_objects = []
     for line in raw_text.strip().split("\n"):  # Split by newlines
         try:
-            json_objects.append(json.loads(line))
+            obj = json.loads(line)
+            if isinstance(obj, dict):
+                # Remove userData from top level
+                obj.pop("userData", None)
+                # Remove userData from direct child dicts (one level deep, no recursion)
+                for key, value in list(obj.items()):
+                    if isinstance(value, dict):
+                        value.pop("userData", None)
+                    elif isinstance(value, list):
+                        for item in value:
+                            if isinstance(item, dict):
+                                item.pop("userData", None)
+            json_objects.append(obj)
         except json.JSONDecodeError:
             logging.error(f"Failed to decode JSON from line: {line}")
     return json_objects
 
 
 def normalize_datetime(date_str: Optional[str]) -> str:
-    """takes a date string as input and attempts to convert it into a standardized ISO 8601 format with UTC timezone (“Z” suffix) and microsecond precision.
-    if any error occurs during parsing or formatting, it returns the original input string.
-    The method specifically handles older Python versions by removing a trailing “Z” and truncating microseconds to 6 digits before parsing.
-    """
     if not date_str:
         return "UNKNOWN_TIMESTAMP"
 
     try:
-        # older versions of python do not support `Z` appendix nor more than 6 digits of microsecond precision
         date_str_no_z = date_str.rstrip("Z")
 
         parts = date_str_no_z.split(".")
@@ -148,61 +142,3 @@ def parse_json_objects(
     logs.sort(key=lambda x: x[0])
 
     return logs
-
-
-def parse_logs(
-    raw_logs: str,
-    labels_config: CoralogixLabelsConfig,
-) -> List[FlattenedLog]:
-    """Processes the HTTP response and extracts only log outputs."""
-    try:
-        json_objects = parse_json_lines(raw_logs)
-        if not json_objects:
-            raise Exception("No valid JSON objects found.")
-        return parse_json_objects(
-            json_objects=json_objects, labels_config=labels_config
-        )
-    except Exception as e:
-        logging.error(
-            f"Unexpected error in format_logs for a coralogix API response: {str(e)}"
-        )
-        raise e
-
-
-def build_coralogix_link_to_logs(
-    config: CoralogixConfig, lucene_query: str, start: str, end: str
-) -> str:
-    query_param = urllib.parse.quote_plus(lucene_query)
-
-    return f"https://{config.team_hostname}.app.{config.domain}/#/query-new/logs?query={query_param}&querySyntax=dataprime&time=from:{start},to:{end}"
-
-
-def merge_log_results(
-    a: CoralogixQueryResult, b: CoralogixQueryResult
-) -> CoralogixQueryResult:
-    """
-    Merges two CoralogixQueryResult objects, deduplicating logs and sorting them by timestamp.
-
-    """
-    if a.error is None and b.error:
-        return a
-    elif b.error is None and a.error:
-        return b
-    elif a.error and b.error:
-        return a
-
-    combined_logs = a.logs + b.logs
-
-    if not combined_logs:
-        deduplicated_logs_set = set()
-    else:
-        deduplicated_logs_set = set(combined_logs)
-
-    # Assumes timestamps are in a format sortable as strings (e.g., ISO 8601)
-    sorted_logs = sorted(list(deduplicated_logs_set), key=lambda log: log.timestamp)
-
-    return CoralogixQueryResult(
-        logs=sorted_logs,
-        http_status=a.http_status if a.http_status is not None else b.http_status,
-        error=a.error,
-    )