from typing import Any, Dict, List, Optional, Union
from uuid import UUID
import mcp.types as types
from prefect import get_client
from .server import mcp
@mcp.tool
async def get_work_queues(
limit: Optional[int] = None,
offset: Optional[int] = None,
name: Optional[str] = None,
is_paused: Optional[bool] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Get a list of work queues with optional filtering.
Args:
limit: Maximum number of work queues to return
offset: Number of work queues to skip
name: Filter by name
is_paused: Filter by paused status
Returns:
A list of work queues with their details
"""
async with get_client() as client:
# Build filter parameters
filters = {}
if name:
filters["name"] = {"like_": f"%{name}%"}
if is_paused is not None:
filters["is_paused"] = {"eq_": is_paused}
work_queues = await client.read_work_queues(
limit=limit,
offset=offset,
**filters
)
work_queues_result = {
"work_queues": [work_queue.dict() for work_queue in work_queues]
}
return [types.TextContent(type="text", text=str(work_queues_result))]
@mcp.tool
async def get_work_queue(
work_queue_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Get details of a specific work queue by ID.
Args:
work_queue_id: The work queue UUID
Returns:
Work queue details
"""
async with get_client() as client:
work_queue = await client.read_work_queue(UUID(work_queue_id))
return [types.TextContent(type="text", text=str(work_queue.dict()))]
@mcp.tool
async def get_work_queue_by_name(
name: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Get a work queue by name.
Args:
name: The work queue name
Returns:
Work queue details
"""
async with get_client() as client:
work_queue = await client.read_work_queue_by_name(name)
return [types.TextContent(type="text", text=str(work_queue.dict()))]
@mcp.tool
async def create_work_queue(
name: str,
description: Optional[str] = None,
is_paused: Optional[bool] = None,
concurrency_limit: Optional[int] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Create a work queue.
Args:
name: The name for the work queue
description: Optional description
is_paused: Whether the queue should be paused upon creation
concurrency_limit: Optional concurrency limit
Returns:
Details of the created work queue
"""
async with get_client() as client:
work_queue = await client.create_work_queue(
name=name,
description=description,
is_paused=is_paused,
concurrency_limit=concurrency_limit,
)
return [types.TextContent(type="text", text=str(work_queue.dict()))]
@mcp.tool
async def update_work_queue(
work_queue_id: str,
name: Optional[str] = None,
description: Optional[str] = None,
is_paused: Optional[bool] = None,
concurrency_limit: Optional[int] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Update a work queue.
Args:
work_queue_id: The work queue UUID
name: New name
description: New description
is_paused: New paused status
concurrency_limit: New concurrency limit
Returns:
Details of the updated work queue
"""
async with get_client() as client:
# Prepare update data
update_data = {}
if name is not None:
update_data["name"] = name
if description is not None:
update_data["description"] = description
if is_paused is not None:
update_data["is_paused"] = is_paused
if concurrency_limit is not None:
update_data["concurrency_limit"] = concurrency_limit
work_queue = await client.update_work_queue(
id=UUID(work_queue_id),
**update_data
)
return [types.TextContent(type="text", text=str(work_queue.dict()))]
@mcp.tool
async def delete_work_queue(
work_queue_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Delete a work queue by ID.
Args:
work_queue_id: The work queue UUID
Returns:
Confirmation message
"""
async with get_client() as client:
await client.delete_work_queue(UUID(work_queue_id))
return [types.TextContent(type="text", text=f"Work queue '{work_queue_id}' deleted successfully.")]
@mcp.tool
async def pause_work_queue(
work_queue_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Pause a work queue.
Args:
work_queue_id: The work queue UUID
Returns:
Details of the updated work queue
"""
async with get_client() as client:
work_queue = await client.update_work_queue(
id=UUID(work_queue_id),
is_paused=True
)
return [types.TextContent(type="text", text=str(work_queue.dict()))]
@mcp.tool
async def resume_work_queue(
work_queue_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
"""
Resume a work queue.
Args:
work_queue_id: The work queue UUID
Returns:
Details of the updated work queue
"""
async with get_client() as client:
work_queue = await client.update_work_queue(
id=UUID(work_queue_id),
is_paused=False
)
return [types.TextContent(type="text", text=str(work_queue.dict()))]
