I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r4002577968
########## superset-frontend/src/features/ai/components/ChatChartEmbed.tsx: ########## @@ -0,0 +1,542 @@ +/** + * 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 A chart rendered inside a chat message. + * + * The assistant does not send a chart, it sends a `form_data_key` it stored — so + * what arrives in the transcript is a reference the client resolves, and the + * rendered chart is the real thing, with the real permissions, rather than an + * image of one. + * + * The awkward part is timing: the key can exist before the query behind it has + * finished. Rather than show the chart's own "No data" state (which reads as a + * broken answer) the component keeps a spinner up and re-renders on a growing + * backoff until rows appear, giving up after a bounded number of attempts. + */ + +import { useCallback, useEffect, useRef, useState } from 'react'; +import { + type QueryFormData, + StatefulChart, + SupersetClient, +} from '@superset-ui/core'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Loading } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { ErrorBoundary } from 'src/components/ErrorBoundary'; + +const VALID_KEY_PATTERN = /^[a-zA-Z0-9_-]+$/; +const MIN_HEIGHT = 100; +const MAX_HEIGHT = 800; +const DEFAULT_HEIGHT = 300; +const FETCH_TIMEOUT_MS = 30_000; +const MAX_RETRIES = 3; +const RETRY_DELAYS_MS = [500, 1500, 3000]; + +/** Width used until the container has been measured. */ +const FALLBACK_CHART_WIDTH = 600; + +// Backoff for the "waiting for chart data" poll. The delay grows exponentially +// per attempt up to a ceiling, so a slow query is waited out without hammering +// the backend. +const POLL_BASE_DELAY_MS = 1000; +const POLL_MAX_DELAY_MS = 30_000; + +/** + * Polling stops after this many attempts. + * + * A retry re-issues the chart's data request, which will use the results cache if + * the query has landed but will otherwise execute it. Polling forever would keep + * re-issuing it, so an unfinished query surfaces the retry control instead. + */ +const MAX_POLL_ATTEMPTS = 6; + +const getPollDelayMs = (attempt: number): number => + Math.min(POLL_BASE_DELAY_MS * 2 ** attempt, POLL_MAX_DELAY_MS); + +export interface ChartEmbedParams { + formDataKey: string | null; + height: number; + title: string | null; +} + +/** + * Parse key=value lines from the content of a ```superset-chart fenced block. + * + * Rules are strict on purpose: the block is model output, so `form_data_key` must + * match `/^[a-zA-Z0-9_-]+$/` before it reaches a URL, and `height` is clamped. + * Unknown keys are ignored so a newer backend can add some without breaking an + * older client. + */ +export function parseChartEmbedParams(codeText: string): ChartEmbedParams { + const result: ChartEmbedParams = { + formDataKey: null, + height: DEFAULT_HEIGHT, + title: null, + }; + + const lines = codeText + .split('\n') + .map(line => line.trim()) + .filter(Boolean); + + lines.forEach(line => { + const eqIndex = line.indexOf('='); + if (eqIndex <= 0) { + return; + } + + const key = line.slice(0, eqIndex).trim().toLowerCase(); + const value = line.slice(eqIndex + 1).trim(); + + if (key === 'form_data_key') { + if (value && VALID_KEY_PATTERN.test(value)) { + result.formDataKey = value; + } + return; + } + if (key === 'height') { + const parsed = parseInt(value, 10); + if (!Number.isNaN(parsed)) { + result.height = Math.max(MIN_HEIGHT, Math.min(MAX_HEIGHT, parsed)); + } + return; + } + if (key === 'title' && value) { + result.title = value; + } + }); + + return result; +} + +const ChartContainer = styled.div` + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow: hidden; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + background: ${({ theme }) => theme.colorBgContainer}; +`; + +const ChartHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 2}px + ${({ theme }) => theme.sizeUnit * 3}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgLayout}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const ChartTitle = styled.span` + font-weight: ${({ theme }) => theme.fontWeightStrong}; + color: ${({ theme }) => theme.colorText}; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + flex: 1; + min-width: 0; +`; + +const ChartActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: center; + flex-shrink: 0; + margin-left: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ActionLink = styled.a` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorPrimary}; + cursor: pointer; + text-decoration: none; + + &:hover { + text-decoration: underline; + } +`; + +const ActionButton = styled.button` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + cursor: pointer; + background: none; + border: none; + padding: 2px ${({ theme }) => theme.sizeUnit / 2}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover { + color: ${({ theme }) => theme.colorPrimary}; + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const ChartBody = styled.div<{ height: number }>` + height: ${({ height }) => height}px; + position: relative; +`; + +// Covers the chart while it reports no data, hiding the underlying "No data" +// state (which looks broken) behind a spinner while refreshing continues. +const ChartDataOverlay = styled.div` + position: absolute; + inset: 0; + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + background: ${({ theme }) => theme.colorBgContainer}; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + z-index: 2; +`; + +const CenteredMessage = styled.div<{ height: number }>` + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + height: ${({ height }) => height}px; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + text-align: center; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +interface ChatChartEmbedProps { + formDataKey: string; + height?: number; + title?: string; +} + +type FetchState = + | { status: 'loading' } + | { status: 'loaded'; formData: QueryFormData } + | { status: 'error'; message: string }; + +const exploreUrlFor = (formDataKey: string): string => + `/explore/?form_data_key=${encodeURIComponent(formDataKey)}`; + +export function ChatChartEmbedInner({ + formDataKey, + height = DEFAULT_HEIGHT, + title, +}: ChatChartEmbedProps) { + const [fetchState, setFetchState] = useState<FetchState>({ + status: 'loading', + }); + const [chartWidth, setChartWidth] = useState(0); + const [chartRenderKey, setChartRenderKey] = useState(0); + // True while the chart's query results are still unavailable (either an empty + // result or a cache miss), which is what keeps the overlay up. + const [isAwaitingData, setIsAwaitingData] = useState(false); + const chartBodyRef = useRef<HTMLDivElement>(null); + const pollTimeoutRef = useRef<ReturnType<typeof setTimeout>>(); + + // Bumping `chartRenderKey` remounts the chart, which re-issues its data + // request. `force` is left off so the results cache is preferred. + const scheduleNextPoll = useCallback(() => { + if (chartRenderKey >= MAX_POLL_ATTEMPTS) { + setIsAwaitingData(false); + return; + } + setIsAwaitingData(true); + const delay = getPollDelayMs(chartRenderKey); + if (pollTimeoutRef.current) { + clearTimeout(pollTimeoutRef.current); + } + pollTimeoutRef.current = setTimeout( + () => setChartRenderKey(key => key + 1), + delay, + ); + }, [chartRenderKey]); + + useEffect( + () => () => { + if (pollTimeoutRef.current) { + clearTimeout(pollTimeoutRef.current); + } + }, + [], + ); + + useEffect(() => { + const element = chartBodyRef.current; + if (!element) { + return undefined; + } + + const initialWidth = Math.floor(element.getBoundingClientRect().width); + if (initialWidth > 0) { + setChartWidth(initialWidth); + } + + // The panel is resizable by the host, so the chart is measured rather than + // given a fixed width. + const observer = new ResizeObserver(entries => { + const [entry] = entries; + if (entry) { + const width = Math.floor(entry.contentRect.width); + if (width > 0) { + setChartWidth(width); + } + } + }); + observer.observe(element); + return () => observer.disconnect(); + }, [fetchState.status]); + + const exploreUrl = exploreUrlFor(formDataKey); + + const fetchFormData = useCallback( + async (attempt: number = 0) => { + setFetchState({ status: 'loading' }); + setIsAwaitingData(false); + setChartRenderKey(0); + + const retry = (): boolean => { + if (attempt >= MAX_RETRIES) { + return false; + } + const delay = RETRY_DELAYS_MS[attempt] ?? 1000; + setTimeout(() => { Review Comment: The embed retry timer and request cleanup are still an open base UI issue. This update does not cancel the stale callback; the unmount-before-backoff fake-timer test you describe is the right regression for that path. Leaving it open. ########## superset/models/ai.py: ########## @@ -0,0 +1,270 @@ +# 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. +""" +Persistence for AI assistant conversations. + +Conversations live in Superset's own metadata database, which keeps the +feature deployable with no extra infrastructure and makes ownership +enforceable with the same DAO filters used everywhere else. + +These models live under ``superset/models/`` rather than inside +``superset/ai/`` because background workers need them without importing the +API module. +""" + +from __future__ import annotations + +import uuid as uuid_module +from typing import Any + +import sqlalchemy as sa +from flask_appbuilder import Model +from sqlalchemy.orm import relationship, validates +from sqlalchemy_utils import UUIDType + +from superset.ai.types import ( + MessageExtra, + MessageRole, + MessageStatus, + ThreadStatus, +) +from superset.models.helpers import AuditMixinNullable +from superset.utils import json +from superset.utils.core import MediumText + +#: Bumped when the meaning of keys inside an ``extra_json`` blob changes, so a +#: reader can tell an old row from a new one instead of guessing. +EXTRA_JSON_VERSION = 1 + + +class AIChatThread(AuditMixinNullable, Model): + """ + One conversation between a user and the assistant. + + Ownership is expressed through ``created_by_fk`` (supplied by + :class:`AuditMixinNullable`) and enforced by the DAO's base filter, so a + thread identifier is not by itself a capability. + """ + + __tablename__ = "ai_chat_threads" + __table_args__ = ( + # Serves the "my threads, most recent first" list query. + sa.Index("ix_ai_chat_threads_owner_recent", "created_by_fk", "changed_on"), + ) + + id = sa.Column(sa.Integer, primary_key=True) + #: The only identifier exposed over HTTP. Integer ids stay internal. + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + title = sa.Column(sa.String(512), nullable=True) + status = sa.Column( + sa.String(32), + nullable=False, + default=ThreadStatus.ACTIVE.value, + server_default=ThreadStatus.ACTIVE.value, + ) + #: Which agent profile this thread was last run with. + agent_key = sa.Column(sa.String(64), nullable=True) + extra_json = sa.Column(MediumText(), nullable=True) + + # No ``passive_deletes``: SQLite does not enforce foreign keys unless + # ``PRAGMA foreign_keys=ON``, so deferring the cascade to the database + # would orphan messages there. The ORM deletes children itself, and the + # ``ON DELETE CASCADE`` in the DDL remains a backstop for direct SQL. + messages = relationship( + "AIChatMessage", + back_populates="thread", + cascade="all, delete-orphan", + order_by="AIChatMessage.created_on", + ) + + def __repr__(self) -> str: + return f"<AIChatThread {self.uuid} [{self.status}]>" + + @validates("status") + def _validate_status(self, _key: str, value: Any) -> str: + """Reject unknown lifecycle values at assignment time.""" + return ThreadStatus(value).value + + @property + def message_count(self) -> int: + """ + Number of stored messages. + + Derived rather than denormalised: a counter column would have to be kept + in step with cascade deletes and retention pruning, and a counter that + drifts is worse than one query — it reports a conversation length nobody + can reconcile against the rows. + """ + return len(self.messages) + + @property + def extra(self) -> dict[str, Any]: + """Parsed ``extra_json``, or an empty dict when absent or corrupt.""" + return _load_json_object(self.extra_json) + + +class AIChatMessage(AuditMixinNullable, Model): + """ + A single turn in a conversation. + + Assistant messages are inserted before inference begins so a client that + reconnects mid-run has a row to attach to, then transition through + ``streaming`` to a terminal status. + """ + + __tablename__ = "ai_chat_messages" + __table_args__ = ( + sa.Index("ix_ai_chat_messages_thread_created", "thread_id", "created_on"), + # Makes client-supplied idempotency real: replaying a request cannot + # create a second row for the same turn. The constraint carries the role + # as well, because one request legitimately produces both a user message + # and the assistant message answering it. + sa.UniqueConstraint( + "thread_id", + "request_id", + "role", + name="uq_ai_chat_messages_thread_request_role", + ), + ) + + id = sa.Column(sa.Integer, primary_key=True) + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + thread_id = sa.Column( + sa.Integer, + sa.ForeignKey("ai_chat_threads.id", ondelete="CASCADE"), + nullable=False, + ) + role = sa.Column(sa.String(16), nullable=False) + #: Unbounded model or user text; deliberately not a plain ``Text`` column, + #: which caps at 64 KB on MySQL. + content = sa.Column(MediumText(), nullable=False, default="") + status = sa.Column( + sa.String(32), + nullable=False, + default=MessageStatus.COMPLETE.value, + server_default=MessageStatus.COMPLETE.value, + ) + #: Client-generated idempotency key for the turn. + request_id = sa.Column(sa.String(96), nullable=True) + #: Serialised :class:`~superset.ai.types.MessageExtra`. + extra_json = sa.Column(MediumText(), nullable=True) + + thread = relationship("AIChatThread", back_populates="messages") + + def __repr__(self) -> str: + return f"<AIChatMessage {self.uuid} {self.role} [{self.status}]>" + + @validates("role") + def _validate_role(self, _key: str, value: Any) -> str: + """Reject unknown authors at assignment time.""" + return MessageRole(value).value + + @validates("status") + def _validate_status(self, _key: str, value: Any) -> str: + """Reject unknown lifecycle values at assignment time.""" + return MessageStatus(value).value + + @property + def is_terminal(self) -> bool: + """Whether this message will never change again.""" + return MessageStatus(self.status) in MessageStatus.terminal() + + @property + def extra(self) -> MessageExtra: + """Parsed ``extra_json``, or an empty dict when absent or corrupt.""" + return _load_json_object(self.extra_json) # type: ignore[return-value] + + def update_extra(self, updates: MessageExtra) -> None: + """ + Merge keys into ``extra_json``. + + Merge rather than replace, because a run writes tool calls and token + usage at different moments. + """ + merged: dict[str, Any] = dict(self.extra) + merged.update(updates) + merged["version"] = EXTRA_JSON_VERSION + self.extra_json = json.dumps(merged) + + +class AIChatFeedback(AuditMixinNullable, Model): + """ + A thumbs up or down on an assistant message. + + A first-class table rather than a log line, so the signal can actually be + aggregated and joined back to the conversation that produced it. + """ + + __tablename__ = "ai_chat_feedback" + __table_args__ = ( + # One verdict per user per message; a repeat vote updates in place. + sa.UniqueConstraint( + "message_id", + "created_by_fk", + name="uq_ai_chat_feedback_message_user", + ), + ) + + id = sa.Column(sa.Integer, primary_key=True) + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + message_id = sa.Column( + sa.Integer, + sa.ForeignKey("ai_chat_messages.id", ondelete="CASCADE"), + nullable=False, + ) + liked = sa.Column(sa.Boolean, nullable=False) + comment = sa.Column(MediumText(), nullable=True) + + message = relationship("AIChatMessage") Review Comment: The feedback cascade remains unresolved in this branch. I am tracking the matching new base review at https://github.com/apache/superset/pull/42805#discussion_r4000933698 so the relationship and SQLite deletion regression can be fixed once in the base. -- 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]
