Copilot commented on code in PR #43237: URL: https://github.com/apache/superset/pull/43237#discussion_r4002625804
########## superset/ai/page_context.py: ########## @@ -0,0 +1,372 @@ +# 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. +""" +What the user is looking at, rendered for the model. + +This is what lets someone ask "why is this number lower than last week?" while +looking at a dashboard and get an answer about *that* chart. Without it the +assistant is a search box that happens to live in Superset. + +The client gathers the context — it is the only party that knows which tab is +open, what is typed in the editor, and which filters are applied — and this +module turns it into prose. Everything here is treated as untrusted: a dashboard +title, a chart description or a markdown block is authored by a user, so it is +data and never instruction. +""" + +from __future__ import annotations + +from typing import Any + +#: Ceiling on the whole rendered block. Page context competes with conversation +#: history for the same budget, so an enormous dashboard cannot crowd out the +#: question being asked. +MAX_CONTEXT_CHARS = 20_000 + +#: Ceiling on the editor SQL specifically. A pasted migration script should not +#: consume the entire context, and the useful part is near the top. +MAX_SQL_CHARS = 10_000 + +#: Markdown authored on a dashboard is how a team explains its own data, so it +#: is worth real space — but bounded, and only a handful of blocks. +MAX_MARKDOWN_BLOCKS = 10 +MAX_MARKDOWN_BLOCK_CHARS = 4_000 + +#: Lists that could otherwise be unbounded. +MAX_CHARTS = 50 +MAX_FILTERS = 25 +MAX_TABLES = 20 + +#: Page types the client may report. An unknown value renders as "other" rather +#: than being echoed back into the prompt. +KNOWN_PAGE_TYPES = frozenset( + {"sqllab", "explore", "dashboard", "chart", "home", "other"} +) + + +def render_page_context(context: Any) -> str: + """ + Render the client's page context as a prompt section. + + Returns an empty string when there is nothing useful, so the caller can + append unconditionally. Never raises: a malformed payload from a stale + client costs the model some context, and should not cost the user an answer. + """ + if not isinstance(context, dict): + return "" + + try: + return _render(context)[:MAX_CONTEXT_CHARS] + except Exception: # pylint: disable=broad-except Review Comment: The client-supplied dashboard titles, markdown, filter values, and SQL are interpolated directly into the system prompt by `_render`; the prose warning that follows the header does not stop one of those values from being interpreted as instructions. This path should apply `sanitize_for_llm_context` to untrusted leaves (while keeping operational IDs/page type usable) before interpolation, with a regression test for an injection string, so page context has the same framing guarantee as tool data. ########## 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) + Review Comment: `AppendAIChatMessageCommand` already makes replayed `request_id` calls return the existing row and exposes whether an insert occurred, but this path ignores that result and unconditionally allocates a new run and overwrites the assistant row's run context. Retrying the same payload can therefore return the same `message_uuid` with a different `run_id` and start a second inference, defeating the idempotency contract and potentially double-charging. Reuse the existing run (or make the replay wait for/return its recorded run) instead of starting another one. ########## 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( + f"{_CANCEL_PREFIX}{run_id}", True, timeout=_CANCEL_TTL_SECONDS + ) + except Exception: # pylint: disable=broad-except + logger.warning("Could not record cancellation for AI run %s", run_id) + + +def is_cancelled(run_id: str) -> bool: + """Whether a stop has been requested for this run.""" + from superset.extensions import cache_manager + + if run_id in _CANCELLED_LOCALLY: + return True + try: + return bool(cache_manager.cache.get(f"{_CANCEL_PREFIX}{run_id}")) + except Exception: # pylint: disable=broad-except + # A cache that cannot be read must not make every run appear cancelled; + # that would stop all inference the moment the cache went away. + return False + + +def clear_cancel(run_id: str) -> None: + """Drop a run's cancellation flag.""" + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.discard(run_id) + try: + cache_manager.cache.delete(f"{_CANCEL_PREFIX}{run_id}") + except Exception: # pylint: disable=broad-except + logger.debug("Could not clear cancellation flag for AI run %s", run_id) + + +def stream_turn(request: TurnRequest) -> Iterator[StreamEvent]: + """ + Answer a turn, yielding events as they happen. + + This is the primary entry point. Inline execution consumes it directly from + inside the streaming response, which means the producer and the reader are + the same process by construction — important because Superset runs several + web workers, and a turn that published to one process's in-memory queue + while the browser's stream landed on another would appear to hang forever. + + Never raises for an operational failure: a failure is an ``error`` event and + an ``error`` message status, because the caller may already have flushed + response headers or may be a worker with no one to report to. + """ + recorder = start_run( + run_id=request.run_id, + thread_uuid=request.thread_uuid, + user_id=request.user_id, + ) + # Shared with ``_run`` so the ``finally`` below can see the runtime's + # partial result and whether the message was already written. + state: dict[str, Any] = {} + try: + # Bound here rather than inside ``_run`` so that a run which fails before + # it has resolved a profile still produces a start and an end, and so + # that the runtime can report its own spans without the runtime contract + # growing a telemetry parameter. + with bind_run(recorder): + recorder.run_started() + yield from _run(request, state) + except Exception as ex: # pylint: disable=broad-except + logger.exception("AI turn failed for run %s", request.run_id) + recorder.error(ex) + recorder.run_ended(outcome=RunOutcome.ERROR) + answer, extra = _partial_from_state(state) + extra["outcome"] = RunOutcome.ERROR.value + _finalise_message( + request.assistant_message_uuid, + # The generic text rather than the exception: this is persisted and + # served back to the browser, so it must not carry internals. The + # detail is in the log line above, keyed by run id. + content=answer or GENERIC_ERROR_MESSAGE, + status=MessageStatus.ERROR, + extra=extra, + ) + state["finalised"] = True + yield error_event() + yield done_event(ok=False) + finally: + clear_cancel(request.run_id) + # A client that stops the run, or simply navigates away, abandons this + # generator part-way through. Nothing above will have written the + # message, so it would otherwise sit in ``streaming`` with no content + # for ever — the user loses both the partial answer and any record that + # the turn happened. Persist whatever was produced. + _abandon_message(request.assistant_message_uuid, state) + # Idempotent, so the ordinary paths above win. + recorder.run_ended(outcome=RunOutcome.CANCELLED) + + +def execute_turn(request: TurnRequest) -> RunOutcome: + """ + Answer a turn, publishing events to the event bus. + + Used by worker execution, where the reader is in another process. Shares its + whole body with :func:`stream_turn` so the two execution modes cannot drift + apart in behaviour. + """ + from superset.ai.eventbus import get_event_bus + + bus = get_event_bus() + outcome = RunOutcome.SUCCESS + + for event in stream_turn(request): + bus.publish(request.run_id, event) + if event.type is StreamEventType.ERROR: + outcome = RunOutcome.ERROR + elif event.type is StreamEventType.CANCELLED: + outcome = RunOutcome.CANCELLED + elif event.type is StreamEventType.DONE and not event.payload.get("ok"): + # A run that ended un-ok without an explicit error frame timed out. + if outcome is RunOutcome.SUCCESS: + outcome = RunOutcome.TIMEOUT + + return outcome + + +def _run(request: TurnRequest, state: dict[str, Any]) -> Iterator[StreamEvent]: + """Assemble and drive the run. See :func:`stream_turn` for error policy.""" + from superset.ai.factories import ( + get_profiles, + get_provider, + get_runtime, + get_tools_for_profile, + ) + from superset.ai.policy import load_policy_chain + from superset.ai.runtime.base import RunRequest + from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO + + recorder = current_run() + + thread = AIChatThreadDAO.find_by_uuid_for_user(request.thread_uuid, request.user_id) + if thread is None: + # The thread vanished between accepting the message and running it. + recorder.run_ended(outcome=RunOutcome.ERROR) + yield error_event("That conversation is no longer available.") + yield done_event(ok=False) + return + + profile = get_profiles().get(request.profile_key) + tools = get_tools_for_profile(profile) + provider = get_provider() + runtime = get_runtime(provider) + state["runtime"] = runtime + + yield session_event(request.thread_uuid, request.assistant_message_uuid) + + from superset.ai.page_context import render_page_context + + history = _build_history(AIChatMessageDAO.find_for_thread(thread)) + # Recorded as well as prompted with, so the transcript can show what the + # assistant was told about the user's screen. An answer that looks wrong is + # usually an answer to a different question than the reader assumed, and the + # page context is where that difference lives. + rendered_context = render_page_context(request.page_context) + state["page_context"] = rendered_context + system_prompt = _build_system_prompt(tools, rendered_context) + model = _resolved_model(provider, request.model, profile) + + recorder.describe( + agent_key=profile.key, + model=model, + question=_latest_question(history), + ) + + run_request = RunRequest( + messages=history, + system_prompt=system_prompt, + tools=tools, + policies=load_policy_chain(), + model_alias=profile.model_alias, Review Comment: `request.model` is resolved only for telemetry here; `RunRequest` carries just `model_alias`, and `MessagesApiRuntime` builds `CompletionRequest` without the concrete model. Consequently an explicit model in the API request—and the profile's pinned `model` used by `_resolved_model`—never reaches the provider, so every run silently uses the alias default. Add the resolved concrete model to the runtime request and pass it through to the completion request. ########## 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"), Review Comment: The thread model documents `agent_key` as the profile used for the conversation, but this path passes only the per-message payload value and never falls back to the thread's stored profile or updates it after a run. A thread created with a non-default profile therefore silently uses the default on a later message that omits `agent_key`, and the returned thread metadata becomes stale. Resolve the effective profile from the thread and persist the selected profile when accepting the run. ########## 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: + return self.response_422(message=str(ex)) Review Comment: The command used here raises `AIChatFeedbackInvalidError` for a non-assistant message, but this handler catches `AIChatMessageInvalidError`. Such a target therefore bypasses the intended 422 response and is treated as an unexpected server error instead. Catch the feedback-specific exception (and import it) so the documented validation response is returned. -- 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]
