improve agregate with status

src/fgai/graylog_aggregate.py

@@ -48,6 +48,21 @@ def _count_from_records(records: list[dict[str, object]]) -> int:
     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
@@ -56,7 +71,27 @@ class GraylogAggregateSource:
 
     def fetch_count(self, *, range_seconds: int = 300) -> dict[str, object]:
         status = self.client.probe()
-        variants: list[dict[str, object]] = [
+        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 ("group_by", "groups", "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 [],
@@ -104,17 +139,23 @@ class GraylogAggregateSource:
                 "range_seconds": max(1, int(range_seconds)),
             },
         ]
+        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(str(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(detail)
+                errors.append(f"{arguments}: {detail}")
                 continue
             records: list[dict[str, object]] = []
             for item in content if isinstance(content, list) else []:
@@ -130,6 +171,7 @@ class GraylogAggregateSource:
                 "aggregate_events": _count_from_records(records),
                 "aggregate_records": len(records),
                 "aggregate_arguments": arguments,
+                "aggregate_schema_properties": sorted(properties),
             }
         return {
             **status,
@@ -138,4 +180,6 @@ class GraylogAggregateSource:
             "aggregate_events": 0,
             "aggregate_records": 0,
             "aggregate_error": "; ".join(errors[-3:]) or "aggregate_messages failed",
+            "aggregate_errors": errors[-6:],
+            "aggregate_schema_properties": sorted(properties),
         }