from codeocean import CodeOcean
from codeocean.computation import (
Computation,
FileURLs,
Folder,
RunParams,
)
from mcp.server.fastmcp import FastMCP
from codeocean_mcp_server.file_utils import download_and_read_file
from codeocean_mcp_server.models import dataclass_to_pydantic
RunParamsModel = dataclass_to_pydantic(RunParams)
def add_tools(mcp: FastMCP, client: CodeOcean):
"""Add capsule tools to the MCP server."""
@mcp.tool(description=client.computations.get_computation.__doc__)
def get_computation(computation_id: str) -> Computation:
"""Retrieve a specific computation by its unique identifier."""
return client.computations.get_computation(computation_id)
@mcp.tool(
description=(
str(client.computations.run_capsule.__doc__) + "Typical workflow: 1) run_capsule() to start execution "
"2) wait_until_completed() to monitor progress "
"3) list_computation_results() and get_result_file_urls() "
"to retrieve outputs."
)
)
def run_capsule(run_params: RunParamsModel) -> Computation:
"""Execute a capsule or a pipeline in Code Ocean and don't wait."""
params = RunParams(**run_params.model_dump(exclude_none=True))
return client.computations.run_capsule(params)
@mcp.tool(description=client.computations.wait_until_completed.__doc__)
def wait_until_completed(computation_id: str) -> Computation:
"""Wait until a computation completes and return its details."""
computation = client.computations.get_computation(computation_id)
return client.computations.wait_until_completed(computation)
@mcp.tool(
description=(
str(client.computations.list_computation_results.__doc__) + " computation_id is required as string"
)
)
def list_computation_results(computation_id: str) -> Folder:
"""List the output files generated by a completed computation."""
return client.computations.list_computation_results(computation_id)
@mcp.tool(description=(client.computations.get_result_file_urls.__doc__))
def get_result_file_urls(computation_id: str, file_path: str) -> FileURLs:
"""Get view and download URLs for a specific result file from computation."""
return client.computations.get_result_file_urls(computation_id, file_path)
@mcp.tool(description=("Use when you want to read the content of a file from a computation"))
def download_and_read_a_file_from_computation(computation_id: str, file_path: str) -> str:
"""Download a file using the provided URL and return its content."""
file_urls = client.computations.get_result_file_urls(computation_id, file_path)
return download_and_read_file(file_urls.download_url)
