sadpandajoe commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r3870345550
########## 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, + max_turns=profile.max_turns or _config("AI_AGENT_MAX_TURNS", 20), + timeout_seconds=profile.timeout_seconds + or _config("AI_AGENT_TIMEOUT_SECONDS", 300), + should_cancel=lambda: is_cancelled(request.run_id), + ) + + _mark_streaming(request.assistant_message_uuid) + + # The runtime is async and this is a synchronous generator, so the async + # events are drained into a list per batch rather than bridged with a + # thread. Collecting the whole run before yielding would defeat streaming, + # so the loop pulls one event at a time from a dedicated event loop. + yield from _drain(runtime.run(run_request)) + + result = runtime.result + outcome = _outcome_of(result) + if result.error is not None: + # The only place the provider's own words are recorded. They do not go on + # the message: that is served back to the browser, and a transport error + # can name internal hosts. + logger.warning( + "AI run %s failed: %s", + request.run_id, + result.error, + ) + _finalise_message( + request.assistant_message_uuid, + content=_terminal_content(result, outcome), + status=_status_of(outcome), + extra={ Review Comment: Successful turns discard the reasoning that was streamed to the browser, so reopening a completed conversation shows a different thought process than the live run. Could this persist the bounded `result.thoughts` here, as the abnormal-run path does? ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,508 @@ +--- +title: AI Assistant +hide_title: true +sidebar_position: 17 +version: 1 +--- + +# AI Assistant + +The AI Assistant is a conversational interface for exploring your data. A user +asks a question in plain language; the assistant finds relevant datasets, +inspects their schema, writes and runs read-only SQL, and answers with both the +result and the query it used. The shipped profiles are read-only. A deployment +can explicitly add chart and dashboard authoring tools to a gated profile. + +Superset ships **no model provider and talks to no model vendor by default**. +The feature is disabled, and even when enabled it returns `404` until you point +it at a provider you control. Nothing is sent anywhere until you configure it. + +## Enabling it + +Two things are required: the feature flag, and a provider. + +```python +# superset_config.py +FEATURE_FLAGS = { + "AI_ASSISTANT": True, +} + +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.anthropic.AnthropicProvider" +AI_LLM_PROVIDER_CONFIG = { + "api_key": os.environ["ANTHROPIC_API_KEY"], + "models": { + "default": "claude-sonnet-4-5", + "fast": "claude-haiku-4-5", + "reasoning": "claude-opus-4-1", + }, +} +``` + +Install the matching extra: + +```bash +pip install "apache-superset[ai-anthropic]" # or [ai-openai] +``` + +Then run `superset init` so the assistant's permissions are created and assigned Review Comment: Following these enablement steps on an existing deployment leaves the new `ai_chat_*` tables absent, because `superset init` does not apply the migration. Could this also require `superset db upgrade` before the assistant is used? -- 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]
