1
1
import type { UseChatHelpers } from '@ai-sdk/react' ;
2
- import React , { type JSX , useState , useEffect , useMemo } from 'react' ;
2
+ import React , { type JSX , useMemo , useState , useEffect } from 'react' ;
3
3
4
4
import { AggregatedSearchBlock } from './AggregatedSearchBlock' ;
5
5
import { AlertIcon , LoadingIcon , SearchIcon } from './icons' ;
6
6
import { MemoizedMarkdown } from './MemoizedMarkdown' ;
7
7
import type { ScreenStateProps } from './ScreenState' ;
8
8
import type { StoredSearchPlugin } from './stored-searches' ;
9
9
import type { InternalDocSearchHit , StoredAskAiState } from './types' ;
10
- import { extractLinksFromText } from './utils/ai' ;
10
+ import type { AIMessage } from './types/AskiAi' ;
11
+ import { extractLinksFromMessage , getMessageContent } from './utils/ai' ;
11
12
import { groupConsecutiveToolResults } from './utils/groupConsecutiveToolResults' ;
12
13
13
14
export type AskAiScreenTranslations = Partial < {
@@ -46,8 +47,8 @@ export type AskAiScreenTranslations = Partial<{
46
47
} > ;
47
48
48
49
type AskAiScreenProps = Omit < ScreenStateProps < InternalDocSearchHit > , 'translations' > & {
49
- messages : UseChatHelpers [ 'messages' ] ;
50
- status : UseChatHelpers [ 'status' ] ;
50
+ messages : AIMessage [ ] ;
51
+ status : UseChatHelpers < AIMessage > [ 'status' ] ;
51
52
askAiStreamError : Error | null ;
52
53
askAiFetchError : Error | undefined ;
53
54
translations ?: AskAiScreenTranslations ;
@@ -59,8 +60,8 @@ interface AskAiScreenHeaderProps {
59
60
60
61
interface Exchange {
61
62
id : string ;
62
- userMessage : UseChatHelpers [ 'messages' ] [ number ] ;
63
- assistantMessage : UseChatHelpers [ 'messages' ] [ number ] | null ;
63
+ userMessage : AIMessage ;
64
+ assistantMessage : AIMessage | null ;
64
65
}
65
66
66
67
function AskAiScreenHeader ( { disclaimerText } : AskAiScreenHeaderProps ) : JSX . Element {
@@ -71,7 +72,7 @@ interface AskAiExchangeCardProps {
71
72
exchange : Exchange ;
72
73
askAiStreamError : Error | null ;
73
74
isLastExchange : boolean ;
74
- loadingStatus : UseChatHelpers [ 'status' ] ;
75
+ loadingStatus : UseChatHelpers < AIMessage > [ 'status' ] ;
75
76
onSearchQueryClick : ( query : string ) => void ;
76
77
translations : AskAiScreenTranslations ;
77
78
conversations : StoredSearchPlugin < StoredAskAiState > ;
@@ -92,20 +93,25 @@ function AskAiExchangeCard({
92
93
93
94
const showActions = ! isLastExchange || ( isLastExchange && loadingStatus === 'ready' && Boolean ( assistantMessage ) ) ;
94
95
95
- const urlsToDisplay = React . useMemo ( ( ) => extractLinksFromText ( assistantMessage ?. content || '' ) , [ assistantMessage ] ) ;
96
+ const assistantContent = useMemo ( ( ) => getMessageContent ( assistantMessage ) , [ assistantMessage ] ) ;
97
+ const userContent = useMemo ( ( ) => getMessageContent ( userMessage ) , [ userMessage ] ) ;
98
+
99
+ const urlsToDisplay = React . useMemo ( ( ) => extractLinksFromMessage ( assistantMessage ) , [ assistantMessage ] ) ;
96
100
97
101
const displayParts = React . useMemo ( ( ) => {
98
- if ( ! Array . isArray ( assistantMessage ?. parts ) ) {
99
- return assistantMessage ?. content ? [ assistantMessage ?. content ] : [ ] ;
100
- }
101
102
return groupConsecutiveToolResults ( assistantMessage ?. parts || [ ] ) ;
102
103
} , [ assistantMessage ] ) ;
103
104
105
+ const isThinking =
106
+ [ 'submitted' , 'streaming' ] . includes ( loadingStatus ) &&
107
+ isLastExchange &&
108
+ ! displayParts . some ( ( part ) => part . type !== 'step-start' ) ;
109
+
104
110
return (
105
111
< div className = "DocSearch-AskAiScreen-Response-Container" >
106
112
< div className = "DocSearch-AskAiScreen-Response" >
107
113
< div className = "DocSearch-AskAiScreen-Message DocSearch-AskAiScreen-Message--user" >
108
- < p className = "DocSearch-AskAiScreen-Query" > { userMessage . content } </ p >
114
+ < p className = "DocSearch-AskAiScreen-Query" > { userContent ?. text ?? '' } </ p >
109
115
</ div >
110
116
< div className = "DocSearch-AskAiScreen-Message DocSearch-AskAiScreen-Message--assistant" >
111
117
< div className = "DocSearch-AskAiScreen-MessageContent" >
@@ -120,127 +126,113 @@ function AskAiExchangeCard({
120
126
/>
121
127
</ div >
122
128
) }
123
- { loadingStatus === 'submitted' && isLastExchange && (
129
+ { isThinking && (
124
130
< div className = "DocSearch-AskAiScreen-MessageContent-Reasoning" >
125
131
< span className = "shimmer" > { translations . thinkingText || 'Thinking...' } </ span >
126
132
</ div >
127
133
) }
128
- { Array . isArray ( displayParts )
129
- ? displayParts . map ( ( part , idx ) => {
130
- const index = idx ;
131
-
132
- if ( typeof part === 'string' ) {
133
- return (
134
- < MemoizedMarkdown
135
- key = { index }
136
- content = { part }
137
- copyButtonText = { translations . copyButtonText || 'Copy' }
138
- copyButtonCopiedText = { translations . copyButtonCopiedText || 'Copied!' }
139
- isStreaming = { loadingStatus === 'streaming' }
140
- />
141
- ) ;
142
- }
143
-
144
- // aggregated tool call rendering
145
- if ( part && ( part as any ) . type === 'aggregated-tool-call' ) {
134
+ { displayParts . map ( ( part , idx ) => {
135
+ const index = idx ;
136
+
137
+ if ( typeof part === 'string' ) {
138
+ return (
139
+ < MemoizedMarkdown
140
+ key = { index }
141
+ content = { part }
142
+ copyButtonText = { translations . copyButtonText || 'Copy' }
143
+ copyButtonCopiedText = { translations . copyButtonCopiedText || 'Copied!' }
144
+ isStreaming = { loadingStatus === 'streaming' }
145
+ />
146
+ ) ;
147
+ }
148
+
149
+ if ( part . type === 'aggregated-tool-call' ) {
150
+ return (
151
+ < AggregatedSearchBlock
152
+ key = { index }
153
+ queries = { part . queries }
154
+ translations = { translations }
155
+ onSearchQueryClick = { onSearchQueryClick }
156
+ />
157
+ ) ;
158
+ }
159
+
160
+ if ( part . type === 'reasoning' && part . state === 'streaming' ) {
161
+ return (
162
+ < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Reasoning shimmer" >
163
+ < LoadingIcon className = "DocSearch-AskAiScreen-SmallerLoadingIcon" />
164
+ < span className = "shimmer" > Reasoning...</ span >
165
+ </ div >
166
+ ) ;
167
+ }
168
+
169
+ if ( part . type === 'text' ) {
170
+ return (
171
+ < MemoizedMarkdown
172
+ key = { index }
173
+ content = { part . text }
174
+ copyButtonText = { translations . copyButtonText || 'Copy' }
175
+ copyButtonCopiedText = { translations . copyButtonCopiedText || 'Copied!' }
176
+ isStreaming = { part . state === 'streaming' }
177
+ />
178
+ ) ;
179
+ }
180
+ if ( part . type === 'tool-searchIndex' ) {
181
+ switch ( part . state ) {
182
+ case 'input-streaming' :
146
183
return (
147
- < AggregatedSearchBlock
148
- key = { index }
149
- queries = { ( part as any ) . queries }
150
- translations = { translations }
151
- onSearchQueryClick = { onSearchQueryClick }
152
- />
153
- ) ;
154
- }
155
-
156
- if ( part . type === 'reasoning' && assistantMessage ?. parts ?. length === 1 ) {
157
- return (
158
- < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Reasoning shimmer" >
159
- < span className = "shimmer" > Reasoning...</ span >
184
+ < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--PartialCall shimmer" >
185
+ < LoadingIcon className = "DocSearch-AskAiScreen-SmallerLoadingIcon" />
186
+ < span > { translations . preToolCallText || 'Searching...' } </ span >
160
187
</ div >
161
188
) ;
162
- }
163
-
164
- if ( part . type === 'text' ) {
189
+ case 'input-available' :
165
190
return (
166
- < MemoizedMarkdown
167
- key = { index }
168
- content = { part . text }
169
- copyButtonText = { translations . copyButtonText || 'Copy' }
170
- copyButtonCopiedText = { translations . copyButtonCopiedText || 'Copied!' }
171
- isStreaming = { loadingStatus === 'streaming' }
172
- />
191
+ < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--Call shimmer" >
192
+ < LoadingIcon className = "DocSearch-AskAiScreen-SmallerLoadingIcon" />
193
+ < span >
194
+ { `${ translations . duringToolCallText || 'Searching for ' } "${ part . input || '' } " ...` }
195
+ </ span >
196
+ </ div >
173
197
) ;
174
- }
175
- if ( part . type === 'tool-invocation' ) {
176
- const { toolInvocation } = part ;
177
- if ( toolInvocation . toolName === 'searchIndex' ) {
178
- switch ( toolInvocation . state ) {
179
- case 'partial-call' :
180
- return (
181
- < div
182
- key = { index }
183
- className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--PartialCall shimmer"
184
- >
185
- < LoadingIcon className = "DocSearch-AskAiScreen-SmallerLoadingIcon" />
186
- < span > { translations . preToolCallText || 'Searching...' } </ span >
187
- </ div >
188
- ) ;
189
- case 'call' :
190
- return (
191
- < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--Call shimmer" >
192
- < LoadingIcon className = "DocSearch-AskAiScreen-SmallerLoadingIcon" />
193
- < span >
194
- { `${ translations . duringToolCallText || 'Searching for ' } "${ toolInvocation . args ?. query || '' } " ...` }
195
- </ span >
196
- </ div >
197
- ) ;
198
- case 'result' :
199
- return (
200
- < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--Result" >
201
- < SearchIcon size = { 18 } />
202
- < span >
203
- { `${ translations . afterToolCallText || 'Searched for' } ` } { ' ' }
204
- < span
205
- role = "button"
206
- tabIndex = { 0 }
207
- className = "DocSearch-AskAiScreen-MessageContent-Tool-Query"
208
- onKeyDown = { ( e ) => {
209
- if ( e . key === 'Enter' || e . key === ' ' ) {
210
- e . preventDefault ( ) ;
211
- onSearchQueryClick ( toolInvocation . args ?. query || '' ) ;
212
- }
213
- } }
214
- onClick = { ( ) => onSearchQueryClick ( toolInvocation . args ?. query || '' ) }
215
- >
216
- { ' ' }
217
- "{ toolInvocation . args ?. query || '' } "
218
- </ span >
219
- </ span >
220
- </ div >
221
- ) ;
222
- default :
223
- return null ;
224
- }
225
- }
226
- // fallback for unknown tool, should never happen in theory. :shrug:
198
+ case 'output-available' :
227
199
return (
228
- < span key = { index } className = "text-sm italic shimmer" >
229
- { translations . thinkingText || 'Thinking...' }
230
- </ span >
200
+ < div key = { index } className = "DocSearch-AskAiScreen-MessageContent-Tool Tool--Result" >
201
+ < SearchIcon />
202
+ < span >
203
+ { `${ translations . afterToolCallText || 'Searched for' } ` } { ' ' }
204
+ < span
205
+ role = "button"
206
+ tabIndex = { 0 }
207
+ className = "DocSearch-AskAiScreen-MessageContent-Tool-Query"
208
+ onKeyDown = { ( e ) => {
209
+ if ( e . key === 'Enter' || e . key === ' ' ) {
210
+ e . preventDefault ( ) ;
211
+ onSearchQueryClick ( part . output . query || '' ) ;
212
+ }
213
+ } }
214
+ onClick = { ( ) => onSearchQueryClick ( part . output . query || '' ) }
215
+ >
216
+ { ' ' }
217
+ "{ part . output . query || '' } "
218
+ </ span >
219
+ </ span >
220
+ </ div >
231
221
) ;
232
- }
233
- // fallback for unknown part type
234
- return null ;
235
- } )
236
- : assistantMessage ?. content }
222
+ default :
223
+ break ;
224
+ }
225
+ }
226
+ // fallback for unknown part type
227
+ return null ;
228
+ } ) }
237
229
</ div >
238
230
</ div >
239
231
< div className = "DocSearch-AskAiScreen-Answer-Footer" >
240
232
< AskAiScreenFooterActions
241
233
id = { userMessage ?. id || exchange . id }
242
234
showActions = { showActions }
243
- latestAssistantMessageContent = { assistantMessage ?. content || null }
235
+ latestAssistantMessageContent = { assistantContent ?. text || null }
244
236
translations = { translations }
245
237
conversations = { conversations }
246
238
onFeedback = { onFeedback }
0 commit comments