/** * Copyright (c) 2026-present, Goldman Sachs * * Licensed 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. */ import { useMemo, useCallback, useState, useRef, useEffect } from 'react'; import { SparkleStarsIcon, CodeIcon, TableIcon, CopyIcon, CheckIcon, TimesIcon, MinusIcon, PlusIcon, CaretDownIcon, CaretRightIcon, DotIcon, LoadingIcon, LikeIcon, DislikeIcon, ExternalLinkIcon, MarkdownTextViewer, } from '@finos/legend-art'; import { noop } from '@finos/legend-shared'; import { type LegendAIChatProps, type LegendAIAssistantMessage, type LegendAIMessageFeedback, type LegendAIThinkingStep, type LegendAIScopeItem, type LegendAIQuestionIntent, LegendAIMessageFeedbackRating, LegendAIThinkingStepStatus, LegendAIMessageRole, LegendAIErrorType, classifyQuestionIntentFast, } from '../LegendAITypes.js'; import { useLegendAIChatState } from '../stores/LegendAIChatState.js'; import { LegendAIResultGrid } from './LegendAIResultGrid.js'; import { LegendAIAnalysisPanel } from './LegendAIAnalysisPanel.js'; import { LegendAIChatInput } from './LegendAIChatInput.js'; import { buildSuggestedQueries } from './LegendAIChatHelpers.js'; export const LEGEND_AI_ANCHOR_ID = 'legend-ai-anchor'; const COPY_FEEDBACK_DURATION_MS = 2000; const METADATA_CONTEXT_HEADING = '### Metadata context'; const QUERY_ANALYSIS_HEADING = '### Query analysis'; function toUserFacingThinkingLabel(label: string): string { const normalized = label.toLowerCase(); if ( normalized.includes('analyzing your question') || normalized.includes('intent is ambiguous') ) { return 'Understanding your request'; } if ( normalized.includes('building metadata context') || normalized.includes('answering from product metadata') ) { return 'Checking product capabilities and services'; } if ( normalized.includes('found relevant services') || normalized.includes('selecting best service') || normalized.includes('building context from service schemas') || normalized.includes('preparing data query') || normalized.includes('generating sql query') || normalized.includes('verifying query correctness') || normalized.includes('query corrected') || normalized.includes('max verification attempts reached') || normalized.includes('judge approved a non-sql draft') ) { return 'Trying a data query when helpful'; } if ( normalized.includes('retrieved ') || normalized.includes('executing') || normalized.includes('analyzing results') || normalized.includes('verifying answer coverage') ) { return 'Summarizing what matters for your question'; } if (normalized.includes('error')) { return 'Hit an issue while preparing the answer'; } return label; } function formatThinkingSteps( thinkingSteps: LegendAIThinkingStep[], ): LegendAIThinkingStep[] { const formatted: LegendAIThinkingStep[] = []; for (const step of thinkingSteps) { const userLabel = toUserFacingThinkingLabel(step.label); const last = formatted[formatted.length - 1]; if (last?.label === userLabel) { formatted[formatted.length - 1] = { ...last, status: step.status, }; } else { formatted.push({ ...step, label: userLabel, }); } } return formatted; } function splitCombinedAnswer(textAnswer: string | null): { metadataContext: string | null; queryAnalysis: string | null; } { if (!textAnswer) { return { metadataContext: null, queryAnalysis: null }; } const metadataIndex = textAnswer.indexOf(METADATA_CONTEXT_HEADING); if (metadataIndex < 0) { return { metadataContext: null, queryAnalysis: textAnswer }; } const metadataStart = metadataIndex + METADATA_CONTEXT_HEADING.length; const queryIndex = textAnswer.indexOf(QUERY_ANALYSIS_HEADING, metadataStart); const metadataContext = queryIndex >= 0 ? textAnswer.slice(metadataStart, queryIndex).trim() : textAnswer.slice(metadataStart).trim(); const queryAnalysis = queryIndex >= 0 ? textAnswer.slice(queryIndex + QUERY_ANALYSIS_HEADING.length).trim() || null : null; return { metadataContext: metadataContext.length > 0 ? metadataContext : null, queryAnalysis, }; } const AISummaryRenderer = ({ value }: { value: string }): React.ReactNode => ( ); const DEFAULT_SCOPES: LegendAIScopeItem[] = [ { id: 'legend-ai-mcp', label: 'Legend AI MCP', description: 'Model Context Protocol via Marketplace /mcp proxy', }, ]; export function renderStepStatusIcon( status: LegendAIThinkingStepStatus, ): React.ReactNode { if (status === LegendAIThinkingStepStatus.ACTIVE) { return ; } return status === LegendAIThinkingStepStatus.DONE ? ( ) : ( ); } const AssistantMessageView = (props: { msg: LegendAIAssistantMessage; questionText: string; isThinkingVisible: boolean; onToggleThinking: () => void; onMessageFeedback?: ( feedback: LegendAIMessageFeedback, ) => Promise | void; selectedFeedbackRating: LegendAIMessageFeedbackRating | undefined; feedbackSubmitting: boolean; onSuggestedQueryClick?: (query: string) => void; onFallbackAction?: (messageId: string) => void; enghubDocUrl?: string; enthubRequestAccessUrl?: string; }): React.ReactNode => { const { msg, questionText, isThinkingVisible, onToggleThinking, onMessageFeedback, selectedFeedbackRating, feedbackSubmitting, onSuggestedQueryClick, onFallbackAction, enghubDocUrl, enthubRequestAccessUrl, } = props; const hasPermissionAccessLinks = enghubDocUrl !== undefined || enthubRequestAccessUrl !== undefined; const [sqlCopied, setSqlCopied] = useState(false); const copyTimerRef = useRef | undefined>( undefined, ); useEffect( () => () => { if (copyTimerRef.current !== undefined) { clearTimeout(copyTimerRef.current); } }, [], ); const handleCopySql = useCallback(() => { if (msg.sql) { navigator.clipboard.writeText(msg.sql).catch(noop); setSqlCopied(true); if (copyTimerRef.current !== undefined) { clearTimeout(copyTimerRef.current); } copyTimerRef.current = setTimeout(() => { setSqlCopied(false); copyTimerRef.current = undefined; }, COPY_FEEDBACK_DURATION_MS); } }, [msg.sql]); const canShowFeedback = !msg.isProcessing && (msg.textAnswer !== null || msg.gridData !== null || msg.error !== null); const visibleThinkingSteps = formatThinkingSteps(msg.thinkingSteps); const { metadataContext, queryAnalysis } = splitCombinedAnswer( msg.textAnswer, ); const analysisSummary = (() => { if (msg.gridData === null) { return null; } return queryAnalysis ?? (metadataContext === null ? msg.textAnswer : null); })(); const plainAnswer = msg.gridData === null ? (metadataContext ?? msg.textAnswer) : null; const submitFeedback = useCallback( (rating: LegendAIMessageFeedbackRating): void => { const result = onMessageFeedback?.({ messageId: msg.id, rating, question: questionText, ...(msg.textAnswer === null ? {} : { answer: msg.textAnswer }), ...(msg.sql === null ? {} : { sql: msg.sql }), ...(msg.gridData === null ? {} : { rowCount: msg.gridData.rowData.length }), }); if (result instanceof Promise) { result.catch(noop); } }, [msg, onMessageFeedback, questionText], ); return (
{visibleThinkingSteps.length > 0 && (
{!msg.isProcessing && ( )} {isThinkingVisible && (
{visibleThinkingSteps.map((step) => (
{renderStepStatusIcon(step.status)} {step.label}
))}
)}
)} {metadataContext && msg.gridData && (
)} {msg.sql && (
Generated SQL {msg.sqlGenTime && ( {msg.sqlGenTime}s )}
{msg.sql}
)} {msg.isExecuting && (
Executing query...
)} {msg.error && (
{msg.error} {msg.errorType === LegendAIErrorType.PERMISSION && hasPermissionAccessLinks && (
Need access?
{enghubDocUrl && ( View Documentation )} {enthubRequestAccessUrl && ( Request Access )}
)}
)} {msg.gridData && (
Results {msg.gridData.rowData.length} row {msg.gridData.rowData.length === 1 ? '' : 's'} {msg.execTime ? ( <> {' '} {' '} {msg.execTime}s ) : ( '' )}
)} {plainAnswer && (
)} {analysisSummary && msg.gridData && ( )} {msg.isProcessing && !msg.isExecuting && msg.gridData && (
Analyzing results...
)} {!msg.isProcessing && msg.suggestedQueries.length > 0 && onSuggestedQueryClick && (
Try a data query: {msg.suggestedQueries.map((q) => ( ))}
)} {msg.fallbackAction && !msg.isProcessing && onFallbackAction && ( )} {canShowFeedback && (
Did this answer your question?
)}
); }; export const LegendAIChat = (props: LegendAIChatProps): React.ReactNode => { const { services, coordinates, config, metadata, title, plugin, dataProductCoordinates, pureExecutionContext, availableScopes, onMessageFeedback, onClose, onMinimize, } = props; const state = useLegendAIChatState( services, coordinates, config, metadata, plugin, dataProductCoordinates, pureExecutionContext, ); const suggestedQueries = useMemo( () => buildSuggestedQueries(services, metadata), [services, metadata], ); const hasServices = services.length > 0; const inferSuggestedQueryIntent = useCallback( (query: string): LegendAIQuestionIntent => classifyQuestionIntentFast(query, hasServices).intent, [hasServices], ); const hasMessages = state.messages.length > 0; const scopes = availableScopes ?? DEFAULT_SCOPES; const [feedbackByMessageId, setFeedbackByMessageId] = useState< Map >(new Map()); const [pendingFeedbackByMessageId, setPendingFeedbackByMessageId] = useState< Set >(new Set()); const handleMessageFeedback = useCallback( async (feedback: LegendAIMessageFeedback): Promise => { setFeedbackByMessageId((prev) => { const next = new Map(prev); next.set(feedback.messageId, feedback.rating); return next; }); if (!onMessageFeedback) { return; } setPendingFeedbackByMessageId((prev) => { const next = new Set(prev); next.add(feedback.messageId); return next; }); try { await onMessageFeedback(feedback); } catch { setFeedbackByMessageId((prev) => { const next = new Map(prev); next.delete(feedback.messageId); return next; }); } finally { setPendingFeedbackByMessageId((prev) => { const next = new Set(prev); next.delete(feedback.messageId); return next; }); } }, [onMessageFeedback], ); return (
{title ?? 'Legend AI'}
{onMinimize && ( )} {onClose && ( )}
{!hasMessages && (
Ask a question about your data
{suggestedQueries.map((q) => ( ))}
)} {state.messages.map((msg, msgIndex) => { if (msg.role === LegendAIMessageRole.USER) { return (
{msg.text}
); } const isThinkingVisible = msg.isProcessing || state.expandedThinking.has(msgIndex); const previousMessage = msgIndex > 0 ? state.messages[msgIndex - 1] : null; const questionText = previousMessage?.role === LegendAIMessageRole.USER ? previousMessage.text : ''; return ( state.toggleThinking(msgIndex)} onMessageFeedback={handleMessageFeedback} selectedFeedbackRating={feedbackByMessageId.get(msg.id)} feedbackSubmitting={pendingFeedbackByMessageId.has(msg.id)} {...(config.enghubDocUrl === undefined ? {} : { enghubDocUrl: config.enghubDocUrl })} {...(config.enthubRequestAccessUrl === undefined ? {} : { enthubRequestAccessUrl: config.enthubRequestAccessUrl })} onFallbackAction={(messageId): void => state.runFallbackAction(messageId) } onSuggestedQueryClick={(q): void => state.askQuestionWithIntent(q, inferSuggestedQueryIntent(q)) } /> ); })}
); };