Rag search ui improvements

Author: Veikkosuhonen

src/client/components/Rag/Rag.tsx

@@ -84,15 +84,15 @@ const Rag: React.FC = () => {
                 <TableRow>
                   <TableCell>ID</TableCell>
                   <TableCell>Name</TableCell>
-                  <TableCell>Vector Dimensions</TableCell>
+                  <TableCell>Language</TableCell>
                   <TableCell>Number of files</TableCell>
                 </TableRow>
               </TableHead>
               <TableBody>
                 <TableRow>
                   <TableCell>{index.id}</TableCell>
                   <TableCell>{index.metadata?.name}</TableCell>
-                  <TableCell>{index.metadata?.dim}</TableCell>
+                  <TableCell>{index.metadata?.language}</TableCell>
                   <TableCell>{index.ragFileCount}</TableCell>
                 </TableRow>
               </TableBody>

src/client/components/Rag/Search.tsx

@@ -1,34 +1,102 @@
 import { useState } from 'react'
 import { RagIndexAttributes } from '../../../shared/types'
-import { RagChunk } from '../../../shared/rag'
+import { RagChunk, SearchInputParams, SearchParams } from '../../../shared/rag'
 import apiClient from '../../util/apiClient'
+import { Box, Checkbox, Collapse, Fade, FormControl, FormControlLabel, Grow, LinearProgress, TextField, Typography, Zoom } from '@mui/material'
+import { OutlineButtonBlue } from '../ChatV2/general/Buttons'
+import { amber, green, blue } from '@mui/material/colors'
+
+const TimeLineColors = [blue[200], amber[300], green[300]]
 
 export const Search = ({ ragIndex }: { ragIndex: RagIndexAttributes }) => {
   const [query, setQuery] = useState('')
-  const [results, setResults] = useState<RagChunk[]>([])
+  const [vector, setVector] = useState(true)
+  const [ft, setFt] = useState(true)
+  const [rerank, setRerank] = useState(true)
+  const [isLoading, setIsLoading] = useState(false)
+  const [results, setResults] = useState<{ results: RagChunk[]; timings: Record<string, number> }>()
 
   const handleSubmit = async (event: React.FormEvent<HTMLFormElement>) => {
     event.preventDefault()
-    const response = await apiClient.post<RagChunk[]>(`/rag/indices/${ragIndex.id}/search`, {
+    const searchParams: SearchInputParams = {
       query,
-    })
+      vector,
+      ft,
+      rerank,
+    }
+    setIsLoading(true)
+    setResults(undefined)
+    const response = await apiClient.post<{ results: RagChunk[]; timings: Record<string, number> }>(`/rag/indices/${ragIndex.id}/search`, searchParams)
 
-    setResults(response.data ?? [])
+    setResults(response.data)
+    setIsLoading(false)
   }
 
   const handleInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     setQuery(event.target.value)
   }
 
+  const totalTime = Object.values(results?.timings ?? {}).reduce((acc, curr) => acc + curr, 0)
+
   return (
-    <form onSubmit={handleSubmit}>
-      <input type="text" value={query} onChange={handleInputChange} />
-      <button type="submit">Search</button>
-      <ul>
-        {results.map((chunk) => (
-          <li key={chunk.id}>{chunk.content}</li>
-        ))}
-      </ul>
-    </form>
+    <Box my="2rem">
+      <form onSubmit={handleSubmit}>
+        <FormControl>
+          <TextField type="text" value={query} onChange={handleInputChange} label="Search Query" />
+          <FormControlLabel control={<Checkbox checked={vector} onChange={(e) => setVector(e.target.checked)} />} label="Use semantic search" />
+          <FormControlLabel control={<Checkbox checked={ft} onChange={(e) => setFt(e.target.checked)} />} label="Use keyword search" />
+          <FormControlLabel control={<Checkbox checked={rerank} onChange={(e) => setRerank(e.target.checked)} />} label="Use reranking" />
+        </FormControl>
+        <OutlineButtonBlue type="submit">Search</OutlineButtonBlue>
+      </form>
+      {isLoading && <LinearProgress />}
+      <Fade in={!!results?.timings}>
+        {results?.timings ? (
+          <Box my="2rem">
+            <Typography variant="h6">Timings</Typography>
+            <Box display="flex" width="70vw">
+              {Object.entries(results?.timings ?? {}).map(([key, value], idx) => (
+                <Zoom
+                  in={true}
+                  key={key}
+                  style={{ width: `${(value / totalTime) * 100}%`, minWidth: '1rem', whiteSpace: 'nowrap' }}
+                  timeout={{ enter: 1000 + idx * 1000 }}
+                >
+                  <div>
+                    {key}: {value} ms
+                    <Box width="100%" height="1rem" bgcolor={TimeLineColors[idx % TimeLineColors.length]} border="1px solid black" borderRadius="1rem" />
+                  </div>
+                </Zoom>
+              ))}
+            </Box>
+          </Box>
+        ) : (
+          <div />
+        )}
+      </Fade>
+      <Fade in={!!results?.results}>
+        {results?.results ? (
+          <Box my="1rem">
+            <Typography variant="h6" mb="1rem">
+              Results
+            </Typography>
+            <Box>
+              {results?.results?.map((chunk) => (
+                <Box key={chunk.id} my="1rem">
+                  <Typography variant="subtitle2" color="text.secondary">
+                    Source: {chunk.metadata?.ragFileName ?? 'unknown'}
+                  </Typography>
+                  <Typography whiteSpace="pre-line" variant="body2">
+                    {chunk.content}
+                  </Typography>
+                </Box>
+              ))}
+            </Box>
+          </Box>
+        ) : (
+          <div />
+        )}
+      </Fade>
+    </Box>
   )
 }

