konnosDashboard
Message.Parts.tsx273 lines · main
1import { UIMessage as VercelMessage } from '@ai-sdk/react'
2import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai'
3import { BrainIcon, CheckIcon, Loader2 } from 'lucide-react'
4import { cn } from 'ui'
5
6import { DisplayBlockRenderer } from './DisplayBlockRenderer'
7import { EdgeFunctionRenderer } from './EdgeFunctionRenderer'
8import { Tool } from './elements/Tool'
9import { useMessageActionsContext, useMessageInfoContext } from './Message.Context'
10import {
11 deployEdgeFunctionInputSchema,
12 deployEdgeFunctionOutputSchema,
13 parseExecuteSqlChartResult,
14} from './Message.utils'
15import { MessageMarkdown } from './MessageMarkdown'
16
17function MessagePartText({ textPart }: { textPart: TextUIPart }) {
18 const { id, isLoading, readOnly, isUserMessage, state } = useMessageInfoContext()
19
20 return (
21 <MessageMarkdown
22 id={id}
23 isLoading={isLoading}
24 readOnly={readOnly}
25 className={cn(
26 'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h2:font-medium prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 wrap-break-word [&>p:not(:last-child)]:mb-2! [&>*>p:first-child]:mt-0! [&>*>p:last-child]:mb-0! [&>*>*>p:first-child]:mt-0! [&>*>*>p:last-child]:mb-0! [&>ol>li]:pl-4!',
27 isUserMessage && 'text-foreground [&>p]:font-medium',
28 state === 'editing' && 'animate-pulse'
29 )}
30 >
31 {textPart.text}
32 </MessageMarkdown>
33 )
34}
35
36function MessagePartDynamicTool({ toolPart }: { toolPart: DynamicToolUIPart }) {
37 return (
38 <Tool
39 icon={
40 toolPart.state === 'input-streaming' ? (
41 <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
42 ) : (
43 <CheckIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
44 )
45 }
46 label={
47 <div>
48 {toolPart.state === 'input-streaming' ? 'Running ' : 'Ran '}
49 <span className="text-foreground-lighter">{`${toolPart.toolName}`}</span>
50 </div>
51 }
52 />
53 )
54}
55
56function MessagePartTool({ toolPart }: { toolPart: ToolUIPart }) {
57 return (
58 <Tool
59 icon={
60 toolPart.state === 'input-streaming' ? (
61 <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
62 ) : (
63 <CheckIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
64 )
65 }
66 label={
67 <div>
68 {toolPart.state === 'input-streaming' ? 'Running ' : 'Ran '}
69 <span className="text-foreground-lighter">{`${toolPart.type.replace('tool-', '')}`}</span>
70 </div>
71 }
72 />
73 )
74}
75
76function MessagePartReasoning({ reasoningPart }: { reasoningPart: ReasoningUIPart }) {
77 return (
78 <Tool
79 icon={
80 reasoningPart.state === 'streaming' ? (
81 <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
82 ) : (
83 <BrainIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
84 )
85 }
86 label={reasoningPart.state === 'streaming' ? 'Thinking...' : 'Reasoned'}
87 >
88 {reasoningPart.text}
89 </Tool>
90 )
91}
92
93function ToolDisplayExecuteSqlLoading({ label = 'Writing SQL...' }: { label?: string }) {
94 return (
95 <div className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2">
96 <Loader2 className="w-4 h-4 animate-spin" />
97 {label}
98 </div>
99 )
100}
101
102function ToolDisplayExecuteSqlFailure() {
103 return <div className="text-xs text-danger">Failed to execute SQL.</div>
104}
105
106function MessagePartExecuteSql({
107 toolPart,
108 isLastPart,
109}: {
110 toolPart: ToolUIPart
111 isLastPart?: boolean
112}) {
113 const { id, isLastMessage } = useMessageInfoContext()
114 const { addToolApprovalResponse } = useMessageActionsContext()
115
116 const { toolCallId, state, input, output } = toolPart
117
118 if (state === 'input-streaming') {
119 return <ToolDisplayExecuteSqlLoading />
120 }
121
122 if (state === 'output-error') {
123 return <ToolDisplayExecuteSqlFailure />
124 }
125
126 const { data: chart, success } = parseExecuteSqlChartResult(input)
127 if (!success) return null
128
129 if (
130 state === 'input-available' ||
131 state === 'approval-requested' ||
132 state === 'approval-responded' ||
133 state === 'output-denied' ||
134 state === 'output-available'
135 ) {
136 const approvalId = state === 'approval-requested' ? toolPart.approval?.id : undefined
137 return (
138 <div className="w-auto overflow-x-hidden my-4 space-y-2">
139 <DisplayBlockRenderer
140 messageId={id}
141 toolCallId={toolCallId}
142 initialArgs={{
143 sql: chart.sql,
144 label: chart.label,
145 isWriteQuery: chart.isWriteQuery,
146 view: chart.view,
147 xAxis: chart.xAxis,
148 yAxis: chart.yAxis,
149 }}
150 initialResults={output}
151 toolState={state}
152 toolApprovalRespondedApproved={toolPart.approval?.approved}
153 isLastPart={isLastPart}
154 isLastMessage={isLastMessage}
155 onApprove={
156 approvalId
157 ? () => addToolApprovalResponse?.({ id: approvalId, approved: true })
158 : undefined
159 }
160 onDeny={
161 approvalId
162 ? () => addToolApprovalResponse?.({ id: approvalId, approved: false })
163 : undefined
164 }
165 />
166 </div>
167 )
168 }
169
170 return null
171}
172
173const TOOL_DEPLOY_EDGE_FUNCTION_STATES_WITH_INPUT = new Set([
174 'input-available',
175 'approval-requested',
176 'approval-responded',
177 'output-denied',
178 'output-available',
179])
180
181function MessagePartDeployEdgeFunction({ toolPart }: { toolPart: ToolUIPart }) {
182 const { state, input, output } = toolPart
183 const { addToolApprovalResponse } = useMessageActionsContext()
184
185 if (state === 'input-streaming') {
186 return (
187 <div className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2">
188 <Loader2 className="w-4 h-4 animate-spin" />
189 Writing Edge Function...
190 </div>
191 )
192 }
193
194 if (state === 'output-error') {
195 return <p className="text-xs text-danger">Failed to deploy Edge Function.</p>
196 }
197
198 if (!TOOL_DEPLOY_EDGE_FUNCTION_STATES_WITH_INPUT.has(state)) return null
199
200 const parsedInput = deployEdgeFunctionInputSchema.safeParse(input)
201 if (!parsedInput.success) return null
202
203 const parsedOutput = deployEdgeFunctionOutputSchema.safeParse(output)
204 const isInitiallyDeployed =
205 state === 'output-available' && parsedOutput.success && parsedOutput.data.success === true
206
207 const approvalId = state === 'approval-requested' ? toolPart.approval?.id : undefined
208
209 return (
210 <EdgeFunctionRenderer
211 label={parsedInput.data.label}
212 code={parsedInput.data.code}
213 functionName={parsedInput.data.functionName}
214 showConfirmFooter={state === 'approval-requested'}
215 isDeploying={state === 'approval-responded' && toolPart.approval?.approved !== false}
216 initialIsDeployed={isInitiallyDeployed}
217 onApprove={
218 approvalId ? () => addToolApprovalResponse?.({ id: approvalId, approved: true }) : undefined
219 }
220 onDeny={
221 approvalId
222 ? () => addToolApprovalResponse?.({ id: approvalId, approved: false })
223 : undefined
224 }
225 />
226 )
227}
228
229const MessagePart = {
230 Text: MessagePartText,
231 Dynamic: MessagePartDynamicTool,
232 Tool: MessagePartTool,
233 Reasoning: MessagePartReasoning,
234 ExecuteSql: MessagePartExecuteSql,
235 DeployEdgeFunction: MessagePartDeployEdgeFunction,
236} as const
237
238export function MessagePartSwitcher({
239 part,
240 isLastPart,
241}: {
242 part: NonNullable<VercelMessage['parts']>[number]
243 isLastPart?: boolean
244}) {
245 switch (part.type) {
246 case 'dynamic-tool': {
247 return <MessagePart.Dynamic toolPart={part} />
248 }
249 case 'tool-list_policies':
250 case 'tool-search_docs':
251 case 'tool-get_active_incidents':
252 case 'tool-load_knowledge': {
253 return <MessagePart.Tool toolPart={part} />
254 }
255 case 'reasoning':
256 return <MessagePart.Reasoning reasoningPart={part} />
257 case 'text':
258 return <MessagePart.Text textPart={part} />
259
260 case 'tool-execute_sql': {
261 return <MessagePart.ExecuteSql toolPart={part} isLastPart={isLastPart} />
262 }
263 case 'tool-deploy_edge_function': {
264 return <MessagePart.DeployEdgeFunction toolPart={part} />
265 }
266
267 case 'source-url':
268 case 'source-document':
269 case 'file':
270 default:
271 return null
272 }
273}