I3eka commented on code in PR #43136: URL: https://github.com/apache/superset/pull/43136#discussion_r4061195021
########## 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) + except Exception: # pylint: disable=broad-except + logger.warning("Could not publish AI event for run %s", run_id) + + def consume( + self, + run_id: str, + timeout_seconds: float, + poll_seconds: float = 1.0, + ) -> Iterator[StreamEvent | None]: + import time + + stream = self._stream(run_id) + deadline = time.monotonic() + timeout_seconds + last_id = "-" + + while time.monotonic() < deadline: + try: + entries = self._cache.xrange(stream, last_id, "+", 100) + except Exception: # pylint: disable=broad-except + logger.warning("Could not read AI events for run %s", run_id) + yield IDLE + time.sleep(poll_seconds) + continue + + fresh = [entry for entry in entries if _entry_id(entry) != last_id] + if not fresh: + yield IDLE + time.sleep(poll_seconds) + continue + + for entry in fresh: + last_id = _entry_id(entry) + event = _decode(entry) + if event is None: + continue + yield event + if event.type in _TERMINAL: + return + + def close(self, run_id: str) -> None: + # The stream is left to expire rather than deleted, so a reader that is + # still catching up is not cut off mid-replay. + try: + self._cache.expire(self._stream(run_id), self._ttl) + except Exception: # pylint: disable=broad-except + logger.debug("Could not set TTL on AI event stream for %s", run_id) + + +def _entry_id(entry: Any) -> str: + """Stream entry id, tolerating bytes from the Redis client.""" + raw = entry[0] + return raw.decode() if isinstance(raw, bytes) else str(raw) + + +def _decode(entry: Any) -> StreamEvent | None: + """Rebuild an event from a stream entry, skipping anything malformed.""" + fields = entry[1] + raw = fields.get(b"data") or fields.get("data") + if raw is None: + return None + if isinstance(raw, bytes): + raw = raw.decode() + try: + decoded = json.loads(raw) + return StreamEvent(StreamEventType(decoded["type"]), decoded["payload"]) + except (json.JSONDecodeError, KeyError, ValueError, TypeError): + logger.warning("Discarding malformed AI event") + return None + + +def get_event_bus() -> BaseEventBus: + """ + Build the configured bus, refusing combinations that cannot work. + + An in-memory bus with worker execution is a silent failure — every stream + would sit empty while the run completed elsewhere — so it is rejected at + construction rather than discovered in production. + """ + from flask import current_app + + mode = current_app.config.get("AI_ASSISTANT_EXECUTION_MODE", "inline") + kind = current_app.config.get("AI_ASSISTANT_EVENT_BUS", "memory") + + if mode == "worker" and kind == "memory": + raise RuntimeError( + "AI_ASSISTANT_EXECUTION_MODE='worker' requires " + "AI_ASSISTANT_EVENT_BUS='redis': an in-process bus cannot carry " + "events from a Celery worker to the web process." + ) + + if kind == "memory": + return _memory_bus() + + return RedisStreamEventBus( + cache=_stream_backend(), + prefix=current_app.config.get("AI_ASSISTANT_EVENT_STREAM_PREFIX", "ai-events-"), + ttl_seconds=current_app.config.get("AI_ASSISTANT_EVENT_TTL_SECONDS", 900), + ) + + +def _stream_backend() -> Any: + """ + Build a cache client that can speak Redis streams. + + Deliberately not ``cache_manager.cache``: the general-purpose cache is a + Flask-Caching client with no stream commands, so publishing through it would + fail with an ``AttributeError`` on the first event. The stream methods live + on the same backend classes the async-query channel uses, and those are + constructed from a config dict rather than taken from the extension. + """ + from flask import current_app + + from superset.async_events.cache_backend import ( Review Comment: Thanks, confirmed. Both Redis and Sentinel construction failed with `ModuleNotFoundError` once I removed the test's `_stream_backend` stub. Fixed in ddaaed36f4 using `superset.coordination.cache_backend`. The regression constructs the actual configured backend and stubs only stream I/O; both cases now pass. The combined AI/migration suite is 699 passing. This repairs bus construction, not worker-loss recovery or cross-process cancellation. ########## superset/migrations/versions/2026-09-14_00-00_e84f26b1c903_merge_ai_chat_with_master.py: ########## @@ -0,0 +1,34 @@ +# 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. +"""Merge AI chat tables with the main migration chain. + +Revision ID: e84f26b1c903 +Revises: a1c4f7e29b31, 7e2c9a4f1b83 +Create Date: 2026-09-14 00:00:00.000000 + +""" + +revision = "e84f26b1c903" +down_revision = ("a1c4f7e29b31", "7e2c9a4f1b83") Review Comment: Confirmed against current master: the upstream single-head regression failed with `e84f26b1c903` and `60f94cd6cd11`. Maintenance merge 8e6321fa05 incorporates master 1f1e09c9 without rewriting the published branch and reuses the same no-op merge revision `73c5b8e2a194` as #43135. The graph regression now passes; both new merge functions execute as no-ops. I did not run a full fresh-database migration or touch a running database. The SIP hold remains in place. ########## tests/unit_tests/ai/test_suggestions.py: ########## @@ -0,0 +1,179 @@ +# 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. +""" +Model-generated opening suggestions. + +The client keeps its own locally derived list, so every failure mode here has to +end in an empty list rather than an exception — an empty return is what makes the +fallback take over. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +import pytest + + +def test_parsing_a_plain_json_array() -> None: + """The shape the instruction asks for.""" + from superset.ai.suggestions import _parse + + assert _parse('["How many orders?", "Revenue by region?"]', 3) == [ + "How many orders?", + "Revenue by region?", + ] + + +def test_parsing_tolerates_a_code_fence_and_preamble() -> None: + """ + A model that wraps the array has still done the useful part. + + Being strict here would throw away a good answer over presentation. + """ + from superset.ai.suggestions import _parse + + fenced = 'Here you go:\n```json\n["One?", "Two?"]\n```' + assert _parse(fenced, 3) == ["One?", "Two?"] + + +def test_parsing_caps_to_the_requested_count() -> None: + """A model that ignores the limit does not get to overflow the row.""" + from superset.ai.suggestions import _parse + + assert _parse('["a", "b", "c", "d", "e"]', 2) == ["a", "b"] + + +def test_parsing_drops_duplicates_case_insensitively() -> None: + """Variety was asked for; near-duplicates are not variety.""" + from superset.ai.suggestions import _parse + + assert _parse('["Revenue?", "revenue?", "Orders?"]', 3) == [ + "Revenue?", + "Orders?", + ] + + +def test_parsing_drops_non_strings_and_blanks() -> None: + """A ragged array yields only its usable entries.""" + from superset.ai.suggestions import _parse + + assert _parse('["Good?", "", 5, null, {"a": 1}, " ", "Also good?"]', 3) == [ + "Good?", + "Also good?", + ] + + +def test_parsing_clips_an_overlong_suggestion() -> None: + """A chip cannot show an essay, so one is not stored.""" + from superset.ai.suggestions import _parse, MAX_SUGGESTION_CHARS + + [only] = _parse(f'["{"x" * 500}"]', 1) + assert len(only) == MAX_SUGGESTION_CHARS + + +def test_parsing_salvages_an_array_wrapped_in_an_object() -> None: + """ + A model that replies ``{"prompts": [...]}`` has still answered. + + The array is taken from wherever it is, which is the same tolerance that + handles a code fence: the instruction asks for a bare array, but refusing a + near-miss would drop a perfectly good answer. + """ + from superset.ai.suggestions import _parse + + assert _parse('{"prompts": ["How many orders?"]}', 3) == ["How many orders?"] Review Comment: Rechecked on ddaaed36f4: the exact object-wrapped example still returns the expected list, and `test_parsing_salvages_an_array_wrapped_in_an_object` passes in the full 698-test AI suite. The array is extracted before JSON decoding, so this particular failure report is incorrect. Resolving this test-failure thread; the separate provider-cleanup issue in the review overview is real and remains open. ########## 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") Review Comment: There is now a focused fix in #44497 (68836a2738), based on #42805: the cancel route resolves the run message inside the owned thread before setting any flag. Its real SQLite route/DAO tests cover both execution-mode configurations and reject missing, foreign-thread and other-owned-thread run IDs without calling cancellation. That PR is not yet inherited here, so this thread remains open. It does not fix the separate shared-cache/cross-process cancellation issue. -- 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]
