Prefect MCP Server

from typing import Any, Callable, Dict, List, Optional, Union from uuid import UUID import mcp.types as types from prefect import get_client def get_all_functions() -> list[tuple[Callable, str, str]]: return [ (get_concurrency_limits, "get_concurrency_limits", "Get all concurrency limits"), (get_concurrency_limit, "get_concurrency_limit", "Get a concurrency limit by ID"), (get_concurrency_limit_by_tag, "get_concurrency_limit_by_tag", "Get a concurrency limit by tag"), (create_concurrency_limit, "create_concurrency_limit", "Create a concurrency limit"), (update_concurrency_limit, "update_concurrency_limit", "Update a concurrency limit"), (delete_concurrency_limit, "delete_concurrency_limit", "Delete a concurrency limit"), (increment_concurrency_limit, "increment_concurrency_limit", "Increment a concurrency limit"), (decrement_concurrency_limit, "decrement_concurrency_limit", "Decrement a concurrency limit"), (reset_concurrency_limit, "reset_concurrency_limit", "Reset a concurrency limit by tag"), ] async def get_concurrency_limits() -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get all concurrency limits. Returns: A list of concurrency limits """ async with get_client() as client: try: # Prefect API doesn't support limit/offset on concurrency limits endpoint directly # Use the filter endpoint instead concurrency_limits = await client.read_concurrency_limits_filter() limits_result = { "concurrency_limits": [limit.dict() for limit in concurrency_limits] } return [types.TextContent(type="text", text=str(limits_result))] except Exception as e: error_message = f"Error fetching concurrency limits: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_concurrency_limit( limit_id: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get a concurrency limit by ID. Args: limit_id: The concurrency limit UUID Returns: Concurrency limit details """ async with get_client() as client: try: concurrency_limit = await client.read_concurrency_limit(UUID(limit_id)) return [types.TextContent(type="text", text=str(concurrency_limit.dict()))] except Exception as e: error_message = f"Error fetching concurrency limit {limit_id}: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_concurrency_limit_by_tag( tag: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get a concurrency limit by tag. Args: tag: The concurrency limit tag Returns: Concurrency limit details """ async with get_client() as client: try: concurrency_limit = await client.read_concurrency_limit_by_tag(tag) return [types.TextContent(type="text", text=str(concurrency_limit.dict()))] except Exception as e: error_message = f"Error fetching concurrency limit for tag '{tag}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def create_concurrency_limit( tag: str, concurrency_limit: int, active: Optional[bool] = True, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Create a concurrency limit. Args: tag: The tag to limit concurrency_limit: The maximum allowed concurrency active: Whether the limit is active Returns: Details of the created concurrency limit """ async with get_client() as client: try: limit = await client.create_concurrency_limit( tag=tag, concurrency_limit=concurrency_limit, active=active ) return [types.TextContent(type="text", text=str(limit.dict()))] except Exception as e: error_message = f"Error creating concurrency limit: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def update_concurrency_limit( limit_id: str, concurrency_limit: Optional[int] = None, active: Optional[bool] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Update a concurrency limit. Args: limit_id: The concurrency limit UUID concurrency_limit: The new maximum allowed concurrency active: Whether the limit is active Returns: Details of the updated concurrency limit """ async with get_client() as client: try: # Prepare update data update_data = {} if concurrency_limit is not None: update_data["concurrency_limit"] = concurrency_limit if active is not None: update_data["active"] = active updated_limit = await client.update_concurrency_limit( id=UUID(limit_id), **update_data ) return [types.TextContent(type="text", text=str(updated_limit.dict()))] except Exception as e: error_message = f"Error updating concurrency limit {limit_id}: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def delete_concurrency_limit( limit_id: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Delete a concurrency limit by ID. Args: limit_id: The concurrency limit UUID Returns: Confirmation message """ async with get_client() as client: try: await client.delete_concurrency_limit(UUID(limit_id)) return [types.TextContent(type="text", text=f"Concurrency limit '{limit_id}' deleted successfully.")] except Exception as e: error_message = f"Error deleting concurrency limit {limit_id}: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def increment_concurrency_limit( tag: str, delta: int = 1, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Increment a concurrency limit by tag. Args: tag: The concurrency limit tag delta: Amount to increment by (default 1) Returns: Updated concurrency limit details """ async with get_client() as client: try: updated_limit = await client.increment_concurrency_limit( tag=tag, delta=delta ) return [types.TextContent(type="text", text=str(updated_limit.dict()))] except Exception as e: error_message = f"Error incrementing concurrency limit for tag '{tag}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def decrement_concurrency_limit( tag: str, delta: int = 1, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Decrement a concurrency limit by tag. Args: tag: The concurrency limit tag delta: Amount to decrement by (default 1) Returns: Updated concurrency limit details """ async with get_client() as client: try: updated_limit = await client.decrement_concurrency_limit( tag=tag, delta=delta ) return [types.TextContent(type="text", text=str(updated_limit.dict()))] except Exception as e: error_message = f"Error decrementing concurrency limit for tag '{tag}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def reset_concurrency_limit( tag: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Reset a concurrency limit by tag, setting its current count to 0. Args: tag: The concurrency limit tag Returns: Updated concurrency limit details """ async with get_client() as client: try: updated_limit = await client.reset_concurrency_limit(tag=tag) return [types.TextContent(type="text", text=str(updated_limit.dict()))] except Exception as e: error_message = f"Error resetting concurrency limit for tag '{tag}': {str(e)}" return [types.TextContent(type="text", text=error_message)]