math_in_python.py

import asyncio  # noqa -- swapping to trio would be beneficial, but not blocking atm
import json
import os
from dataclasses import dataclass
from typing import Annotated, Any, Literal

from mcp.server.fastmcp import FastMCP
from pydantic import Field

from taiga.spec import Grade
from taiga.tools.base import ToolResult

# This is part of the reference impl
# ----------------------------------

mcp = FastMCP("taiga")

TEST_MODE = os.environ.get("MCP_TESTING_MODE", "1") in ["1", "true"]

if TEST_MODE:
    # Note, these tools are only available in testing mode for the purpose of testing
    # If the enviroment performs well with these tools, it will also work with our internal
    # implementation

    from taiga.tools.bash import BashTool
    from taiga.tools.edit import Command, EditTool
    from taiga.tools.tmux import TmuxTool

    bash_tool = BashTool()

    @mcp.tool()
    async def bash(command: str | None = None, restart: bool = False) -> ToolResult:
        return await bash_tool(command, restart)

    edit_tool = EditTool()

    @mcp.tool(
        name="str_replace_editor",
        description="Create and edit files using str_replace_editor.  Please use absolute paths for all file names.  When writing files please work within /workdir.",
    )
    async def str_replace_editor(
        *,
        command: Command,
        path: str,
        file_text: str | None = None,
        view_range: list[int] | None = None,
        old_str: str | None = None,
        new_str: str | None = None,
        insert_line: int | None = None,
        insert_text: str | None = None,
    ) -> ToolResult:
        return await edit_tool(
            command=command,
            path=path,
            file_text=file_text,
            view_range=view_range,
            old_str=old_str,
            new_str=new_str,
            insert_line=insert_line,
            insert_text=insert_text,
        )

    tmux_tool = TmuxTool()

    @mcp.tool(
        description=(
            "Run commands with tmux. The `args` are passed directly to the tmux binary "
            '(e.g. ["send-keys", "-t", "base", "python3 server.py", "Enter"]). '
            "A detached session named 'base' is created automatically on first use. "
            'Passing a "-t" target with anything other than a direct capture-pane call automatically '
            "captures that pane in a loop until its content settles, a `patterns` regex matches, or "
            "`timeout` seconds elapse. Captured output is a sequence of timestamped frames; when "
            "nothing changes between frames, no new frame output is shown. "
            "For multiline text, pass it via `text` and put a literal $text "
            "placeholder in `args`. Use tmux for long-running or interactive processes; use bash for "
            "one-shot commands."
        )
    )
    async def tmux(
        args: list[str] | None = None,
        text: str | None = None,
        timeout: float | None = None,
        frame_rate: float | None = None,
        patterns: list[str] | None = None,
        restart: bool = False,
    ) -> ToolResult:
        return await tmux_tool(
            args=args,
            text=text,
            timeout=timeout,
            frame_rate=frame_rate,
            patterns=patterns,
            restart=restart,
        )

# This is the contractor provided environment
# -------------------------------------------


@dataclass
class Problem:
    id: str
    statement: str
    solution: str


template = """
Write a python script into a file using str_replace_editor to
<STATEMENT>
and run it using bashtool to get the results.
"""

hinted_template = """
Write a python script into a file using str_replace_editor to
<STATEMENT>
and run it using bashtool to get the results.

HINT: you should be able to solve the problem using standard python libraries.
"""

current_problem: Problem | None = None


def _get_problem(problem_id: str, extra_fields: dict[str, Any] | None) -> Problem:
    """Resolve the problem from extra_fields supplied by Taiga / problems-metadata.json."""
    global current_problem
    assert current_problem is None or current_problem.id == problem_id, (
        f"setup_problem can only be called once with id: {problem_id=}, current_problem.id={current_problem.id}"
    )
    if not extra_fields or "statement" not in extra_fields or "solution" not in extra_fields:
        raise ValueError(
            f"No problem data for id={problem_id!r}: supply extra_fields.statement and extra_fields.solution"
        )
    current_problem = Problem(
        id=problem_id,
        statement=extra_fields["statement"],
        solution=str(extra_fields["solution"]),
    )
    return current_problem


