I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r4002577053
########## superset/ai/runtime/messages.py: ########## @@ -0,0 +1,574 @@ +# 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. +""" +The default runtime: a plain tool-use loop over the provider's message API. + +Chosen as the default because it needs nothing beyond an HTTP call — no agent +engine subprocess, no working directory, no bundled binary — so it works with +whatever provider a deployment configures. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import AsyncIterator +from typing import Any + +from superset.ai.events import ( + assistant_delta_event, + checkpoint_event, + error_event, + final_event, + GENERIC_ERROR_MESSAGE, + StreamEvent, + thinking_event, + thoughts_event, +) +from superset.ai.llm.base import ( + CompletionRequest, + LLMError, + LLMResponse, + Message, + StreamEventKind, + ToolCall, + ToolResult, +) +from superset.ai.runtime.base import BaseAgentRuntime, RunRequest, RunResult +from superset.ai.telemetry import ( + current_run, + POLICY_DENIED, + RunRecorder, + TOOL_UNAVAILABLE, +) +from superset.ai.types import MessageRole, ProgressStage, TokenUsage + +logger = logging.getLogger(__name__) + +#: How much of a tool's output is kept on the persisted message. The model +#: still sees the whole thing; this is the audit copy. +_RECORDED_OUTPUT_LIMIT = 2_000 + +#: Size of the chunks the finished answer is delivered in. +_DELIVERY_CHUNK_SIZE = 512 + +#: How much reasoning is kept on the result. Reasoning can run several times +#: longer than the answer, and this is persisted next to it. +_RECORDED_THOUGHTS_LIMIT = 8_000 + +_NO_ANSWER = ( + "I wasn't able to reach an answer for that. Try narrowing the question, " + "or naming the dataset you have in mind." +) + + +class MessagesApiRuntime(BaseAgentRuntime): + """ + Alternates model calls and tool calls until the model stops asking. + + Two behaviours are worth understanding before changing this class. + + First, prose the model emits *before* a tool call is treated as reasoning, + not answer: it becomes a ``thoughts`` event and is dropped from the answer. + A model narrating "the orders table looks right, let me check" is stating a + hypothesis it may abandon, and appending that to the answer produces a + reply that contradicts itself. + + Second, the loop always terminates and never raises for an operational + failure. By the time it runs, response headers have been flushed and an + exception can no longer become an HTTP status, so every failure is an event. + """ + + def __init__(self, provider: Any) -> None: + super().__init__(provider) + self._result = RunResult() + #: Set when the model signals it has finished answering. + self._finished = False + #: The most recent round trip's response, or ``None`` if it failed. The + #: turn methods are generators and cannot return a value. + self._last_response: LLMResponse | None = None + #: Whether any answer text has already been sent as it was generated. The + #: finished answer is only replayed in chunks when it has not. + self._streamed_text = False + + @property + def result(self) -> RunResult: + return self._result + + async def run(self, request: RunRequest) -> AsyncIterator[StreamEvent]: + self._result = RunResult() + self._finished = False + self._last_response = None + self._streamed_text = False + answer_parts: list[str] = [] + + yield thinking_event(ProgressStage.START, "Working on your question") + + # The provider's connection pool belongs to the loop this run is driven + # on, and the caller closes that loop as soon as the run ends. Closing + # here — inside the loop, however the run finishes, including when the + # generator is abandoned mid-way by a user pressing stop — is what keeps + # a client from being finalised against a dead loop. + try: + async for event in self._turn_loop(request, answer_parts): + yield event + + # A run that failed or was abandoned has already said so; emitting an + # answer as well would contradict it. + if self._result.error is not None or self._result.cancelled: + return + + answer = "\n\n".join(part for part in answer_parts if part).strip() + self._result.answer = answer or _NO_ANSWER + + # Only replayed when nothing was streamed — a provider without + # streaming support still gets to deliver its answer progressively. + # Replaying after live text would show the answer twice. + if not self._streamed_text: + for chunk in _chunk(self._result.answer): + yield assistant_delta_event(chunk) + yield final_event(self._result.answer) + finally: + await self.provider.aclose() + + async def _turn_loop( + self, + request: RunRequest, + answer_parts: list[str], + ) -> AsyncIterator[StreamEvent]: + """ + Alternate model and tool calls until the model stops or a budget runs out. + + Appends to ``answer_parts`` rather than returning the answer, because an + async generator cannot both yield events and return a value. + """ + deadline = time.monotonic() + request.timeout_seconds + conversation = list(request.messages) + + for turn in range(1, request.max_turns + 1): + self._result.turns = turn + + if self._should_stop(request, deadline): + if self._result.timed_out: + yield thinking_event( + ProgressStage.FALLBACK, + "Taking longer than expected — answering with what I have", + ) + return + + async for event in self._safe_turn(request, conversation, turn): + yield event + response = self._last_response + if response is None: + yield error_event() + return + + async for event in self._consume( + request, response, conversation, answer_parts + ): + yield event + + if self._finished or self._result.cancelled: + return + + # Budget exhausted without the model choosing to stop. + yield thinking_event( + ProgressStage.FALLBACK, + "Reached the step limit — answering with what I have", + ) + + async def _consume( + self, + request: RunRequest, + response: LLMResponse, + conversation: list[Message], + answer_parts: list[str], + ) -> AsyncIterator[StreamEvent]: + """Act on one model response, running any tools it asked for.""" + if response.thinking: + self._record_thoughts(response.thinking) + yield thoughts_event(response.thinking) + + if not response.wants_tools: + self._finished = True + if response.text: + answer_parts.append(response.text) + # Recorded as it arrives, not just at the end, so a run stopped + # after this point still persists what the user already saw. + self._result.answer = "\n\n".join( + part for part in answer_parts if part + ).strip() + return + + # Prose accompanying a tool call is reasoning, not answer. + if response.text: + self._record_thoughts(response.text) + yield thoughts_event(response.text) + + conversation.append( + Message( + role=MessageRole.ASSISTANT, + content=response.text, + tool_calls=list(response.tool_calls), + ) + ) + + results: list[ToolResult] = [] + async for event in self._run_tools(request, response.tool_calls, results): + yield event + + conversation.append(Message(role=MessageRole.USER, tool_results=results)) + + async def _run_tools( + self, + request: RunRequest, + calls: list[ToolCall], + results: list[ToolResult], + ) -> AsyncIterator[StreamEvent]: + """Execute this turn's tool calls, appending outcomes to ``results``.""" + for call in calls: + if self._cancelled(request): + self._result.cancelled = True + return + + yield thinking_event( + ProgressStage.TOOL, + f"Running {call.name}", + {"tool_name": call.name}, + ) + result, detail = self._invoke_tool(request, call) + results.append(result) + record = self._record_call(call, result, detail) + + # The frame carries the same record that is persisted, rather than a + # subset assembled separately. The subset was missing the arguments + # and the output, so a step expanded during a run showed nothing at + # all unless its tool happened to supply a display — and then filled + # itself in on reload, which looked like the detail arrived late. + # Sharing one record makes that class of drift impossible. + yield checkpoint_event( + f"{'Failed' if result.is_error else 'Finished'} {call.name}", + # ``tool_name`` as well as ``name``: the progress frames use that + # key, so a consumer reading either finds what it expects. + {"tool_name": call.name, **record}, + ) + + async def _safe_turn( + self, + request: RunRequest, + conversation: list[Message], + turn: int, + ) -> AsyncIterator[StreamEvent]: + """ + One model round trip, converting failure into a ``None`` response. + + A generator rather than a coroutine so the answer can reach the client as + the model produces it. The response is handed back on + :attr:`_last_response` because an async generator cannot both yield events + and return a value — the same reason ``_turn_loop`` writes into + ``answer_parts``. + + The failure detail goes to the log; the caller emits a message that cannot + leak a URL, a credential or a fragment of someone else's query. + """ + recorder = current_run() + started = time.monotonic() + self._last_response = None + try: + async for event in self._one_turn(request, conversation): + yield event + except LLMError as ex: + logger.warning("AI provider error on turn %s: %s", turn, ex) + self._result.error = str(ex) + self._trace_model_call(recorder, request, turn, started, error=ex) + self._last_response = None + return + except Exception as ex: # pylint: disable=broad-except + logger.exception("Unexpected error in AI runtime on turn %s", turn) + self._result.error = GENERIC_ERROR_MESSAGE + self._trace_model_call(recorder, request, turn, started, error=ex) + self._last_response = None + return + self._trace_model_call( + recorder, request, turn, started, response=self._last_response + ) + + def _trace_model_call( + self, + recorder: RunRecorder, + request: RunRequest, + turn: int, + started: float, + response: LLMResponse | None = None, + error: BaseException | None = None, + ) -> None: + """ + Report one round trip to telemetry. + + Content is passed as-is; whether any of it survives into a trace is the + redaction policy's decision, made in one place rather than here. + """ + if not recorder.enabled: + return + usage = response.usage if response is not None else TokenUsage() + recorder.model_call( + turn=turn, + # The concrete identifier when the provider reported one, and the + # capability tier otherwise, so a trace can always be grouped by + # what the run asked for. + model=usage.get("model") or request.model_alias.value, + duration_ms=int((time.monotonic() - started) * 1000), + input_tokens=usage.get("input_tokens"), + output_tokens=usage.get("output_tokens"), + stop_reason=response.stop_reason if response is not None else None, + error_type=type(error).__name__ if error is not None else None, + system_prompt=request.system_prompt, + response_text=response.text if response is not None else None, + ) + if error is not None: + recorder.error(error) + + async def _one_turn( + self, + request: RunRequest, + conversation: list[Message], + ) -> AsyncIterator[StreamEvent]: + """ + Call the model once, yielding answer text as the model produces it. + + Streaming is used when the provider supports it. The assembled response + is left on :attr:`_last_response` rather than returned, because a + generator cannot do both; it has the same shape either way, so callers do + not branch on which path ran. + """ + completion = CompletionRequest( Review Comment: The concrete-model propagation gap remains in the base runtime; the refreshed authoring branch does not fix it. The resolved model needs to reach `CompletionRequest`, with a test where an alias maps to a different default. Keeping this open rather than treating the recorded model name as proof. ########## superset/ai/runtime/messages.py: ########## @@ -0,0 +1,574 @@ +# 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. +""" +The default runtime: a plain tool-use loop over the provider's message API. + +Chosen as the default because it needs nothing beyond an HTTP call — no agent +engine subprocess, no working directory, no bundled binary — so it works with +whatever provider a deployment configures. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import AsyncIterator +from typing import Any + +from superset.ai.events import ( + assistant_delta_event, + checkpoint_event, + error_event, + final_event, + GENERIC_ERROR_MESSAGE, + StreamEvent, + thinking_event, + thoughts_event, +) +from superset.ai.llm.base import ( + CompletionRequest, + LLMError, + LLMResponse, + Message, + StreamEventKind, + ToolCall, + ToolResult, +) +from superset.ai.runtime.base import BaseAgentRuntime, RunRequest, RunResult +from superset.ai.telemetry import ( + current_run, + POLICY_DENIED, + RunRecorder, + TOOL_UNAVAILABLE, +) +from superset.ai.types import MessageRole, ProgressStage, TokenUsage + +logger = logging.getLogger(__name__) + +#: How much of a tool's output is kept on the persisted message. The model +#: still sees the whole thing; this is the audit copy. +_RECORDED_OUTPUT_LIMIT = 2_000 + +#: Size of the chunks the finished answer is delivered in. +_DELIVERY_CHUNK_SIZE = 512 + +#: How much reasoning is kept on the result. Reasoning can run several times +#: longer than the answer, and this is persisted next to it. +_RECORDED_THOUGHTS_LIMIT = 8_000 + +_NO_ANSWER = ( + "I wasn't able to reach an answer for that. Try narrowing the question, " + "or naming the dataset you have in mind." +) + + +class MessagesApiRuntime(BaseAgentRuntime): + """ + Alternates model calls and tool calls until the model stops asking. + + Two behaviours are worth understanding before changing this class. + + First, prose the model emits *before* a tool call is treated as reasoning, + not answer: it becomes a ``thoughts`` event and is dropped from the answer. + A model narrating "the orders table looks right, let me check" is stating a + hypothesis it may abandon, and appending that to the answer produces a + reply that contradicts itself. + + Second, the loop always terminates and never raises for an operational + failure. By the time it runs, response headers have been flushed and an + exception can no longer become an HTTP status, so every failure is an event. + """ + + def __init__(self, provider: Any) -> None: + super().__init__(provider) + self._result = RunResult() + #: Set when the model signals it has finished answering. + self._finished = False + #: The most recent round trip's response, or ``None`` if it failed. The + #: turn methods are generators and cannot return a value. + self._last_response: LLMResponse | None = None + #: Whether any answer text has already been sent as it was generated. The + #: finished answer is only replayed in chunks when it has not. + self._streamed_text = False + + @property + def result(self) -> RunResult: + return self._result + + async def run(self, request: RunRequest) -> AsyncIterator[StreamEvent]: + self._result = RunResult() + self._finished = False + self._last_response = None + self._streamed_text = False + answer_parts: list[str] = [] + + yield thinking_event(ProgressStage.START, "Working on your question") + + # The provider's connection pool belongs to the loop this run is driven + # on, and the caller closes that loop as soon as the run ends. Closing + # here — inside the loop, however the run finishes, including when the + # generator is abandoned mid-way by a user pressing stop — is what keeps + # a client from being finalised against a dead loop. + try: + async for event in self._turn_loop(request, answer_parts): + yield event + + # A run that failed or was abandoned has already said so; emitting an + # answer as well would contradict it. + if self._result.error is not None or self._result.cancelled: + return + + answer = "\n\n".join(part for part in answer_parts if part).strip() + self._result.answer = answer or _NO_ANSWER + + # Only replayed when nothing was streamed — a provider without + # streaming support still gets to deliver its answer progressively. + # Replaying after live text would show the answer twice. + if not self._streamed_text: + for chunk in _chunk(self._result.answer): + yield assistant_delta_event(chunk) + yield final_event(self._result.answer) + finally: + await self.provider.aclose() + + async def _turn_loop( + self, + request: RunRequest, + answer_parts: list[str], + ) -> AsyncIterator[StreamEvent]: + """ + Alternate model and tool calls until the model stops or a budget runs out. + + Appends to ``answer_parts`` rather than returning the answer, because an + async generator cannot both yield events and return a value. + """ + deadline = time.monotonic() + request.timeout_seconds + conversation = list(request.messages) + + for turn in range(1, request.max_turns + 1): + self._result.turns = turn + + if self._should_stop(request, deadline): + if self._result.timed_out: + yield thinking_event( + ProgressStage.FALLBACK, + "Taking longer than expected — answering with what I have", + ) + return + + async for event in self._safe_turn(request, conversation, turn): + yield event + response = self._last_response + if response is None: + yield error_event() + return + + async for event in self._consume( + request, response, conversation, answer_parts + ): + yield event + + if self._finished or self._result.cancelled: + return + + # Budget exhausted without the model choosing to stop. + yield thinking_event( + ProgressStage.FALLBACK, + "Reached the step limit — answering with what I have", + ) + + async def _consume( + self, + request: RunRequest, + response: LLMResponse, + conversation: list[Message], + answer_parts: list[str], + ) -> AsyncIterator[StreamEvent]: + """Act on one model response, running any tools it asked for.""" + if response.thinking: + self._record_thoughts(response.thinking) + yield thoughts_event(response.thinking) + + if not response.wants_tools: + self._finished = True + if response.text: + answer_parts.append(response.text) + # Recorded as it arrives, not just at the end, so a run stopped + # after this point still persists what the user already saw. + self._result.answer = "\n\n".join( + part for part in answer_parts if part + ).strip() + return + + # Prose accompanying a tool call is reasoning, not answer. + if response.text: + self._record_thoughts(response.text) + yield thoughts_event(response.text) + + conversation.append( + Message( + role=MessageRole.ASSISTANT, + content=response.text, + tool_calls=list(response.tool_calls), + ) + ) + + results: list[ToolResult] = [] + async for event in self._run_tools(request, response.tool_calls, results): + yield event + + conversation.append(Message(role=MessageRole.USER, tool_results=results)) + + async def _run_tools( + self, + request: RunRequest, + calls: list[ToolCall], + results: list[ToolResult], + ) -> AsyncIterator[StreamEvent]: + """Execute this turn's tool calls, appending outcomes to ``results``.""" + for call in calls: + if self._cancelled(request): + self._result.cancelled = True + return + + yield thinking_event( + ProgressStage.TOOL, + f"Running {call.name}", + {"tool_name": call.name}, + ) + result, detail = self._invoke_tool(request, call) + results.append(result) + record = self._record_call(call, result, detail) + + # The frame carries the same record that is persisted, rather than a + # subset assembled separately. The subset was missing the arguments + # and the output, so a step expanded during a run showed nothing at + # all unless its tool happened to supply a display — and then filled + # itself in on reload, which looked like the detail arrived late. + # Sharing one record makes that class of drift impossible. + yield checkpoint_event( + f"{'Failed' if result.is_error else 'Finished'} {call.name}", + # ``tool_name`` as well as ``name``: the progress frames use that + # key, so a consumer reading either finds what it expects. + {"tool_name": call.name, **record}, + ) + + async def _safe_turn( + self, + request: RunRequest, + conversation: list[Message], + turn: int, + ) -> AsyncIterator[StreamEvent]: + """ + One model round trip, converting failure into a ``None`` response. + + A generator rather than a coroutine so the answer can reach the client as + the model produces it. The response is handed back on + :attr:`_last_response` because an async generator cannot both yield events + and return a value — the same reason ``_turn_loop`` writes into + ``answer_parts``. + + The failure detail goes to the log; the caller emits a message that cannot + leak a URL, a credential or a fragment of someone else's query. + """ + recorder = current_run() + started = time.monotonic() + self._last_response = None + try: + async for event in self._one_turn(request, conversation): + yield event + except LLMError as ex: + logger.warning("AI provider error on turn %s: %s", turn, ex) + self._result.error = str(ex) + self._trace_model_call(recorder, request, turn, started, error=ex) + self._last_response = None + return + except Exception as ex: # pylint: disable=broad-except + logger.exception("Unexpected error in AI runtime on turn %s", turn) + self._result.error = GENERIC_ERROR_MESSAGE + self._trace_model_call(recorder, request, turn, started, error=ex) + self._last_response = None + return + self._trace_model_call( + recorder, request, turn, started, response=self._last_response + ) + + def _trace_model_call( + self, + recorder: RunRecorder, + request: RunRequest, + turn: int, + started: float, + response: LLMResponse | None = None, + error: BaseException | None = None, + ) -> None: + """ + Report one round trip to telemetry. + + Content is passed as-is; whether any of it survives into a trace is the + redaction policy's decision, made in one place rather than here. + """ + if not recorder.enabled: + return + usage = response.usage if response is not None else TokenUsage() + recorder.model_call( + turn=turn, + # The concrete identifier when the provider reported one, and the + # capability tier otherwise, so a trace can always be grouped by + # what the run asked for. + model=usage.get("model") or request.model_alias.value, + duration_ms=int((time.monotonic() - started) * 1000), + input_tokens=usage.get("input_tokens"), + output_tokens=usage.get("output_tokens"), + stop_reason=response.stop_reason if response is not None else None, + error_type=type(error).__name__ if error is not None else None, + system_prompt=request.system_prompt, + response_text=response.text if response is not None else None, + ) + if error is not None: + recorder.error(error) + + async def _one_turn( + self, + request: RunRequest, + conversation: list[Message], + ) -> AsyncIterator[StreamEvent]: + """ + Call the model once, yielding answer text as the model produces it. + + Streaming is used when the provider supports it. The assembled response + is left on :attr:`_last_response` rather than returned, because a + generator cannot do both; it has the same shape either way, so callers do + not branch on which path ran. + """ + completion = CompletionRequest( + messages=conversation, + system=request.system_prompt, + model_alias=request.model_alias, + tools=tuple(request.tools.definitions()) if request.tools else (), + ) + + if not self.provider.supports_streaming: Review Comment: Still open: the existing pre-call deadline check does not bound a hung provider await, and the normal runtime calls do not use the retry wrapper. This needs a single remaining-budget contract in the base runtime, not independent timeout guesses in the authoring adapter. ########## 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: You are right that the warning text is not containment. I reopened this thread: the mutating daemon thread may outlive the reported timeout, and the shorter profile deadline is not propagated. This remains a blocker on this authoring PR, not something I am asking the reviewer to overlook. A safe fix needs the base run/deadline and ambiguous-write reconciliation contract; no claim of cancellation or safe retry is made by this update. ########## superset-frontend/src/features/ai/hooks/useAIAction.ts: ########## @@ -0,0 +1,102 @@ +/** + * 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. + */ + +/** + * @fileoverview How the rest of the app asks the assistant something. + * + * A "Debug this query" button in SQL Lab, or a chart menu item, needs to open the + * assistant with a question already asked. It does that by dispatching one DOM + * event that the panel listens for. + * + * A DOM event rather than a shared module or a global: the caller may be mounted + * in a different React tree from the panel (the panel is mounted by the chat host) + * and must not import it, or the assistant's code would be pulled into every + * bundle that offers an action. The panel is the only listener, and if it is not + * mounted the event is simply unheard. + */ + +import { useCallback } from 'react'; +import { chat } from 'src/core/chat'; +import type { AIActionPayload } from '../types'; + +/** The event the panel listens for. */ +export const AI_ACTION_EVENT = 'superset-ai-action'; + +export type AiActionEvent = CustomEvent<AIActionPayload>; + +/** + * Asks the assistant a question in a new conversation. + * + * Safe to call whether or not the panel is mounted: the chat host is asked to open + * first, which mounts it, and the event is dispatched after so a panel that has + * just mounted has its listener attached. + */ +export const triggerAIAction = (payload: AIActionPayload): void => { + if (!payload.prompt.trim()) { + return; + } + chat.open(); + window.dispatchEvent( + new CustomEvent<AIActionPayload>(AI_ACTION_EVENT, { detail: payload }), + ); +}; Review Comment: The mount-before-dispatch race is still present in the inherited action/panel handoff. No fix is included in this refresh. The base needs a queued payload plus acknowledgement from the mounted listener, with a closed-panel action test; keeping this open. ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,508 @@ +--- +title: AI Assistant +hide_title: true +sidebar_position: 17 +version: 1 +--- + +# AI Assistant + +The AI Assistant is a conversational interface for exploring your data. A user +asks a question in plain language; the assistant finds relevant datasets, +inspects their schema, writes and runs read-only SQL, and answers with both the +result and the query it used. The shipped profiles are read-only. A deployment +can explicitly add chart and dashboard authoring tools to a gated profile. + +Superset ships **no model provider and talks to no model vendor by default**. +The feature is disabled, and even when enabled it returns `404` until you point +it at a provider you control. Nothing is sent anywhere until you configure it. + +## Enabling it + +Two things are required: the feature flag, and a provider. + +```python +# superset_config.py +FEATURE_FLAGS = { + "AI_ASSISTANT": True, +} + +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.anthropic.AnthropicProvider" +AI_LLM_PROVIDER_CONFIG = { + "api_key": os.environ["ANTHROPIC_API_KEY"], + "models": { + "default": "claude-sonnet-4-5", + "fast": "claude-haiku-4-5", + "reasoning": "claude-opus-4-1", + }, +} +``` + +Install the matching extra: + +```bash +pip install "apache-superset[ai-anthropic]" # or [ai-openai] +``` + +Then run `superset init` so the assistant's permissions are created and assigned +to roles. Without this the endpoints return `403`. + +Conversations are stored in Superset's metadata database, so no extra +infrastructure is needed for the default configuration. + +### Which roles get access + +`superset init` grants `can_read`/`can_write` on `AIAssistant` to **Admin** and +**Alpha** only. "Write" here means writing one's own conversation. The shipped +profiles remain read-only; an operator who enables asset-authoring tools must +also gate that profile, and each tool enforces the current user's normal asset +and dataset permissions. + +**Gamma does not get it by default.** The assistant runs queries and costs +money per question, so it is granted deliberately rather than inherited. To +give it to Gamma users, add `can_read`/`can_write` on `AIAssistant` to Gamma or +to a custom role. + +Every query the assistant runs is subject to the *user's own* database and +dataset permissions. It cannot read anything the person chatting with it could +not read themselves. + +Because it is not in Gamma, it is also not inherited by the Public role when +`PUBLIC_ROLE_LIKE = "Gamma"` — an anonymous visitor cannot reach the assistant +unless you grant it explicitly. + +## Choosing a provider + +`AI_LLM_PROVIDER_CLASS` is a dotted path to a +`superset.ai.llm.base.BaseLLMProvider` subclass. Two are bundled: + +| Class | Use for | +| --- | --- | +| `superset.ai.llm.anthropic.AnthropicProvider` | The Anthropic Messages API | +| `superset.ai.llm.openai_compatible.OpenAICompatibleProvider` | OpenAI, and anything exposing an OpenAI-compatible endpoint — vLLM, Ollama, a private gateway | + +`AI_LLM_PROVIDER_CONFIG` is passed to the provider's constructor and its +contents are provider-defined. For the OpenAI-compatible provider, `base_url` +points it anywhere: + +```python +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.openai_compatible.OpenAICompatibleProvider" +AI_LLM_PROVIDER_CONFIG = { + "base_url": "https://llm.internal.example.com/v1";, + "api_key": os.environ["MY_GATEWAY_KEY"], + "models": {"default": "our-hosted-model"}, +} +``` + +Everything vendor-specific — URLs, authentication, model naming — lives in the +provider. Superset core contains none of it, so a self-hosted model or a private +gateway needs configuration rather than a fork. + +### Model tiers and selection + +Profiles and prompts refer to capability *tiers* (`default`, `fast`, +`reasoning`), never to a vendor's model names. The provider maps tiers to +concrete models via the `models` dict. A tier you do not configure is an error +when requested, never a silent substitution — so cost and answer quality stay +attributable to the model actually used. + +Users may also pin a specific model per turn. Only models present in your +`models` mapping are accepted; anything else is rejected. + +## Agent profiles + +A profile bundles the decisions that differ between a quick answer and a careful +investigation: which tools are available, which model tier, and how many steps. +Two ship by default — `default` and `analyst`. + +**Which tools a model may invoke is a decision each deployment makes**, so +profiles are fully configurable. `AI_AGENT_PROFILES` maps a profile key to the +fields you want to override, leaving the rest alone: + +```python +AI_AGENT_PROFILES = { + # Let the assistant search and inspect, but never run SQL. + "default": {"tools": ["search_assets", "list_databases", "get_schema"]}, + + # Let the analyst profile think harder and longer. + "analyst": {"model_alias": "reasoning", "max_turns": 60}, + + # Add a profile only some users may select. + "deep": { + "name": "Deep analysis", + "description": "Slow, thorough, multi-step.", + "tools": ["search_assets", "get_schema", "execute_sql"], + "required_permission": ("can_write", "AIAssistant"), + }, + + # Opt-in authoring. The tools still enforce normal asset and dataset RBAC. + "builder": { + "name": "Dashboard builder", + "tools": [ + "search_assets", + "get_schema", + "create_virtual_dataset", + "generate_chart", + "generate_dashboard", + ], + "required_permission": ("can_write", "Dashboard"), + }, +} +``` + +A tool name that does not exist is an error naming the typo and listing the +valid names, rather than an assistant that quietly lacks a capability. An empty +`tools` list is valid and means conversation with no data access. + +`required_permission` is enforced on both the listing *and* the run path, so a +profile a user cannot see is also one they cannot invoke by posting its key. + +### Available tools + +| Tool | What it does | +| --- | --- | +| `search_assets` | Finds datasets, charts and dashboards the user can see | +| `list_databases` | Lists database connections exposed to SQL Lab | +| `get_schema` | Lists schemas, tables and columns | +| `execute_sql` | Runs a **read-only** query | +| `validate_sql` | Checks a query without running it | +| `get_chart_context` | Reads a chart's definition | +| `get_dashboard_context` | Reads a dashboard's definition | +| `create_virtual_dataset` | Saves a read-only SQL query as a chartable dataset; opt-in only | +| `generate_chart` | Previews or saves a native chart; opt-in only | +| `generate_dashboard` | Creates a dashboard from saved chart IDs; opt-in only | + +## Customising the prompt + +Three levers, in increasing order of bluntness. + +**Add to it.** `AI_EXTRA_PROMPT_SECTIONS` appends your own sections. This is +where deployment-specific knowledge belongs — your table conventions, your +warehouse's dialect quirks, how your business defines a metric. The shipped +prompt is deliberately generic and mentions no particular database engine. + +**Remove from it.** `AI_DISABLED_PROMPT_SECTIONS` drops a shipped section by +key, for when you disagree with one. The safety section cannot be disabled. + +**Replace it.** `AI_SYSTEM_PROMPT` substitutes the whole thing. + +:::warning +Setting `AI_SYSTEM_PROMPT` discards the shipped safety and prompt-injection +rules along with everything else. Your deployment then owns them. +::: + +`AI_SYSTEM_PROMPT_MUTATOR` is a last-mile callable applied after assembly, +mirroring `SQL_QUERY_MUTATOR`. + +## Where turns execute + +`AI_ASSISTANT_EXECUTION_MODE` decides where the work happens. + +**`"inline"`** (default) runs the turn in the web process. Nothing extra to +deploy. + +**`"worker"`** hands it to Celery. Web workers stay free, and a browser that +loses its connection can rejoin a run in progress. It requires Celery and a +Redis event bus: + +```python +AI_ASSISTANT_EXECUTION_MODE = "worker" +AI_ASSISTANT_EVENT_BUS = "redis" +AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG = { + "CACHE_TYPE": "RedisCache", + "CACHE_REDIS_HOST": "redis", + "CACHE_REDIS_PORT": 6379, + "CACHE_REDIS_DB": 0, +} + +class CeleryConfig: + imports = ( + # ... your existing imports ... + "superset.ai.tasks", + ) +``` + +Streams need Redis commands the general-purpose cache client does not expose, +which is why the bus is configured separately rather than reusing `CACHE_CONFIG`. + +Selecting `"worker"` with the in-memory event bus raises rather than leaving +every stream silently empty, and so does selecting the Redis bus without a +usable connection. + +A turn is deliberately **not** retried after a worker crash: inference costs +money, and re-running a turn the user may already have partly seen would charge +twice. The message records that it failed and the user can ask again. + +## Safety and limits + +Guards are applied before any tool runs, configured via +`AI_AGENT_TOOL_POLICIES`: + +- **Read-only SQL.** Enforced using Superset's own SQL parser, not pattern + matching — so a write hidden behind a comment, a CTE, a second statement, or + an unparseable construct is refused. `EXPLAIN`, `SHOW` and `DESCRIBE` are + permitted; everything the parser cannot vouch for is not. +- **Identifier safety.** Table and column names are resolved against metadata + the user may see rather than interpolated into SQL. + +These bound blast radius; they do not replace authorization. Every tool that +touches a data-bearing object performs the same permission check the REST API +does. + +Result sizes are capped by `AI_AGENT_MAX_RESULT_ROWS` and +`AI_AGENT_MAX_RESULT_BYTES`, and truncation is reported rather than hidden. Turn +length is bounded by `AI_AGENT_MAX_TURNS` and `AI_AGENT_TIMEOUT_SECONDS`; a run +that exhausts either answers with what it has. + +Content that arrives from your warehouse or asset metadata — table comments, +chart titles, column labels — is marked as untrusted in the prompt, because a +value in a database is data and not an instruction. + +### Cancellation + +Cancellation is cooperative: a run stops at its next step boundary. A run inside +a single long model call or a single long query will not stop until that call +returns. + +## Monitoring and tracing + +Superset bundles **no integration with any AI monitoring product**. Instead it +exposes a small sink interface, `AITelemetry`, and calls it once per run, once +per model round trip and once per tool call. Whatever you already use — +Braintrust, LangSmith, Langfuse, Arize Phoenix, an OpenTelemetry collector, a +self-hosted alternative, or a table in your own warehouse — you connect by +implementing that interface and listing it in `AI_TELEMETRY`. + +Entries are instances or dotted paths, exactly as for `EVENT_LOGGER` and +`STATS_LOGGER`. Two sinks ship in-tree and depend on nothing external: + +```python +# superset_config.py +import logging + +from superset.ai.telemetry import LoggingAITelemetry, StatsLoggerAITelemetry + +AI_TELEMETRY = [ + # One structured line per span, at the level you choose. + LoggingAITelemetry(level=logging.INFO), + # Counters and timings through your configured STATS_LOGGER. + StatsLoggerAITelemetry(), +] +``` + +`StatsLoggerAITelemetry` emits under a `superset.ai.` prefix: `run.start`, +`run.end`, `run.outcome.<outcome>`, `run.duration_ms`, `run.turns`, +`run.tokens.input`, `run.tokens.output`, `model_call`, +`model_call.duration_ms`, `model_call.error`, `error`, and per tool +`tool_call.<tool>`, `tool_call.<tool>.duration_ms`, `tool_call.<tool>.error` +and `tool_call.<tool>.truncated`. User, run and thread identifiers deliberately +never appear in a metric name — a metric per user is how a metrics backend gets +brought down. That detail belongs in a trace, which is what a custom sink is +for. + +### The content trade-off + +`AI_TELEMETRY_REDACT_CONTENT` defaults to `True`, and telemetry then carries +**structure and measurements only**: durations, token counts, model names, tool +names, outcomes, error classes, and the run, thread and user identifiers. No +question, no answer, no SQL, no row of data. Redaction is applied where the +trace is built, so a sink cannot receive content by accident even if it looks +for it. + +Setting it to `False` is what makes a trace genuinely useful for debugging +answer quality — you can read the prompt that produced a wrong answer and the +statement it ran. It also means the text of business questions and values from +your warehouse leave Superset for whichever service your sinks talk to. In many +organisations that is a decision for someone other than the person editing the +config file. `AI_TELEMETRY_MAX_CONTENT_CHARS` (default 10,000) caps any single +content field so one large result cannot dominate a payload. + +### A custom sink + +Every method has a no-op default, so implement only the ones you need — a sink +that only wants token counts overrides `on_model_call` and nothing else. + +```python +from superset.ai.telemetry import AITelemetry, ModelCallTrace, RunTrace + + +class TracingServiceTelemetry(AITelemetry): + """Forwards runs to an external tracing service.""" + + def __init__(self, client): + self._client = client + + def on_run_start(self, run: RunTrace) -> None: + self._client.start_span(run.run_id, name="superset.ai.run", attributes={ + "thread": run.thread_uuid, + "user": run.user_id, + }) + + def on_model_call(self, run: RunTrace, call: ModelCallTrace) -> None: + self._client.event(run.run_id, "model_call", { + "turn": call.turn, + "model": call.model, + "input_tokens": call.input_tokens, + "output_tokens": call.output_tokens, + # None unless you have turned redaction off. + "prompt": call.system_prompt, + }) + + def on_run_end(self, run: RunTrace) -> None: + self._client.end_span(run.run_id, status=str(run.outcome), attributes={ + "duration_ms": run.duration_ms, + "turns": run.turns, + "usage": run.usage, + }) + + +AI_TELEMETRY = [TracingServiceTelemetry(client=my_tracing_client)] +``` + +Three things to know before you write one: + +- **Sinks are called on the thread answering the user.** Anything that makes a + network call should hand off to a queue or a background thread; otherwise a + slow monitoring backend becomes slow answers. +- **A sink that raises cannot break a run.** Failures are logged once and + ignored, and the other configured sinks still receive everything. The same + applies to a dotted path that will not import: it is skipped with a warning + rather than taking the assistant down, because a missing observer loses the + record of a run and not the run itself. +- **`agent_key`, `model` and `question` are resolved after the run starts**, so + a `RunTrace` passed to `on_run_start` may carry less than the one passed to + the later hooks. Read those on `on_run_end`. + +## Connecting your own MCP servers + +The assistant's built-in tools cover Superset itself. To let it reach anything +else — your data catalog, a metrics service, a ticketing system — attach an +[MCP](https://modelcontextprotocol.io) server. Superset bundles no third-party +integration and connects to nothing by default; you name the servers. + +```bash +pip install "apache-superset[ai-mcp]" +``` + +```python +AI_AGENT_MCP_SERVERS = { + "acme_catalog": { + "url": "https://mcp.acme.internal/mcp";, + "transport": "streamable_http", # or "sse" + "headers": {"Authorization": f"Bearer {os.environ['ACME_MCP_TOKEN']}"}, + "timeout_seconds": 30, + "tool_allowlist": ["search_tables"], # omit to offer every tool + }, +} + +# Then let a profile use it. +AI_AGENT_PROFILES = { + "default": {"mcp_servers": ["acme_catalog"]}, +} +``` + +Its tools appear to the model as `mcp__acme_catalog__search_tables`. The +namespace means a foreign tool can never shadow a built-in one, and it is the +name to use in `tool_allowlist` and `tool_denylist`. + +### What Superset does to keep a foreign server contained + +A third-party server is untrusted input, and possibly untrusted intent: + +- **Everything it returns is marked as untrusted** before the model sees it, so + text in a tool result is treated as data rather than instructions. Tool + *descriptions* get the same treatment, since they enter the prompt every turn. +- **No Superset credential is ever forwarded.** Only the headers you configured + for that server are sent — never the user's session cookie, CSRF token, or an + inbound authorization header. +- **SQL execution through a foreign server is refused by default.** Superset's + read-only enforcement and per-dataset authorization cannot apply to a query + another system runs, so allowing it would silently bypass both. Set + `AI_AGENT_MCP_DENY_FOREIGN_SQL = False` to accept that trade deliberately. +- **Results obey the same size cap** as built-in tools, and the cap is applied + while reading, so a hostile server cannot exhaust memory before truncation. +- **A server being down does not break the assistant.** Discovery failure means + that server contributes no tools for the turn; the built-ins keep working. + +A profile naming a server you have not configured is an error, because a typo +there is indistinguishable at runtime from an agent that has quietly lost a +capability. Note that discovery happens per turn, so a slow server adds its +latency to every turn that uses it. + +## Retention + +Conversations are kept for `AI_ASSISTANT_MESSAGE_RETENTION_DAYS` (default 30). +Pruning is not automatic — schedule it if you want it enforced. + +## Trying it locally + +The development `docker compose` stack can bring the assistant up against the +example data. Put the settings in `docker/.env-local`, which is untracked: + +```bash +# docker/.env-local + +# Point at any OpenAI-compatible endpoint, including a private gateway. +SUPERSET_AI_LLM_BASE_URL=https://your-gateway/v1 +SUPERSET_AI_LLM_API_KEY=your-token +SUPERSET_AI_MODEL_DEFAULT=your-model-name +``` + +```bash +docker compose up +``` + +The assistant appears once both a URL and a key are present; with neither set the +stack behaves exactly as it did before. The model providers are optional extras, +so add whichever one you need to `docker/requirements-local.txt`: + +``` +openai>=1.60.0,<2 +``` + +The local stack also logs traces to the container output and, unlike the +production default, includes prompts and SQL in them. + +## Full configuration reference Review Comment: Fixed in 0d5be1cba7: the reference now includes all six `AI_SUGGESTED_PROMPTS_*` settings and `AI_ASSISTANT_EVENT_STREAM_PREFIX`, with defaults checked against `superset/config.py` and the suggestions consumer. ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,508 @@ +--- +title: AI Assistant +hide_title: true +sidebar_position: 17 +version: 1 +--- + +# AI Assistant + +The AI Assistant is a conversational interface for exploring your data. A user +asks a question in plain language; the assistant finds relevant datasets, +inspects their schema, writes and runs read-only SQL, and answers with both the +result and the query it used. The shipped profiles are read-only. A deployment +can explicitly add chart and dashboard authoring tools to a gated profile. + +Superset ships **no model provider and talks to no model vendor by default**. +The feature is disabled, and even when enabled it returns `404` until you point +it at a provider you control. Nothing is sent anywhere until you configure it. + +## Enabling it + +Two things are required: the feature flag, and a provider. Review Comment: Fixed the setup instructions in 0d5be1cba7 by explicitly enabling `ENABLE_EXTENSIONS` alongside `AI_ASSISTANT`. I verified that the refreshed `AppContent` still gates chat hosts through `FeatureFlag.EnableExtensions`; no UI decoupling is implied. ########## superset/ai/orchestrator.py: ########## @@ -0,0 +1,647 @@ +# 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. +""" +Runs one assistant turn end to end. + +Sits between the HTTP layer and the runtime: loads the conversation, assembles +the prompt, resolves the tools the chosen profile allows, drives the runtime, +publishes every event to the bus, and records the outcome on the assistant +message. + +Deliberately independent of *where* it runs. The same function body serves the +inline path and the Celery path, which is what makes the execution mode a +configuration choice rather than two implementations that drift apart. +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid as uuid_module +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass +from typing import Any + +from superset.ai.events import ( + cancelled_event, + done_event, + error_event, + GENERIC_ERROR_MESSAGE, + session_event, + StreamEvent, +) +from superset.ai.llm.base import Message +from superset.ai.telemetry import bind_run, current_run, start_run +from superset.ai.types import MessageRole, MessageStatus, RunOutcome, StreamEventType +from superset.utils.decorators import transaction + +logger = logging.getLogger(__name__) + +#: Cache key prefix for a run's cancellation flag. A flag rather than a signal +#: because a worker cannot be interrupted mid-call reliably; the runtime checks +#: this between steps. +_CANCEL_PREFIX = "ai-cancel-" + +#: How long a cancellation request stays meaningful. +_CANCEL_TTL_SECONDS = 900 + +#: Stored when a run is stopped before it produced any answer, so the +#: transcript still records that the turn happened. +_STOPPED_WITHOUT_ANSWER = "_Stopped before an answer was produced._" + +#: Stored when a run exhausted its time budget without saying anything. Phrased +#: as something the user can act on, because retrying is usually the right move. +_TIMED_OUT_WITHOUT_ANSWER = ( + "The assistant ran out of time before it could answer. Please try again." +) + +#: Ceiling on the page context recorded on a message. Well below the prompt's own +#: limit: this is stored per turn and read back with the whole transcript. +_RECORDED_CONTEXT_LIMIT = 4_000 + + +@dataclass +class TurnRequest: + """One unit of work: answer the latest message on a thread.""" + + thread_uuid: str + user_id: int + run_id: str + #: Assistant message row to fill in. Created before the run starts so a + #: client that reconnects has something to attach to. + assistant_message_uuid: str + profile_key: str | None = None + #: Concrete model to pin, overriding the profile's tier. + model: str | None = None + #: What the user had on screen when they asked. Supplied by the client, + #: which is the only party that knows which tab is open, what is typed in + #: the editor and which filters are applied. + page_context: dict[str, Any] | None = None + + def to_payload(self) -> dict[str, Any]: + """Serialise for the task broker.""" + return { + "thread_uuid": self.thread_uuid, + "user_id": self.user_id, + "run_id": self.run_id, + "assistant_message_uuid": self.assistant_message_uuid, + "profile_key": self.profile_key, + "model": self.model, + "page_context": self.page_context, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> TurnRequest: + """Rebuild from a broker payload.""" + return cls(**payload) + + +def new_run_id() -> str: + """Identifier for one run, used as the event-stream key.""" + return str(uuid_module.uuid4()) + + +#: Runs cancelled in this process. +#: +#: Held alongside the cache rather than instead of it. Superset's default cache +#: is a null cache, which accepts a write and discards it — so a cache-only +#: implementation would leave cancellation silently broken on a default install, +#: with the button appearing to work and nothing stopping. This set makes inline +#: execution correct with no cache at all; the cache is what carries a +#: cancellation across processes for worker execution. +_CANCELLED_LOCALLY: set[str] = set() + + +def request_cancel(run_id: str) -> None: + """ + Ask a run to stop. + + Cooperative by design: the flag is recorded here and observed by the runtime + between steps. A run blocked inside a single long model call or query will + not notice until that call returns, which is a real limit worth documenting + rather than hiding. + """ + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.add(run_id) + try: + cache_manager.cache.set( + f"{_CANCEL_PREFIX}{run_id}", True, timeout=_CANCEL_TTL_SECONDS + ) + except Exception: # pylint: disable=broad-except + logger.warning("Could not record cancellation for AI run %s", run_id) + + +def is_cancelled(run_id: str) -> bool: + """Whether a stop has been requested for this run.""" + from superset.extensions import cache_manager + + if run_id in _CANCELLED_LOCALLY: + return True + try: + return bool(cache_manager.cache.get(f"{_CANCEL_PREFIX}{run_id}")) + except Exception: # pylint: disable=broad-except + # A cache that cannot be read must not make every run appear cancelled; + # that would stop all inference the moment the cache went away. + return False + + +def clear_cancel(run_id: str) -> None: + """Drop a run's cancellation flag.""" + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.discard(run_id) + try: + cache_manager.cache.delete(f"{_CANCEL_PREFIX}{run_id}") + except Exception: # pylint: disable=broad-except + logger.debug("Could not clear cancellation flag for AI run %s", run_id) + + +def stream_turn(request: TurnRequest) -> Iterator[StreamEvent]: + """ + Answer a turn, yielding events as they happen. + + This is the primary entry point. Inline execution consumes it directly from + inside the streaming response, which means the producer and the reader are + the same process by construction — important because Superset runs several + web workers, and a turn that published to one process's in-memory queue + while the browser's stream landed on another would appear to hang forever. + + Never raises for an operational failure: a failure is an ``error`` event and + an ``error`` message status, because the caller may already have flushed + response headers or may be a worker with no one to report to. + """ + recorder = start_run( + run_id=request.run_id, + thread_uuid=request.thread_uuid, + user_id=request.user_id, + ) + # Shared with ``_run`` so the ``finally`` below can see the runtime's + # partial result and whether the message was already written. + state: dict[str, Any] = {} + try: + # Bound here rather than inside ``_run`` so that a run which fails before + # it has resolved a profile still produces a start and an end, and so + # that the runtime can report its own spans without the runtime contract + # growing a telemetry parameter. + with bind_run(recorder): + recorder.run_started() + yield from _run(request, state) + except Exception as ex: # pylint: disable=broad-except + logger.exception("AI turn failed for run %s", request.run_id) + recorder.error(ex) + recorder.run_ended(outcome=RunOutcome.ERROR) + answer, extra = _partial_from_state(state) + extra["outcome"] = RunOutcome.ERROR.value + _finalise_message( + request.assistant_message_uuid, + # The generic text rather than the exception: this is persisted and + # served back to the browser, so it must not carry internals. The + # detail is in the log line above, keyed by run id. + content=answer or GENERIC_ERROR_MESSAGE, + status=MessageStatus.ERROR, + extra=extra, + ) + state["finalised"] = True + yield error_event() + yield done_event(ok=False) + finally: + clear_cancel(request.run_id) + # A client that stops the run, or simply navigates away, abandons this + # generator part-way through. Nothing above will have written the + # message, so it would otherwise sit in ``streaming`` with no content + # for ever — the user loses both the partial answer and any record that + # the turn happened. Persist whatever was produced. + _abandon_message(request.assistant_message_uuid, state) + # Idempotent, so the ordinary paths above win. + recorder.run_ended(outcome=RunOutcome.CANCELLED) + + +def execute_turn(request: TurnRequest) -> RunOutcome: + """ + Answer a turn, publishing events to the event bus. + + Used by worker execution, where the reader is in another process. Shares its + whole body with :func:`stream_turn` so the two execution modes cannot drift + apart in behaviour. + """ + from superset.ai.eventbus import get_event_bus + + bus = get_event_bus() + outcome = RunOutcome.SUCCESS + + for event in stream_turn(request): + bus.publish(request.run_id, event) + if event.type is StreamEventType.ERROR: + outcome = RunOutcome.ERROR + elif event.type is StreamEventType.CANCELLED: + outcome = RunOutcome.CANCELLED + elif event.type is StreamEventType.DONE and not event.payload.get("ok"): + # A run that ended un-ok without an explicit error frame timed out. + if outcome is RunOutcome.SUCCESS: + outcome = RunOutcome.TIMEOUT + + return outcome + + +def _run(request: TurnRequest, state: dict[str, Any]) -> Iterator[StreamEvent]: + """Assemble and drive the run. See :func:`stream_turn` for error policy.""" + from superset.ai.factories import ( + get_profiles, + get_provider, + get_runtime, + get_tools_for_profile, + ) + from superset.ai.policy import load_policy_chain + from superset.ai.runtime.base import RunRequest + from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO + + recorder = current_run() + + thread = AIChatThreadDAO.find_by_uuid_for_user(request.thread_uuid, request.user_id) + if thread is None: + # The thread vanished between accepting the message and running it. + recorder.run_ended(outcome=RunOutcome.ERROR) + yield error_event("That conversation is no longer available.") + yield done_event(ok=False) + return + + profile = get_profiles().get(request.profile_key) + tools = get_tools_for_profile(profile) + provider = get_provider() + runtime = get_runtime(provider) + state["runtime"] = runtime + + yield session_event(request.thread_uuid, request.assistant_message_uuid) + + from superset.ai.page_context import render_page_context + + history = _build_history(AIChatMessageDAO.find_for_thread(thread)) + # Recorded as well as prompted with, so the transcript can show what the + # assistant was told about the user's screen. An answer that looks wrong is + # usually an answer to a different question than the reader assumed, and the + # page context is where that difference lives. + rendered_context = render_page_context(request.page_context) + state["page_context"] = rendered_context + system_prompt = _build_system_prompt(tools, rendered_context) + model = _resolved_model(provider, request.model, profile) + + recorder.describe( + agent_key=profile.key, + model=model, + question=_latest_question(history), + ) + + run_request = RunRequest( + messages=history, + system_prompt=system_prompt, + tools=tools, + policies=load_policy_chain(), + model_alias=profile.model_alias, + max_turns=profile.max_turns or _config("AI_AGENT_MAX_TURNS", 20), + timeout_seconds=profile.timeout_seconds + or _config("AI_AGENT_TIMEOUT_SECONDS", 300), + should_cancel=lambda: is_cancelled(request.run_id), + ) + + _mark_streaming(request.assistant_message_uuid) + + # The runtime is async and this is a synchronous generator, so the async + # events are drained into a list per batch rather than bridged with a + # thread. Collecting the whole run before yielding would defeat streaming, + # so the loop pulls one event at a time from a dedicated event loop. + yield from _drain(runtime.run(run_request)) + + result = runtime.result + outcome = _outcome_of(result) + if result.error is not None: + # The only place the provider's own words are recorded. They do not go on + # the message: that is served back to the browser, and a transport error + # can name internal hosts. + logger.warning( + "AI run %s failed: %s", + request.run_id, + result.error, + ) + _finalise_message( + request.assistant_message_uuid, + content=_terminal_content(result, outcome), + status=_status_of(outcome), + extra={ + "outcome": outcome.value, + "agent_key": profile.key, + "model": model, + "tool_calls": result.tool_calls, + "turns": result.turns, + **_recorded_context(rendered_context), + }, + ) + + recorder.run_ended( + outcome=outcome, + turns=result.turns, + answer=result.answer, + ) + + state["finalised"] = True + + if outcome is RunOutcome.CANCELLED: + yield cancelled_event() + yield done_event(ok=outcome is RunOutcome.SUCCESS) + + +def _drain(source: AsyncIterator[StreamEvent]) -> Iterator[StreamEvent]: + """ + Pull an async iterator one item at a time from a synchronous caller. + + A single event loop is kept for the whole run and stepped with + ``__anext__``, so each event reaches the client as it is produced rather + than after the run completes. + """ + loop = asyncio.new_event_loop() + try: + iterator = source.__aiter__() + while True: + try: + yield loop.run_until_complete(iterator.__anext__()) + except StopAsyncIteration: + return + finally: + loop.close() + + +def _build_history(messages: list[Any]) -> list[Message]: + """ + Convert stored rows into provider messages, trimmed to the configured budget. + + Trimming is newest-first by count and then by total characters, because an + old turn is less useful than a recent one and an oversized request is + rejected outright by every provider. + """ + max_messages = _config("AI_ASSISTANT_MAX_HISTORY_MESSAGES", 25) + max_chars = _config("AI_ASSISTANT_MAX_HISTORY_CHARS", 100_000) + + usable = [ + message + for message in messages + if message.content and message.role != MessageRole.SYSTEM.value + ] + window = usable[-max_messages:] + + # Trimmed to the budget, but never to nothing: a single over-budget message + # is still sent, because the provider's own error about it is more useful + # than a request with no question in it. + total = sum(len(message.content) for message in window) + while len(window) > 1 and total > max_chars: + dropped = window.pop(0) + total -= len(dropped.content) + + return [ + Message(role=MessageRole(message.role), content=message.content) + for message in window + ] + + +def _latest_question(history: list[Message]) -> str | None: + """ + The question this turn is answering. + + Offered to telemetry, which drops it unless a deployment has turned + redaction off. ``None`` when the turn was somehow queued with no user + message, which is a state worth being able to see rather than crash on. + """ + for message in reversed(history): + if message.role is MessageRole.USER and message.content: + return message.content + return None + + +def _build_system_prompt(tools: Any, rendered_context: str) -> str: + """ + Assemble the system prompt for the tools actually on offer. + + The page context arrives already rendered, and is appended after assembly + rather than joining the section list, because it is per-request data rather + than a configured section — and because the layering rules deliberately + refuse content from anywhere but ``superset.core`` in the prompt's own + sections. Rendering happens in the caller so the same text can be recorded on + the message without rendering it twice. + """ + from flask import current_app + + from superset.ai.prompts import assemble_system_prompt + from superset.ai.prompts.core import core_sections + + prompt = assemble_system_prompt( Review Comment: The knowledge-provider integration is still an open base feature gap; updating master does not connect configured providers to prompt assembly or a knowledge tool. I am not treating the presence of the configuration key as working functionality. This needs the implementation or a reduced public contract in #42805. ########## superset/ai/api.py: ########## @@ -0,0 +1,989 @@ +# 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. +""" +REST API for the AI assistant. + +Every route carries ``@protect()`` and is reached through ``@expose`` on a +``BaseSupersetApi`` subclass, which is what makes Flask-AppBuilder's +authorization actually run. Ownership is enforced a second time in the command +and DAO layers, so a conversation identifier is never on its own a capability. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Generator +from typing import Any, cast + +from flask import current_app, request, Response, stream_with_context +from flask_appbuilder.api import expose, permission_name, protect, safe +from marshmallow import ValidationError + +from superset.ai.events import ( + error_event, + KEEPALIVE_FRAME, + KEEPALIVE_INTERVAL_SECONDS, +) +from superset.ai.schemas import ( + AgentResponseSchema, + CancelPostSchema, + FeedbackPostSchema, + MessagePostSchema, + RunAcceptedResponseSchema, + SuggestedPromptsPostSchema, + ThreadDetailResponseSchema, + ThreadPostSchema, + ThreadPutSchema, + ThreadResponseSchema, +) +from superset.ai.types import MessageRole, MessageStatus +from superset.commands.ai.exceptions import ( + AIChatMessageInvalidError, + AIChatMessageNotFoundError, + AIChatThreadInvalidError, + AIChatThreadNotFoundError, +) +from superset.extensions import event_logger +from superset.utils.core import get_user_id +from superset.utils.decorators import transaction +from superset.views.base_api import BaseSupersetApi, statsd_metrics + +logger = logging.getLogger(__name__) + +#: Upper bound on how long a client may hold a stream open, so an abandoned +#: browser tab cannot pin a worker indefinitely. +_STREAM_TIMEOUT_SECONDS = 900 + +#: How often a reader checks the event bus for new frames. +#: +#: Deliberately separate from ``KEEPALIVE_INTERVAL_SECONDS``. Passing the +#: keep-alive interval as the poll interval made the reader sleep fifteen seconds +#: between checks and then deliver everything that had accumulated in one batch — +#: so a worker-mode run showed no streaming at all: the answer and every tool call +#: appeared in fifteen-second lumps. One controls responsiveness, the other how +#: often an idle connection is reassured; they are not the same number. +_EVENT_POLL_SECONDS = 0.1 + + +class AIRestApi(BaseSupersetApi): + """Conversations with the AI assistant.""" + + resource_name = "ai" + openapi_spec_tag = "AI Assistant" + allow_browser_login = True + class_permission_name = "AIAssistant" + + openapi_spec_component_schemas = ( + AgentResponseSchema, + CancelPostSchema, + FeedbackPostSchema, + MessagePostSchema, + RunAcceptedResponseSchema, + SuggestedPromptsPostSchema, + ThreadDetailResponseSchema, + ThreadPostSchema, + ThreadPutSchema, + ThreadResponseSchema, + ) + + @expose("/agent/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def agents(self) -> Response: + """List agent profiles the current user may select. + --- + get: + summary: List available agent profiles + responses: + 200: + description: Available profiles + content: + application/json: + schema: + type: object + properties: + result: + type: array + items: + $ref: '#/components/schemas/AgentResponseSchema' + 401: + $ref: '#/components/responses/401' + 403: + $ref: '#/components/responses/403' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.factories import get_profiles + + profiles = get_profiles().visible_to_current_user() + return self.response(200, result=[p.to_public_dict() for p in profiles]) + + @expose("/model/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def models(self) -> Response: + """List models this deployment has configured. + --- + get: + summary: List selectable models + responses: + 200: + description: Configured model identifiers + content: + application/json: + schema: + type: object + properties: + result: + type: array + items: + type: string + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.factories import get_provider + + return self.response(200, result=get_provider().available_models()) + + @expose("/thread/", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post_thread", + log_to_statsd=False, + ) + def post_thread(self) -> Response: + """Create a conversation. + --- + post: + summary: Create a conversation + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ThreadPostSchema' + responses: + 201: + description: Conversation created + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/ThreadResponseSchema' + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import CreateAIChatThreadCommand + + try: + payload = ThreadPostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + try: + thread = CreateAIChatThreadCommand( + user_id=self._user_id(), + title=payload.get("title"), + agent_key=payload.get("agent_key"), + ).run() + except AIChatThreadInvalidError as ex: + return self.response_422(message=str(ex)) + return self.response(201, result=_thread_dict(thread)) + + @expose("/thread/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def get_threads(self) -> Response: + """List the current user's conversations. + --- + get: + summary: List conversations + parameters: + - in: query + name: limit + schema: + type: integer + - in: query + name: offset + schema: + type: integer + responses: + 200: + description: Conversations + content: + application/json: + schema: + type: object + properties: + count: + type: integer + result: + type: array + items: + $ref: '#/components/schemas/ThreadResponseSchema' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import AIChatThreadDAO + + limit = request.args.get("limit", type=int) or 50 + offset = request.args.get("offset", type=int) or 0 + threads = AIChatThreadDAO.find_all_for_user( + self._user_id(), limit=limit, offset=offset + ) + return self.response( + 200, + count=len(threads), + result=[_thread_dict(thread) for thread in threads], + ) + + @expose("/thread/<thread_uuid>", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def get_thread(self, thread_uuid: str) -> Response: + """Fetch a conversation and its messages. + --- + get: + summary: Get a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + responses: + 200: + description: Conversation with messages + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/ThreadDetailResponseSchema' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import ( + AIChatFeedbackDAO, + AIChatMessageDAO, + AIChatThreadDAO, + ) + + user_id = self._user_id() + thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, user_id) + if thread is None: + return self.response_404() + + messages = AIChatMessageDAO.find_for_thread(thread) + # Resolved for the whole transcript at once so the panel can show which + # replies this user already rated; without it a reload loses the verdict + # and the message looks unrated. + verdicts = AIChatFeedbackDAO.find_verdicts_for_user( + [message.id for message in messages], user_id + ) + detail = _thread_dict(thread) + detail["messages"] = [ + _message_dict(message, liked=verdicts.get(message.id)) + for message in messages + ] + return self.response(200, result=detail) + + @expose("/thread/<thread_uuid>", methods=("PUT",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.put_thread", + log_to_statsd=False, + ) + def put_thread(self, thread_uuid: str) -> Response: + """Rename or archive a conversation. + --- + put: + summary: Update a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ThreadPutSchema' + responses: + 200: + description: Conversation updated + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 422: + $ref: '#/components/responses/422' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import UpdateAIChatThreadCommand + + try: + payload = ThreadPutSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + try: + thread = UpdateAIChatThreadCommand( + thread_uuid, + self._user_id(), + title=payload.get("title"), + status=payload.get("status"), + ).run() + except AIChatThreadNotFoundError: + return self.response_404() + except AIChatThreadInvalidError as ex: + return self.response_422(message=str(ex)) + return self.response(200, result=_thread_dict(thread)) + + @expose("/thread/<thread_uuid>", methods=("DELETE",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.delete_thread", + log_to_statsd=False, + ) + def delete_thread(self, thread_uuid: str) -> Response: + """Delete a conversation and its messages. + --- + delete: + summary: Delete a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + responses: + 200: + description: Conversation deleted + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import DeleteAIChatThreadCommand + + try: + DeleteAIChatThreadCommand(thread_uuid, self._user_id()).run() + except AIChatThreadNotFoundError: + return self.response_404() + return self.response(200, message="OK") + + @expose("/thread/<thread_uuid>/message", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post_message", + log_to_statsd=False, + ) + def post_message(self, thread_uuid: str) -> Response: + """Post a user message and start a run. + --- + post: + summary: Post a message + description: > + Stores the user's message, creates a placeholder assistant message, + and starts a run. Returns immediately; consume the answer from the + stream endpoint using the returned run identifier. + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/MessagePostSchema' + responses: + 202: + description: Run accepted + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/RunAcceptedResponseSchema' + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 422: + $ref: '#/components/responses/422' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.orchestrator import new_run_id + from superset.commands.ai import AppendAIChatMessageCommand + + try: + payload = MessagePostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + user_id = self._user_id() + + try: + user_message = AppendAIChatMessageCommand( + thread_uuid, + user_id, + MessageRole.USER, + payload["content"], + request_id=payload.get("request_id"), + ).run() + # Created up front so a client that reconnects before any token + # arrives still has a row to attach its stream to. + assistant_message = AppendAIChatMessageCommand( + thread_uuid, + user_id, + MessageRole.ASSISTANT, + "", + request_id=payload.get("request_id"), + status=MessageStatus.PENDING, + ).run() + except AIChatThreadNotFoundError: + return self.response_404() + except (AIChatMessageInvalidError, AIChatThreadInvalidError) as ex: + return self.response_422(message=str(ex)) + + run_id = new_run_id() + _record_run_context(assistant_message, run_id, payload) + + self._start_run( + thread_uuid=thread_uuid, + user_id=user_id, + run_id=run_id, + assistant_message_uuid=str(assistant_message.uuid), + agent_key=payload.get("agent_key"), Review Comment: The persisted conversation profile is still not the per-turn fallback in this branch, and reopening tabs is part of the same contract. This is unchanged base behavior in the refresh. It needs a create/reopen/submit-without-agent-key regression before the thread can close. -- 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]
