Prefect MCP Server

"""Logs module for MCP Prefect.""" from datetime import datetime from typing import List, Optional, Union import mcp.types as types from prefect import get_client from .server import mcp @mcp.tool async def get_logs( limit: Optional[int] = None, offset: Optional[int] = None, flow_run_id: Optional[str] = None, task_run_id: Optional[str] = None, level: Optional[int] = None, timestamp_after: Optional[str] = None, timestamp_before: Optional[str] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get logs with optional filtering. Args: limit: Maximum number of logs to return offset: Number of logs to skip flow_run_id: Filter by flow run ID task_run_id: Filter by task run ID level: Filter by log level (10=DEBUG, 20=INFO, 30=WARNING, 40=ERROR, 50=CRITICAL) timestamp_after: ISO formatted datetime string for logs after this time timestamp_before: ISO formatted datetime string for logs before this time Returns: A list of logs with their details """ try: async with get_client() as client: # Build filter parameters filter_dict = {} if flow_run_id: filter_dict["flow_run_id"] = {"any_": [flow_run_id]} if task_run_id: filter_dict["task_run_id"] = {"any_": [task_run_id]} if level is not None: filter_dict["level"] = {"ge_": level} if timestamp_after: filter_dict["timestamp"] = {"after_": timestamp_after} if timestamp_before: if "timestamp" in filter_dict: filter_dict["timestamp"]["before_"] = timestamp_before else: filter_dict["timestamp"] = {"before_": timestamp_before} response = await client._client.post( "/logs/filter", json={ "logs": filter_dict if filter_dict else None, "limit": limit, "offset": offset, "sort": "TIMESTAMP_DESC" } ) logs = response.json() return [types.TextContent(type="text", text=str({"logs": logs}))] except Exception as e: error_message = f"Error fetching logs: {str(e)}" return [types.TextContent(type="text", text=error_message)] @mcp.tool async def create_log( message: str, level: int = 20, flow_run_id: Optional[str] = None, task_run_id: Optional[str] = None, name: Optional[str] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Create a log entry. Args: message: The log message level: Log level (10=DEBUG, 20=INFO, 30=WARNING, 40=ERROR, 50=CRITICAL) flow_run_id: Associate with a flow run task_run_id: Associate with a task run name: Logger name Returns: Confirmation of log creation """ try: async with get_client() as client: log_data = { "name": name or "prefect.mcp", "level": level, "message": message, "timestamp": datetime.utcnow().isoformat(), } if flow_run_id: log_data["flow_run_id"] = flow_run_id if task_run_id: log_data["task_run_id"] = task_run_id response = await client._client.post( "/logs/", json=[log_data] # Logs API expects a list ) return [types.TextContent(type="text", text="Log created successfully.")] except Exception as e: error_message = f"Error creating log: {str(e)}" return [types.TextContent(type="text", text=error_message)] @mcp.tool async def get_flow_run_logs( flow_run_id: str, limit: Optional[int] = None, offset: Optional[int] = None, level: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get logs for a specific flow run. Args: flow_run_id: The flow run UUID limit: Maximum number of logs to return offset: Number of logs to skip level: Minimum log level to return Returns: Logs for the flow run """ try: async with get_client() as client: filter_dict = {"flow_run_id": {"any_": [flow_run_id]}} if level is not None: filter_dict["level"] = {"ge_": level} response = await client._client.post( "/logs/filter", json={ "logs": filter_dict, "limit": limit, "offset": offset, "sort": "TIMESTAMP_ASC" } ) logs = response.json() # Format logs for better readability formatted_logs = [] for log in logs: formatted_logs.append({ "timestamp": log.get("timestamp"), "level": log.get("level"), "message": log.get("message"), "name": log.get("name") }) return [types.TextContent(type="text", text=str({"flow_run_logs": formatted_logs}))] except Exception as e: error_message = f"Error fetching flow run logs: {str(e)}" return [types.TextContent(type="text", text=error_message)] @mcp.tool async def get_task_run_logs( task_run_id: str, limit: Optional[int] = None, offset: Optional[int] = None, level: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get logs for a specific task run. Args: task_run_id: The task run UUID limit: Maximum number of logs to return offset: Number of logs to skip level: Minimum log level to return Returns: Logs for the task run """ try: async with get_client() as client: filter_dict = {"task_run_id": {"any_": [task_run_id]}} if level is not None: filter_dict["level"] = {"ge_": level} response = await client._client.post( "/logs/filter", json={ "logs": filter_dict, "limit": limit, "offset": offset, "sort": "TIMESTAMP_ASC" } ) logs = response.json() # Format logs for better readability formatted_logs = [] for log in logs: formatted_logs.append({ "timestamp": log.get("timestamp"), "level": log.get("level"), "message": log.get("message"), "name": log.get("name") }) return [types.TextContent(type="text", text=str({"task_run_logs": formatted_logs}))] except Exception as e: error_message = f"Error fetching task run logs: {str(e)}" return [types.TextContent(type="text", text=error_message)]