# Implementation notes: setup_problem will only be called once per enviroment instance
@mcp.tool()
async def setup_problem(
    problem_id: str = Field(description="The id of the problem to solve"),
    use_hinted_problem: bool = Field(
        description="If true, setup a 'hinted' verison of the problem which gives more guidance", default=True
    ),
    extra_fields: dict[str, Any] | None = None,
) -> str:
    """Starts the enviroment and returns the problem statement"""
    await asyncio.sleep(0)

    problem = _get_problem(problem_id, extra_fields)

    problem_template = hinted_template if use_hinted_problem else template
    prompt = problem_template.replace("<STATEMENT>", problem.statement)
    if extra_fields is not None and extra_fields.get("prompt_addendum") is not None:
        prompt += "\n" + extra_fields["prompt_addendum"]
    return prompt


# Implementation note: grade_problem will only be called once per enviroment instance
@mcp.tool()
async def grade_problem(
    problem_id: str,
    transcript: str = Field(description="The entire transcript produced by the model and its tool calls"),
    extra_fields: dict[str, Any] | None = None,
) -> Grade:
    """Check your solution for grading. Returns a Grade object making sure to include all components that make up the score as subscores."""
    # Note that this is a temporary signature and a more complete reference will be provided.
    await asyncio.sleep(0)

    answer = _get_problem(problem_id, extra_fields).solution
    score = float(answer in str(transcript))

    # If you would like, you can write anything you want to an output directory. If you set "output_directory" in the problem metadata,
    # it will be stored in the taiga website for your viewing.
    # Here's an example:
    output_directory = "/tmp/out"
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)
    with open(os.path.join(output_directory, "example.txt"), "w") as f:
        f.write(
            "Here are some example logs from the problem run. Because we specified the output_directory as /tmp/out in the problems-metadata file, you will be able to see this in the taiga website. Useful for logs! \n"
        )

    return Grade(
        subscores={"matched_solution": score},
        weights={"matched_solution": 1},
        metadata={"test": "this was graded successfully!"},
    )


# Simple custom tool
# FastMCP uses pydantic to do param validation
@mcp.tool()
async def restricted_echo_tool(
    to_echo: Annotated[Literal["hello", "goodbye", "new"] | None, Field(description="The value to echo")] = "hello",
) -> str | None:
    """Directly returns the value specified by to_echo"""
    await asyncio.sleep(0)

    return to_echo


# More complex custom tool
@dataclass
class TodoItem:
    id: Annotated[str, Field(description="Unique ID for the todo list item")]
    content: Annotated[str, Field(description="Todo list item content")]
    status: Annotated[
        Literal["pending", "in_progress", "completed"], Field(description="Current status of this todo list item")
    ]
    priority: Annotated[
        Literal["high", "medium", "low"], Field(description="The priority level of this todo list item")
    ]

    def to_dict(self):
        return {
            "id": self.id,
            "content": self.content,
            "status": self.status,
            "priority": self.priority,
        }


todo_items: list[TodoItem] = []


# FastMCP uses pydantic to do basic param validation
@mcp.tool()
async def todo_tool(
    operation: Annotated[
        Literal["read", "write"],
        Field(
            description="The operation to perform - either 'read' to get the current todo list or 'write' to replace the entire todo list"
        ),
    ],
    todos: Annotated[
        list[TodoItem] | None,
        Field(
            description="Only required for the 'write' operation. Contains the list of todo items to replace the current todo list",
        ),
    ] = None,
) -> str:
    """Manages todo lists - read and write todo items"""
    await asyncio.sleep(0)

    global todo_items

    if operation == "read":
        if not todo_items:
            raise Exception("The todo list is currently empty")

        todo_dicts = [todo.to_dict() for todo in todo_items]
        return json.dumps(todo_dicts, indent=2)
    elif operation == "write":
        if todos is None:
            raise ValueError("The 'todos' parameter is required for write operations")
        todo_items = todos
        return f"Successfully replaced the todo list with the {len(todos)} provided items"
    else:
        raise ValueError(f"Invalid operation '{operation}'. Must be either 'read' or 'write'")


def main():
    # Initialize and run the server as root; you can use files and services that require root permissions
    # once init is done, the server will run as the model user to prevent it from accessing problem data
    os.chdir("/workdir")
    mcp.run(transport="stdio")