from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
import cellrank as cr
import scvelo as scv
from ...schema.kernel import *
from scmcp_shared.util import filter_args, forward_request, get_ads
from scmcp_shared.logging_config import setup_logger
from scmcp_shared.schema import AdataInfo
logger = setup_logger()
kernel_mcp = FastMCP("CellrankMCP-Kernel-Server")
@kernel_mcp.tool()
def create_kernel(request: KernelModel, adinfo: AdataInfo = AdataInfo()):
"""Create a CellRank kernel based on the specified type and parameters."""
try:
# Check if AnnData object exists in the session
result = forward_request("kernel_create_kernel", request, adinfo)
if result is not None:
return result
ads = get_ads()
adata = ads.get_adata(adinfo=adinfo).copy()
kernel_type = request.kernel
kernel_dic = {
"pseudotime": cr.kernels.PseudotimeKernel,
"cytotrace": cr.kernels.CytoTRACEKernel,
"velocity": cr.kernels.VelocityKernel,
"connectivity": cr.kernels.ConnectivityKernel,
"realtime": cr.kernels.RealTimeKernel,
}
if kernel_type not in kernel_dic:
return {
"status": "error",
"message": f"Unsupported kernel type: {kernel_type}",
}
kernel_class = kernel_dic[kernel_type]
kwargs = filter_args(request, kernel_class)
if kernel_type == "cytotrace":
if "Ms" not in adata.layers or "Mu" not in adata.layers:
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)
kernel = kernel_class(adata=adata, **kwargs).compute_cytotrace()
else:
kernel = kernel_class(adata=adata, **kwargs)
ads.cr_kernel[kernel_type] = kernel
ads.set_adata(adata, adinfo=adinfo)
return {
"status": "success",
"message": f"Successfully created {kernel_type} kernel",
"kernel_type": kernel_type,
}
except ToolError as e:
raise ToolError(e)
except Exception as e:
if hasattr(e, "__context__") and e.__context__:
raise ToolError(e.__context__)
else:
raise ToolError(e)
@kernel_mcp.tool()
def compute_transition_matrix(
request: ComputeTransitionMatrixModel, adinfo: AdataInfo = AdataInfo()
):
"""Compute transition matrix for a specified kernel using appropriate parameters."""
try:
result = forward_request("compute_transition_matrix", request, adinfo)
if result is not None:
return result
kernel_type = request.kernel
ads = get_ads()
kernel = ads.cr_kernel[kernel_type].copy()
kwargs = filter_args(
request, kernel.compute_transition_matrix, show_progress_bar=False
)
kernel.compute_transition_matrix(**kwargs)
ads.cr_kernel[kernel_type] = kernel
return {
"status": "success",
"message": f"Successfully computed transition matrix for {kernel_type} kernel",
}
except ToolError as e:
raise ToolError(e)
except Exception as e:
if hasattr(e, "__context__") and e.__context__:
raise ToolError(e.__context__)
else:
raise ToolError(e)
