show file search results properly

Author: Veikkosuhonen

src/client/App.tsx

@@ -99,7 +99,7 @@ const App = () => {
           <AppContext.Provider value={appRef}>
             <Box minHeight="100vh" display="flex" flexDirection="column" ref={appRef}>
               <NavBar />
-              <Container component="main" sx={{ mt: '4rem', mb: '10rem' }}>
+              <Container component="main" sx={{ mt: '4rem', mb: '10rem' }} maxWidth="xl">
                 <Outlet />
               </Container>
               <Footer />

src/client/components/ChatV2/ChatV2.tsx

@@ -6,7 +6,7 @@ import useLocalStorageState from '../../hooks/useLocalStorageState'
 import { DEFAULT_MODEL } from '../../../config'
 import useInfoTexts from '../../hooks/useInfoTexts'
 import { Message } from '../../types'
-import { FileCitation, ResponseStreamEventData } from '../../../shared/types'
+import { FileCitation, FileSearchResult, ResponseStreamEventData } from '../../../shared/types'
 import useRetryTimeout from '../../hooks/useRetryTimeout'
 import { useTranslation } from 'react-i18next'
 import { handleCompletionStreamError } from './error'
@@ -49,7 +49,7 @@ export const ChatV2 = () => {
   const [activePromptId, setActivePromptId] = useState('')
   const [fileName, setFileName] = useState<string>('')
   const [completion, setCompletion] = useState('')
-  const [citations, setCitations] = useState<FileCitation[]>([])
+  const [fileSearchResult, setFileSearchResult] = useLocalStorageState<FileSearchResult>('last-file-search', null)
   const [streamController, setStreamController] = useState<AbortController>()
   const [alertOpen, setAlertOpen] = useState(false)
   const [disallowedFileType, setDisallowedFileType] = useState('')
@@ -75,7 +75,7 @@ export const ChatV2 = () => {
       const reader = stream.getReader()
 
       let content = ''
-      const citations: FileCitation[] = []
+      let fileSearchResult: FileSearchResult
 
       while (true) {
         const { value, done } = await reader.read()
@@ -86,7 +86,12 @@ export const ChatV2 = () => {
         for (const chunk of data.split('\n')) {
           if (!chunk || chunk.trim().length === 0) continue
 
-          const parsedChunk: ResponseStreamEventData = JSON.parse(chunk)
+          let parsedChunk: ResponseStreamEventData = null
+          try {
+            parsedChunk = JSON.parse(chunk)
+          } catch (_e) {
+            console.error('Could not parse the chunk:', chunk)
+          }
 
           switch (parsedChunk.type) {
             case 'writing':
@@ -96,8 +101,11 @@ export const ChatV2 = () => {
 
             case 'annotation':
               console.log('Received annotation:', parsedChunk.annotation)
-              setCitations((prev) => [...prev, parsedChunk.annotation])
-              citations.push(parsedChunk.annotation)
+              break
+
+            case 'fileSearchDone':
+              fileSearchResult = parsedChunk.fileSearch
+              setFileSearchResult(parsedChunk.fileSearch)
               break
 
             case 'complete':
@@ -115,8 +123,7 @@ export const ChatV2 = () => {
         }
       }
 
-      setMessages((prev: Message[]) => prev.concat({ role: 'assistant', content, citations }))
-      setCitations([])
+      setMessages((prev: Message[]) => prev.concat({ role: 'assistant', content, fileSearchResult }))
     } catch (err: any) {
       handleCompletionStreamError(err, fileName)
     } finally {
@@ -135,7 +142,7 @@ export const ChatV2 = () => {
     setMessage({ content: '' })
     setPrevResponse({ id: '' })
     setCompletion('')
-    setCitations([])
+    setFileSearchResult(undefined)
     setStreamController(new AbortController())
     setRetryTimeout(() => {
       if (streamController) {
@@ -176,6 +183,7 @@ export const ChatV2 = () => {
     setMessage({ content: '' })
     setPrevResponse({ id: '' })
     setCompletion('')
+    setFileSearchResult(null)
     setStreamController(undefined)
     setTokenUsageWarning('')
     setTokenWarningVisible(false)
@@ -218,8 +226,8 @@ export const ChatV2 = () => {
         {courseId ? <Link to={'/v2'}>CurreChat</Link> : <Link to={'/v2/sandbox'}>Ohtu Sandbox</Link>}
       </Box>
       <Box sx={{ display: 'flex' }}>
-        <Box ref={chatContainerRef}>
-          <Conversation messages={messages} completion={completion} citations={citations} />
+        <Box ref={chatContainerRef} flex={1}>
+          <Conversation messages={messages} completion={completion} fileSearchResult={fileSearchResult} />
           <ChatBox
             disabled={false}
             onSubmit={(message) => {
@@ -230,7 +238,11 @@ export const ChatV2 = () => {
             }}
           />
         </Box>
-        {ragIndex && <CitationsBox messages={messages} citations={citations} ragIndex={ragIndex} />}
+        {ragIndex && (
+          <Box flex={1}>
+            <CitationsBox messages={messages} fileSearchResult={fileSearchResult} />
+          </Box>
+        )}
       </Box>
     </Box>
   )

src/client/components/ChatV2/CitationsBox.tsx

@@ -1,76 +1,51 @@
 import { Box, Paper, Typography } from '@mui/material'
-import { FileCitation, RagFileAttributes, RagIndexAttributes } from '../../../shared/types'
+import { FileSearchResult } from '../../../shared/types'
 import { Message } from '../../types'
-import { useQuery } from '@tanstack/react-query'
 import Markdown from 'react-markdown'
 
-const useFileCitationText = (citation: FileCitation, ragIndex: RagIndexAttributes) => {
-  const ragFileId = ragIndex?.ragFiles?.find((file) => file.filename === citation.filename)?.id
-
-  const { data, ...rest } = useQuery<RagFileAttributes & { fileContent: string }>({
-    queryKey: ['file-citation-content', ragIndex.id, citation.filename],
-    queryFn: async () => {
-      const response = await fetch(`/api/rag/indices/${ragIndex.id}/files/${ragFileId}`)
-      if (!response.ok) {
-        throw new Error('Failed to fetch file text')
-      }
-      return response.json()
-    },
-    retry: false,
-    enabled: !!ragFileId,
-  })
-
-  return {
-    ragFile: data,
-    ...rest,
-  }
-}
-
-const Citation = ({ citation, ragIndex }: { citation: FileCitation; ragIndex: RagIndexAttributes }) => {
-  const { ragFile, isLoading, error } = useFileCitationText(citation, ragIndex)
-
-  const contentAtIdx = ragFile?.fileContent?.substring(citation.index, citation.index + 800) || ''
+type FileItem = FileSearchResult['results'][number]
 
+const FileItemComponent = ({ fileItem }: { fileItem: FileItem }) => {
   return (
-    <Box sx={{ mb: 2, pl: 1, borderLeft: '2px solid #ccc' }}>
+    <Paper sx={{ p: 1, mt: 2 }}>
       <Typography variant="body2" color="textSecondary">
-        {citation.filename} (Index: {citation.index})
+        {fileItem.filename} (score: {fileItem.score})
       </Typography>
-      {isLoading ? (
-        <Typography variant="body2" color="textSecondary">
-          Loading citation text...
-        </Typography>
-      ) : error ? (
-        <Typography variant="body2" color="error">
-          Error loading citation text: {error.message}
-        </Typography>
-      ) : (
-        ragFile && <Markdown>{contentAtIdx}</Markdown>
-      )}
-    </Box>
+      <Markdown>{fileItem.text}</Markdown>
+    </Paper>
   )
 }
 
-const MessageCitations = ({ citations, ragIndex }: { citations: FileCitation[]; ragIndex: RagIndexAttributes }) => {
+const MessageFileSearchResult = ({ fileSearchResult }: { fileSearchResult: FileSearchResult }) => {
   return (
     <Box sx={{ mt: 1 }}>
-      {citations.map((citation, index) => (
-        <Citation key={index} citation={citation} ragIndex={ragIndex} />
+      <Typography>Searched for:</Typography>
+      <Box display="flex" gap={2}>
+        {fileSearchResult.queries.map((q, idx) => (
+          <Typography variant="body2" key={idx}>
+            "{q}"
+          </Typography>
+        ))}
+      </Box>
+
+      {fileSearchResult.results.map((result, key) => (
+        <FileItemComponent key={key} fileItem={result} />
       ))}
     </Box>
   )
 }
 
-export const CitationsBox = ({ messages, citations, ragIndex }: { messages: Message[]; citations: FileCitation[]; ragIndex: RagIndexAttributes }) => {
-  const messageCitations = [...messages.map((message) => (Array.isArray(message.citations) ? message.citations : [])), citations]
+export const CitationsBox = ({ messages, fileSearchResult }: { messages: Message[]; fileSearchResult?: FileSearchResult }) => {
+  const messageCitations = [...messages.map((m) => m.fileSearchResult).filter(Boolean)]
+  if (fileSearchResult) {
+    messageCitations.push(fileSearchResult)
+  }
 
   return (
-    <Paper>
-      <Box sx={{ width: 400, padding: 2 }}>
-        {messageCitations.map((c, index) => (
-          <MessageCitations key={index} citations={c} ragIndex={ragIndex} />
-        ))}
-      </Box>
-    </Paper>
+    <Box p={2}>
+      {messageCitations.map((c, index) => (
+        <MessageFileSearchResult key={index} fileSearchResult={fileSearchResult} />
+      ))}
+    </Box>
   )
 }

src/client/components/ChatV2/Conversation.tsx

@@ -3,31 +3,19 @@ import { Message } from '../../types'
 import ReactMarkdown from 'react-markdown'
 import remarkGfm from 'remark-gfm'
 import { Assistant } from '@mui/icons-material'
-import { FileCitation } from '../../../shared/types'
+import { FileCitation, FileSearchResult } from '../../../shared/types'
 
 const MessageItem = ({ message }: { message: Message }) => (
   <Paper
-    elevation={3}
+    elevation={message.role === 'assistant' ? 0 : 3}
     sx={{
       my: '2rem',
       ml: message.role === 'assistant' ? '0' : '2rem',
       mr: message.role === 'assistant' ? '2rem' : '0',
       p: '1rem',
-      backgroundColor: message.role === 'user' ? '#ffffff' : '#e0f7fa',
-      borderRadius: message.role === 'assistant' ? '0 1rem 1rem 1rem' : '1rem 0 1rem 1rem',
     }}
   >
     <ReactMarkdown remarkPlugins={[remarkGfm]}>{message.content}</ReactMarkdown>
-    {message.citations && message.citations.length > 0 && (
-      <Box sx={{ mt: 1, fontSize: '0.875rem', color: 'gray' }}>
-        Citations:
-        {message.citations.map((citation, index) => (
-          <Box key={index} sx={{ display: 'block' }}>
-            {citation.filename}
-          </Box>
-        ))}
-      </Box>
-    )}
   </Paper>
 )
 
@@ -52,12 +40,12 @@ const PöhinäLogo = () => (
   </Box>
 )
 
-export const Conversation = ({ messages, completion, citations }: { messages: Message[]; completion: string; citations: FileCitation[] }) => (
+export const Conversation = ({ messages, completion, fileSearchResult }: { messages: Message[]; completion: string; fileSearchResult: FileSearchResult }) => (
   <Box sx={{ flex: 1, overflowY: 'auto', gap: 2 }}>
     {messages.map((message, idx) => (
       <MessageItem key={idx} message={message} />
     ))}
-    {completion && <MessageItem message={{ role: 'assistant', content: completion, citations }} />}
+    {completion && <MessageItem message={{ role: 'assistant', content: completion, fileSearchResult }} />}
     {messages.length === 0 && <PöhinäLogo />}
   </Box>
 )

src/client/types.ts

@@ -1,4 +1,4 @@
-import { FileCitation } from '../shared/types'
+import { FileSearchResult } from "../shared/types"
 
 export type SetState<T> = React.Dispatch<React.SetStateAction<T>>
 
@@ -7,7 +7,7 @@ export type Role = 'system' | 'assistant' | 'user'
 export interface Message {
   role: Role
   content: string
-  citations?: FileCitation[]
+  fileSearchResult?: FileSearchResult
 }
 
 interface Term {

src/server/routes/openai.ts

@@ -197,13 +197,13 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
 
   res.setHeader('content-type', 'text/event-stream')
 
-  const completion = await responsesClient.handleResponse({
+  const result = await responsesClient.handleResponse({
     events,
     encoding,
     res,
   })
 
-  tokenCount += completion.tokenCount
+  tokenCount += result.tokenCount
 
   let userToCharge = user
   if (inProduction && req.hijackedBy) {
@@ -228,10 +228,11 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
   console.log('consentToSave', options.saveConsent, user.username)
 
   if (consentToSave) {
+    // @todo: should file search results also be saved?
     const discussion = {
       userId: user.id,
       courseId,
-      response: completion.response,
+      response: result.response,
       metadata: options,
     }
     await Discussion.create(discussion)

src/server/util/azure/ResponsesAPI.ts

@@ -104,13 +104,20 @@ export class ResponsesClient {
           console.log('file search completed')
           break
 
-        case 'response.output_item.done':
-          console.log('OUTPUT_ITEM DONE???', JSON.stringify(event, null, 2))
-          return {
-            tokenCount,
-            response: contents.join(''),
+        case 'response.output_item.done': {
+          if (event.item.type === 'file_search_call') {
+            this.write(
+              {
+                type: 'fileSearchDone',
+                fileSearch: event.item,
+              },
+              res,
+            )
           }
 
+          break
+        }
+
         case 'response.output_text.annotation.added':
           console.log(event)
           this.write(

src/shared/types.ts

@@ -1,3 +1,5 @@
+import type { ResponseFileSearchToolCall } from 'openai/resources/responses/responses'
+
 export type RagIndexMetadata = {
   name: string
   dim?: number
@@ -26,6 +28,8 @@ export type FileCitation = {
   type: 'file_citation'
 }
 
+export type FileSearchResult = ResponseFileSearchToolCall
+
 export type ResponseStreamEventData =
   | {
       type: 'start'
@@ -39,6 +43,10 @@ export type ResponseStreamEventData =
       type: 'complete'
       prevResponseId: string
     }
+  | {
+      type: 'fileSearchDone'
+      fileSearch: FileSearchResult
+    }
   | {
       type: 'error'
       error: any