rag setup

Author: Veikkosuhonen

package-lock.json

@@ -112,7 +112,8 @@
     "unfuck-utf8-headers-middleware": "^1.0.1",
     "vite": "^5.0.12",
     "winston": "^3.11.0",
-    "winston-gelf-transporter": "^1.0.2"
+    "winston-gelf-transporter": "^1.0.2",
+    "zod": "^3.24.4"
   },
   "nodemonConfig": {
     "ignore": [

package.json

@@ -26,6 +26,9 @@ export const DEFAUL_CONTEXT_LIMIT =
 export const DEFAULT_RESET_CRON =
   process.env.DEFAULT_RESET_CRON || '0 0 1 */3 *'
 
+export const EMBED_MODEL = 'text-embedding-small'
+export const EMBED_DIM = 1024
+
 export const validModels = [
   {
     name: 'gpt-4',

src/config.ts

@@ -0,0 +1,36 @@
+import { DataTypes } from 'sequelize'
+
+import { Migration } from '../connection'
+
+export const up: Migration = ({ context: queryInterface }) =>
+  queryInterface.createTable('rag_indices', {
+    id: {
+      type: DataTypes.INTEGER,
+      allowNull: false,
+      primaryKey: true,
+      autoIncrement: true,
+    },
+    user_id: {
+      type: DataTypes.STRING,
+      allowNull: false,
+    },
+    course_id: {
+      type: DataTypes.STRING,
+      allowNull: true,
+    },
+    metadata: {
+      type: DataTypes.JSONB,
+      allowNull: true,
+    },
+    created_at: {
+      type: DataTypes.DATE,
+      allowNull: false,
+    },
+    updated_at: {
+      type: DataTypes.DATE,
+      allowNull: false,
+    },
+  })
+
+export const down: Migration = ({ context: queryInterface }) =>
+  queryInterface.dropTable('rag_indices')

src/server/db/migrations/20250516_00_init_rag_indices.ts

@@ -5,6 +5,7 @@ import Prompt from './prompt'
 import Enrolment from './enrolment'
 import Responsibility from './responsibilities'
 import Discussion from './discussion'
+import RagIndex from './ragIndex'
 
 User.belongsToMany(ChatInstance, {
   through: UserChatInstanceUsage,
@@ -42,6 +43,10 @@ Responsibility.belongsTo(ChatInstance, { as: 'chatInstance' })
 
 ChatInstance.hasMany(Responsibility, { as: 'responsibilities' })
 
+User.hasMany(RagIndex, { as: 'ragIndices' })
+
+RagIndex.belongsTo(User, { as: 'user' })
+
 export {
   User,
   ChatInstance,
@@ -50,4 +55,5 @@ export {
   Enrolment,
   Responsibility,
   Discussion,
+  RagIndex,
 }

src/server/db/models/index.ts

@@ -0,0 +1,56 @@
+import {
+  Model,
+  InferAttributes,
+  InferCreationAttributes,
+  CreationOptional,
+  DataTypes,
+} from 'sequelize'
+
+import { sequelize } from '../connection'
+
+export type RagIndexMetadata = {
+  name: string
+  dim: number
+}
+
+class RagIndex extends Model<
+  InferAttributes<RagIndex>,
+  InferCreationAttributes<RagIndex>
+> {
+  declare id: CreationOptional<number>
+
+  declare userId: string
+
+  declare courseId?: string
+
+  declare metadata: RagIndexMetadata
+}
+
+RagIndex.init(
+  {
+    id: {
+      type: DataTypes.INTEGER,
+      allowNull: false,
+      primaryKey: true,
+      autoIncrement: true,
+    },
+    userId: {
+      type: DataTypes.STRING,
+      allowNull: false,
+    },
+    courseId: {
+      type: DataTypes.STRING,
+      allowNull: true,
+    },
+    metadata: {
+      type: DataTypes.JSONB,
+      allowNull: true,
+    },
+  },
+  {
+    underscored: true,
+    sequelize,
+  }
+)
+
+export default RagIndex

src/server/db/models/ragIndex.ts

@@ -20,7 +20,6 @@ import { connectToDatabase } from './db/connection'
 import seed from './db/seeders'
 import setupCron from './util/cron'
 import { updateLastRestart } from './util/lastRestart'
-import { initRag } from './util/rag'
 
 const app = express()
 
@@ -45,9 +44,6 @@ app.listen(PORT, async () => {
   await connectToDatabase()
   await seed()
   await updateLastRestart()
-  if (RAG_ENABLED) {
-    await initRag()
-  }
   if (inProduction || inStaging) {
     await setupCron()
   }

src/server/index.ts

@@ -1,21 +1,78 @@
 import { Router } from 'express'
-import { searchRag } from '../util/rag'
+import { EMBED_DIM } from '../../config'
+import { createChunkIndex, deleteChunkIndex } from '../services/rag/chunkDb'
+import { RagIndex } from '../db/models'
+import { RequestWithUser } from '../types'
+import z from 'zod'
+import { queryRagIndex } from '../services/rag/query'
 
 const router = Router()
 
+const IndexCreationSchema = z.object({
+  name: z.string().min(1).max(100),
+  dim: z.number().min(1024).max(1024).default(EMBED_DIM),
+})
+
+router.post('/indices', async (req, res) => {
+  const { user } = req as RequestWithUser
+  const { name, dim } = IndexCreationSchema.parse(req.body)
+
+  const ragIndex = await RagIndex.create({
+    userId: user.id,
+    metadata: {
+      name,
+      dim,
+    },
+  })
+
+  await createChunkIndex(ragIndex)
+
+  res.json(ragIndex)
+})
+
+router.delete('/indices/:id', async (req, res) => {
+  const { user } = req as unknown as RequestWithUser // <- fix type
+  const { id } = req.params
+
+  const ragIndex = await RagIndex.findOne({
+    where: { id, userId: user.id },
+  })
+
+  if (!ragIndex) {
+    res.status(404).json({ error: 'Index not found' })
+    return
+  }
+
+  await deleteChunkIndex(ragIndex)
+
+  await ragIndex.destroy()
+
+  res.json({ message: 'Index deleted' })
+})
+
+router.get('/indices', async (_req, res) => {
+  const indices = await RagIndex.findAll()
+
+  res.json(indices)
+})
+
+const RagIndexQuerySchema = z.object({
+  query: z.string().min(1).max(1000),
+  topK: z.number().min(1).max(100).default(5),
+  indexId: z.number(),
+})
 router.post('/query', async (req, res) => {
-  console.log('Received request on /rag/query')
-  console.log('Request body:', req.body)
-  console.log('Request headers:', req.headers)
-  console.log('Request method:', req.method)
-  try {
-    const prompt = req.body.prompt
-    const answer = await searchRag(prompt)
-    res.json(answer)
-  } catch (error) {
-    console.error('Error in /rag/query:', error)
-    res.status(500).json({ error: 'Rag failed' })
+  const { query, topK, indexId } = RagIndexQuerySchema.parse(req.body)
+
+  const ragIndex = await RagIndex.findByPk(indexId)
+
+  if (!ragIndex) {
+    res.status(404).json({ error: 'Index not found' })
+    return
   }
+
+  const results = await queryRagIndex(ragIndex, query, topK)
+  res.json(results)
 })
 
 export default router

src/server/routes/rag.ts

@@ -0,0 +1,123 @@
+import { RagIndex } from '../../db/models'
+import { redisClient } from '../../util/redis'
+
+export const createChunkIndex = async (ragIndex: RagIndex) => {
+  try {
+    await redisClient.ft.create(
+      ragIndex.metadata.name,
+      {
+        metadata: {
+          type: 'TEXT',
+        },
+        content: {
+          type: 'TEXT',
+        },
+        embedding: {
+          type: 'VECTOR',
+          TYPE: 'FLOAT32',
+          ALGORITHM: 'HNSW',
+          DIM: ragIndex.metadata.dim,
+          DISTANCE_METRIC: 'L2',
+        },
+      },
+      {
+        ON: 'HASH',
+        PREFIX: `idx:${ragIndex.metadata.name}`,
+      }
+    )
+
+    console.log(`Index ${ragIndex.metadata.name} created`)
+  } catch (err: any) {
+    if (err.message.includes('Index already exists')) {
+      console.log(`Index ${ragIndex.metadata.name} already exists`)
+    } else {
+      console.error('Error creating index', err)
+    }
+  }
+}
+
+export const deleteChunkIndex = async (ragIndex: RagIndex) => {
+  try {
+    await redisClient.ft.dropIndex(ragIndex.metadata.name)
+    console.log(`Index ${ragIndex.metadata.name} deleted`)
+  } catch (err: any) {
+    if (err.message.includes('Index not found')) {
+      console.log(`Index ${ragIndex.metadata.name} not found`)
+    } else {
+      console.error('Error deleting index', err)
+    }
+  }
+}
+
+export const addChunk = async (
+  ragIndex: RagIndex,
+  {
+    id,
+    metadata,
+    content,
+    embedding,
+  }: {
+    id: string
+    metadata?: { [key: string]: any }
+    content: string
+    embedding: number[]
+  }
+) => {
+  const embeddingBuffer = Buffer.copyBytesFrom(new Float32Array(embedding))
+
+  // Check if the embedding length is correct
+  if (embeddingBuffer.length !== 4 * ragIndex.metadata.dim) {
+    throw new Error(
+      `Embedding length is incorrect, got ${embeddingBuffer.length} bytes`
+    )
+  }
+
+  await redisClient.hSet(`idx:${ragIndex.metadata.name}:${id}`, {
+    metadata: JSON.stringify(metadata || {}),
+    content,
+    embedding: embeddingBuffer,
+  })
+
+  console.log(`Document ${id} added to index ${ragIndex.metadata.name}`)
+}
+
+export const searchKChunks = async (
+  ragIndex: RagIndex,
+  embedding: number[],
+  k: number
+) => {
+  const embeddingBuffer = Buffer.copyBytesFrom(new Float32Array(embedding))
+
+  if (embeddingBuffer.length !== 4 * ragIndex.metadata.dim) {
+    throw new Error(
+      `Embedding length is incorrect, got ${embeddingBuffer.length} bytes`
+    )
+  }
+
+  const queryString = `(*)=>[KNN ${k} @embedding $vec_param AS score]`
+
+  const results = await redisClient.ft.search(
+    ragIndex.metadata.name,
+    queryString,
+    {
+      PARAMS: {
+        vec_param: embeddingBuffer,
+      },
+      DIALECT: 2,
+      RETURN: ['content', 'title', 'score'], // Specify the fields to return
+    }
+  )
+
+  return results as {
+    total: number
+    documents: {
+      id: string
+      value: {
+        content: string
+        title: string
+        score: number
+        metadata: string
+      }
+    }[]
+  }
+}