make work with local embedding & betterments

Author: Veikkosuhonen

src/client/components/Rag/ProgressReporter.tsx

@@ -1,65 +1,97 @@
-import { Box, Typography } from '@mui/material'
+import { Box, Table, TableBody, TableCell, TableHead, TableRow, Typography } from '@mui/material'
+import { orderBy } from 'lodash'
 import { useEffect, useReducer } from 'react'
+import { IngestionPipelineStageKey, IngestionPipelineStageKeys, IngestionPipelineStages } from '../../../shared/constants'
 
 type ProgressEvent = {
   stage: string
-  item?: string
+  items?: string[]
   done?: boolean
   error?: string
 }
 
-type ProgressState = {
-  [fileName: string]: {
-    [stage: string]: {
-      count: number
-      done: boolean
+type ProgressState = Record<
+  IngestionPipelineStageKey,
+  {
+    count: number
+    done: boolean
+    error: boolean
+    files: {
+      [fileName: string]: {
+        count: number
+        done: boolean
+        error: boolean
+      }
     }
   }
-}
+>
+
+const getInitialProgressState = () =>
+  Object.fromEntries(
+    IngestionPipelineStageKeys.map((stage) => [
+      stage,
+      {
+        count: 0,
+        done: false,
+        error: false,
+        files: {},
+      },
+    ]),
+  ) as ProgressState
 
 type Action = { type: 'UPDATE'; payload: ProgressEvent } | { type: 'RESET' }
 
 const progressReducer = (state: ProgressState, action: Action): ProgressState => {
   switch (action.type) {
     case 'UPDATE': {
-      const { stage, item, done } = action.payload
+      const { stage, items, done, error } = action.payload
+
       if (done) {
-        // mark all items at this stage as done
-        const newState = { ...state }
-        for (const file in newState) {
-          if (newState[file][stage]) {
-            newState[file][stage].done = true
-          }
+        return {
+          ...state,
+          [stage]: {
+            ...state[stage],
+            done: true,
+            error: !!error || false,
+          },
         }
-        return newState
       }
-      if (!item) return state
 
-      const fileStages = state[item] || {}
-      const stageData = fileStages[stage] || { count: 0, done: false }
+      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 += 1
+          return acc
+        },
+        { ...state[stage].files },
+      )
 
       return {
         ...state,
-        [item]: {
-          ...fileStages,
-          [stage]: {
-            count: stageData.count + 1,
-            done: stageData.done || false,
-          },
+        [stage]: {
+          ...state[stage],
+          done: done || state[stage]?.done,
+          error: !!error || state[stage]?.error,
+          count: state[stage]?.count + items.length,
+          files: updatedFiles,
         },
       }
     }
 
     case 'RESET':
-      return {}
+      return getInitialProgressState()
 
     default:
       return state
   }
 }
 
-export const ProgressReporter: React.FC<{ stream: ReadableStream | null }> = ({ stream }) => {
-  const [progress, dispatch] = useReducer(progressReducer, {})
+export const ProgressReporter: React.FC<{ filenames: string[]; stream: ReadableStream | null; onError: () => void }> = ({ filenames, stream, onError }) => {
+  const [progress, dispatch] = useReducer(progressReducer, getInitialProgressState())
 
   useEffect(() => {
     let reader: ReadableStreamDefaultReader<Uint8Array> | null = null
@@ -69,6 +101,8 @@ export const ProgressReporter: React.FC<{ stream: ReadableStream | null }> = ({
         reader = stream.getReader()
       } catch (error) {
         console.error('Error getting reader from stream:', error)
+        dispatch({ type: 'RESET' })
+        onError()
         return
       }
 
@@ -90,36 +124,79 @@ export const ProgressReporter: React.FC<{ stream: ReadableStream | null }> = ({
                 dispatch({ type: 'UPDATE', payload: jsonChunk })
               } catch (err) {
                 console.error('Invalid chunk:', line, err)
+                onError()
               }
             }
           }
         }
       }
 
       readStream()
+    } else {
+      dispatch({ type: 'RESET' })
+    }
 
-      return () => {
-        if (reader) {
-          reader.releaseLock()
-        }
+    return () => {
+      if (reader) {
+        reader.releaseLock()
+        dispatch({ type: 'RESET' })
       }
     }
   }, [stream])
 
+  console.log('Progress state:', progress)
+
   return (
-    <Box>
-      {Object.entries(progress).map(([file, stages]) => (
-        <div key={file} className="border p-2 rounded shadow">
-          <h3 className="font-bold text-lg">{file}</h3>
-          <ul className="ml-4 list-disc">
-            {Object.entries(stages).map(([stage, { count, done }]) => (
-              <li key={stage}>
-                <span className="font-medium">{stage}:</span> {done ? "✅ Done" : `🔄 Processing (${count}x)`}
-              </li>
+    <Table size="small">
+      <TableHead>
+        <TableRow>
+          <TableCell>File</TableCell>
+          {IngestionPipelineStageKeys.map((stage) => (
+            <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'}
+              </Typography>
+            </TableCell>
+          ))}
+        </TableRow>
+      </TableHead>
+      <TableBody>
+        {filenames.map((filename) => (
+          <TableRow key={filename}>
+            <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>
             ))}
-          </ul>
-        </div>
-      ))}
-    </Box>
+          </TableRow>
+        ))}
+      </TableBody>
+    </Table>
   )
 }

src/client/components/Rag/Rag.tsx

@@ -106,6 +106,7 @@ const Rag: React.FC = () => {
   const uploadMutation = useUploadMutation(selectedIndex)
   const [modalOpen, setModalOpen] = useState(false)
   const [stream, setStream] = useState<ReadableStream | null>(null)
+  const [filenames, setFilenames] = useState<string[]>([])
 
   const handleSubmit = async (event: React.FormEvent) => {
     event.preventDefault()
@@ -128,9 +129,13 @@ const Rag: React.FC = () => {
     setInputValue('')
   }
 
+  const handleUploadError = () => {
+    setStream(null)
+  }
+
   return (
     <Box sx={{ display: 'flex', gap: 2 }}>
-      <Dialog open={!!selectedIndex && modalOpen} onClose={() => setModalOpen(false)}>
+      <Dialog open={!!selectedIndex && modalOpen} onClose={() => setModalOpen(false)} fullWidth maxWidth="md">
         <DialogTitle>Edit {selectedIndex?.metadata?.name}</DialogTitle>
         <Box sx={{ padding: 2 }}>
           <Box sx={{ display: 'flex', gap: 2 }}>
@@ -143,6 +148,7 @@ const Rag: React.FC = () => {
                   console.log('Files selected:', files)
                   if (files && files.length > 0) {
                     const stream = await uploadMutation.mutateAsync(files)
+                    setFilenames(Array.from(files).map((file) => file.name))
                     setStream(stream)
                   }
                 }}
@@ -164,7 +170,7 @@ const Rag: React.FC = () => {
             </Button>
           </Box>
           <Box mt={2}>
-            <ProgressReporter stream={stream} />
+            <ProgressReporter filenames={filenames} stream={stream} onError={handleUploadError} />
           </Box>
         </Box>
       </Dialog>

src/config.ts

@@ -12,7 +12,7 @@ export const PUBLIC_URL = process.env.PUBLIC_URL || ''
 
 export const UPDATER_CRON_ENABLED = process.env.UPDATER_CRON_ENABLED === 'true'
 
-export const OLLAMA_HOST = process.env.OLLAMA_HOST || 'http://ollama:11434'
+export const OLLAMA_URL = process.env.OLLAMA_HOST || 'http://ollama:11434/v1/'
 export const RAG_ENABLED = process.env.RAG_ENABLED === 'true'
 
 export const DEFAULT_TOKEN_LIMIT = Number(process.env.DEFAULT_TOKEN_LIMIT) || 150_000
@@ -24,7 +24,7 @@ export const DEFAUL_CONTEXT_LIMIT = Number(process.env.DEFAUL_CONTEXT_LIMIT) ||
 export const DEFAULT_RESET_CRON = process.env.DEFAULT_RESET_CRON || '0 0 1 */3 *'
 
 export const EMBED_MODEL = process.env.EMBED_MODEL ?? 'text-embedding-small'
-export const EMBED_DIM = 1024
+export const EMBED_DIM = process.env.EMBED_DIM ? Number(process.env.EMBED_DIM) : 1024
 
 export const validModels = [
   {

src/server/routes/rag.ts

@@ -10,14 +10,15 @@ 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()
 
 const UPLOAD_DIR = 'uploads/rag'
 
 const IndexCreationSchema = z.object({
   name: z.string().min(1).max(100),
-  dim: z.number().min(1024).max(1024).default(EMBED_DIM),
+  dim: z.number().min(EMBED_DIM).max(EMBED_DIM).default(EMBED_DIM),
 })
 
 router.post('/indices', async (req, res) => {
@@ -136,7 +137,7 @@ router.post('/indices/:id/upload', [indexUploadDirMiddleware, uploadMiddleware],
   res.setHeader('Content-Type', 'application/x-ndjson')
   res.setHeader('Transfer-Encoding', 'chunked')
 
-  const openAiClient = getAzureOpenAIClient(EMBED_MODEL)
+  const openAiClient = getOllamaOpenAIClient() // getAzureOpenAIClient(EMBED_MODEL)
 
   const progressReporter = await ingestionPipeline(openAiClient, `uploads/rag/${id}`, ragIndex)
 

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

@@ -35,7 +35,7 @@ export class Chunker extends Transform {
       }),
     )
 
-    this.progressReporter.reportProgress(data.fileName)
+    this.progressReporter.reportProgress([data.fileName])
 
     callback()
   }

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

@@ -47,6 +47,7 @@ export class Embedder extends Transform {
   }
 
   private async embedBatch(callback: (error?: Error | null) => void) {
+    const currentBatchFilenames = this.currentBatch.map((chunk) => chunk.metadata.filename)
     try {
       const chunkContents = this.currentBatch.map((chunk) => chunk.content.join('\n'))
       const startedAt = Date.now()
@@ -65,15 +66,16 @@ export class Embedder extends Transform {
         // Save embedded chunk to cache
         const path = `${this.cachePath}/${embeddedChunk.id}.json`
         await writeFile(path, JSON.stringify(embeddedChunk, null, 2), 'utf-8')
-        this.progressReporter.reportProgress(embeddedChunk.metadata.filename)
       })
+      this.progressReporter.reportProgress(currentBatchFilenames)
 
       // Reset the current batch
       this.currentBatch = []
 
       callback()
     } catch (error) {
-      console.error(`Error saving chunk to cache: ${error}`)
+      console.error(`Embedding stage ${error}`)
+      this.progressReporter.reportError('Embedding chunk failed', currentBatchFilenames)
       callback(error as Error)
     }
   }

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

@@ -63,7 +63,7 @@ export class FileLoader extends Readable {
         this.progressReporter.reportDone()
       } else {
         this.push(file.value)
-        this.progressReporter.reportProgress(file.value.fileName)
+        this.progressReporter.reportProgress([file.value.fileName])
       }
     })
   }

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

@@ -7,8 +7,8 @@ import { RedisStorer } from './storer.ts'
 import type OpenAI from 'openai'
 import RagIndex from '../../../db/models/ragIndex.ts'
 import { TextExtractor } from './textExtractor.ts'
-import { Readable } from 'node:stream'
 import { ProgressReporter } from './progressReporter.ts'
+import { IngestionPipelineStageKeys } from '../../../../shared/constants.ts'
 
 // Pipeline debug cache in pipeline/
 // Check if exists, if not create it.
@@ -37,8 +37,8 @@ export const ingestionPipeline = async (client: OpenAI, loadpath: string, ragInd
     new RedisStorer(ragIndex),
   ]
 
-  stages.forEach((stage) => {
-    stage.progressReporter = progressReporter.getStageReporter(stage.constructor.name)
+  stages.forEach((stage, idx) => {
+    stage.progressReporter = progressReporter.getStageReporter(IngestionPipelineStageKeys[idx], idx)
   })
 
   pipeline(stages)
@@ -47,8 +47,8 @@ export const ingestionPipeline = async (client: OpenAI, loadpath: string, ragInd
       progressReporter.emit('end')
     })
     .catch((error) => {
-      console.error('Pipeline error:', error)
-      progressReporter.emit('error', 'Pipeline error')
+      console.error('Unhandled pipeline error:', error)
+      progressReporter.reportError(error)
     })
 
   return progressReporter