|
| 1 | +from typing import Annotated |
| 2 | + |
| 3 | +from fastmcp import FastMCP |
| 4 | +from fastmcp.utilities.types import Image |
| 5 | +from PIL import Image as PILImage |
| 6 | +from pydantic import BaseModel, Field |
| 7 | + |
| 8 | +from askui.tools.askui.askui_controller import AskUiControllerClient |
| 9 | +from askui.tools.computer import Action20250124, Computer20250124Tool, ScrollDirection |
| 10 | +from askui.utils.image_utils import ImageSource |
| 11 | + |
| 12 | +mcp = FastMCP(name="AskUI Computer MCP") |
| 13 | + |
| 14 | + |
| 15 | +active_display = 1 |
| 16 | + |
| 17 | + |
| 18 | +@mcp.tool( |
| 19 | + description="Interact with your computer", |
| 20 | +) |
| 21 | +def computer( |
| 22 | + action: Action20250124, |
| 23 | + text: str | None = None, |
| 24 | + coordinate: tuple[Annotated[int, Field(ge=0)], Annotated[int, Field(ge=0)]] |
| 25 | + | None = None, |
| 26 | + scroll_direction: ScrollDirection | None = None, |
| 27 | + scroll_amount: Annotated[int, Field(ge=0)] | None = None, |
| 28 | + duration: Annotated[float, Field(ge=0.0, le=100.0)] | None = None, |
| 29 | + key: str | None = None, |
| 30 | +) -> Image | None | str: |
| 31 | + with AskUiControllerClient(display=active_display) as agent_os: |
| 32 | + result = Computer20250124Tool(agent_os=agent_os)( |
| 33 | + action=action, |
| 34 | + text=text, |
| 35 | + coordinate=coordinate, |
| 36 | + scroll_direction=scroll_direction, |
| 37 | + scroll_amount=scroll_amount, |
| 38 | + duration=duration, |
| 39 | + key=key, |
| 40 | + ) |
| 41 | + if isinstance(result, PILImage.Image): |
| 42 | + src = ImageSource(result) |
| 43 | + return Image(data=src.to_bytes(), format="png") |
| 44 | + return result |
| 45 | + |
| 46 | + |
| 47 | +class Display(BaseModel): |
| 48 | + id: int |
| 49 | + |
| 50 | + |
| 51 | +class DisplayListResponse(BaseModel): |
| 52 | + data: list[Display] |
| 53 | + |
| 54 | + |
| 55 | +@mcp.tool(description="List all available displays") |
| 56 | +def list_displays() -> DisplayListResponse: |
| 57 | + with AskUiControllerClient(display=active_display) as agent_os: |
| 58 | + return DisplayListResponse( |
| 59 | + data=[Display(id=display.id) for display in agent_os.list_displays().data], |
| 60 | + ) |
| 61 | + |
| 62 | + |
| 63 | +@mcp.tool( |
| 64 | + description="Set the active display from which screenshots are taken / on which actions are performed (coordinates are relative to the active display)" |
| 65 | +) |
| 66 | +def set_active_display( |
| 67 | + display_id: Annotated[int, Field(ge=1)], |
| 68 | +) -> None: |
| 69 | + global active_display |
| 70 | + active_display = display_id |
| 71 | + |
| 72 | + |
| 73 | +@mcp.tool( |
| 74 | + description="Retrieve the active display from which screenshots are taken / on which actions are performed (coordinates are relative to the active display)" |
| 75 | +) |
| 76 | +def retrieve_active_display() -> Display: |
| 77 | + return Display(id=active_display) |
0 commit comments