Skip to content

Commit ac06d0a

Browse files
VeikkosuhonenVeikkosuhonen
committed
refactor chat v2 process stream into hook
1 parent 6d3251b commit ac06d0a

File tree

2 files changed

+28
-127
lines changed

2 files changed

+28
-127
lines changed

‎src/client/components/ChatV2/ChatV2.tsx‎

Lines changed: 27 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { useContext, useEffect, useRef, useState } from 'react'
77
import { useTranslation } from 'react-i18next'
88
import { useParams } from 'react-router-dom'
99
import { ALLOWED_FILE_TYPES, DEFAULT_ASSISTANT_INSTRUCTIONS, DEFAULT_MODEL, DEFAULT_MODEL_TEMPERATURE, FREE_MODEL, validModels } from '../../../config'
10-
import type { FileSearchCompletedData, ResponseStreamEventData } from '../../../shared/types'
10+
import type { FileSearchCompletedData } from '../../../shared/types'
1111
import { getLanguageValue } from '../../../shared/utils'
1212
import useCourse from '../../hooks/useCourse'
1313
import useInfoTexts from '../../hooks/useInfoTexts'
@@ -18,7 +18,6 @@ import useUserStatus from '../../hooks/useUserStatus'
1818
import type { Message } from '../../types'
1919
import { AppContext } from '../../util/AppContext'
2020
import { ChatBox } from './ChatBox'
21-
import { FileSearchInfo } from './CitationsBox'
2221
import { Conversation } from './Conversation'
2322
import { DisclaimerModal } from './Disclaimer'
2423
import { handleCompletionStreamError } from './error'
@@ -28,16 +27,16 @@ import { SettingsModal } from './SettingsModal'
2827
import { getCompletionStream } from './util'
2928
import { OutlineButtonBlack } from './generics/Buttons'
3029
import useCurrentUser from '../../hooks/useCurrentUser'
30+
import { useChatStream } from './useChatStream'
3131
import Annotations from './Annotations'
32-
import { enqueueSnackbar } from 'notistack'
3332

