sadpandajoe commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r3870381181
########## superset/ai/tools/authoring.py: ########## @@ -0,0 +1,307 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Native AI adapters for Superset's existing MCP authoring tools.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from importlib import import_module +from threading import Thread +from typing import Any, ClassVar, TypeVar + +from pydantic import BaseModel, ValidationError + +from superset.ai.tools.base import AITool, ToolError, ToolOutput +from superset.mcp_service.chart.schemas import GenerateChartRequest +from superset.mcp_service.dashboard.schemas import GenerateDashboardRequest +from superset.mcp_service.dataset.schemas import CreateVirtualDatasetRequest +from superset.utils import json + +ModelT = TypeVar("ModelT", bound=BaseModel) +ToolCaller = Callable[[BaseModel], Any] + +_MCP_TOOL_MODULES = { + "create_virtual_dataset": ( + "superset.mcp_service.dataset.tool.create_virtual_dataset" + ), + "generate_chart": "superset.mcp_service.chart.tool.generate_chart", + "generate_dashboard": ("superset.mcp_service.dashboard.tool.generate_dashboard"), +} + + +def _tool_schema(model: type[BaseModel]) -> dict[str, Any]: + """Expose a request model without its server-only warning field.""" + schema = model.model_json_schema() + properties = dict(schema.get("properties", {})) + properties.pop("sanitization_warnings", None) + schema["properties"] = properties + if required := schema.get("required"): + schema["required"] = [ + name for name in required if name != "sanitization_warnings" + ] + return schema + + +def _validate(model: type[ModelT], payload: dict[str, Any], label: str) -> ModelT: + """Turn Pydantic errors into a correction the model can act on.""" + try: + return model.model_validate(payload) + except ValidationError as ex: + issues = [] + for error in ex.errors(include_url=False)[:3]: + location = ".".join(str(part) for part in error["loc"]) + issues.append(f"{location}: {error['msg']}") + raise ToolError(f"Invalid {label} request: {'; '.join(issues)}.") from ex + + +def _payload(response: Any) -> dict[str, Any]: + if isinstance(response, BaseModel): + return response.model_dump(mode="json", exclude_none=True) + if isinstance(response, dict): + return response + raise ToolError("Superset returned an unexpected authoring response.") + + +async def _call_mcp_tool(tool_name: str, request: BaseModel) -> Any: + """Call the registered tool through FastMCP so it gets a real context.""" + import_module(_MCP_TOOL_MODULES[tool_name]) + + from fastmcp import Client + + from superset.mcp_service.app import mcp + + arguments = { + "request": request.model_dump( + mode="json", + exclude={"sanitization_warnings"}, + exclude_none=True, + ) + } + async with Client(mcp) as client: + result = await client.call_tool(tool_name, arguments) + + if result.is_error: + raise ToolError(f"Superset could not run {tool_name}.") + return ( + result.structured_content + if result.structured_content is not None + else result.data + ) + + +def _run_mcp_tool(tool_name: str, request: BaseModel) -> dict[str, Any]: + """Run FastMCP off the agent loop with isolated Flask request state.""" + from flask import current_app, g + + try: + app = current_app._get_current_object() + user = getattr(g, "user", None) + except RuntimeError as ex: + raise ToolError("Authoring requires an authenticated request.") from ex + + username = getattr(user, "username", None) + email = getattr(user, "email", None) + if not username and not email: + raise ToolError("Authoring requires an authenticated user.") + + outcome: dict[str, Any] = {} + + def run() -> None: + try: + from flask import g as worker_g + + from superset.mcp_service.auth import load_user_with_relationships + + with app.test_request_context(): + worker_g.user = load_user_with_relationships( + username=str(username) if username else None, + email=str(email) if email else None, + ) + if worker_g.user is None: + raise ToolError("The authenticated user could not be reloaded.") + outcome["value"] = asyncio.run(_call_mcp_tool(tool_name, request)) + except BaseException as ex: # noqa: BLE001 + outcome["error"] = ex + + worker = Thread(target=run, name="superset-ai-authoring", daemon=True) + worker.start() + worker.join() Review Comment: A stalled MCP/database call leaves this unbounded join waiting forever, so the turn cannot reach its timeout or cancellation checks and a web/worker slot remains occupied. Could this wait be bounded and surface a controlled tool failure? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
