fix mcp search
This commit is contained in:
@@ -44,8 +44,17 @@ class GraylogStreamSource:
|
||||
|
||||
def fetch(self) -> tuple[list[LogEvent], dict[str, object]]:
|
||||
status = self.client.probe()
|
||||
result = self.client.call_tool("search_messages", {"query": self.query, "stream": self.stream, "limit": 1000})
|
||||
arguments: dict[str, object] = {"query": self.query, "limit": 1000, "range_seconds": 300}
|
||||
if self.stream:
|
||||
arguments["streams"] = [self.stream]
|
||||
result = self.client.call_tool("search_messages", arguments)
|
||||
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 search failed",
|
||||
)
|
||||
raise RuntimeError(f"graylog_search_error: {detail}")
|
||||
records: list[dict[str, object]] = []
|
||||
for item in content if isinstance(content, list) else []:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
|
||||
@@ -4,22 +4,29 @@ from fgai.graylog_source import GraylogStreamSource
|
||||
|
||||
|
||||
class _Client:
|
||||
def __init__(self):
|
||||
self.arguments = None
|
||||
|
||||
def probe(self):
|
||||
return {"status": "connected"}
|
||||
|
||||
def call_tool(self, _name, _arguments):
|
||||
def call_tool(self, _name, arguments):
|
||||
self.arguments = arguments
|
||||
return {"result": {"content": [{"type": "text", "text": '{"messages":[{"client":"10.0.0.1","server":"8.8.8.8","result":"deny"}]}' }]}}
|
||||
|
||||
|
||||
class GraylogSourceTests(unittest.TestCase):
|
||||
def test_applies_custom_mapping_to_generic_stream_message(self):
|
||||
client = _Client()
|
||||
events, status = GraylogStreamSource(
|
||||
_Client(), "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}'
|
||||
client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}'
|
||||
).fetch()
|
||||
self.assertEqual(status["events_fetched"], 1)
|
||||
self.assertEqual(events[0].src_ip, "10.0.0.1")
|
||||
self.assertEqual(events[0].dst_ip, "8.8.8.8")
|
||||
self.assertEqual(events[0].action, "deny")
|
||||
self.assertEqual(client.arguments["streams"], ["vpn"])
|
||||
self.assertEqual(client.arguments["range_seconds"], 300)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user