sadpandajoe commented on code in PR #42805:
URL: https://github.com/apache/superset/pull/42805#discussion_r4033192956
##########
docs/admin_docs/configuration/ai-assistant.mdx:
##########
@@ -0,0 +1,489 @@
+---
+title: AI Assistant
+hide_title: true
+sidebar_position: 17
+version: 1
+---
+
+# AI Assistant
+
+The AI Assistant is a conversational interface for exploring your data. A user
+asks a question in plain language; the assistant finds relevant datasets,
+inspects their schema, writes and runs read-only SQL, and answers with both the
+result and the query it used.
+
+Superset ships **no model provider and talks to no model vendor by default**.
+The feature is disabled, and even when enabled it returns `404` until you point
+it at a provider you control. Nothing is sent anywhere until you configure it.
+
+## Enabling it
+
+Two things are required: the feature flag, and a provider.
+
+```python
+# superset_config.py
+FEATURE_FLAGS = {
+ "AI_ASSISTANT": True,
+}
+
+AI_LLM_PROVIDER_CLASS = "superset.ai.llm.anthropic.AnthropicProvider"
+AI_LLM_PROVIDER_CONFIG = {
+ "api_key": os.environ["ANTHROPIC_API_KEY"],
+ "models": {
+ "default": "claude-sonnet-4-5",
+ "fast": "claude-haiku-4-5",
+ "reasoning": "claude-opus-4-1",
+ },
+}
+```
+
+Install the matching extra:
+
+```bash
+pip install "apache-superset[ai-anthropic]" # or [ai-openai]
+```
+
+Then run `superset init` so the assistant's permissions are created and
assigned
Review Comment:
The setup path creates permissions but never applies the new chat-table
migration. On an existing deployment, following these steps leaves the metadata
tables absent and the first conversation fails. Could the guide require
`superset db upgrade` before `superset init`?
##########
superset/ai/api.py:
##########
@@ -0,0 +1,989 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+REST API for the AI assistant.
+
+Every route carries ``@protect()`` and is reached through ``@expose`` on a
+``BaseSupersetApi`` subclass, which is what makes Flask-AppBuilder's
+authorization actually run. Ownership is enforced a second time in the command
+and DAO layers, so a conversation identifier is never on its own a capability.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Generator
+from typing import Any, cast
+
+from flask import current_app, request, Response, stream_with_context
+from flask_appbuilder.api import expose, permission_name, protect, safe
+from marshmallow import ValidationError
+
+from superset.ai.events import (
+ error_event,
+ KEEPALIVE_FRAME,
+ KEEPALIVE_INTERVAL_SECONDS,
+)
+from superset.ai.schemas import (
+ AgentResponseSchema,
+ CancelPostSchema,
+ FeedbackPostSchema,
+ MessagePostSchema,
+ RunAcceptedResponseSchema,
+ SuggestedPromptsPostSchema,
+ ThreadDetailResponseSchema,
+ ThreadPostSchema,
+ ThreadPutSchema,
+ ThreadResponseSchema,
+)
+from superset.ai.types import MessageRole, MessageStatus
+from superset.commands.ai.exceptions import (
+ AIChatMessageInvalidError,
+ AIChatMessageNotFoundError,
+ AIChatThreadInvalidError,
+ AIChatThreadNotFoundError,
+)
+from superset.extensions import event_logger
+from superset.utils.core import get_user_id
+from superset.utils.decorators import transaction
+from superset.views.base_api import BaseSupersetApi, statsd_metrics
+
+logger = logging.getLogger(__name__)
+
+#: Upper bound on how long a client may hold a stream open, so an abandoned
+#: browser tab cannot pin a worker indefinitely.
+_STREAM_TIMEOUT_SECONDS = 900
+
+#: How often a reader checks the event bus for new frames.
+#:
+#: Deliberately separate from ``KEEPALIVE_INTERVAL_SECONDS``. Passing the
+#: keep-alive interval as the poll interval made the reader sleep fifteen
seconds
+#: between checks and then deliver everything that had accumulated in one
batch —
+#: so a worker-mode run showed no streaming at all: the answer and every tool
call
+#: appeared in fifteen-second lumps. One controls responsiveness, the other how
+#: often an idle connection is reassured; they are not the same number.
+_EVENT_POLL_SECONDS = 0.1
+
+
+class AIRestApi(BaseSupersetApi):
+ """Conversations with the AI assistant."""
+
+ resource_name = "ai"
+ openapi_spec_tag = "AI Assistant"
+ allow_browser_login = True
+ class_permission_name = "AIAssistant"
+
+ openapi_spec_component_schemas = (
+ AgentResponseSchema,
+ CancelPostSchema,
+ FeedbackPostSchema,
+ MessagePostSchema,
+ RunAcceptedResponseSchema,
+ SuggestedPromptsPostSchema,
+ ThreadDetailResponseSchema,
+ ThreadPostSchema,
+ ThreadPutSchema,
+ ThreadResponseSchema,
+ )
+
+ @expose("/agent/", methods=("GET",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("read")
+ def agents(self) -> Response:
+ """List agent profiles the current user may select.
+ ---
+ get:
+ summary: List available agent profiles
+ responses:
+ 200:
+ description: Available profiles
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ type: array
+ items:
+ $ref: '#/components/schemas/AgentResponseSchema'
+ 401:
+ $ref: '#/components/responses/401'
+ 403:
+ $ref: '#/components/responses/403'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.ai.factories import get_profiles
+
+ profiles = get_profiles().visible_to_current_user()
+ return self.response(200, result=[p.to_public_dict() for p in
profiles])
+
+ @expose("/model/", methods=("GET",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("read")
+ def models(self) -> Response:
+ """List models this deployment has configured.
+ ---
+ get:
+ summary: List selectable models
+ responses:
+ 200:
+ description: Configured model identifiers
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ type: array
+ items:
+ type: string
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.ai.factories import get_provider
+
+ return self.response(200, result=get_provider().available_models())
+
+ @expose("/thread/", methods=("POST",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs:
f"{self.__class__.__name__}.post_thread",
+ log_to_statsd=False,
+ )
+ def post_thread(self) -> Response:
+ """Create a conversation.
+ ---
+ post:
+ summary: Create a conversation
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/ThreadPostSchema'
+ responses:
+ 201:
+ description: Conversation created
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ $ref: '#/components/schemas/ThreadResponseSchema'
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.commands.ai import CreateAIChatThreadCommand
+
+ try:
+ payload = ThreadPostSchema().load(request.json or {})
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+ try:
+ thread = CreateAIChatThreadCommand(
+ user_id=self._user_id(),
+ title=payload.get("title"),
+ agent_key=payload.get("agent_key"),
+ ).run()
+ except AIChatThreadInvalidError as ex:
+ return self.response_422(message=str(ex))
+ return self.response(201, result=_thread_dict(thread))
+
+ @expose("/thread/", methods=("GET",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("read")
+ def get_threads(self) -> Response:
+ """List the current user's conversations.
+ ---
+ get:
+ summary: List conversations
+ parameters:
+ - in: query
+ name: limit
+ schema:
+ type: integer
+ - in: query
+ name: offset
+ schema:
+ type: integer
+ responses:
+ 200:
+ description: Conversations
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ count:
+ type: integer
+ result:
+ type: array
+ items:
+ $ref: '#/components/schemas/ThreadResponseSchema'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.daos.ai import AIChatThreadDAO
+
+ limit = request.args.get("limit", type=int) or 50
+ offset = request.args.get("offset", type=int) or 0
+ threads = AIChatThreadDAO.find_all_for_user(
+ self._user_id(), limit=limit, offset=offset
+ )
+ return self.response(
+ 200,
+ count=len(threads),
+ result=[_thread_dict(thread) for thread in threads],
+ )
+
+ @expose("/thread/<thread_uuid>", methods=("GET",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("read")
+ def get_thread(self, thread_uuid: str) -> Response:
+ """Fetch a conversation and its messages.
+ ---
+ get:
+ summary: Get a conversation
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ responses:
+ 200:
+ description: Conversation with messages
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ $ref: '#/components/schemas/ThreadDetailResponseSchema'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.daos.ai import (
+ AIChatFeedbackDAO,
+ AIChatMessageDAO,
+ AIChatThreadDAO,
+ )
+
+ user_id = self._user_id()
+ thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, user_id)
+ if thread is None:
+ return self.response_404()
+
+ messages = AIChatMessageDAO.find_for_thread(thread)
+ # Resolved for the whole transcript at once so the panel can show which
+ # replies this user already rated; without it a reload loses the
verdict
+ # and the message looks unrated.
+ verdicts = AIChatFeedbackDAO.find_verdicts_for_user(
+ [message.id for message in messages], user_id
+ )
+ detail = _thread_dict(thread)
+ detail["messages"] = [
+ _message_dict(message, liked=verdicts.get(message.id))
+ for message in messages
+ ]
+ return self.response(200, result=detail)
+
+ @expose("/thread/<thread_uuid>", methods=("PUT",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs:
f"{self.__class__.__name__}.put_thread",
+ log_to_statsd=False,
+ )
+ def put_thread(self, thread_uuid: str) -> Response:
+ """Rename or archive a conversation.
+ ---
+ put:
+ summary: Update a conversation
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/ThreadPutSchema'
+ responses:
+ 200:
+ description: Conversation updated
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ 422:
+ $ref: '#/components/responses/422'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.commands.ai import UpdateAIChatThreadCommand
+
+ try:
+ payload = ThreadPutSchema().load(request.json or {})
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+ try:
+ thread = UpdateAIChatThreadCommand(
+ thread_uuid,
+ self._user_id(),
+ title=payload.get("title"),
+ status=payload.get("status"),
+ ).run()
+ except AIChatThreadNotFoundError:
+ return self.response_404()
+ except AIChatThreadInvalidError as ex:
+ return self.response_422(message=str(ex))
+ return self.response(200, result=_thread_dict(thread))
+
+ @expose("/thread/<thread_uuid>", methods=("DELETE",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs:
f"{self.__class__.__name__}.delete_thread",
+ log_to_statsd=False,
+ )
+ def delete_thread(self, thread_uuid: str) -> Response:
+ """Delete a conversation and its messages.
+ ---
+ delete:
+ summary: Delete a conversation
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ responses:
+ 200:
+ description: Conversation deleted
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.commands.ai import DeleteAIChatThreadCommand
+
+ try:
+ DeleteAIChatThreadCommand(thread_uuid, self._user_id()).run()
+ except AIChatThreadNotFoundError:
+ return self.response_404()
+ return self.response(200, message="OK")
+
+ @expose("/thread/<thread_uuid>/message", methods=("POST",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs:
f"{self.__class__.__name__}.post_message",
+ log_to_statsd=False,
+ )
+ def post_message(self, thread_uuid: str) -> Response:
+ """Post a user message and start a run.
+ ---
+ post:
+ summary: Post a message
+ description: >
+ Stores the user's message, creates a placeholder assistant message,
+ and starts a run. Returns immediately; consume the answer from the
+ stream endpoint using the returned run identifier.
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/MessagePostSchema'
+ responses:
+ 202:
+ description: Run accepted
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ $ref: '#/components/schemas/RunAcceptedResponseSchema'
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ 422:
+ $ref: '#/components/responses/422'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.ai.orchestrator import new_run_id
+ from superset.commands.ai import AppendAIChatMessageCommand
+
+ try:
+ payload = MessagePostSchema().load(request.json or {})
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+ user_id = self._user_id()
+
+ try:
+ user_message = AppendAIChatMessageCommand(
+ thread_uuid,
+ user_id,
+ MessageRole.USER,
+ payload["content"],
+ request_id=payload.get("request_id"),
+ ).run()
+ # Created up front so a client that reconnects before any token
+ # arrives still has a row to attach its stream to.
+ assistant_message = AppendAIChatMessageCommand(
+ thread_uuid,
+ user_id,
+ MessageRole.ASSISTANT,
+ "",
+ request_id=payload.get("request_id"),
+ status=MessageStatus.PENDING,
+ ).run()
+ except AIChatThreadNotFoundError:
+ return self.response_404()
+ except (AIChatMessageInvalidError, AIChatThreadInvalidError) as ex:
+ return self.response_422(message=str(ex))
+
+ run_id = new_run_id()
+ _record_run_context(assistant_message, run_id, payload)
+
+ self._start_run(
+ thread_uuid=thread_uuid,
+ user_id=user_id,
+ run_id=run_id,
+ assistant_message_uuid=str(assistant_message.uuid),
+ agent_key=payload.get("agent_key"),
+ model=payload.get("model"),
+ page_context=payload.get("page_context"),
+ )
+
+ return self.response(
+ 202,
+ result={
+ "message_uuid": str(user_message.uuid),
+ "assistant_message_uuid": str(assistant_message.uuid),
+ "run_id": run_id,
+ },
+ )
+
+ @expose("/thread/<thread_uuid>/stream", methods=("GET",))
+ @protect()
+ @statsd_metrics
+ @permission_name("read")
+ def stream(self, thread_uuid: str) -> Response:
+ """Stream a run's events.
+ ---
+ get:
+ summary: Stream assistant events
+ description: >
+ Server-sent events for one run. Frame names are session, thinking,
+ thoughts, checkpoint, assistant_delta, final, error, cancelled and
+ done. The done frame is always last and reports whether the run
+ succeeded.
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ - in: query
+ name: run_id
+ required: true
+ schema:
+ type: string
+ responses:
+ 200:
+ description: An event stream
+ content:
+ text/event-stream:
+ schema:
+ type: string
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ # No @safe here: once headers are flushed an exception can no longer
+ # become a status code, so failures are reported as in-band error
frames.
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO
+
+ run_id = request.args.get("run_id")
+ if not run_id:
+ return self.response_400(message="run_id is required")
+
+ # Ownership is checked before the stream opens; the run identifier
alone
+ # must not grant access to another user's conversation.
+ thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid,
self._user_id())
+ if thread is None:
+ return self.response_404()
+
+ pending = _find_run_message(AIChatMessageDAO.find_for_thread(thread),
run_id)
+ if pending is None:
+ return self.response_404()
+
+ turn = None
+ if current_app.config.get("AI_ASSISTANT_EXECUTION_MODE") != "worker":
+ from superset.ai.orchestrator import TurnRequest
+
+ extra = pending.extra
+ turn = TurnRequest(
+ thread_uuid=thread_uuid,
+ user_id=self._user_id(),
+ run_id=run_id,
+ assistant_message_uuid=str(pending.uuid),
+ profile_key=extra.get("agent_key"),
+ model=extra.get("model"),
+ page_context=extra.get("page_context"),
+ )
+
+ generator = self._build_stream(run_id, turn)
+ response = Response(
+ generator,
+ content_type="text/event-stream; charset=utf-8",
+ headers={
+ "Cache-Control": "no-cache, no-transform",
+ "Connection": "keep-alive",
+ # Defeats proxy buffering, which otherwise holds frames until
+ # the response completes and makes streaming pointless.
+ "X-Accel-Buffering": "no",
+ "Content-Encoding": "identity",
+ },
+ direct_passthrough=False,
+ )
+ response.implicit_sequence_conversion = False
+ return response
+
+ @expose("/thread/<thread_uuid>/cancel", methods=("POST",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ def cancel(self, thread_uuid: str) -> Response:
+ """Ask a run to stop.
+ ---
+ post:
+ summary: Cancel a run
+ description: >
+ Cancellation is cooperative: the run stops at its next step
+ boundary. A run inside a single long model call or query will not
+ stop until that call returns.
+ parameters:
+ - in: path
+ name: thread_uuid
+ required: true
+ schema:
+ type: string
+ format: uuid
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/CancelPostSchema'
+ responses:
+ 200:
+ description: Cancellation recorded
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ """
+ if (unavailable := self._reject_if_unconfigured()) is not None:
+ return unavailable
+
+ from superset.ai.orchestrator import request_cancel
+ from superset.daos.ai import AIChatThreadDAO
+
+ try:
+ payload = CancelPostSchema().load(request.json or {})
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+ if AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, self._user_id())
is None:
+ return self.response_404()
+
+ request_cancel(payload["run_id"])
+ return self.response(200, message="OK")
+
+ @expose("/feedback", methods=("POST",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("write")
+ def feedback(self) -> Response:
+ """Rate an assistant message.
+ ---
+ post:
+ summary: Submit feedback
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/FeedbackPostSchema'
+ responses:
+ 200:
+ description: Feedback recorded
+ 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 SubmitAIChatFeedbackCommand
+
+ try:
+ payload = FeedbackPostSchema().load(request.json or {})
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+ try:
+ SubmitAIChatFeedbackCommand(
+ payload["message_uuid"],
+ self._user_id(),
+ liked=payload["liked"],
+ comment=payload.get("comment"),
+ ).run()
+ except AIChatMessageNotFoundError:
+ return self.response_404()
+ except AIChatMessageInvalidError as ex:
Review Comment:
Rating a user or system message raises `AIChatFeedbackInvalidError`, but
this handler only catches `AIChatMessageInvalidError`, so the documented
invalid-feedback case escapes through `@safe` as a 500 instead of a 422. Could
this catch the exception the command actually raises?
##########
superset/config.py:
##########
@@ -1780,6 +1789,13 @@ class CeleryConfig: # pylint:
disable=too-few-public-methods
"superset.tasks.slack",
"superset.tasks.export_dashboard_excel",
"superset.tasks.version_history_retention",
+ # Imported unconditionally, like the rest: importing the module only
+ # registers the task, and nothing enqueues it unless the AI assistant
is
+ # switched on and configured for worker execution. Leaving it out
meant a
+ # deployment that did switch worker mode on got
+ # "Received unregistered task of type 'ai.run_turn'" and a stream that
+ # never produced a frame.
+ "superset.ai.tasks",
Review Comment:
This registers `ai.run_turn` only in the built-in Celery config. The Helm
chart generates its own `CeleryConfig.imports` without `superset.ai.tasks`, so
an official worker-mode deployment discards the queued task as unregistered and
leaves the assistant row pending. Could the chart configuration import this
task too?
--
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]
