sadpandajoe commented on code in PR #42805: URL: https://github.com/apache/superset/pull/42805#discussion_r3821577552
########## superset/ai/tools/sql.py: ########## @@ -0,0 +1,540 @@ +# 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. +""" +Running and checking SQL. + +Both tools are reads. ``execute_sql`` decides that with Superset's own parser +rather than a keyword match, because a prefix match is defeated by a leading +comment, a CTE wrapping a mutation, and a second statement smuggled after a +legitimate ``SELECT`` — all of which the parser already handles for the rest of +Superset. + +Authorization is layered deliberately: + +1. ``DatabaseDAO.find_by_id`` applies ``DatabaseFilter``, so a database the user + has no grant on is indistinguishable from one that does not exist. +2. ``expose_in_sqllab`` is honoured, so an operator who withheld a connection + from ad-hoc querying has also withheld it here. +3. The database's ``allow_dml`` setting is checked. +4. Every statement must be non-mutating. +5. ``security_manager.raise_for_access`` with ``force_dataset_match=True`` — the + same strictness SQL Lab applies — so the tables referenced must resolve to + datasets the user may read. +""" + +from __future__ import annotations + +import logging +import time +from decimal import Decimal +from typing import Any, Callable, ClassVar + +from superset.ai.tools.base import AITool, ToolError, ToolOutput + +logger = logging.getLogger(__name__) + +#: Rows returned when the caller does not ask for a specific limit. Small on +#: purpose: an agent inspecting data needs a shape, not a dump, and the model +#: pays for every row in context. +DEFAULT_ROW_LIMIT = 100 + +#: Rows kept in the UI summary. The summary is persisted on the message and sent +#: to the browser, so it carries a sample rather than the result set. +DISPLAY_SAMPLE_ROWS = 20 + +#: Characters of executed SQL shown in the UI summary. +DISPLAY_SQL_CHARS = 4000 + + +def _database_or_refuse(database_id: Any) -> Any: + """ + Resolve a database the current user is allowed to query, or refuse. + + Uses the DAO so that ``DatabaseFilter`` — the same filter the database REST + API applies — decides visibility. A database the user cannot see is reported + as "not found" rather than "forbidden", so the tool cannot be used to probe + for the existence of connections. + """ + from superset.daos.database import DatabaseDAO + + if not isinstance(database_id, int) or isinstance(database_id, bool): + raise ToolError("'database_id' must be an integer. Use list_databases first.") + + database = DatabaseDAO.find_by_id(database_id) + if database is None: + raise ToolError( + f"No database with id {database_id} is available to you. " + f"Call list_databases to see the ones you can query." + ) + if not database.expose_in_sqllab: + raise ToolError( + f"Database {database.database_name!r} is not available for ad-hoc " + f"queries. Call list_databases to see the ones that are." + ) + return database + + +def _parse_or_refuse(sql: str, database: Any) -> Any: + """ + Parse ``sql`` for the database's engine, or refuse. + + Fails closed: SQL that will not parse cannot be shown to be read-only, so it + is refused rather than executed. The parser error is logged rather than + returned, because its message quotes the offending query back and that text + is not ours to echo. + """ + from superset.sql.parse import SQLScript + + if not isinstance(sql, str) or not sql.strip(): + raise ToolError("'sql' must be a non-empty string.") + + try: + return SQLScript(sql, database.db_engine_spec.engine) + except Exception: # pylint: disable=broad-except + logger.info("Refusing unparseable SQL for database %s", database.id) + raise ToolError( + "That SQL could not be parsed. Send a single, syntactically valid " + "read-only statement." + ) from None + + +def _assert_read_only(script: Any, database: Any) -> None: + """ + Refuse anything that is not a read. + + The ``allow_dml`` check comes first because it is the more specific refusal: + telling the model the *deployment* forbids writes on this connection is more + useful than a generic "read-only tool" message. The blanket mutation check + then applies even where ``allow_dml`` is enabled, since this tool is a read + regardless of what the connection would otherwise permit. + """ + has_mutation = script.has_mutation() + + if has_mutation and not database.allow_dml: + raise ToolError( + f"Writes are disabled on database " + f"{database.database_name!r} (allow_dml is off). Rewrite this as a " + f"SELECT." + ) + if has_mutation: + raise ToolError( + "This tool only runs read-only SQL. Rewrite this as a SELECT — " + "statements that modify data or schema are refused." + ) + # A statement the parser could not fully model has no enumerable table + # references, so neither the mutation check above nor the per-table + # authorization below can vouch for it. Refused rather than guessed at. + if script.has_unparseable_statement: + raise ToolError( + "That SQL contains a statement this tool cannot verify as " + "read-only. Send a plain SELECT." + ) + + +def _raise_for_sql_access( + database: Any, + sql: str, + catalog: str | None, + schema: str | None, +) -> None: + """ + Check the user may read every table the query touches. + + ``force_dataset_match=True`` matches what SQL Lab's own pre-execute + validator uses: each referenced table must resolve to a registered dataset + the user has access to, rather than falling through to a broader + catalog- or schema-level grant. + """ + from superset import security_manager + from superset.exceptions import SupersetSecurityException + + try: + security_manager.raise_for_access( + database=database, + sql=sql, + catalog=catalog, + schema=schema, + force_dataset_match=True, + ) + except SupersetSecurityException as ex: + # The exception message names the tables that were denied, which is + # exactly what lets the model pick a different source. + raise ToolError(str(ex.error.message)) from None + + +def _row_limit(requested: Any) -> int: + """ + Clamp the caller's limit to the configured ceiling. + + The model may not raise the cap by asking for more; ``limit`` narrows only. + """ + # A misconfigured or absent ceiling falls back to the default rather than + # becoming unbounded: sending a query with no limit is the one outcome this + # function exists to prevent. + ceiling = DEFAULT_ROW_LIMIT + try: + from flask import current_app + + configured = current_app.config.get("AI_AGENT_MAX_RESULT_ROWS") + if configured is not None: + ceiling = int(configured) + except Exception: # pylint: disable=broad-except + ceiling = DEFAULT_ROW_LIMIT + if ceiling < 1: + ceiling = DEFAULT_ROW_LIMIT + + if requested is None: + return min(DEFAULT_ROW_LIMIT, ceiling) + if not isinstance(requested, int) or isinstance(requested, bool) or requested < 1: + raise ToolError("'limit' must be a positive integer.") + return min(requested, ceiling) + + +def _columns_and_records( + data: Any, +) -> tuple[list[dict[str, str]], list[dict[str, Any]]]: + """ + Normalise a statement's result rows. + + ``Database.execute`` returns a ``DataFrame`` when a row limit was supplied + and a plain list of mappings when one was not, so both shapes are handled + rather than relying on the caller always producing the first. Assuming the + frame is how this tool previously raised ``AttributeError`` on a result it + had asked for perfectly legitimately. + """ + if hasattr(data, "columns") and hasattr(data, "to_dict"): + columns = [ + {"name": str(name), "type": str(data[name].dtype)} for name in data.columns + ] + return columns, list(data.to_dict(orient="records")) + + records = [dict(row) for row in (data or [])] + names: list[str] = [] + for record in records: + for key in record: + if key not in names: + names.append(key) + # Without a frame there are no dtypes to report; the values themselves still + # carry their types through ``_json_safe``. + return [{"name": str(name), "type": "unknown"} for name in names], records + + +def _json_safe(value: Any) -> Any: + """ + Coerce one warehouse value into something JSON can carry. + + ``Decimal`` becomes a float and binary becomes text (or hex when it is not + text at all); everything else exotic — dates, intervals, UUIDs, driver + types — becomes its string form. Lossy by design: the model reads these, it + does not compute on them. + """ + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, Decimal): + return float(value) + if isinstance(value, (bytes, memoryview)): + raw = bytes(value) + try: + return raw.decode("utf-8") + except UnicodeDecodeError: + return raw.hex() + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + return str(value) + + +def _execute_via_database( + database: Any, + sql: str, + catalog: str | None, + schema: str | None, + limit: int, +) -> Any: + """ + The single point at which a warehouse is touched. + + Isolated as a module-level function so a test can substitute it and exercise + every guard above without a live connection. ``Database.execute`` is the + same entry point SQL Lab and the MCP service use, so this inherits Jinja + rendering, ``SQL_QUERY_MUTATOR``, disallowed-function and disallowed-table + checks, row-level security, and the executor's own ``allow_dml`` gate. + """ + from superset_core.queries.types import QueryOptions + + return database.execute( + sql, + QueryOptions(catalog=catalog, schema=schema, limit=limit), + ) + + +def _result_to_payload(result: Any, limit: int) -> dict[str, Any]: + """ + Flatten a ``QueryResult`` into rows the model can read. + + Only the last data-bearing statement is returned. A read-only script with + several ``SELECT``s is unusual, and returning every result set multiplies + the context cost for a case the model can trivially split into two calls. + """ + from superset_core.queries.types import QueryStatus + + if result.status != QueryStatus.SUCCESS: + raise ToolError( + f"The query did not complete: {result.error_message or result.status}" + ) + + statement = next( + (item for item in reversed(result.statements) if item.data is not None), + None, + ) + if statement is None: + return { + "rows": [], + "row_count": 0, + "columns": [], + "note": "No rows returned.", + "executed_sql": None, + } + + columns, records = _columns_and_records(statement.data) + rows = [ + {str(key): _json_safe(value) for key, value in record.items()} + for record in records[:limit] + ] + + payload: dict[str, Any] = { + "rows": rows, + "row_count": len(rows), + "columns": columns, + # The SQL the warehouse actually ran, after the limit and any row-level + # security rewrite. This is what the user needs to see to trust the + # answer, and what they would paste into SQL Lab to check it. + "executed_sql": getattr(statement, "executed_sql", None), + } + if len(records) > limit: + payload["truncated"] = True + payload["note"] = ( + f"Showing {limit} of {len(records)} rows. Add a tighter WHERE clause " + f"or aggregate to see the rest." + ) + return payload + + +def _sql_display( + database: Any, + payload: dict[str, Any], + duration_ms: int, +) -> dict[str, Any]: + """ + Build the UI summary for one query. + + Deliberately not the model-facing payload: the row sample is smaller, and + only the connection's name and id appear — never its URI or credentials. + """ + executed = payload.get("executed_sql") or "" + return { + "kind": "sql_result", + "database_id": database.id, + "database_name": database.database_name, + "executed_sql": str(executed)[:DISPLAY_SQL_CHARS], + "executed_sql_truncated": len(str(executed)) > DISPLAY_SQL_CHARS, + "columns": [column["name"] for column in payload.get("columns", [])], + "rows": payload.get("rows", [])[:DISPLAY_SAMPLE_ROWS], + "row_count": payload.get("row_count", 0), + "sample_only": len(payload.get("rows", [])) > DISPLAY_SAMPLE_ROWS, + "truncated": bool(payload.get("truncated", False)), + "duration_ms": duration_ms, + } + + +class ExecuteSqlTool(AITool): + """Run a read-only query and return its rows.""" + + name: ClassVar[str] = "execute_sql" + description: ClassVar[str] = ( + "Run a read-only SQL query against a Superset database connection and " + "return the rows. Only SELECT-style statements are accepted; anything " + "that writes data or changes schema is refused. Results are capped, so " + "aggregate or filter in SQL rather than asking for everything. Call " + "list_databases for a database_id and get_schema for table and column " + "names before writing the query." + ) + input_schema: ClassVar[dict[str, Any]] = { + "type": "object", + "properties": { + "database_id": { + "type": "integer", + "description": "Database connection to query, from list_databases.", + }, + "sql": { + "type": "string", + "description": "A single read-only SQL statement.", + }, + "schema": { + "type": "string", + "description": ( + "Schema unqualified table names resolve to. Optional; the " + "database's default is used when omitted." + ), + }, + "catalog": { + "type": "string", + "description": "Catalog to query, for engines that have them.", + }, + "limit": { + "type": "integer", + "description": ( + "Maximum rows to return. Narrows the default only; it " + "cannot raise the configured ceiling." + ), + }, + }, + "required": ["database_id", "sql"], + } + + def __init__( + self, + executor: Callable[[Any, str, str | None, str | None, int], Any] | None = None, + ) -> None: + # Injectable so the guards are testable without a warehouse. + self._executor = executor or _execute_via_database + + def run( # pylint: disable=too-many-arguments + self, + database_id: Any = None, + sql: Any = None, + schema: Any = None, + catalog: Any = None, + limit: Any = None, + **_ignored: Any, + ) -> ToolOutput: + database = _database_or_refuse(database_id) + script = _parse_or_refuse(sql, database) + _assert_read_only(script, database) + + schema_name = str(schema) if schema else None + catalog_name = str(catalog) if catalog else None + _raise_for_sql_access(database, sql, catalog_name, schema_name) Review Comment: This lets an Alpha user who has access to a registered dataset execute SQL through the assistant even without the additive SQL Lab permission: the built-in profiles expose `execute_sql`, while this path checks dataset access but never `can_execute_sql_query` on `SQLLab`. Could the tool require that permission before calling the database executor? ########## 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) Review Comment: When a worker-mode client disconnects before opening the stream, no consumer calls `close()`, so this Redis stream never receives the configured TTL and can retain up to 10,000 events indefinitely. Could the producer set or refresh the expiry when publishing (or after its terminal event)? ########## superset-frontend/src/features/ai/components/ChatTabsMenu.tsx: ########## @@ -0,0 +1,356 @@ +/** + * 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. + */ + +/** + * @fileoverview The conversation list. + * + * Conversations live behind one menu rather than a tab strip: the panel is narrow + * enough in floating mode that a strip would truncate every name, and the list + * doubles as the history of past conversations, which a strip cannot be. + */ + +import { useCallback, useState } from 'react'; +import type { MouseEvent as ReactMouseEvent } from 'react'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Button, Dropdown, Popconfirm } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import type { ChatTab } from '../types'; + +const MenuContainer = styled.div` + background: ${({ theme }) => theme.colorBgElevated}; + border-radius: ${({ theme }) => theme.borderRadius}px; + box-shadow: ${({ theme }) => theme.boxShadowSecondary}; + min-width: ${({ theme }) => theme.sizeUnit * 65}px; + max-height: ${({ theme }) => theme.sizeUnit * 100}px; + overflow-y: auto; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; +`; + +const MenuHeader = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + text-transform: uppercase; + letter-spacing: 0.5px; +`; + +const NewChatButton = styled.button` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + width: 100%; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + color: ${({ theme }) => theme.colorPrimary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + background: none; + border: none; + text-align: left; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const TabItem = styled.div<{ isActive: boolean }>` + display: flex; + align-items: center; + justify-content: space-between; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + background: ${({ theme, isActive }) => + isActive ? theme.colorFillSecondary : 'transparent'}; + border-left: 3px solid + ${({ theme, isActive }) => (isActive ? theme.colorPrimary : 'transparent')}; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + + .action-btn { + opacity: 1; + } + } +`; + +const TabInfo = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex: 1; + overflow: hidden; +`; + +const TabName = styled.span` + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + max-width: ${({ theme }) => theme.sizeUnit * 35}px; +`; + +const TabNameInput = styled.input` + width: 100%; + max-width: ${({ theme }) => theme.sizeUnit * 40}px; + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + background: ${({ theme }) => theme.colorBgContainer}; + border: 1px solid ${({ theme }) => theme.colorBorder}; + border-radius: ${({ theme }) => theme.borderRadius}px; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const TabTimestamp = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextQuaternary}; + white-space: nowrap; + flex-shrink: 0; +`; + +const ActionButtons = styled.div` + display: flex; + align-items: center; + gap: 2px; +`; + +const ActionButton = styled.button` + background: none; + border: none; + padding: ${({ theme }) => theme.sizeUnit}px; + cursor: pointer; + color: ${({ theme }) => theme.colorTextSecondary}; + opacity: 0; + transition: all ${({ theme }) => theme.motionDurationMid}; + display: flex; + align-items: center; + justify-content: center; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover, + &:focus-visible { + opacity: 1; + color: ${({ theme }) => theme.colorError}; + background: ${({ theme }) => theme.colorErrorBg}; + } +`; + +const Divider = styled.div` + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + margin: ${({ theme }) => theme.sizeUnit}px 0; +`; + +const EmptyState = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 5}px + ${({ theme }) => theme.sizeUnit * 4}px; + text-align: center; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const MINUTE_SECONDS = 60; +const HOUR_MINUTES = 60; +const DAY_HOURS = 24; +const WEEK_DAYS = 7; + +export const formatRelativeTime = (timestamp: number): string => { + const seconds = Math.floor((Date.now() - timestamp) / 1000); + if (seconds < MINUTE_SECONDS) { + return t('just now'); + } + const minutes = Math.floor(seconds / MINUTE_SECONDS); + if (minutes < HOUR_MINUTES) { + return t('%sm', String(minutes)); + } + const hours = Math.floor(minutes / HOUR_MINUTES); + if (hours < DAY_HOURS) { + return t('%sh', String(hours)); + } + const days = Math.floor(hours / DAY_HOURS); + if (days < WEEK_DAYS) { + return t('%sd', String(days)); + } + return new Date(timestamp).toLocaleDateString(undefined, { + month: 'short', + day: 'numeric', + }); +}; + +interface ChatTabsMenuProps { + tabs: ChatTab[]; + activeTabId: string; + onSelectTab: (tabId: string) => void; + onNewChat: () => void; + onDeleteTab: (tabId: string) => void; + onRenameTab: (tabId: string, name: string) => void; +} + +export const ChatTabsMenu = ({ + tabs, + activeTabId, + onSelectTab, + onNewChat, + onDeleteTab, + onRenameTab, +}: ChatTabsMenuProps) => { + const [editingTabId, setEditingTabId] = useState<string | null>(null); + const [editingName, setEditingName] = useState(''); + + const startEditing = useCallback((event: ReactMouseEvent, tab: ChatTab) => { + event.stopPropagation(); + setEditingTabId(tab.id); + setEditingName(tab.name); + }, []); + + const cancelEditing = useCallback(() => { + setEditingTabId(null); + setEditingName(''); + }, []); + + const commitRename = useCallback( + (tabId: string) => { + const trimmedName = editingName.trim(); + if (trimmedName) { + onRenameTab(tabId, trimmedName); + } + cancelEditing(); + }, + [cancelEditing, editingName, onRenameTab], + ); + + const menuContent = ( + <MenuContainer data-test="chat-tabs-menu"> + <MenuHeader>{t('Conversations')}</MenuHeader> + <NewChatButton type="button" onClick={onNewChat}> + <Icons.PlusOutlined iconSize="s" /> + <span>{t('New Chat')}</span> + </NewChatButton> + <Divider /> + {tabs.length === 0 ? ( + <EmptyState>{t('No conversations yet')}</EmptyState> + ) : ( + tabs.map(tab => ( + <TabItem + key={tab.id} + isActive={tab.id === activeTabId} + onClick={() => onSelectTab(tab.id)} + > + <TabInfo> + <Icons.MessageOutlined iconSize="s" /> + {editingTabId === tab.id ? ( + <TabNameInput + autoFocus + value={editingName} + onChange={event => setEditingName(event.target.value)} + onClick={event => event.stopPropagation()} + onBlur={() => commitRename(tab.id)} + onKeyDown={event => { + event.stopPropagation(); + if (event.key === 'Enter') { + commitRename(tab.id); + } else if (event.key === 'Escape') { + cancelEditing(); + } + }} + aria-label={t('Conversation name')} + /> + ) : ( + <TabName>{tab.name}</TabName> + )} + {tab.updatedAt !== undefined && ( + <TabTimestamp>{formatRelativeTime(tab.updatedAt)}</TabTimestamp> + )} + </TabInfo> + <ActionButtons> + <ActionButton + type="button" + className="action-btn" + onClick={event => startEditing(event, tab)} + title={t('Rename conversation')} + aria-label={t('Rename conversation')} + > + <Icons.EditOutlined iconSize="s" /> + </ActionButton> + {/* A conversation with messages is confirmed before deletion; an + empty one is discarded without a prompt. */} + {tab.messages.length > 0 ? ( Review Comment: Unopened threads intentionally have `messages: []` even when they contain persisted messages, so deleting one follows this empty-thread branch and skips confirmation. Could the tab retain `message_count` from the list API and require confirmation whenever it is nonzero or unknown? ########## 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( Review Comment: The user and pending assistant rows have already been committed when broker submission fails here, so an unavailable Celery broker leaves a permanent empty pending message with no worker that can finalize it. Could this failure atomically terminalize/remove those accepted rows, or use a durable outbox boundary? -- 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]