3433
export const ChatV2 = () => {
3534
const { courseId } = useParams()
3635

3736
const { data: course } = useCourse(courseId)
3837

3938
const { ragIndices } = useRagIndices(courseId)
40-
const { infoTexts, isLoading: infoTextsLoading } = useInfoTexts()
39+
const { infoTexts } = useInfoTexts()
4140

4241
const { userStatus, isLoading: statusLoading, refetch: refetchStatus } = useUserStatus(courseId)
4342

@@ -65,19 +64,13 @@ export const ChatV2 = () => {
6564
const [fileSearch, setFileSearch] = useLocalStorageState<FileSearchCompletedData>(`${localStoragePrefix}-last-file-search`)
6665

6766
// App States
68-
const [isFileSearching, setIsFileSearching] = useState<boolean>(false)
6967
const [settingsModalOpen, setSettingsModalOpen] = useState<boolean>(false)
7068
const [fileName, setFileName] = useState<string>('')
7169
const [tokenUsageWarning, setTokenUsageWarning] = useState<string>('')
7270
const [tokenUsageAlertOpen, setTokenUsageAlertOpen] = useState<boolean>(false)
7371
const [allowedModels, setAllowedModels] = useState<string[]>([])
7472
const [saveConsent, setSaveConsent] = useState<boolean>(false)
7573

76-
// Chat Streaming states
77-
const [completion, setCompletion] = useState<string>('')
78-
const [isCompletionDone, setIsCompletionDone] = useState<boolean>(true)
79-
const [streamController, setStreamController] = useState<AbortController>()
80-
8174
// RAG states
8275
const [ragIndexId, setRagIndexId] = useState<number | undefined>()
8376
const [ragDisplay, setRagDisplay] = useState<boolean>(true)
@@ -93,106 +86,27 @@ export const ChatV2 = () => {
9386

9487
const [setRetryTimeout, clearRetryTimeout] = useRetryTimeout()
9588

96-
const decoder = new TextDecoder()
9789
const { t, i18n } = useTranslation()
98-
const { language } = i18n
99-
100-
const disclaimerInfo = infoTexts?.find((infoText) => infoText.name === 'disclaimer')?.text[language] ?? null
10190

102-
const processStream = async (stream: ReadableStream) => {
103-
let content = ''
104-
let error = ''
105-
let fileSearch: FileSearchCompletedData
91+
const disclaimerInfo = infoTexts?.find((infoText) => infoText.name === 'disclaimer')?.text[i18n.language] ?? null
10692

107-
try {
108-
const reader = stream.getReader()
109-
110-
while (true) {
111-
const { value, done } = await reader.read()
112-
if (done) break
113-
114-
const data = decoder.decode(value)
115-
116-
let accumulatedChunk = ''
117-
for (const chunk of data.split('\n')) {
118-
if (!chunk || chunk.trim().length === 0) continue
119-
120-
let parsedChunk: ResponseStreamEventData | undefined
121-
try {
122-
parsedChunk = JSON.parse(chunk)
123-
} catch (e: any) {
124-
console.error('Error', e)
125-
console.error('Could not parse the chunk:', chunk)
126-
accumulatedChunk += chunk
127-
128-
try {
129-
parsedChunk = JSON.parse(accumulatedChunk)
130-
accumulatedChunk = ''
131-
} catch (e: any) {
132-
console.error('Error', e)
133-
console.error('Could not parse the accumulated chunk:', accumulatedChunk)
134-
}
135-
}
136-
137-
if (!parsedChunk) continue
138-
139-
switch (parsedChunk.type) {
140-
case 'writing':
141-
setCompletion((prev) => prev + parsedChunk.text)
142-
content += parsedChunk.text
143-
break
144-
145-
case 'annotation':
146-
console.log('Received annotation:', parsedChunk.annotation)
147-
break
148-
149-
case 'fileSearchStarted':
150-
setIsFileSearching(true)
151-
break
152-
153-
case 'fileSearchDone':
154-
fileSearch = parsedChunk.fileSearch
155-
setFileSearch(parsedChunk.fileSearch)
156-
setIsFileSearching(false)
157-
break
158-
159-
case 'error':
160-
error += parsedChunk.error
161-
break
162-
163-
case 'complete':
164-
setPrevResponse({ id: parsedChunk.prevResponseId })
165-
break
166-
167-
default:
168-
break
169-
}
170-
}
93+
const { processStream, completion, isStreaming, isFileSearching, streamController } = useChatStream({
94+
onComplete: ({ message, previousResponseId }) => {
95+
if (previousResponseId) {
96+
setPrevResponse({ id: previousResponseId })
17197
}
172-
} catch (err: any) {
173-
handleCompletionStreamError(err, fileName)
174-
error += '\nResponse stream was interrupted'
175-
} finally {
176-
if (content.length > 0) {
177-
setMessages((prev: Message[]) =>
178-
prev.concat({
179-
role: 'assistant',
180-
content,
181-
error: error.length > 0 ? error : undefined,
182-
fileSearchResult: fileSearch,
183-
}),
184-
)
98+
if (message.content.length > 0) {
99+
setMessages((prev: Message[]) => prev.concat(message))
100+
refetchStatus()
185101
}
186-
187-
setStreamController(undefined)
188-
setCompletion('')
189-
setIsCompletionDone(true)
190-
refetchStatus()
191-
setFileName('')
192-
setIsFileSearching(false)
193-
clearRetryTimeout()
194-
}
195-
}
102+
},
103+
onError: (error) => {
104+
handleCompletionStreamError(error, fileName)
105+
},
106+
onFileSearchComplete: (fileSearch) => {
107+
setFileSearch(fileSearch)
108+
},
109+
})
196110

197111
const handleSubmit = async (message: string, ignoreTokenUsageWarning: boolean) => {
198112
const formData = new FormData()
@@ -215,15 +129,11 @@ export const ChatV2 = () => {
215129

216130
setMessages(newMessages)
217131
setPrevResponse({ id: '' })
218-
setCompletion('')
219-
setIsCompletionDone(false)
220132
if (fileInputRef.current) {
221133
fileInputRef.current.value = ''
222134
}
223135
setFileName('')
224136
setFileSearch(undefined)
225-
setIsFileSearching(false)
226-
setStreamController(new AbortController())
227137
setRetryTimeout(() => {
228138
if (streamController) {
229139
streamController.abort()
@@ -277,8 +187,6 @@ export const ChatV2 = () => {
277187
if (window.confirm('Are you sure you want to empty this conversation?')) {
278188
setMessages([])
279189
setPrevResponse({ id: '' })
280-
setCompletion('')
281-
setIsCompletionDone(true)
282190
if (!ragDisplay) {
283191
handleRagDisplay()
284192
}
@@ -287,7 +195,6 @@ export const ChatV2 = () => {
287195
}
288196
setFileName('')
289197
setFileSearch(undefined)
290-
setStreamController(undefined)
291198
setTokenUsageWarning('')
292199
setTokenUsageAlertOpen(false)
293200
setRetryTimeout(() => {
@@ -300,8 +207,6 @@ export const ChatV2 = () => {
300207
}
301208

302209
const handleCancel = () => {
303-
setIsCompletionDone(true)
304-
setStreamController(undefined)
305210
setTokenUsageWarning('')
306211
setTokenUsageAlertOpen(false)
307212
clearRetryTimeout()
@@ -310,7 +215,7 @@ export const ChatV2 = () => {
310215
useEffect(() => {
311216
// Scrolls to bottom on initial load only
312217
if (!appContainerRef?.current || !conversationRef.current || messages.length === 0) return
313-
if (isCompletionDone) {
218+
if (!isStreaming) {
314219
const container = appContainerRef?.current
315220
if (container) {
316221
container.scrollTo({
@@ -327,7 +232,7 @@ export const ChatV2 = () => {
327232

328233
const lastNode = conversationRef.current.lastElementChild as HTMLElement
329234

330-
if (lastNode.classList.contains('message-role-assistant') && !isCompletionDone) {
235+
if (lastNode.classList.contains('message-role-assistant') && isStreaming) {
331236
const container = appContainerRef.current
332237

333238
const containerRect = container.getBoundingClientRect()
@@ -341,7 +246,7 @@ export const ChatV2 = () => {
341246
behavior: 'smooth',
342247
})
343248
}
344-
}, [isCompletionDone])
249+
}, [isStreaming])
345250

346251
useEffect(() => {
347252
if (!userStatus) return
@@ -369,7 +274,7 @@ export const ChatV2 = () => {
369274
}
370275
}, [userStatus, course])
371276

372-
if (statusLoading || infoTextsLoading) return null
277+
if (statusLoading) return null
373278

374279
if (course && course.usageLimit === 0) {
375280
return (
@@ -499,13 +404,13 @@ export const ChatV2 = () => {
499404
>
500405
<Alert severity="info">{t('chat:testUseInfo')}</Alert>
501406
<Conversation
502-
courseName={course && getLanguageValue(course.name, language)}
407+
courseName={course && getLanguageValue(course.name, i18n.language)}
503408
courseDate={course?.activityPeriod}
504409
conversationRef={conversationRef}
505410
expandedNodeHeight={window.innerHeight - (inputFieldRef.current?.clientHeight ?? 0) - 300}
506411
messages={messages}
507412
completion={completion}
508-
isCompletionDone={isCompletionDone}
413+
isCompletionDone={!isStreaming}
509414
setActiveFileSearchResult={setActiveFileSearchResult}
510415
/>
511416
</Box>
@@ -523,7 +428,7 @@ export const ChatV2 = () => {
523428
}}
524429
>
525430
<ChatBox
526-
disabled={!isCompletionDone}
431+
disabled={isStreaming}
527432
currentModel={activeModel.name}
528433
availableModels={allowedModels}
529434
fileInputRef={fileInputRef}
@@ -545,8 +450,6 @@ export const ChatV2 = () => {
545450

546451
{/* Annotations columns ----------------------------------------------------------------------------------------------------- */}
547452

548-
549-
{/* LEGACY - keep here for legacy when new annotations flow is work in progres */}
550453
{/* {showFileSearch && (
551454
<FileSearchInfo
552455
isFileSearching={isFileSearching}
@@ -568,9 +471,7 @@ export const ChatV2 = () => {
568471
borderLeft: activeFileSearchResult ? '1px solid rgba(0, 0, 0, 0.12)' : 'none',
569472
}}
570473
>
571-
<Box sx={{ position: 'sticky', top: 65, padding: '2rem' }}>
572-
{activeFileSearchResult && <Annotations fileSearchResult={activeFileSearchResult} />}
573-
</Box>
474+
<Box sx={{ position: 'sticky', top: 65, padding: '2rem' }}>{activeFileSearchResult && <Annotations fileSearchResult={activeFileSearchResult} />}</Box>
574475
</Box>
575476

576477
{/* Modals --------------------------------------*/}

‎src/server/util/azure/ResponsesAPI.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { createMockStream } from './mocks/MockStream'
1111
import { createFileSearchTool } from './util'
1212
import { FileSearchResultsStore } from '../../services/azureFileSearch/fileSearchResultsStore'
1313
import { getAzureOpenAIClient } from './client'
14-
import { RagIndex } from '../../db/models'
14+
import type { RagIndex } from '../../db/models'
1515
import OpenAI from 'openai'
1616
import { ApplicationError } from '../ApplicationError'
1717

0 commit comments

Comments
 (0)