rag index mayheming

Author: Veikkosuhonen

src/client/components/Rag/ProgressReporter.tsx

@@ -1,25 +1,21 @@
-import { Box, Table, TableBody, TableCell, TableHead, TableRow, Typography } from '@mui/material'
-import { orderBy } from 'lodash'
+import { LinearProgress, Table, TableBody, TableCell, TableHead, TableRow, Typography } from '@mui/material'
 import { useEffect, useReducer } from 'react'
 import { IngestionPipelineStageKey, IngestionPipelineStageKeys, IngestionPipelineStages } from '../../../shared/constants'
 
 type ProgressEvent = {
   stage: string
   items?: string[]
-  done?: boolean
   error?: string
 }
 
 type ProgressState = Record<
   IngestionPipelineStageKey,
   {
     count: number
-    done: boolean
     error: boolean
     files: {
       [fileName: string]: {
         count: number
-        done: boolean
         error: boolean
       }
     }
@@ -32,7 +28,6 @@ const getInitialProgressState = () =>
       stage,
       {
         count: 0,
-        done: false,
         error: false,
         files: {},
       },
@@ -44,25 +39,14 @@ type Action = { type: 'UPDATE'; payload: ProgressEvent } | { type: 'RESET' }
 const progressReducer = (state: ProgressState, action: Action): ProgressState => {
   switch (action.type) {
     case 'UPDATE': {
-      const { stage, items, done, error } = action.payload
-
-      if (done) {
-        return {
-          ...state,
-          [stage]: {
-            ...state[stage],
-            done: true,
-            error: !!error || false,
-          },
-        }
-      }
+      const { stage, items, error } = action.payload
 
       if (!items) return state
 
       const updatedFiles = items.reduce(
         (acc, fileName) => {
           if (!acc[fileName]) {
-            acc[fileName] = { count: 0, done: done || false, error: !!error || false }
+            acc[fileName] = { count: 0, error: !!error || false }
           }
           acc[fileName].count += 1
           return acc
@@ -74,7 +58,6 @@ const progressReducer = (state: ProgressState, action: Action): ProgressState =>
         ...state,
         [stage]: {
           ...state[stage],
-          done: done || state[stage]?.done,
           error: !!error || state[stage]?.error,
           count: state[stage]?.count + items.length,
           files: updatedFiles,
@@ -144,8 +127,6 @@ export const ProgressReporter: React.FC<{ filenames: string[]; stream: ReadableS
     }
   }, [stream])
 
-  console.log('Progress state:', progress)
-
   return (
     <Table size="small">
       <TableHead>
@@ -155,13 +136,16 @@ export const ProgressReporter: React.FC<{ filenames: string[]; stream: ReadableS
             <TableCell key={stage}>
               <Typography variant="body2">{IngestionPipelineStages[stage].name}</Typography>
               <Typography variant="caption" color="textSecondary">
-                {progress[stage].count}{' '}
-              </Typography>
-              <Typography variant="caption" color="textSecondary">
-                {progress[stage].done ? 'Done' : progress[stage].error ? 'Error' : 'In Progress'}
+                {progress[stage].count}/{filenames.length}
               </Typography>
+              {progress[stage].error && (
+                <Typography variant="caption" color="error">
+                  Error
+                </Typography>
+              )}
             </TableCell>
           ))}
+          <TableCell>Status</TableCell>
         </TableRow>
       </TableHead>
       <TableBody>
@@ -170,30 +154,31 @@ export const ProgressReporter: React.FC<{ filenames: string[]; stream: ReadableS
             <TableCell component="th" scope="row">
               {filename}
             </TableCell>
-            {IngestionPipelineStageKeys.map((stage) => (
-              <TableCell
-                key={stage}
-                sx={{
-                  transition: 'background-color 0.3s',
-                  backgroundColor: progress[stage]?.error
-                    ? 'error.light'
-                    : progress[stage]?.done
-                      ? 'success.light'
-                      : progress[stage]?.files[filename]?.error
-                        ? 'error.light'
-                        : progress[stage]?.files[filename]?.count
-                          ? 'info.light'
-                          : 'inherit',
-                }}
-              >
-                <Box display="flex" gap={2}>
-                  <Typography variant="body2">{progress[stage].files[filename]?.count > 1 || ''}</Typography>
-                  <Typography variant="caption" color="textSecondary">
-                    {progress[stage]?.files[filename]?.done ? 'Done' : progress[stage].files[filename]?.error ? 'Error' : ''}
-                  </Typography>
-                </Box>
-              </TableCell>
-            ))}
+            <TableCell colSpan={IngestionPipelineStageKeys.length}>
+              <LinearProgress
+                variant="determinate"
+                value={(IngestionPipelineStageKeys.reduce((acc, stage) => acc + (progress[stage].files[filename]?.count ? 1 : 0), 0) / IngestionPipelineStageKeys.length) * 100}
+              />
+            </TableCell>
+            <TableCell>
+              {IngestionPipelineStageKeys.some((stage) => progress[stage].files[filename]?.error) ? (
+                <Typography variant="body2" color="error">
+                  Error
+                </Typography>
+              ) : progress['store']?.files[filename]?.count > 0 ? (
+                <Typography variant="body2" color="textSecondary">
+                  Completed
+                </Typography>
+              ) : IngestionPipelineStageKeys.some((stage) => progress[stage].files[filename]?.count) ? (
+                <Typography variant="body2" color="textSecondary">
+                  In Progress
+                </Typography>
+              ) : (
+                <Typography variant="body2" color="textSecondary">
+                  Not Started
+                </Typography>
+              )}
+            </TableCell>
           </TableRow>
         ))}
       </TableBody>

src/client/components/Rag/Rag.tsx

@@ -24,6 +24,7 @@ type RagIndexAttributes = {
     dim: number
   }
   numOfChunks: number
+  filenames: string[]
 }
 
 const useRagIndices = () => {
@@ -135,7 +136,7 @@ const Rag: React.FC = () => {
 
   return (
     <Box sx={{ display: 'flex', gap: 2 }}>
-      <Dialog open={!!selectedIndex && modalOpen} onClose={() => setModalOpen(false)} fullWidth maxWidth="md">
+      <Dialog open={!!selectedIndex && modalOpen} onClose={() => { setModalOpen(false); refetch(); }} fullWidth maxWidth="md">
         <DialogTitle>Edit {selectedIndex?.metadata?.name}</DialogTitle>
         <Box sx={{ padding: 2 }}>
           <Box sx={{ display: 'flex', gap: 2 }}>
@@ -207,6 +208,7 @@ const Rag: React.FC = () => {
                 <TableRow>
                   <TableCell>ID</TableCell>
                   <TableCell>Name</TableCell>
+                  <TableCell>Files</TableCell>
                   <TableCell>Dim</TableCell>
                   <TableCell>Num chunks</TableCell>
                 </TableRow>
@@ -215,6 +217,17 @@ const Rag: React.FC = () => {
                 <TableRow>
                   <TableCell>{index.id}</TableCell>
                   <TableCell>{index.metadata.name}</TableCell>
+                  <TableCell>
+                    {index.filenames.length ? (
+                      <>
+                        {index.filenames.map((filename) => (
+                          <p key={filename}>{filename}</p>
+                        ))}
+                      </>
+                    ) : (
+                      'No files'
+                    )}
+                  </TableCell>
                   <TableCell>{index.metadata.dim}</TableCell>
                   <TableCell>{index.numOfChunks}</TableCell>
                 </TableRow>

src/server/db/migrations/20250521_00_add_files_to_rag_indexes.ts

@@ -0,0 +1,13 @@
+import { QueryInterface, DataTypes } from 'sequelize'
+
+export async function up({ context: queryInterface }) {
+  await queryInterface.addColumn('rag_indices', 'filenames', {
+    type: DataTypes.ARRAY(DataTypes.STRING),
+    allowNull: false,
+    comment: 'Original filenames of the files uploaded for this index',
+  })
+}
+
+export async function down({ context: queryInterface }) {
+  await queryInterface.removeColumn('rag_indices', 'filenames')
+}

src/server/db/models/ragIndex.ts

@@ -11,6 +11,8 @@ class RagIndex extends Model<InferAttributes<RagIndex>, InferCreationAttributes<
   declare courseId?: string
 
   declare metadata: RagIndexMetadata
+
+  declare filenames: string[]
 }
 
 RagIndex.init(
@@ -33,6 +35,10 @@ RagIndex.init(
       type: DataTypes.JSONB,
       allowNull: true,
     },
+    filenames: {
+      type: DataTypes.ARRAY(DataTypes.STRING),
+      allowNull: false,
+    },
   },
   {
     underscored: true,

src/server/routes/rag.ts

@@ -1,15 +1,13 @@
 import { NextFunction, Request, Response, Router } from 'express'
-import { EMBED_DIM, EMBED_MODEL } from '../../config'
+import { EMBED_DIM } from '../../config'
 import { createChunkIndex, deleteChunkIndex, getNumberOfChunks } from '../services/rag/chunkDb'
 import { RagIndex } from '../db/models'
 import { RequestWithUser } from '../types'
 import z from 'zod'
 import { queryRagIndex } from '../services/rag/query'
 import { ingestionPipeline } from '../services/rag/ingestion/pipeline'
-import { getAzureOpenAIClient } from '../util/azure'
 import multer from 'multer'
 import { mkdir, rm, stat } from 'fs/promises'
-import { Readable } from 'stream'
 import { getOllamaOpenAIClient } from '../util/ollama'
 
 const router = Router()
@@ -30,7 +28,9 @@ router.post('/indices', async (req, res) => {
     metadata: {
       name,
       dim,
+      numOfChunks: 0,
     },
+    filenames: [],
   })
 
   await createChunkIndex(ragIndex)

src/server/services/rag/ingestion/chunker.ts

@@ -1,47 +1,21 @@
-import { Transform } from 'node:stream'
 import { chunkingAlgorithms } from './chunkingAlgorithms.ts'
 import { mkdirSync } from 'node:fs'
 import { writeFile } from 'node:fs/promises'
 import { TextData } from './textExtractor.ts'
-import { StageReporter } from './progressReporter.ts'
 
-export class Chunker extends Transform {
-  private cachePath: string
-  public progressReporter: StageReporter
+export async function createChunks(data: TextData, cachePath: string) {
+  const chunksDir = `${cachePath}/chunks`
+  mkdirSync(chunksDir, { recursive: true })
 
-  constructor(cachePath: string) {
-    super({ objectMode: true })
+  const chunkingAlgorithm = chunkingAlgorithms[data.chunkingStrategy]
+  const chunks = chunkingAlgorithm(data)
 
-    this.cachePath = cachePath + '/chunks'
+  await Promise.all(
+    chunks.map((chunk) => {
+      const chunkPath = `${chunksDir}/${chunk.id}.json`
+      return writeFile(chunkPath, JSON.stringify(chunk, null, 2), 'utf-8')
+    }),
+  )
 
-    // Make sure the cache path exists
-    mkdirSync(this.cachePath, { recursive: true })
-  }
-
-  async _transform(data: TextData, _encoding: BufferEncoding, callback: (error?: Error | null) => void) {
-    const chunkingAlgorithm = chunkingAlgorithms[data.chunkingStrategy]
-
-    const chunks = chunkingAlgorithm(data)
-    for (const chunk of chunks) {
-      this.push(chunk)
-    }
-
-    // Save chunks to cache
-
-    await Promise.all(
-      chunks.map((chunk) => {
-        const chunkPath = `${this.cachePath}/${chunk.id}.json`
-        return writeFile(chunkPath, JSON.stringify(chunk, null, 2), 'utf-8')
-      }),
-    )
-
-    this.progressReporter.reportProgress([data.fileName])
-
-    callback()
-  }
-
-  _flush(callback: (error?: Error | null) => void) {
-    this.progressReporter.reportDone()
-    callback()
-  }
+  return chunks
 }