Public/fgAI
SHA256
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add 5000 events

Browse Source
This commit is contained in:
larssand
2026-06-21 22:31:20 +02:00
parent 64a7996b65
commit f99e699da6
2 changed files with 24 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
src/fgai/graylog_source.py
View File

@@ -52,34 +52,38 @@ class GraylogStreamSource:
if not isinstance(self.mapping, dict): if not isinstance(self.mapping, dict):
raise RuntimeError("invalid_graylog_field_mapping") raise RuntimeError("invalid_graylog_field_mapping")
def fetch(self) -> tuple[list[LogEvent], dict[str, object]]: def fetch(self, *, max_events: int = 5_000) -> tuple[list[LogEvent], dict[str, object]]:
status = self.client.probe() status = self.client.probe()
mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)] mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
arguments: dict[str, object] = { arguments: dict[str, object] = {
"query": self.query, "query": self.query,
"limit": 1000,
"range_seconds": 300, "range_seconds": 300,
"fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])), "fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])),
} }
if self.stream: if self.stream:
arguments["streams"] = [self.stream] arguments["streams"] = [self.stream]
result = self.client.call_tool("search_messages", arguments) events: list[LogEvent] = []
content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else [] page_size = 1_000
if isinstance(result.get("result"), dict) and result["result"].get("isError"): pages = 0
detail = next( while len(events) < max_events:
(str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), result = self.client.call_tool("search_messages", {**arguments, "limit": page_size, "offset": len(events)})
"Graylog search failed", content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else []
) if isinstance(result.get("result"), dict) and result["result"].get("isError"):
raise RuntimeError(f"graylog_search_error: {detail}") detail = next((str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), "Graylog search failed")
records: list[dict[str, object]] = [] raise RuntimeError(f"graylog_search_error: {detail}")
for item in content if isinstance(content, list) else []: records: list[dict[str, object]] = []
if isinstance(item, dict) and item.get("type") == "text": for item in content if isinstance(content, list) else []:
try: if isinstance(item, dict) and item.get("type") == "text":
records.extend(_records(json.loads(str(item.get("text", ""))))) try:
except json.JSONDecodeError: records.extend(_records(json.loads(str(item.get("text", "")))))
continue except json.JSONDecodeError:
events = [self._event(record) for record in records] continue
status.update({"source": "graylog_mcp", "events_fetched": len(events)}) events.extend(self._event(record) for record in records)
pages += 1
if len(records) < page_size:
break
latest = max((event.fields.get("eventtime", "") for event in events), default="")
status.update({"source": "graylog_mcp", "events_fetched": len(events), "pages": pages, "truncated": len(events) >= max_events, "latest_event_time": latest})
return events, status return events, status
def _event(self, record: dict[str, object]) -> LogEvent: def _event(self, record: dict[str, object]) -> LogEvent:

1
tests/test_graylog_source.py
View File

@@ -28,6 +28,7 @@ class GraylogSourceTests(unittest.TestCase):
self.assertEqual(client.arguments["streams"], ["vpn"]) self.assertEqual(client.arguments["streams"], ["vpn"])
self.assertEqual(client.arguments["range_seconds"], 300) self.assertEqual(client.arguments["range_seconds"], 300)
self.assertIn("client", client.arguments["fields"]) self.assertIn("client", client.arguments["fields"])
self.assertEqual(client.arguments["offset"], 0)
if __name__ == "__main__": if __name__ == "__main__":