I3eka commented on code in PR #43134:
URL: https://github.com/apache/superset/pull/43134#discussion_r4003213008
##########
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.
+AI_ASSISTANT_MESSAGE_RETENTION_DAYS = 30
Review Comment:
Retention is still not implemented in this telemetry branch; the presence of
`AI_ASSISTANT_MESSAGE_RETENTION_DAYS` in configuration/docs is not enforcement.
This needs the base pruning task/scheduling contract or removal of the promise.
Keeping this open.
##########
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:
The Helm task-registration path is still a base deployment gap. This refresh
does not add `superset.ai.tasks` to Helm's overriding Celery imports; it needs
a worker registration test as well as the default source configuration change.
Leaving this open.
##########
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)
+def run_turn(self: Any, payload: dict[str, Any]) -> str: # noqa: ARG001
+ """
+ Answer one assistant turn.
+
+ ``acks_late`` is deliberately not set: a turn costs money to run, so
Review Comment:
The inherited acknowledgement policy and hard-worker-loss recovery remain
unresolved. Explicit early acknowledgement alone would prevent redelivery but
can still strand a message; this update does not claim crash-safe execution. It
needs the common base run-lifecycle design and regression.
##########
superset/ai/api.py:
##########
@@ -0,0 +1,996 @@
+# 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_command = AppendAIChatMessageCommand(
+ thread_uuid,
+ user_id,
+ MessageRole.ASSISTANT,
+ "",
+ request_id=payload.get("request_id"),
+ status=MessageStatus.PENDING,
+ )
+ assistant_message = assistant_message_command.run()
+ except AIChatThreadNotFoundError:
+ return self.response_404()
+ except (AIChatMessageInvalidError, AIChatThreadInvalidError) as ex:
+ return self.response_422(message=str(ex))
+
+ if assistant_message_command.created:
Review Comment:
Correct — the `created` guard added in this follow-up can strand the
existing pending row when submission fails. This is not resolved by the
telemetry/TTL fixes, and I am keeping it open as a blocker on our side. Simply
re-enqueueing every pending row is unsafe after an ambiguous broker response;
this needs a durable claim/submission state shared with the base lifecycle work.
##########
superset/ai/orchestrator.py:
##########
@@ -0,0 +1,647 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Runs one assistant turn end to end.
+
+Sits between the HTTP layer and the runtime: loads the conversation, assembles
+the prompt, resolves the tools the chosen profile allows, drives the runtime,
+publishes every event to the bus, and records the outcome on the assistant
+message.
+
+Deliberately independent of *where* it runs. The same function body serves the
+inline path and the Celery path, which is what makes the execution mode a
+configuration choice rather than two implementations that drift apart.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid as uuid_module
+from collections.abc import AsyncIterator, Iterator
+from dataclasses import dataclass
+from typing import Any
+
+from superset.ai.events import (
+ cancelled_event,
+ done_event,
+ error_event,
+ GENERIC_ERROR_MESSAGE,
+ session_event,
+ StreamEvent,
+)
+from superset.ai.llm.base import Message
+from superset.ai.telemetry import bind_run, current_run, start_run
+from superset.ai.types import MessageRole, MessageStatus, RunOutcome,
StreamEventType
+from superset.utils.decorators import transaction
+
+logger = logging.getLogger(__name__)
+
+#: Cache key prefix for a run's cancellation flag. A flag rather than a signal
+#: because a worker cannot be interrupted mid-call reliably; the runtime checks
+#: this between steps.
+_CANCEL_PREFIX = "ai-cancel-"
+
+#: How long a cancellation request stays meaningful.
+_CANCEL_TTL_SECONDS = 900
+
+#: Stored when a run is stopped before it produced any answer, so the
+#: transcript still records that the turn happened.
+_STOPPED_WITHOUT_ANSWER = "_Stopped before an answer was produced._"
+
+#: Stored when a run exhausted its time budget without saying anything. Phrased
+#: as something the user can act on, because retrying is usually the right
move.
+_TIMED_OUT_WITHOUT_ANSWER = (
+ "The assistant ran out of time before it could answer. Please try again."
+)
+
+#: Ceiling on the page context recorded on a message. Well below the prompt's
own
+#: limit: this is stored per turn and read back with the whole transcript.
+_RECORDED_CONTEXT_LIMIT = 4_000
+
+
+@dataclass
+class TurnRequest:
+ """One unit of work: answer the latest message on a thread."""
+
+ thread_uuid: str
+ user_id: int
+ run_id: str
+ #: Assistant message row to fill in. Created before the run starts so a
+ #: client that reconnects has something to attach to.
+ assistant_message_uuid: str
+ profile_key: str | None = None
+ #: Concrete model to pin, overriding the profile's tier.
+ model: str | None = None
+ #: What the user had on screen when they asked. Supplied by the client,
+ #: which is the only party that knows which tab is open, what is typed in
+ #: the editor and which filters are applied.
+ page_context: dict[str, Any] | None = None
+
+ def to_payload(self) -> dict[str, Any]:
+ """Serialise for the task broker."""
+ return {
+ "thread_uuid": self.thread_uuid,
+ "user_id": self.user_id,
+ "run_id": self.run_id,
+ "assistant_message_uuid": self.assistant_message_uuid,
+ "profile_key": self.profile_key,
+ "model": self.model,
+ "page_context": self.page_context,
+ }
+
+ @classmethod
+ def from_payload(cls, payload: dict[str, Any]) -> TurnRequest:
+ """Rebuild from a broker payload."""
+ return cls(**payload)
+
+
+def new_run_id() -> str:
+ """Identifier for one run, used as the event-stream key."""
+ return str(uuid_module.uuid4())
+
+
+#: Runs cancelled in this process.
+#:
+#: Held alongside the cache rather than instead of it. Superset's default cache
+#: is a null cache, which accepts a write and discards it — so a cache-only
+#: implementation would leave cancellation silently broken on a default
install,
+#: with the button appearing to work and nothing stopping. This set makes
inline
+#: execution correct with no cache at all; the cache is what carries a
+#: cancellation across processes for worker execution.
+_CANCELLED_LOCALLY: set[str] = set()
+
+
+def request_cancel(run_id: str) -> None:
+ """
+ Ask a run to stop.
+
+ Cooperative by design: the flag is recorded here and observed by the
runtime
+ between steps. A run blocked inside a single long model call or query will
+ not notice until that call returns, which is a real limit worth documenting
+ rather than hiding.
+ """
+ from superset.extensions import cache_manager
+
+ _CANCELLED_LOCALLY.add(run_id)
+ try:
+ cache_manager.cache.set(
Review Comment:
This still requires a guaranteed shared cancellation backend; the dedicated
event-bus configuration does not configure the general cache. Tracking the
matching base review at
https://github.com/apache/superset/pull/42805#discussion_r4000933695. The
refresh does not claim a process-separated cancellation 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]
