add agregated search

src/fgai/graylog_aggregate.py

@@ -0,0 +1,114 @@
+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)
+
+
+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) -> dict[str, object]:
+        status = self.client.probe()
+        variants: list[dict[str, object]] = [
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "group_by": [],
+                "metrics": [{"function": "count"}],
+            },
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "groups": [],
+                "series": [{"function": "count"}],
+            },
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+            },
+        ]
+        errors: list[str] = []
+        for arguments in variants:
+            try:
+                result = self.client.call_tool("aggregate_messages", arguments)
+            except RuntimeError as exc:
+                errors.append(str(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(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,
+            }
+        return {
+            **status,
+            "source": "graylog_mcp_aggregate",
+            "aggregate_status": "error",
+            "aggregate_events": 0,
+            "aggregate_records": 0,
+            "aggregate_error": "; ".join(errors[-3:]) or "aggregate_messages failed",
+        }