3ce187da02Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""MCP helper endpoints for data recipe."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
|
|
from fastapi import APIRouter
|
|
|
|
from core.data_recipe.service import build_mcp_providers
|
|
from loggers import get_logger
|
|
from models.data_recipe import (
|
|
McpToolsListRequest,
|
|
McpToolsListResponse,
|
|
McpToolsProviderResult,
|
|
)
|
|
from utils.utils import safe_error_detail
|
|
|
|
logger = get_logger(__name__)
|
|
router = APIRouter()
|
|
|
|
|
|
@router.post("/mcp/tools", response_model = McpToolsListResponse)
|
|
def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
|
|
try:
|
|
from data_designer.engine.mcp import io as mcp_io
|
|
except ImportError as exc:
|
|
logger.error(
|
|
"data_recipe.mcp.dependencies_unavailable",
|
|
error = str(exc),
|
|
exc_info = True,
|
|
)
|
|
return McpToolsListResponse(
|
|
providers = [
|
|
McpToolsProviderResult(
|
|
name = "",
|
|
error = "MCP dependencies unavailable.",
|
|
)
|
|
]
|
|
)
|
|
|
|
providers: list[McpToolsProviderResult] = []
|
|
tool_to_providers: dict[str, list[str]] = defaultdict(list)
|
|
|
|
from core.inference.mcp_client import stdio_mcp_enabled
|
|
|
|
for provider_payload in payload.mcp_providers:
|
|
provider_name = str(provider_payload.get("name", "")).strip()
|
|
if provider_payload.get("provider_type") == "stdio" and not stdio_mcp_enabled():
|
|
providers.append(
|
|
McpToolsProviderResult(
|
|
name = provider_name,
|
|
error = "Local (stdio) MCP servers are disabled on this host.",
|
|
)
|
|
)
|
|
continue
|
|
built = build_mcp_providers({"mcp_providers": [provider_payload]})
|
|
if len(built) != 1:
|
|
providers.append(
|
|
McpToolsProviderResult(
|
|
name = provider_name,
|
|
error = "Unsupported MCP provider config.",
|
|
)
|
|
)
|
|
continue
|
|
|
|
provider = built[0]
|
|
try:
|
|
tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec)
|
|
tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")})
|
|
for tool_name in tool_names:
|
|
tool_to_providers[tool_name].append(provider.name)
|
|
providers.append(
|
|
McpToolsProviderResult(
|
|
name = provider.name,
|
|
tools = tool_names,
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
logger.error(
|
|
"data_recipe.mcp.list_tools_failed",
|
|
error = str(exc),
|
|
exc_info = True,
|
|
)
|
|
providers.append(
|
|
McpToolsProviderResult(
|
|
name = provider.name or provider_name,
|
|
error = safe_error_detail(exc, fallback = "Failed to load tools."),
|
|
)
|
|
)
|
|
|
|
duplicate_tools = {
|
|
tool_name: provider_names
|
|
for tool_name, provider_names in sorted(tool_to_providers.items())
|
|
if len(provider_names) > 1
|
|
}
|
|
|
|
return McpToolsListResponse(
|
|
providers = providers,
|
|
duplicate_tools = duplicate_tools,
|
|
)
|