I3eka commented on code in PR #43132:
URL: https://github.com/apache/superset/pull/43132#discussion_r4384599612
##########
superset/config.py:
##########
@@ -2948,6 +2964,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
+# ``ai.prune_conversations`` Celery task, which must be scheduled to run.
Review Comment:
Fixed here in `b9a2abec6a` by carrying over the implementation from #43135.
`ai.prune_conversations` removes expired threads and their messages/feedback;
the default beat schedule includes it. Tests verify dependent-row deletion
while retaining a newer conversation, plus task registration. Custom schedules
still need the documented entry; nonpositive retention disables pruning.
##########
superset/ai/tasks.py:
##########
@@ -0,0 +1,89 @@
+# 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.
+"""
+Background execution of assistant turns.
+
+Used when ``AI_ASSISTANT_EXECUTION_MODE`` is ``"worker"``. The task body is a
+thin wrapper: all the work lives in
+:func:`superset.ai.orchestrator.execute_turn`, so the two execution modes
cannot
+diverge in behaviour.
+
+To enable, add ``"superset.ai.tasks"`` to ``CeleryConfig.imports`` and set the
+execution mode and a Redis event bus.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from superset.ai.orchestrator import execute_turn, TurnRequest
+from superset.extensions import celery_app
+
+logger = logging.getLogger(__name__)
+
+
+@celery_app.task(name="ai.run_turn", bind=True, soft_time_limit=None)
Review Comment:
Fixed the observed Celery child-failure path in `c510dfb724`. Parent-process
Request callbacks handle worker loss and hard time limits, commit the terminal
error and publish final/error/done without rerunning the turn. Early
acknowledgement is explicit. The atomic claim records the winning delivery, so
a failed duplicate cannot terminate another worker's run; completed/cancelled
rows are preserved.
The regressions invoke Celery's real parent callbacks outside a Flask
context and exercise concurrent database claims. A loss of the entire Celery
parent/host cannot invoke those callbacks; that separate limitation is now
explicit in the guide.
##########
superset/ai/eventbus.py:
##########
@@ -0,0 +1,314 @@
+# 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.
+"""
+Carries streamed events from whatever produced them to the HTTP response.
+
+Two implementations, matching the two execution modes. Inline execution needs
+nothing more than an in-process queue. Worker execution needs a shared,
+*replayable* channel — replayable because a browser that loses its connection
+must be able to rejoin a run already in progress, which rules out
+publish/subscribe: a subscriber that was absent when an event was published
+never sees it.
+
+The Redis implementation therefore uses streams, and reuses the cache backend
+that Superset's async-query channel already configures rather than introducing
+a second Redis client to operate.
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from typing import Any
+
+from superset.ai.events import StreamEvent
+from superset.ai.types import StreamEventType
+from superset.utils import json
+
+logger = logging.getLogger(__name__)
+
+#: Yielded by :meth:`BaseEventBus.consume` when nothing arrived within the poll
+#: interval, so a caller can emit a keep-alive rather than block indefinitely.
+IDLE = None
+
+#: Terminal event types. Seeing one ends consumption, so a reader does not hang
+#: waiting for a producer that has already finished.
+_TERMINAL = frozenset(
+ {StreamEventType.DONE, StreamEventType.ERROR, StreamEventType.CANCELLED}
+)
+
+
+class BaseEventBus(ABC):
+ """A per-run channel of events."""
+
+ @abstractmethod
+ def publish(self, run_id: str, event: StreamEvent) -> None:
+ """Append an event to a run's channel."""
+
+ @abstractmethod
+ def consume(
+ self,
+ run_id: str,
+ timeout_seconds: float,
+ poll_seconds: float = 1.0,
+ ) -> Iterator[StreamEvent | None]:
+ """
+ Yield a run's events until a terminal one arrives or time runs out.
+
+ Yields :data:`IDLE` when a poll interval passes with nothing new, which
+ is the caller's cue to send a keep-alive frame.
+ """
+
+ @abstractmethod
+ def close(self, run_id: str) -> None:
+ """Release any resources held for a run."""
+
+
+class MemoryEventBus(BaseEventBus):
+ """
+ An in-process queue per run.
+
+ Correct only when the producer and the streaming request share a process.
+ Selecting this alongside worker execution would leave every stream silent,
+ which :func:`get_event_bus` refuses to allow.
+ """
+
+ def __init__(self) -> None:
+ self._queues: dict[str, queue.SimpleQueue[StreamEvent]] = {}
+
+ def _queue_for(self, run_id: str) -> queue.SimpleQueue[StreamEvent]:
+ return self._queues.setdefault(run_id, queue.SimpleQueue())
+
+ def publish(self, run_id: str, event: StreamEvent) -> None:
+ self._queue_for(run_id).put(event)
+
+ def consume(
+ self,
+ run_id: str,
+ timeout_seconds: float,
+ poll_seconds: float = 1.0,
+ ) -> Iterator[StreamEvent | None]:
+ import time
+
+ # Deliberately not ``_queue_for``: reading must not create a channel.
+ # This bus lives for the life of the process, so a client polling
+ # unknown run identifiers would otherwise grow the dict without bound.
+ channel = self._queues.get(run_id)
+ deadline = time.monotonic() + timeout_seconds
+
+ while True:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return
+ if channel is None:
+ # The producer may not have published yet; look again rather
+ # than deciding the run does not exist. Only report idle if it
+ # is still absent, so a channel that appeared during the wait
+ # is drained on this pass instead of costing an extra tick.
+ channel = self._queues.get(run_id)
+ if channel is None:
+ yield IDLE
+ time.sleep(min(poll_seconds, remaining))
+ continue
+ try:
+ # Bounded by whichever is sooner, so a generous poll interval
+ # cannot overshoot the caller's deadline.
+ event = channel.get(timeout=min(poll_seconds, remaining))
+ except queue.Empty:
+ yield IDLE
+ continue
+ yield event
+ if event.type in _TERMINAL:
+ return
+
+ def close(self, run_id: str) -> None:
+ self._queues.pop(run_id, None)
+
+
+class RedisStreamEventBus(BaseEventBus):
+ """
+ A Redis stream per run.
+
+ Replayable by construction: a reconnecting reader starts from the beginning
+ of the stream and catches up, which is what makes worker execution usable
+ from a browser on a flaky connection.
+ """
+
+ def __init__(
+ self,
+ cache: Any,
+ prefix: str = "ai-events-",
+ ttl_seconds: int = 900,
+ ) -> None:
+ self._cache = cache
+ self._prefix = prefix
+ self._ttl = ttl_seconds
+
+ def _stream(self, run_id: str) -> str:
+ return f"{self._prefix}{run_id}"
+
+ def publish(self, run_id: str, event: StreamEvent) -> None:
+ payload = {
+ "data": json.dumps({"type": event.type.value, "payload":
event.payload})
+ }
+ # A failure to publish must not kill the run that is producing useful
+ # work; the reader will time out and the answer is still persisted.
+ try:
+ self._cache.xadd(self._stream(run_id), payload, "*", 10_000)
Review Comment:
Confirmed and fixed here in `c8b4233c59`, carrying over #43134's event-bus
fix. Every successful publish establishes/refreshes expiry, including the first
progress event with no reader. Those regressions failed on this branch before
the change and pass afterward. The same update preserves the final `done` after
error/cancelled frames.
##########
docs/admin_docs/configuration/ai-assistant.mdx:
##########
@@ -0,0 +1,491 @@
+---
+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`.
+Worker cancellation flags use this same connection, so no second cache setting
+is required.
+
+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
Review Comment:
Fixed in `90c16210c2`. The guide adds the provider to
`docker/requirements-local.txt` before `docker compose up`, matching the
bootstrap's dependency-install step. The main enablement instructions also
include the extension UI flag and `superset db upgrade` before permission
initialization.
##########
superset-frontend/src/features/ai/hooks/useChatBot.ts:
##########
@@ -0,0 +1,1323 @@
+/**
+ * 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 Conversation state and the send loop.
+ *
+ * Runs are tracked per conversation, not globally. That is the point of the
+ * structure: a user can start something slow in one conversation, switch to
+ * another and keep working, and come back to find the first still going. A
single
+ * `isLoading` flag would have made switching away cancel or corrupt the run.
+ *
+ * The server owns the transcript. A finished run is re-read from it rather
than
+ * assembled from the frames, so the tool calls persisted on the message are
what
+ * the user sees, and what they see survives a reload.
+ */
+
+import { useCallback, useEffect, useRef, useState } from 'react';
+import type { TextAreaRef } from 'antd/es/input/TextArea';
+import { logging } from '@apache-superset/core/utils';
+import { t } from '@apache-superset/core/translation';
+import {
+ type AiAgent,
+ type AiToolCall,
+ type ChatMessageWithMeta,
+ type ChatTab,
+ type CheckpointPayload,
+} from '../types';
+import {
+ AGENT_STORAGE_KEY,
+ ChatRequestAbortedError,
+ ChatStreamEventError,
+ ChatStreamTimeoutError,
+ DEFAULT_AGENT_KEY,
+ DEFAULT_CHAT_AGENT,
+ cancelChatRun,
+ describeRequestError,
+ fetchAgents,
+ fetchSuggestedPrompts,
+ loadStoredAgentKey,
+ normalizeChatAgents,
+ startRun,
+ streamRun,
+ submitFeedback,
+} from './chatRequest';
+import {
+ NEW_CHAT_NAME,
+ createThread,
+ deleteThread as deleteThreadApi,
+ getThread,
+ listThreads,
+ threadToTab,
+ updateThread,
+} from './chatThreadsApi';
+import { buildQuickPrompts } from './quickPrompts';
+import {
+ buildPageContextPayload,
+ usePageContext,
+ type PageContext,
+} from './usePageContext';
+
+/** Cache of the conversation list, so the menu renders before the list
arrives. */
+export const CHAT_TABS_STORAGE_KEY = 'superset-chat-tabs';
+
+/** Which conversation was last open. */
+export const ACTIVE_TAB_STORAGE_KEY = 'superset-chat-active-tab';
+
+/** Recent inputs, recalled with the arrow keys. */
+export const HISTORY_STORAGE_KEY = 'superset-chat-history';
+
+export { AGENT_STORAGE_KEY } from './chatRequest';
+
+/** How many inputs the arrow-key history keeps. */
+const MAX_INPUT_HISTORY = 50;
+
+/** A conversation title derived from a message is clipped to this. */
+const MAX_TAB_NAME_LENGTH = 30;
+
+export type ChatRunStatus = 'running' | 'cancelling';
+
+/** Shared empty list, so a render with no steps yet keeps a stable identity.
*/
+const EMPTY_TOOL_CALLS: AiToolCall[] = [];
+
+interface ActiveChatRun {
+ requestId: string;
+ tabId: string;
+ threadId: string;
+ runId?: string;
+ controller: AbortController;
+ isStreaming: boolean;
+ liveThoughts: string;
+ liveToolLog: string;
+ /**
+ * Steps taken so far, as structured records rather than log lines.
+ *
+ * Carried alongside `liveToolLog` so a run in flight can be rendered the
same
+ * way a finished one is — expandable per step, with the SQL and the rows it
+ * returned — instead of as a wall of text that only becomes legible once the
+ * transcript is re-read from the server.
+ */
+ liveToolCalls: AiToolCall[];
+ /** The page context this run was given, so the live view can show it too. */
+ livePageContext?: string;
+ /**
+ * The answer so far, as the model produces it.
+ *
+ * Rendered directly: the deltas used to be folded into `liveThinking`, which
+ * nothing displayed, so an answer appeared in one piece the moment the run
+ * ended however long it had taken to generate.
+ */
+ liveAnswer: string;
+ liveThinking: string;
+ status: ChatRunStatus;
+ startedAt: number;
+ checkpoint: CheckpointPayload | null;
+}
+
+/**
+ * An identifier for a turn.
+ *
+ * Drawn from `crypto`, not `Math.random`. These become the idempotency key on
a
+ * turn and the handle used to cancel one, so a value another session could
guess
+ * is a correctness and a security problem rather than merely a collision risk.
+ */
+const generateId = (): string => {
+ if (typeof crypto.randomUUID === 'function') {
+ return crypto.randomUUID();
+ }
+ // Older engines expose the entropy source without the convenience wrapper.
+ const bytes = new Uint8Array(16);
+ crypto.getRandomValues(bytes);
+ return Array.from(bytes, byte => byte.toString(16).padStart(2,
'0')).join('');
+};
+
+const createNewTab = (name: string = NEW_CHAT_NAME): ChatTab => ({
+ id: generateId(),
+ name,
+ messages: [],
+ createdAt: Date.now(),
+});
+
+const truncateTabName = (
+ name: string,
+ maxLength: number = MAX_TAB_NAME_LENGTH,
+): string =>
+ name.length <= maxLength ? name : `${name.substring(0, maxLength)}...`;
+
+const readJson = <T>(key: string, fallback: T): T => {
+ try {
+ const stored = localStorage.getItem(key);
+ return stored ? (JSON.parse(stored) as T) : fallback;
+ } catch (caught) {
+ logging.warn(`[ai] could not read ${key}`, caught);
+ return fallback;
+ }
+};
+
+const writeJson = (key: string, value: unknown): void => {
+ try {
+ localStorage.setItem(key, JSON.stringify(value));
+ } catch (caught) {
+ logging.warn(`[ai] could not write ${key}`, caught);
+ }
+};
+
+/**
+ * Reconciles the server's transcript with what is already on screen.
+ *
+ * The server's copy is authoritative — it carries the tool calls — but it is
not
+ * necessarily complete the moment a run ends, and replacing outright would
then
+ * erase an answer the user has just read. So anything local that the server
has
+ * not accounted for is kept, matched by identity first and by role and content
+ * second, which is how a locally-appended turn is recognised once the server
+ * returns its own copy of it under a real uuid.
+ */
+export const mergeMessages = (
+ fromServer: ChatMessageWithMeta[],
+ local: ChatMessageWithMeta[],
+): ChatMessageWithMeta[] => {
+ const serverIds = new Set(fromServer.map(message => message.id));
+ const serverTurns = new Set(
+ fromServer.map(message => `${message.role}:${message.content}`),
+ );
+ const unaccounted = local.filter(
+ message =>
+ !serverIds.has(message.id) &&
+ !serverTurns.has(`${message.role}:${message.content}`),
+ );
+ return [...fromServer, ...unaccounted];
+};
+
+/**
+ * The `page_context` body for one turn.
+ *
+ * Returns undefined when there is nothing to send, so an omitted field is
+ * distinguishable from an empty one.
+ */
+export const buildRequestPageContext = (
+ context: PageContext | undefined,
+ directive?: string,
+): Record<string, unknown> | undefined => {
+ const payload = context ? buildPageContextPayload(context) : undefined;
+ if (!directive) {
+ return payload;
+ }
+ const existing = payload?.helper_directives;
+ return {
+ ...payload,
+ helper_directives: [
+ directive,
+ ...(Array.isArray(existing) ? existing : []),
+ ],
+ };
+};
+
+export interface UseChatBotReturn {
+ // Conversations
+ chatTabs: ChatTab[];
+ activeTabId: string;
+ activeTab: ChatTab | undefined;
+ threadsLoaded: boolean;
+ handleNewChat: () => Promise<string>;
+ handleSelectTab: (tabId: string) => Promise<void>;
+ handleDeleteTab: (tabId: string) => Promise<void>;
+ handleRenameTab: (tabId: string, newName: string) => void;
+ // Messages of the active conversation
+ messages: ChatMessageWithMeta[];
+ // Input
+ inputValue: string;
+ setInputValue: (value: string) => void;
+ handleKeyDown: (event: React.KeyboardEvent) => void;
+ inputRef: React.RefObject<TextAreaRef>;
+ messagesEndRef: React.RefObject<HTMLDivElement>;
+ // The run in flight, if any, for the active conversation
+ isLoading: boolean;
+ isStreamingResponse: boolean;
+ liveThoughts: string;
+ liveToolLog: string;
+ /** Steps taken so far in the run in flight, for the structured live view. */
+ liveToolCalls: AiToolCall[];
+ /** The page context the run in flight was given. */
+ livePageContext?: string;
+ /** The answer so far for the run in flight. */
+ liveAnswer: string;
+ checkpoint: CheckpointPayload | null;
+ activeRunStatus: ChatRunStatus | null;
+ error?: string;
+ // Actions
+ sendMessage: (
+ messageOverride?: string,
+ systemPromptOverride?: string,
+ ) => Promise<void>;
+ handleCancelRun: () => Promise<void>;
+ handleCheckpointContinue: () => void;
+ handleFeedback: (messageId: string, feedback: 'like' | 'dislike') => void;
+ messageFeedback: Record<string, 'like' | 'dislike'>;
+ // Suggestions
+ /** The message whose run just ended; its thought process stays open. */
+ justCompletedId?: string;
+ quickPrompts: string[];
+ loadQuickPrompts: () => void;
+ applyQuickPrompt: (prompt: string) => Promise<void>;
+ // Agent profiles
+ agents: AiAgent[];
+ selectedAgent: string;
+ setSelectedAgent: (key: string) => void;
+ // Page context
+ pageContext: PageContext;
+ includePageContext: boolean;
+ toggleIncludePageContext: () => void;
+}
+
+export const useChatBot = (): UseChatBotReturn => {
+ const [chatTabs, setChatTabs] = useState<ChatTab[]>(() =>
+ readJson<ChatTab[]>(CHAT_TABS_STORAGE_KEY, []).map(tab => ({
+ // The cache is a placeholder for the menu; message bodies are re-read
from
+ // the server so a stale cache cannot show a conversation that has moved
on.
+ ...tab,
+ messages: [],
+ })),
+ );
+ const [activeTabId, setActiveTabId] = useState<string>(() => {
+ try {
+ return localStorage.getItem(ACTIVE_TAB_STORAGE_KEY) ?? '';
+ } catch {
+ return '';
+ }
+ });
+ const [threadsLoaded, setThreadsLoaded] = useState(false);
+ const [error, setError] = useState<string | undefined>(undefined);
+
+ const [inputValue, setInputValue] = useState('');
+ const [activeRunsByTab, setActiveRunsByTab] = useState<
+ Record<string, ActiveChatRun>
+ >({});
+ const [quickPrompts, setQuickPrompts] = useState<string[]>([]);
+ const [messageFeedback, setMessageFeedback] = useState<
+ Record<string, 'like' | 'dislike'>
+ >({});
+ const [includePageContext, setIncludePageContext] = useState(true);
+ /**
+ * The assistant message whose run has only just ended.
+ *
+ * Its thought process stays open, because collapsing it the instant the
answer
+ * lands moves everything below it — the answer the user is mid-sentence
through
+ * jumps up the panel. Older messages start closed.
+ */
+ const [justCompletedId, setJustCompletedId] = useState<string | undefined>();
+ const [agents, setAgents] = useState<AiAgent[]>([DEFAULT_CHAT_AGENT]);
+ const [selectedAgent, setSelectedAgent] = useState<string>(() =>
+ loadStoredAgentKey(AGENT_STORAGE_KEY),
+ );
+
+ const [messageHistory, setMessageHistory] = useState<string[]>(() =>
+ readJson<string[]>(HISTORY_STORAGE_KEY, []),
+ );
+ const [historyIndex, setHistoryIndex] = useState(-1);
+ const [currentDraft, setCurrentDraft] = useState('');
+
+ const messagesEndRef = useRef<HTMLDivElement>(null);
+ const inputRef = useRef<TextAreaRef>(null);
+
+ /**
+ * The run map and the conversation list are also held in refs, and the refs
are
+ * the authority.
+ *
+ * The send loop has to ask "is this still my run?" between awaits, and it
cannot
+ * ask React: a run that starts and fails inside one batch never causes a
render,
+ * so a ref synced at render time would still be empty and the loop would
discard
+ * its own result as stale. Writing the ref at the point of mutation removes
that
+ * window. The callbacks read the refs rather than the state so their
identities
+ * do not churn on every streamed frame, which would restart effects mid-run.
+ */
+ const activeRunsByTabRef = useRef<Record<string, ActiveChatRun>>({});
+ const chatTabsRef = useRef<ChatTab[]>(chatTabs);
+ const activeTabIdRef = useRef(activeTabId);
+ activeTabIdRef.current = activeTabId;
+
+ const updateRuns = useCallback(
+ (
+ updater: (
+ previous: Record<string, ActiveChatRun>,
+ ) => Record<string, ActiveChatRun>,
+ ) => {
+ activeRunsByTabRef.current = updater(activeRunsByTabRef.current);
+ setActiveRunsByTab(activeRunsByTabRef.current);
+ },
+ [],
+ );
+
+ const updateTabs = useCallback(
+ (updater: (previous: ChatTab[]) => ChatTab[]) => {
+ chatTabsRef.current = updater(chatTabsRef.current);
+ setChatTabs(chatTabsRef.current);
+ },
+ [],
+ );
+
+ /** Resolved when the user answers a checkpoint; see `streamRun`. */
+ const checkpointGateRef = useRef<{ resolve: () => void } | null>(null);
+ const mountedRef = useRef(true);
+ useEffect(
+ () => () => {
+ mountedRef.current = false;
+ },
+ [],
+ );
+
+ const activeTab = chatTabs.find(tab => tab.id === activeTabId);
+ const messages = activeTab?.messages ?? [];
+
+ const activeRun = activeRunsByTab[activeTabId];
+ const isLoading = Boolean(activeRun);
+ const isStreamingResponse = activeRun?.isStreaming ?? false;
+ const liveThoughts = activeRun?.liveThoughts ?? '';
+ const liveToolLog = activeRun?.liveToolLog ?? '';
+ const liveToolCalls = activeRun?.liveToolCalls ?? EMPTY_TOOL_CALLS;
+ const livePageContext = activeRun?.livePageContext;
+ const liveAnswer = activeRun?.liveAnswer ?? '';
+ const checkpoint = activeRun?.checkpoint ?? null;
+ const activeRunStatus = activeRun?.status ?? null;
+
+ const pageContext = usePageContext();
+ const pageContextRef = useRef(pageContext);
+ pageContextRef.current = pageContext;
+
+ const fail = useCallback(async (caught: unknown, fallback: string) => {
+ const message = await describeRequestError(caught, fallback);
+ logging.error('[ai] assistant request failed', caught);
+ if (mountedRef.current) {
+ setError(message);
+ }
+ }, []);
+
+ // -----------------------------------------------------------------------
+ // Persistence of the small things
+ // -----------------------------------------------------------------------
+
+ useEffect(() => {
+ // Only the shell of each conversation is cached; see the initialiser.
+ writeJson(
+ CHAT_TABS_STORAGE_KEY,
+ chatTabs.map(tab => ({ ...tab, messages: [] })),
+ );
+ }, [chatTabs]);
+
+ useEffect(() => {
+ try {
+ localStorage.setItem(ACTIVE_TAB_STORAGE_KEY, activeTabId);
+ } catch (caught) {
+ logging.warn('[ai] could not remember the active conversation', caught);
+ }
+ }, [activeTabId]);
+
+ useEffect(() => {
+ if (messageHistory.length > 0) {
+ writeJson(HISTORY_STORAGE_KEY, messageHistory);
+ }
+ }, [messageHistory]);
+
+ useEffect(() => {
+ try {
+ localStorage.setItem(AGENT_STORAGE_KEY, selectedAgent);
+ } catch (caught) {
+ logging.warn('[ai] could not remember the selected agent', caught);
+ }
+ }, [selectedAgent]);
+
+ // Follows the transcript as it grows, including while a run streams. Guarded
+ // because `scrollIntoView` is absent in environments without a layout
engine,
+ // and failing to scroll must not take the panel down.
+ useEffect(() => {
+ messagesEndRef.current?.scrollIntoView?.({ behavior: 'smooth' });
+ }, [messages, liveToolLog, liveThoughts]);
+
+ // -----------------------------------------------------------------------
+ // Conversation management
+ // -----------------------------------------------------------------------
+
+ const setMessagesForTab = useCallback(
+ (
+ tabId: string,
+ updater: (previous: ChatMessageWithMeta[]) => ChatMessageWithMeta[],
+ ) => {
+ updateTabs(previous =>
+ previous.map(tab =>
+ tab.id === tabId ? { ...tab, messages: updater(tab.messages) } : tab,
+ ),
+ );
+ },
+ [updateTabs],
+ );
+
+ const refreshThreadMessages = useCallback(
+ async (threadId: string) => {
+ const { thread, messages: threadMessages } = await getThread(threadId);
+ if (!mountedRef.current) {
+ return;
+ }
+ const refreshed = threadToTab(thread, threadMessages);
+ // A locally set title wins: the user may have renamed the conversation
+ // while the request was in flight.
+ updateTabs(previous =>
+ previous.map(tab =>
+ tab.threadId === threadId
+ ? {
+ ...refreshed,
+ name: tab.name || refreshed.name,
+ messages: mergeMessages(refreshed.messages, tab.messages),
+ }
+ : tab,
+ ),
+ );
+ },
+ [updateTabs],
+ );
+
+ const handleNewChat = useCallback(async (): Promise<string> => {
+ try {
+ const thread = await createThread(
+ undefined,
+ selectedAgent === DEFAULT_AGENT_KEY ? undefined : selectedAgent,
+ );
+ const tab = threadToTab(thread);
+ updateTabs(previous => [tab, ...previous]);
+ setActiveTabId(tab.id);
+ activeTabIdRef.current = tab.id;
+ setError(undefined);
+ return tab.id;
+ } catch (caught) {
+ await fail(caught, t('The conversation could not be created.'));
+ // A local tab still lets the user type; the thread is created on send.
+ const tab = createNewTab();
+ updateTabs(previous => [tab, ...previous]);
+ setActiveTabId(tab.id);
+ activeTabIdRef.current = tab.id;
+ return tab.id;
+ }
+ }, [fail, selectedAgent, updateTabs]);
+
+ const handleSelectTab = useCallback(
+ async (tabId: string) => {
+ setActiveTabId(tabId);
+ const tab = chatTabsRef.current.find(candidate => candidate.id ===
tabId);
+ // Messages are fetched on first view, not up front: a user with fifty
+ // conversations should not pay for forty-nine of them.
+ if (tab?.threadId && tab.messages.length === 0) {
+ try {
+ await refreshThreadMessages(tab.threadId);
+ } catch (caught) {
+ await fail(caught, t('The conversation could not be loaded.'));
+ }
+ }
+ },
+ [fail, refreshThreadMessages],
+ );
+
+ const handleDeleteTab = useCallback(
+ async (tabId: string) => {
+ const tab = chatTabsRef.current.find(candidate => candidate.id ===
tabId);
+ if (tab?.threadId) {
+ try {
+ await deleteThreadApi(tab.threadId);
+ } catch (caught) {
+ await fail(caught, t('The conversation could not be deleted.'));
+ return;
+ }
+ }
+ updateTabs(previous => {
+ const remaining = previous.filter(candidate => candidate.id !== tabId);
+ if (tabId === activeTabIdRef.current) {
+ setActiveTabId(remaining[0]?.id ?? '');
+ activeTabIdRef.current = remaining[0]?.id ?? '';
+ }
+ return remaining;
+ });
+ },
+ [fail, updateTabs],
+ );
+
+ const handleRenameTab = useCallback(
+ (tabId: string, newName: string) => {
+ const trimmedName = newName.trim();
+ if (!trimmedName) {
+ return;
+ }
+ const tab = chatTabsRef.current.find(candidate => candidate.id ===
tabId);
+ updateTabs(previous =>
+ previous.map(candidate =>
+ candidate.id === tabId
+ ? { ...candidate, name: trimmedName }
+ : candidate,
+ ),
+ );
+ // Renamed locally first, then on the server: the menu should not wait
for a
+ // round trip to show what the user just typed.
+ if (tab?.threadId) {
+ updateThread(tab.threadId, { title: trimmedName }).catch(caught => {
+ logging.warn('[ai] could not rename the conversation', caught);
+ });
+ }
+ },
+ [updateTabs],
+ );
+
+ /** Titles an untitled conversation after its first message. */
+ const nameTabFromMessage = useCallback(
+ (tabId: string, threadId: string, message: string) => {
+ const title = truncateTabName(message);
+ updateTabs(previous =>
+ previous.map(tab =>
+ tab.id === tabId && tab.name === NEW_CHAT_NAME
+ ? { ...tab, name: title }
+ : tab,
+ ),
+ );
+ updateThread(threadId, { title }).catch(caught => {
+ logging.warn('[ai] could not title the conversation', caught);
+ });
+ },
+ [updateTabs],
+ );
+
+ /**
+ * The thread backing a tab, creating it if the tab is only local.
+ *
+ * A tab's id becomes its thread uuid once it has one, so the id is rewritten
+ * here and the new one returned — callers must use it from then on.
+ */
+ const ensureThread = useCallback(
+ async (tabId: string): Promise<{ tabId: string; threadId: string }> => {
+ const tab = chatTabsRef.current.find(candidate => candidate.id ===
tabId);
+ if (tab?.threadId) {
+ return { tabId, threadId: tab.threadId };
+ }
+ const thread = await createThread(
+ undefined,
+ selectedAgent === DEFAULT_AGENT_KEY ? undefined : selectedAgent,
+ );
+ updateTabs(previous =>
+ previous.map(candidate =>
+ candidate.id === tabId
+ ? { ...candidate, id: thread.uuid, threadId: thread.uuid }
+ : candidate,
+ ),
+ );
+ if (activeTabIdRef.current === tabId) {
+ setActiveTabId(thread.uuid);
+ activeTabIdRef.current = thread.uuid;
+ }
+ return { tabId: thread.uuid, threadId: thread.uuid };
+ },
+ [selectedAgent, updateTabs],
+ );
+
+ // -----------------------------------------------------------------------
+ // Run bookkeeping
+ // -----------------------------------------------------------------------
+
+ const isRunCurrent = useCallback(
+ (tabId: string, requestId: string): boolean => {
+ const run = activeRunsByTabRef.current[tabId];
+ return Boolean(run && run.requestId === requestId);
+ },
+ [],
+ );
+
+ const updateRunState = useCallback(
+ (
+ tabId: string,
+ requestId: string,
+ updater: (run: ActiveChatRun) => ActiveChatRun,
+ ) => {
+ updateRuns(previous => {
+ const current = previous[tabId];
+ if (!current || current.requestId !== requestId) {
+ return previous;
+ }
+ return { ...previous, [tabId]: updater(current) };
+ });
+ },
+ [updateRuns],
+ );
+
+ const clearRunIfMatches = useCallback(
+ (tabId: string, requestId: string) => {
+ updateRuns(previous => {
+ const current = previous[tabId];
+ if (!current || current.requestId !== requestId) {
+ return previous;
+ }
+ const { [tabId]: _removed, ...rest } = previous;
+ return rest;
+ });
+ },
+ [updateRuns],
+ );
+
+ const releaseCheckpointGate = useCallback(() => {
+ checkpointGateRef.current?.resolve();
+ checkpointGateRef.current = null;
+ }, []);
+
+ const handleCheckpointContinue = useCallback(() => {
+ releaseCheckpointGate();
+ const tabId = activeTabIdRef.current;
+ const run = activeRunsByTabRef.current[tabId];
+ if (run) {
+ updateRunState(tabId, run.requestId, current => ({
+ ...current,
+ checkpoint: null,
+ }));
+ }
+ }, [releaseCheckpointGate, updateRunState]);
+
+ const handleCancelRun = useCallback(async () => {
+ const tabId = activeTabIdRef.current;
+ const run = activeRunsByTabRef.current[tabId];
+ if (!run) {
+ return;
+ }
+ // A run paused at a checkpoint is not reading, so the gate is released
first
+ // or the abort would not be noticed until the user pressed Continue.
+ releaseCheckpointGate();
+ updateRunState(tabId, run.requestId, current => ({
+ ...current,
+ status: 'cancelling',
+ }));
+ run.controller.abort();
+ if (run.runId) {
Review Comment:
Confirmed and fixed in `1a5d4bc056`. Stop keeps the run in cancelling state
while the start POST is pending. When its handle arrives, the client sends
cancellation for that exact run instead of attaching a stream; another send
cannot overlap that interval. The deferred-POST regression failed before the
fix and passes now.
##########
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(
Review Comment:
Fixed here in `5a5010f26e`, using the tested #44561 implementation. Run ID
and accepted context are saved with the new placeholder, not overwritten on
replay. Queue resubmission is limited to an unclaimed run, and an atomic claim
prevents duplicate inference. The concurrent-session test also verifies the
winning delivery's ownership after the worker-loss fix.
--
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]