src/server/routes/rag/rag.ts

@@ -68,9 +68,6 @@ router.post('/indices', async (req, res) => {
     userId: user.id,
     metadata: {
       name,
-      dim,
-      // azureVectorStoreId: vectorStore.id,
-      ragIndexFilterValue: randomUUID(),
       language,
     },
   })

src/server/routes/rag/ragIndex.ts

@@ -9,6 +9,7 @@ import { ApplicationError } from '../../util/ApplicationError'
 import { ingestRagFiles } from '../../services/rag/ingestion'
 import { search } from '../../services/rag/search'
 import { getRedisVectorStore } from '../../services/rag/vectorStore'
+import { SearchSchema } from '../../../shared/rag'
 
 const ragIndexRouter = Router()
 
@@ -213,18 +214,14 @@ ragIndexRouter.post('/upload', [indexUploadDirMiddleware, uploadMiddleware], asy
   res.json({ message: 'Files uploaded successfully' })
 })
 
-const RagIndexSearchSchema = z.object({
-  query: z.string().min(1).max(1000),
-})
-
 ragIndexRouter.post('/search', async (req, res) => {
   const ragIndexRequest = req as unknown as RagIndexRequest
   const { ragIndex } = ragIndexRequest
-  const { query } = RagIndexSearchSchema.parse(req.body)
+  const searchParams = SearchSchema.parse(req.body)
 
-  const results = await search(query, ragIndex)
+  const { results, timings } = await search(ragIndex, searchParams)
 
-  res.json(results)
+  res.json({ results, timings })
 })
 
 export default ragIndexRouter

src/server/routes/testUtils.ts

@@ -66,9 +66,8 @@ router.post('/reset-test-data', async (req, res) => {
     defaults: {
       userId,
       metadata: {
-        ragIndexFilterValue: 'mock',
         name: `rag-${testUserIdx}`,
-        azureVectorStoreId: 'mock',
+        language: 'English',
       },
     },
   })

src/server/services/rag/ingestion.ts

