sadpandajoe commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r3911479703
########## superset/ai/tools/authoring.py: ########## @@ -0,0 +1,310 @@ +# 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( Review Comment: This hands the chat user's identity to the MCP call only through `g.user`, but the MCP auth hook prefers `MCP_DEV_USERNAME` before that fallback. A deployment that still has that documented dev setting will therefore run `generate_chart`/`generate_dashboard` as the configured user rather than the requester, bypassing the requester's dataset and write permissions. Could this call pin the reloaded requester as the authenticated principal instead of relying on the lowest-priority fallback? ########## superset/ai/tools/authoring.py: ########## @@ -0,0 +1,310 @@ +# 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(float(app.config.get("AI_AGENT_TIMEOUT_SECONDS", 300))) + + if worker.is_alive(): + raise ToolError("Superset authoring timed out.") Review Comment: The timeout reports the authoring operation as failed but leaves the daemon thread running, so a slow create can still commit after the user or model retries and produce a duplicate chart, dashboard, or dataset. Could the call be cancelled/contained before returning a retryable failure, or at least report that the operation may still complete? ########## docs/static/feature-flags.json: ########## @@ -9,6 +9,12 @@ "lifecycle": "development", "description": "Enables Table V2 (AG Grid) viz plugin" }, + { + "name": "AI_ASSISTANT", + "default": false, + "lifecycle": "development", + "description": "Enables the conversational AI assistant: an experimental feature that answers questions about your data by running read-only queries. Off by default and inert until a model provider is configured via AI_LLM_PROVIDER_CLASS \u2014 with the flag on but no provider the endpoints still return 404, so enabling the flag alone sends nothing anywhere. The assistant's own tools are read-only; it cannot create or modify assets." Review Comment: This says the assistant cannot create or modify assets, but this PR adds profile-opt-in tools that create virtual datasets, charts, and dashboards under the same flag. Could the feature-flag description mirror the admin guide's qualified wording and be regenerated from the corrected config comment? -- 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]
