fgAI/src/fgai/graylog_aggregate.py

from __future__ import annotations

import json
from collections.abc import Iterable

from .graylog_mcp import GraylogMcpClient


def _records(value: object) -> Iterable[dict[str, object]]:
    if isinstance(value, dict):
        schema = value.get("schema")
        datarows = value.get("datarows")
        if isinstance(schema, list) and isinstance(datarows, list):
            fields = [str(column.get("field", column.get("name", ""))) for column in schema if isinstance(column, dict)]
            for row in datarows:
                if isinstance(row, list):
                    yield {field: row[index] for index, field in enumerate(fields) if field and index < len(row)}
            return
        for key in ("rows", "data", "results", "messages"):
            if isinstance(value.get(key), list):
                yield from (item for item in value[key] if isinstance(item, dict))
                return
        if value:
            yield value


def _number(value: object) -> int:
    try:
        return int(float(str(value)))
    except (TypeError, ValueError):
        return 0


def _count_from_records(records: list[dict[str, object]]) -> int:
    if not records:
        return 0
    keys = ("count", "event_count", "events", "total", "COUNT()", "count()")
    for record in records:
        for key, value in record.items():
            normalized_key = str(key).lower()
            if normalized_key in {item.lower() for item in keys} or "count" in normalized_key:
                count = _number(value)
                if count:
                    return count
    if len(records) == 1:
        numeric_values = [_number(value) for value in records[0].values()]
        return max(numeric_values or [0])
    return len(records)


def _schema_properties(schema: object) -> set[str]:
    if not isinstance(schema, dict):
        return set()
    properties = schema.get("properties")
    if not isinstance(properties, dict):
        return set()
    return {str(key) for key in properties}


def _filter_supported(arguments: dict[str, object], properties: set[str]) -> dict[str, object]:
    if not properties:
        return arguments
    return {key: value for key, value in arguments.items() if key in properties}


class GraylogAggregateSource:
    def __init__(self, client: GraylogMcpClient, stream: str, query: str = "*") -> None:
        self.client = client
        self.stream = stream
        self.query = query or "*"

    def fetch_count(self, *, range_seconds: int = 300, probe_status: dict[str, object] | None = None) -> dict[str, object]:
        if probe_status is not None:
            status = dict(probe_status)
        else:
            try:
                status = self.client.probe()
            except RuntimeError as exc:
                return {
                    "status": "error",
                    "source": "graylog_mcp_aggregate",
                    "aggregate_status": "error",
                    "aggregate_events": 0,
                    "aggregate_records": 0,
                    "aggregate_error": f"probe_error: {exc}",
                    "aggregate_errors": [f"probe_error: {exc}"],
                    "aggregate_schema_properties": [],
                }
        tool_schemas = status.get("tool_schemas", {}) if isinstance(status.get("tool_schemas"), dict) else {}
        aggregate_schema = tool_schemas.get("aggregate_messages", {}) if isinstance(tool_schemas, dict) else {}
        properties = _schema_properties(aggregate_schema)
        base = {
            "query": self.query,
            "streams": [self.stream] if self.stream else [],
            "range_seconds": max(1, int(range_seconds)),
        }
        schema_variants: list[dict[str, object]] = []
        metric_keys = [key for key in ("metrics", "series") if key in properties]
        group_keys = [key for key in ("groupings", "fields") if key in properties]
        if properties and metric_keys:
            for metric_key in metric_keys:
                for metric_value in (["count()"], ["count"], [{"function": "count"}]):
                    candidate = {**base, metric_key: metric_value}
                    for group_key in group_keys:
                        candidate[group_key] = []
                    schema_variants.append(_filter_supported(candidate, properties))
        if properties:
            schema_variants.append(_filter_supported(base, properties))
        fallback_variants: list[dict[str, object]] = [
            {
                "query": self.query,
                "streams": [self.stream] if self.stream else [],
                "range_seconds": max(1, int(range_seconds)),
                "groupings": [],
                "metrics": ["count()"],
            },
            {
                "query": self.query,
                "streams": [self.stream] if self.stream else [],
                "range_seconds": max(1, int(range_seconds)),
                "groupings": [],
                "metrics": ["count"],
            },
            {
                "query": self.query,
                "streams": [self.stream] if self.stream else [],
                "range_seconds": max(1, int(range_seconds)),
                "groupings": [],
                "metrics": [{"function": "count"}],
            },
        ]
        variants = [*schema_variants, *fallback_variants]
        seen_variants: set[str] = set()
        errors: list[str] = []
        for arguments in variants:
            variant_key = json.dumps(arguments, sort_keys=True)
            if variant_key in seen_variants:
                continue
            seen_variants.add(variant_key)
            try:
                result = self.client.call_tool("aggregate_messages", arguments)
            except RuntimeError as exc:
                errors.append(f"{arguments}: {exc}")
                continue
            content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else []
            if isinstance(result.get("result"), dict) and result["result"].get("isError"):
                detail = next((str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), "Graylog aggregate failed")
                errors.append(f"{arguments}: {detail}")
                continue
            records: list[dict[str, object]] = []
            for item in content if isinstance(content, list) else []:
                if isinstance(item, dict) and item.get("type") == "text":
                    try:
                        records.extend(_records(json.loads(str(item.get("text", "")))))
                    except json.JSONDecodeError:
                        continue
            return {
                **status,
                "source": "graylog_mcp_aggregate",
                "aggregate_status": "ok",
                "aggregate_events": _count_from_records(records),
                "aggregate_records": len(records),
                "aggregate_arguments": arguments,
                "aggregate_schema_properties": sorted(properties),
            }
        return {
            **status,
            "source": "graylog_mcp_aggregate",
            "aggregate_status": "error",
            "aggregate_events": 0,
            "aggregate_records": 0,
            "aggregate_error": "; ".join(errors[-3:]) or "aggregate_messages failed",
            "aggregate_errors": errors[-6:],
            "aggregate_schema_properties": sorted(properties),
        }