add correct mcp schema

This commit is contained in:
larssand
2026-06-21 22:27:13 +02:00
parent 35c21332f3
commit 64a7996b65
2 changed files with 20 additions and 3 deletions

View File

@@ -17,9 +17,19 @@ DEFAULT_FIELD_MAP = {
"action": ("action", "event_action", "disposition"),
}
DEFAULT_FIELDS = ["timestamp", "source", "srcip", "dstip", "srcport", "dstport", "service", "action", "severity", "policyid", "subtype", "type", "sentbyte", "rcvdbyte", "hitcount"]
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", "")) 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 ("messages", "results", "events", "data"):
if isinstance(value.get(key), list):
yield from (item for item in value[key] if isinstance(item, dict))
@@ -44,7 +54,13 @@ class GraylogStreamSource:
def fetch(self) -> tuple[list[LogEvent], dict[str, object]]:
status = self.client.probe()
arguments: dict[str, object] = {"query": self.query, "limit": 1000, "range_seconds": 300}
mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
arguments: dict[str, object] = {
"query": self.query,
"limit": 1000,
"range_seconds": 300,
"fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])),
}
if self.stream:
arguments["streams"] = [self.stream]
result = self.client.call_tool("search_messages", arguments)