@@ -42,6 +42,10 @@ export const ingestRagFiles = async (ragIndex: RagIndex) => {
 
     chunkDocuments.forEach((chunkDocument, idx) => {
       chunkDocument.id = `ragIndex-${ragFile.ragIndexId}-${ragFile.filename}-${idx}`
+      chunkDocument.metadata = {
+        ...chunkDocument.metadata,
+        ragFileName: ragFile.filename,
+      }
     })
 
     // console.log(await redisClient.ft.info(vectorStore.indexName))

src/server/services/rag/search.ts

@@ -1,31 +1,68 @@
-import { RagChunk } from '../../../shared/rag'
+import { RagChunk, SearchParams } from '../../../shared/rag'
 import type { RagIndex } from '../../db/models'
 import { getRedisVectorStore } from './vectorStore'
 import { FTSearchRetriever } from './retrievers'
 import { EnsembleRetriever } from 'langchain/retrievers/ensemble'
 import { BM25Retriever } from '@langchain/community/retrievers/bm25'
+import { BaseRetriever } from '@langchain/core/retrievers'
+
+export const search = async (ragIndex: RagIndex, searchParams: SearchParams): Promise<{ results: RagChunk[]; timings: Record<string, number> }> => {
+  console.log('Searching', ragIndex.metadata.name, 'for query:', searchParams.query)
+  const timings: Record<string, number> = {}
 
-export const search = async (query: string, ragIndex: RagIndex): Promise<RagChunk[]> => {
-  console.log('Searching', ragIndex.metadata.name, 'for query:', query)
   const vectorStore = getRedisVectorStore(ragIndex.id)
 
-  const vectorstoreRetriever = vectorStore.asRetriever(8)
+  const vectorstoreRetriever = vectorStore.asRetriever(searchParams.vectorK)
   const ftSearchRetriever = new FTSearchRetriever(vectorStore.indexName)
 
+  const retrievers: BaseRetriever[] = []
+  const weights: number[] = []
+
+  if (searchParams.vector) {
+    retrievers.push(vectorstoreRetriever)
+    weights.push(0.3)
+  }
+
+  if (searchParams.ft) {
+    retrievers.push(ftSearchRetriever)
+    weights.push(0.7)
+  }
+
   const retriever = new EnsembleRetriever({
-    retrievers: [vectorstoreRetriever, ftSearchRetriever],
-    weights: [0.3, 0.7],
+    retrievers,
+    weights,
   })
 
-  const results0 = await retriever.invoke(query)
+  timings.search = Date.now()
+  let results = await retriever.invoke(searchParams.query)
+  timings.search = Date.now() - timings.search
 
-  const reranker = BM25Retriever.fromDocuments(results0, { k: 5 })
+  if (searchParams.rerank) {
+    timings.rerank = Date.now()
+    const reranker = BM25Retriever.fromDocuments(results, { k: searchParams.rerankK })
+    results = await reranker.invoke(searchParams.query)
+    timings.rerank = Date.now() - timings.rerank
+  }
 
-  const results = await reranker.invoke(query)
+  return {
+    results: results.map((doc) => ({
+      id: doc.id,
+      content: doc.pageContent,
+      metadata: parseMetadata(doc.metadata),
+    })),
+    timings,
+  }
+}
 
-  return results.map((doc) => ({
-    id: doc.id,
-    content: doc.pageContent,
-    metadata: doc.metadata,
-  }))
+const parseMetadata = (metadata: Record<string, any> | string) => {
+  if (typeof metadata === 'string') {
+    const unescaped = metadata.replace(/\\/g, '')
+    try {
+      return JSON.parse(unescaped)
+    } catch (error) {
+      console.error('Error parsing metadata:', unescaped, error)
+      return {}
+    }
+  }
+  return metadata
 }

src/server/services/rag/searchTool.ts

@@ -2,18 +2,19 @@ import { tool } from '@langchain/core/tools'
 import { z } from 'zod/v4'
 import { RagIndex } from '../../db/models'
 import { search } from './search'
+import { SearchSchema } from '../../../shared/rag'
 
 export const getRagIndexSearchTool = (ragIndex: RagIndex) =>
   tool(
     async ({ query }: { query: string }) => {
       console.log('Search tool invoked with query:', query)
-      const documents = await search(query, ragIndex)
+      const { results: documents } = await search(ragIndex, SearchSchema.parse({ query }))
       // With responseFormat: content_and_artifact, return content and artifact like this:
       return [documents.map((doc) => doc.content).join('\n\n'), documents]
     },
     {
       name: `document_search`, // Gotcha: function name must match '^[a-zA-Z0-9_\.-]+$' at least in AzureOpenAI. This name must satisfy the name in ChatToolDef type
-      description: `Search documents in the materials (titled '${ragIndex.metadata.name}')`,
+      description: `Search documents in the materials (titled '${ragIndex.metadata.name}'). Prefer ${ragIndex.metadata.language}, which is the language used in the documents.`,
       schema: z.object({
         query: z.string().describe('the query to search for'),
       }),

src/shared/rag.ts

@@ -1,6 +1,23 @@
+import z from 'zod/v4'
+
 export type RagChunk = {
   id?: string
   content: string
-  metadata: Record<string, any>
+  metadata: {
+    ragFileName: string
+    [key: string]: any
+  }
   score?: number
 }
+
+export const SearchSchema = z.object({
+  query: z.string().min(1).max(1000),
+  ft: z.boolean().default(true),
+  vector: z.boolean().default(true),
+  vectorK: z.number().min(1).max(20).default(8),
+  rerank: z.boolean().default(true),
+  rerankK: z.number().min(1).max(20).default(5),
+})
+
+export type SearchInputParams = z.input<typeof SearchSchema>
+export type SearchParams = z.infer<typeof SearchSchema>

src/shared/types.ts

@@ -1,20 +1,13 @@
-import type { ResponseFileSearchToolCall } from 'openai/resources/responses/responses'
-import type { VectorStoreFile } from 'openai/resources/vector-stores/files'
 import type { IngestionPipelineStageKey } from './constants'
-import { ChatToolDef } from './tools'
+import type { ChatToolDef } from './tools'
 
 export type RagIndexMetadata = {
   name: string
-  dim?: number
-  azureVectorStoreId?: string
-  ragIndexFilterValue: string
   instructions?: string
   language?: 'Finnish' | 'English'
 }
 
 export type RagFileMetadata = {
-  chunkingStrategy?: NonNullable<VectorStoreFile['chunking_strategy']>['type']
-  vectorStoreFileId?: string
   usageBytes?: number
 }