sadpandajoe commented on code in PR #43135: URL: https://github.com/apache/superset/pull/43135#discussion_r3817977413
########## superset/migrations/versions/2026-08-04_00-00_a1c4f7e29b31_add_ai_chat_tables.py: ########## @@ -0,0 +1,167 @@ +# 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. +"""add_ai_chat_tables + +Revision ID: a1c4f7e29b31 +Revises: e7d93a524ff6 +Create Date: 2026-08-04 00:00:00.000000 + +""" + +# revision identifiers, used by Alembic. +revision = "a1c4f7e29b31" +down_revision = "f3a8c1d2e9b7" Review Comment: This revision points at an ancestor that already has descendants in the current base, creating a second Alembic head. `superset db upgrade` now stops on multiple heads, so could this be rebased onto the current migration head (or merged) before this lands? ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,489 @@ +--- +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. + +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 +assistant's tools are read-only and it cannot create or modify assets. + +**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"), + }, +} +``` + +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 | + +## 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 Review Comment: This says exhausting `AI_AGENT_MAX_TURNS` answers with partial work, but the new runtime marks that case as an error and emits no final answer. Could the documentation distinguish timeout behavior from step-limit exhaustion so operators do not configure a limit expecting a result that is now reported as incomplete? ########## superset/config.py: ########## @@ -2897,6 +2913,326 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq "CACHE_REDIS_SSL_CA_CERTS": None, } +# --------------------------------------------------------- +# AI assistant +# --------------------------------------------------------- +# Requires the AI_ASSISTANT feature flag. Superset ships no model provider and +# talks to no model vendor by default: until AI_LLM_PROVIDER_CLASS names a +# usable provider the assistant's endpoints return 404. +# +# Dotted path to a superset.ai.llm.base.BaseLLMProvider subclass. Point this at +# a vendor provider, an OpenAI-compatible endpoint, a self-hosted model, or a +# private gateway. Everything vendor-specific — base URLs, authentication, +# model naming — belongs in the provider, not here. +AI_LLM_PROVIDER_CLASS: str | None = None + +# Keyword arguments passed to the provider's constructor. Contents are entirely +# provider-defined. Keep credentials out of this file: read them from the +# environment or a secret store in your own config. +# +# AI_LLM_PROVIDER_CONFIG = { +# "api_key": os.environ["MY_LLM_API_KEY"], +# "base_url": "https://llm.internal.example.com/v1";, +# "models": { +# "default": "some-balanced-model", +# "fast": "some-small-model", +# "reasoning": "some-large-model", +# }, +# } +AI_LLM_PROVIDER_CONFIG: dict[str, Any] = {} + +# Dotted path to a superset.ai.runtime.base.BaseAgentRuntime subclass driving +# the tool-use loop. +AI_AGENT_RUNTIME_CLASS = "superset.ai.runtime.messages.MessagesApiRuntime" + +# Where a turn is executed. +# +# "inline" — in the web worker handling the request. No extra infrastructure, +# but a turn occupies a worker for its whole duration. +# "worker" — handed to Celery; the request streams events from the event bus. +# Survives a browser reconnect and keeps web workers free, at the +# cost of requiring Celery and a shared event bus. +AI_ASSISTANT_EXECUTION_MODE: Literal["inline", "worker"] = "inline" + +# How streamed events travel from producer to the HTTP response. +# +# "memory" — an in-process queue. Correct only when the producer and the +# streaming request are the same process, i.e. inline execution. +# "redis" — Redis streams, via the same cache backend the async-query +# channel uses. Required for "worker" execution mode. +AI_ASSISTANT_EVENT_BUS: Literal["memory", "redis"] = "memory" + +# Redis connection for the AI event bus. Required when AI_ASSISTANT_EVENT_BUS is +# "redis". Streams need commands the general-purpose cache client does not +# expose, so this is configured separately rather than borrowed from +# CACHE_CONFIG. The accepted shape matches +# GLOBAL_ASYNC_QUERIES_CACHE_BACKEND; point both at the same Redis if you like. +AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG: dict[str, Any] = { + "CACHE_TYPE": "RedisCache", + "CACHE_REDIS_HOST": "localhost", + "CACHE_REDIS_PORT": 6379, + "CACHE_REDIS_USER": "", + "CACHE_REDIS_PASSWORD": "", + "CACHE_REDIS_DB": 0, + "CACHE_DEFAULT_TIMEOUT": 300, + "CACHE_REDIS_SSL": False, +} + +# Key prefix for AI event streams when the Redis bus is in use. +AI_ASSISTANT_EVENT_STREAM_PREFIX = "ai-events-" + +# How long a run's event stream is retained, in seconds. Bounds how late a +# reconnecting browser can still pick up a run it lost. +AI_ASSISTANT_EVENT_TTL_SECONDS = 900 + +# Named agent profiles, merged over the built-ins by key. Each value is a dict +# of fields to override, so narrowing one profile does not mean restating the +# rest. The most important field is "tools": which tools that profile may +# invoke. An unknown tool name is a startup error, not a silent omission. +# +# AI_AGENT_PROFILES = { +# # Take the shipped default but forbid raw SQL. +# "default": {"tools": ["search_assets", "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", +# "tools": ["search_assets", "get_schema", "execute_sql"], +# "required_permission": ("can_write", "AIAssistant"), +# }, +# } +AI_AGENT_PROFILES: dict[str, Any] = {} + +# Ceiling on model round trips in a single turn. A turn that needs more than +# this is answered with what it has rather than looping indefinitely. +AI_AGENT_MAX_TURNS = 20 + +# Wall-clock budget for one turn, in seconds. +AI_AGENT_TIMEOUT_SECONDS = 300 + +# Pre-tool-use guards, applied in order. Each is a dotted path to a +# superset.ai.policy.ToolPolicy implementation. These bound blast radius; they +# do not replace the per-object authorization checks inside each tool. +AI_AGENT_TOOL_POLICIES: list[str] = [ + "superset.ai.policy.ReadOnlySqlPolicy", + "superset.ai.policy.IdentifierPolicy", + "superset.ai.policy.ForeignToolPolicy", +] + +# Rows and bytes a single tool result may return before it is truncated. +# Model context is finite, and an unbounded result set exhausts it. +AI_AGENT_MAX_RESULT_ROWS = 500 +AI_AGENT_MAX_RESULT_BYTES = 256 * 1024 + +# External MCP servers whose tools may be offered to an agent profile. Superset +# ships none and integrates with no third-party service: with this empty, nothing +# in superset.ai.mcp is ever reached and the assistant behaves exactly as it does +# without it. +# +# A server listed here is only *available*. It is used by an agent profile that +# names it in its "mcp_servers" field, via AI_AGENT_PROFILES. A profile naming a +# server that is not configured here is an error, not a silently shorter tool +# list. +# +# AI_AGENT_MCP_SERVERS = { +# # The key is the server name. It becomes part of every tool name this +# # server contributes, so keep it short: letters, digits, hyphens and +# # underscores, and no double underscore. +# "acme_catalog": { +# # Required. Absolute http:// or https:// endpoint. +# "url": "https://mcp.acme.internal/mcp";, +# # "streamable_http" (default) or "sse". +# "transport": "streamable_http", +# # The ONLY headers sent to this server. Superset never forwards the +# # user's session cookie, CSRF token or any Superset auth header: an +# # external server is not a party to the user's Superset session. +# # Read secrets from the environment rather than writing them here. +# "headers": {"Authorization": f"Bearer {os.environ['ACME_MCP_TOKEN']}"}, +# # Per-call budget. Bounds how long one call may occupy the worker +# # running the turn. Defaults to 30. +# "timeout_seconds": 30, +# # Which of the server's tools to take. Absent or None means every +# # tool it offers, which lets the server decide what the agent can do. +# # Either the server's own name ("search_tables") or the namespaced +# # name Superset assigns ("mcp__acme_catalog__search_tables") matches. +# "tool_allowlist": ["search_tables"], +# # Refused regardless of the allowlist. +# "tool_denylist": [], +# }, +# } +# +# AI_AGENT_PROFILES = { +# "default": {"mcp_servers": ["acme_catalog"]}, +# } +# +# Every tool from a server is namespaced "mcp__<server>__<tool>". The namespace is +# stable, appears in stored conversation history, and is what makes it impossible +# for a server offering "execute_sql" to displace Superset's own tool of that +# name. Foreign results pass through the same AI_AGENT_MAX_RESULT_BYTES bound and +# the same AI_AGENT_TOOL_POLICIES chain as built-in ones, and are wrapped as +# untrusted content before the model sees them. +# +# A server that is unreachable, slow or unreadable contributes no tools and the +# agent keeps working with the built-ins. Discovery happens while assembling the +# registry for a turn, so a slow server costs up to its timeout at the start of +# each turn that uses it. +# +# Requires the 'mcp' package; it is imported only once a server is configured. +AI_AGENT_MCP_SERVERS: dict[str, Any] = {} + +# Refuse any external MCP tool whose name advertises SQL execution — anything +# containing "execute_sql", "run_sql" or "query" by default. Enforced by +# superset.ai.policy.ForeignToolPolicy. +# +# On by default because Superset's read-only enforcement and its per-datasource +# authorization can only apply to SQL Superset itself runs. A third-party server +# executing SQL goes through neither, so permitting it silently removes both +# controls rather than merely widening the surface. Set this False only if you +# have satisfied yourself that the servers you have configured enforce +# equivalent controls of their own. +AI_AGENT_MCP_DENY_FOREIGN_SQL = True + +# Conversation history sent to the model: the most recent N messages, further +# trimmed oldest-first until under the character budget. +AI_ASSISTANT_MAX_HISTORY_MESSAGES = 25 +AI_ASSISTANT_MAX_HISTORY_CHARS = 100_000 + +# Timezone for the authoritative date given to the model, so it never has to +# infer today's date or weekday. +AI_ASSISTANT_TIMEZONE = "UTC" + +# Days a conversation is retained. Pruning is performed by the Review Comment: The configuration tells operators to schedule `ai.prune_conversations`, but the PR only registers `ai.run_turn` and never consumes the retention setting. Following this guidance leaves prompts and messages unpruned indefinitely; could the pruning task be implemented before documenting it as the retention mechanism? ########## superset/ai/runtime/messages.py: ########## @@ -0,0 +1,577 @@ +# 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. Review Comment: Text that accompanies a tool call is streamed before `_consume` reclassifies it as reasoning. On this new exhaustion path there is no final frame to replace those deltas, so the frontend renders that reasoning as a partial answer plus an error until reload; can the terminal error path retract or avoid exposing those provisional deltas? ########## 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() Review Comment: A replayed `request_id` reuses the existing message rows but still generates a new run ID and schedules another run. Retrying after a lost 202 can therefore execute tools and inference twice while both runs race to finalize one assistant message; can this return the original run instead of starting another one? ########## 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"), + model=payload.get("model"), + page_context=payload.get("page_context"), + ) + + return self.response( + 202, + result={ + "message_uuid": str(user_message.uuid), + "assistant_message_uuid": str(assistant_message.uuid), + "run_id": run_id, + }, + ) + + @expose("/thread/<thread_uuid>/stream", methods=("GET",)) + @protect() + @statsd_metrics + @permission_name("read") + def stream(self, thread_uuid: str) -> Response: + """Stream a run's events. + --- + get: + summary: Stream assistant events + description: > + Server-sent events for one run. Frame names are session, thinking, + thoughts, checkpoint, assistant_delta, final, error, cancelled and + done. The done frame is always last and reports whether the run + succeeded. + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + - in: query + name: run_id + required: true + schema: + type: string + responses: + 200: + description: An event stream + content: + text/event-stream: + schema: + type: string + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + # No @safe here: once headers are flushed an exception can no longer + # become a status code, so failures are reported as in-band error frames. + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO + + run_id = request.args.get("run_id") + if not run_id: + return self.response_400(message="run_id is required") + + # Ownership is checked before the stream opens; the run identifier alone + # must not grant access to another user's conversation. + thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, self._user_id()) + if thread is None: + return self.response_404() + + pending = _find_run_message(AIChatMessageDAO.find_for_thread(thread), run_id) + if pending is None: + return self.response_404() + + turn = None Review Comment: Each matching stream GET reconstructs and executes the turn without atomically claiming a pending run. A reconnect, concurrent stream, or later GET after completion can rerun model/tool work and overwrite the stored answer; should this only execute a successfully claimed pending run? ########## superset/ai/tools/sql.py: ########## @@ -0,0 +1,540 @@ +# 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. +""" +Running and checking SQL. + +Both tools are reads. ``execute_sql`` decides that with Superset's own parser +rather than a keyword match, because a prefix match is defeated by a leading +comment, a CTE wrapping a mutation, and a second statement smuggled after a +legitimate ``SELECT`` — all of which the parser already handles for the rest of +Superset. + +Authorization is layered deliberately: + +1. ``DatabaseDAO.find_by_id`` applies ``DatabaseFilter``, so a database the user + has no grant on is indistinguishable from one that does not exist. +2. ``expose_in_sqllab`` is honoured, so an operator who withheld a connection + from ad-hoc querying has also withheld it here. +3. The database's ``allow_dml`` setting is checked. +4. Every statement must be non-mutating. +5. ``security_manager.raise_for_access`` with ``force_dataset_match=True`` — the + same strictness SQL Lab applies — so the tables referenced must resolve to + datasets the user may read. +""" + +from __future__ import annotations + +import logging +import time +from decimal import Decimal +from typing import Any, Callable, ClassVar + +from superset.ai.tools.base import AITool, ToolError, ToolOutput + +logger = logging.getLogger(__name__) + +#: Rows returned when the caller does not ask for a specific limit. Small on +#: purpose: an agent inspecting data needs a shape, not a dump, and the model +#: pays for every row in context. +DEFAULT_ROW_LIMIT = 100 + +#: Rows kept in the UI summary. The summary is persisted on the message and sent +#: to the browser, so it carries a sample rather than the result set. +DISPLAY_SAMPLE_ROWS = 20 + +#: Characters of executed SQL shown in the UI summary. +DISPLAY_SQL_CHARS = 4000 + + +def _database_or_refuse(database_id: Any) -> Any: + """ + Resolve a database the current user is allowed to query, or refuse. + + Uses the DAO so that ``DatabaseFilter`` — the same filter the database REST + API applies — decides visibility. A database the user cannot see is reported + as "not found" rather than "forbidden", so the tool cannot be used to probe + for the existence of connections. + """ + from superset.daos.database import DatabaseDAO + + if not isinstance(database_id, int) or isinstance(database_id, bool): + raise ToolError("'database_id' must be an integer. Use list_databases first.") + + database = DatabaseDAO.find_by_id(database_id) + if database is None: + raise ToolError( + f"No database with id {database_id} is available to you. " + f"Call list_databases to see the ones you can query." + ) + if not database.expose_in_sqllab: + raise ToolError( + f"Database {database.database_name!r} is not available for ad-hoc " + f"queries. Call list_databases to see the ones that are." + ) + return database + + +def _parse_or_refuse(sql: str, database: Any) -> Any: + """ + Parse ``sql`` for the database's engine, or refuse. + + Fails closed: SQL that will not parse cannot be shown to be read-only, so it + is refused rather than executed. The parser error is logged rather than + returned, because its message quotes the offending query back and that text + is not ours to echo. + """ + from superset.sql.parse import SQLScript + + if not isinstance(sql, str) or not sql.strip(): + raise ToolError("'sql' must be a non-empty string.") + + try: + return SQLScript(sql, database.db_engine_spec.engine) + except Exception: # pylint: disable=broad-except + logger.info("Refusing unparseable SQL for database %s", database.id) + raise ToolError( + "That SQL could not be parsed. Send a single, syntactically valid " + "read-only statement." + ) from None + + +def _assert_read_only(script: Any, database: Any) -> None: + """ + Refuse anything that is not a read. + + The ``allow_dml`` check comes first because it is the more specific refusal: + telling the model the *deployment* forbids writes on this connection is more + useful than a generic "read-only tool" message. The blanket mutation check + then applies even where ``allow_dml`` is enabled, since this tool is a read + regardless of what the connection would otherwise permit. + """ + has_mutation = script.has_mutation() + + if has_mutation and not database.allow_dml: + raise ToolError( + f"Writes are disabled on database " + f"{database.database_name!r} (allow_dml is off). Rewrite this as a " + f"SELECT." + ) + if has_mutation: + raise ToolError( + "This tool only runs read-only SQL. Rewrite this as a SELECT — " + "statements that modify data or schema are refused." + ) + # A statement the parser could not fully model has no enumerable table + # references, so neither the mutation check above nor the per-table + # authorization below can vouch for it. Refused rather than guessed at. + if script.has_unparseable_statement: + raise ToolError( + "That SQL contains a statement this tool cannot verify as " + "read-only. Send a plain SELECT." + ) + + +def _raise_for_sql_access( + database: Any, + sql: str, + catalog: str | None, + schema: str | None, +) -> None: + """ + Check the user may read every table the query touches. + + ``force_dataset_match=True`` matches what SQL Lab's own pre-execute + validator uses: each referenced table must resolve to a registered dataset + the user has access to, rather than falling through to a broader + catalog- or schema-level grant. + """ + from superset import security_manager + from superset.exceptions import SupersetSecurityException + + try: + security_manager.raise_for_access( + database=database, + sql=sql, + catalog=catalog, + schema=schema, + force_dataset_match=True, + ) + except SupersetSecurityException as ex: + # The exception message names the tables that were denied, which is + # exactly what lets the model pick a different source. + raise ToolError(str(ex.error.message)) from None + + +def _row_limit(requested: Any) -> int: + """ + Clamp the caller's limit to the configured ceiling. + + The model may not raise the cap by asking for more; ``limit`` narrows only. + """ + # A misconfigured or absent ceiling falls back to the default rather than + # becoming unbounded: sending a query with no limit is the one outcome this + # function exists to prevent. + ceiling = DEFAULT_ROW_LIMIT + try: + from flask import current_app + + configured = current_app.config.get("AI_AGENT_MAX_RESULT_ROWS") + if configured is not None: + ceiling = int(configured) + except Exception: # pylint: disable=broad-except + ceiling = DEFAULT_ROW_LIMIT + if ceiling < 1: + ceiling = DEFAULT_ROW_LIMIT + + if requested is None: + return min(DEFAULT_ROW_LIMIT, ceiling) + if not isinstance(requested, int) or isinstance(requested, bool) or requested < 1: + raise ToolError("'limit' must be a positive integer.") + return min(requested, ceiling) + + +def _columns_and_records( + data: Any, +) -> tuple[list[dict[str, str]], list[dict[str, Any]]]: + """ + Normalise a statement's result rows. + + ``Database.execute`` returns a ``DataFrame`` when a row limit was supplied + and a plain list of mappings when one was not, so both shapes are handled + rather than relying on the caller always producing the first. Assuming the + frame is how this tool previously raised ``AttributeError`` on a result it + had asked for perfectly legitimately. + """ + if hasattr(data, "columns") and hasattr(data, "to_dict"): + columns = [ + {"name": str(name), "type": str(data[name].dtype)} for name in data.columns + ] + return columns, list(data.to_dict(orient="records")) + + records = [dict(row) for row in (data or [])] + names: list[str] = [] + for record in records: + for key in record: + if key not in names: + names.append(key) + # Without a frame there are no dtypes to report; the values themselves still + # carry their types through ``_json_safe``. + return [{"name": str(name), "type": "unknown"} for name in names], records + + +def _json_safe(value: Any) -> Any: + """ + Coerce one warehouse value into something JSON can carry. + + ``Decimal`` becomes a float and binary becomes text (or hex when it is not + text at all); everything else exotic — dates, intervals, UUIDs, driver + types — becomes its string form. Lossy by design: the model reads these, it + does not compute on them. + """ + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, Decimal): + return float(value) + if isinstance(value, (bytes, memoryview)): + raw = bytes(value) + try: + return raw.decode("utf-8") + except UnicodeDecodeError: + return raw.hex() + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + return str(value) + + +def _execute_via_database( + database: Any, + sql: str, + catalog: str | None, + schema: str | None, + limit: int, +) -> Any: + """ + The single point at which a warehouse is touched. + + Isolated as a module-level function so a test can substitute it and exercise + every guard above without a live connection. ``Database.execute`` is the + same entry point SQL Lab and the MCP service use, so this inherits Jinja + rendering, ``SQL_QUERY_MUTATOR``, disallowed-function and disallowed-table + checks, row-level security, and the executor's own ``allow_dml`` gate. + """ + from superset_core.queries.types import QueryOptions + + return database.execute( + sql, + QueryOptions(catalog=catalog, schema=schema, limit=limit), + ) + + +def _result_to_payload(result: Any, limit: int) -> dict[str, Any]: + """ + Flatten a ``QueryResult`` into rows the model can read. + + Only the last data-bearing statement is returned. A read-only script with + several ``SELECT``s is unusual, and returning every result set multiplies + the context cost for a case the model can trivially split into two calls. + """ + from superset_core.queries.types import QueryStatus + + if result.status != QueryStatus.SUCCESS: + raise ToolError( + f"The query did not complete: {result.error_message or result.status}" + ) + + statement = next( + (item for item in reversed(result.statements) if item.data is not None), + None, + ) + if statement is None: + return { + "rows": [], + "row_count": 0, + "columns": [], + "note": "No rows returned.", + "executed_sql": None, + } + + columns, records = _columns_and_records(statement.data) + rows = [ + {str(key): _json_safe(value) for key, value in record.items()} + for record in records[:limit] + ] + + payload: dict[str, Any] = { + "rows": rows, + "row_count": len(rows), + "columns": columns, + # The SQL the warehouse actually ran, after the limit and any row-level + # security rewrite. This is what the user needs to see to trust the + # answer, and what they would paste into SQL Lab to check it. + "executed_sql": getattr(statement, "executed_sql", None), + } + if len(records) > limit: + payload["truncated"] = True + payload["note"] = ( + f"Showing {limit} of {len(records)} rows. Add a tighter WHERE clause " + f"or aggregate to see the rest." + ) + return payload + + +def _sql_display( + database: Any, + payload: dict[str, Any], + duration_ms: int, +) -> dict[str, Any]: + """ + Build the UI summary for one query. + + Deliberately not the model-facing payload: the row sample is smaller, and + only the connection's name and id appear — never its URI or credentials. + """ + executed = payload.get("executed_sql") or "" + return { + "kind": "sql_result", + "database_id": database.id, + "database_name": database.database_name, + "executed_sql": str(executed)[:DISPLAY_SQL_CHARS], + "executed_sql_truncated": len(str(executed)) > DISPLAY_SQL_CHARS, + "columns": [column["name"] for column in payload.get("columns", [])], + "rows": payload.get("rows", [])[:DISPLAY_SAMPLE_ROWS], + "row_count": payload.get("row_count", 0), + "sample_only": len(payload.get("rows", [])) > DISPLAY_SAMPLE_ROWS, + "truncated": bool(payload.get("truncated", False)), + "duration_ms": duration_ms, + } + + +class ExecuteSqlTool(AITool): + """Run a read-only query and return its rows.""" + + name: ClassVar[str] = "execute_sql" + description: ClassVar[str] = ( + "Run a read-only SQL query against a Superset database connection and " + "return the rows. Only SELECT-style statements are accepted; anything " + "that writes data or changes schema is refused. Results are capped, so " + "aggregate or filter in SQL rather than asking for everything. Call " + "list_databases for a database_id and get_schema for table and column " + "names before writing the query." + ) + input_schema: ClassVar[dict[str, Any]] = { + "type": "object", + "properties": { + "database_id": { + "type": "integer", + "description": "Database connection to query, from list_databases.", + }, + "sql": { + "type": "string", + "description": "A single read-only SQL statement.", + }, + "schema": { + "type": "string", + "description": ( + "Schema unqualified table names resolve to. Optional; the " + "database's default is used when omitted." + ), + }, + "catalog": { + "type": "string", + "description": "Catalog to query, for engines that have them.", + }, + "limit": { + "type": "integer", + "description": ( + "Maximum rows to return. Narrows the default only; it " + "cannot raise the configured ceiling." + ), + }, + }, + "required": ["database_id", "sql"], + } + + def __init__( + self, + executor: Callable[[Any, str, str | None, str | None, int], Any] | None = None, + ) -> None: + # Injectable so the guards are testable without a warehouse. + self._executor = executor or _execute_via_database + + def run( # pylint: disable=too-many-arguments + self, + database_id: Any = None, + sql: Any = None, + schema: Any = None, + catalog: Any = None, + limit: Any = None, + **_ignored: Any, + ) -> ToolOutput: + database = _database_or_refuse(database_id) + script = _parse_or_refuse(sql, database) + _assert_read_only(script, database) + + schema_name = str(schema) if schema else None + catalog_name = str(catalog) if catalog else None + _raise_for_sql_access(database, sql, catalog_name, schema_name) + + row_limit = _row_limit(limit) + started = time.monotonic() + result = self._executor(database, sql, catalog_name, schema_name, row_limit) Review Comment: This executes SQL after database/table checks but never requires the SQL Lab capability. Alpha can use the newly granted assistant without the additive `sql_lab` role, so this grants SQL execution outside the documented role boundary; should the tool enforce that permission before executing? -- 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]
