Store rag index chunks in one vector store

Author: Veikkosuhonen

src/server/routes/openai.ts

@@ -6,7 +6,7 @@ import { ChatInstance, Discussion, RagIndex } from '../db/models'
 import { calculateUsage, checkCourseUsage, checkUsage, incrementCourseUsage, incrementUsage } from '../services/chatInstances/usage'
 import type { RequestWithUser } from '../types'
 import { getCompletionEvents, streamCompletion } from '../util/azure/client'
-import { FileSearchResultsStore } from '../util/azure/fileSearchResultsStore'
+import { FileSearchResultsStore } from '../services/azureFileSearch/fileSearchResultsStore'
 import { ResponsesClient } from '../util/azure/ResponsesAPI'
 import { DEFAULT_RAG_SYSTEM_PROMPT } from '../util/config'
 import logger from '../util/logger'
@@ -166,6 +166,7 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
   const responsesClient = new ResponsesClient({
     model: options.model,
     vectorStoreId,
+    ragIndexId,
     instructions,
     temperature: options.modelTemperature,
     user,
@@ -189,7 +190,6 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
     events,
     encoding,
     res,
-    ragIndexId,
   })
 
   tokenCount += result.tokenCount

src/server/routes/rag/rag.ts

@@ -1,11 +1,13 @@
 import { Router } from 'express'
+import z from 'zod/v4'
 import { EMBED_DIM } from '../../../config'
 import { ChatInstance, ChatInstanceRagIndex, RagFile, RagIndex, Responsibility } from '../../db/models'
-import { RequestWithUser, User } from '../../types'
-import z from 'zod/v4'
+import type { RequestWithUser, User } from '../../types'
+import { ApplicationError } from '../../util/ApplicationError'
 import { getAzureOpenAIClient } from '../../util/azure/client'
+import { TEST_COURSES } from '../../util/config'
 import ragIndexRouter, { ragIndexMiddleware } from './ragIndex'
-import { ApplicationError } from '../../util/ApplicationError'
+import { getPrimaryVectorStoreId } from '../../services/azureFileSearch/vectorStore'
 
 const router = Router()
 
@@ -25,11 +27,17 @@ router.post('/indices', async (req, res) => {
   const { name, dim, chatInstanceId } = IndexCreationSchema.parse(req.body)
 
   const chatInstance = await ChatInstance.findByPk(chatInstanceId, {
-    include: {
-      model: Responsibility,
-      as: 'responsibilities',
-      required: true, // Ensure the user is responsible for the course
-    },
+    include: [
+      {
+        model: Responsibility,
+        as: 'responsibilities',
+        required: true, // Ensure the user is responsible for the course
+      },
+      {
+        model: RagIndex,
+        as: 'ragIndices',
+      },
+    ],
   })
 
   if (!chatInstance) {
@@ -40,17 +48,18 @@ router.post('/indices', async (req, res) => {
     throw ApplicationError.Forbidden('Cannot create index, user is not responsible for the course')
   }
 
-  const client = getAzureOpenAIClient()
-  const vectorStore = await client.vectorStores.create({
-    name,
-  })
+  if (chatInstance.courseId !== TEST_COURSES.OTE_SANDBOX.id && (chatInstance.ragIndices ?? []).length > 0) {
+    throw ApplicationError.Forbidden('Cannot create index, index already exists on the course')
+  }
+
+  const vectorStoreId = await getPrimaryVectorStoreId()
 
   const ragIndex = await RagIndex.create({
     userId: user.id,
     metadata: {
       name,
       dim,
-      azureVectorStoreId: vectorStore.id,
+      azureVectorStoreId: vectorStoreId,
     },
   })
 

src/server/routes/rag/ragIndex.ts

@@ -58,14 +58,6 @@ ragIndexRouter.delete('/', async (req, res) => {
   const ragIndexRequest = req as RagIndexRequest
   const ragIndex = ragIndexRequest.ragIndex
 
-  const client = getAzureOpenAIClient()
-  try {
-    await client.vectorStores.del(ragIndex.metadata.azureVectorStoreId)
-  } catch (error) {
-    console.error(`Failed to delete Azure vector store ${ragIndex.metadata.azureVectorStoreId}:`, error)
-    throw ApplicationError.InternalServerError('Failed to delete Azure vector store')
-  }
-
   const uploadPath = `${UPLOAD_DIR}/${ragIndex.id}`
   try {
     await rm(uploadPath, { recursive: true, force: true })
@@ -242,7 +234,18 @@ ragIndexRouter.post('/upload', [indexUploadDirMiddleware, uploadMiddleware], asy
     ragFiles.map(async (ragFile) => {
       const filePath = `${uploadDirPath}/${ragFile.filename}`
       const stream = fs.createReadStream(filePath)
-      const vectorStoreFile = await client.vectorStores.files.upload(ragIndex.metadata.azureVectorStoreId, stream)
+
+      const uploadedFile = await client.files.create({
+        file: stream,
+        purpose: 'user_data',
+      })
+      const vectorStoreFile = await client.vectorStores.files.create(ragIndex.metadata.azureVectorStoreId, {
+        file_id: uploadedFile.id,
+        attributes: {
+          ragIndexId: ragIndex.id,
+        },
+      })
+
       console.log(`File ${filePath} uploaded to vector store`)
       await RagFile.update(
         {

src/server/util/azure/fileSearchResultsStore.ts renamed to src/server/services/azureFileSearch/fileSearchResultsStore.ts

@@ -1,6 +1,6 @@
 import type { FileSearchResultData } from '../../../shared/types'
 import type { User } from '../../types'
-import { redisClient } from '../redis'
+import { redisClient } from '../../util/redis'
 
 export const FileSearchResultsStore = {
   getKey(user: User, fileSearchId: string): string {

src/server/services/azureFileSearch/vectorStore.ts

@@ -0,0 +1,36 @@
+import { getAzureOpenAIClient } from '../../util/azure/client'
+import { redisClient } from '../../util/redis'
+
+// Strategy: have one big vector store
+
+export const getPrimaryVectorStoreId = async () => {
+  const vectorStoreId = await redisClient.get('primaryVectorStoreId')
+  if (vectorStoreId) {
+    return vectorStoreId
+  }
+
+  const client = getAzureOpenAIClient()
+
+  const vectorStores = (await client.vectorStores.list()).getPaginatedItems()
+
+  let primaryVectorStore = vectorStores.find((store) => store.metadata?.primaryRagIndicesVectorStore === 'true')
+
+  if (!primaryVectorStore) {
+    // Create a new vector store if none exist
+    primaryVectorStore = await client.vectorStores.create({
+      name: 'CurreChat file search vector store',
+      expires_after: {
+        anchor: 'last_active_at',
+        days: 60,
+      },
+      metadata: {
+        // Identify this vector store as the primary one for RAG indices
+        primaryRagIndicesVectorStore: 'true',
+      },
+    })
+  }
+
+  await redisClient.set('primaryVectorStoreId', primaryVectorStore.id)
+
+  return primaryVectorStore.id
+}

src/server/util/azure/ResponsesAPI.ts

@@ -1,27 +1,16 @@
 import type { Tiktoken } from '@dqbd/tiktoken'
 import type { Response } from 'express'
-import { AzureOpenAI } from 'openai'
 import type { FileSearchTool, ResponseIncludable, ResponseInput, ResponseItemsPage, ResponseStreamEvent } from 'openai/resources/responses/responses'
 import type { Stream } from 'openai/streaming'
 import { z } from 'zod/v4'
 import { validModels } from '../../../config'
 import type { ResponseStreamEventData } from '../../../shared/types'
 import type { APIError, User } from '../../types'
-import { AZURE_API_KEY, AZURE_RESOURCE } from '../config'
 import logger from '../logger'
 import { createMockStream } from './mocks/MockStream'
 import { createFileSearchTool } from './util'
-import { FileSearchResultsStore } from './fileSearchResultsStore'
-
-const endpoint = `https://${AZURE_RESOURCE}.openai.azure.com/`
-
-export const getAzureOpenAIClient = (deployment: string) =>
-  new AzureOpenAI({
-    apiKey: AZURE_API_KEY,
-    deployment,
-    apiVersion: '2025-03-01-preview',
-    endpoint,
-  })
+import { FileSearchResultsStore } from '../../services/azureFileSearch/fileSearchResultsStore'
+import { getAzureOpenAIClient } from './client'
 
 const client = getAzureOpenAIClient(process.env.GPT_4O_MINI ?? '')
 
@@ -38,37 +27,43 @@ export class ResponsesClient {
   temperature: number
   tools: FileSearchTool[]
   user: User
+  ragIndexId?: number
 
   constructor({
     model,
     temperature,
     vectorStoreId,
     instructions,
     user,
+    ragIndexId,
   }: {
     model: string
     temperature: number
     vectorStoreId?: string
     instructions?: string
     user: User
+    ragIndexId?: number
   }) {
     const selectedModel = validModels.find((m) => m.name === model)?.deployment
 
     if (!selectedModel) throw new Error(`Invalid model: ${model}, not one of ${validModels.map((m) => m.name).join(', ')}`)
 
-    const fileSearchTool = vectorStoreId
-      ? [
-          createFileSearchTool({
-            vectorStoreId,
-          }),
-        ]
-      : [] // needs to retrun empty array for null
+    const fileSearchTool =
+      vectorStoreId && ragIndexId
+        ? [
+            createFileSearchTool({
+              vectorStoreId,
+              filters: { key: 'ragIndexId', value: ragIndexId, type: 'eq' },
+            }),
+          ]
+        : [] // needs to retrun empty array for null
 
     this.model = selectedModel
     this.temperature = temperature
     this.instructions = instructions ?? ''
     this.tools = fileSearchTool
     this.user = user
+    this.ragIndexId = ragIndexId
   }
 
   async createResponse({
@@ -116,7 +111,7 @@ export class ResponsesClient {
     }
   }
 
-  async handleResponse({ events, encoding, res, ragIndexId }: { events: Stream<ResponseStreamEvent>; encoding: Tiktoken; res: Response; ragIndexId?: number }) {
+  async handleResponse({ events, encoding, res }: { events: Stream<ResponseStreamEvent>; encoding: Tiktoken; res: Response }) {
     let tokenCount = 0
     const contents: string[] = []
 
@@ -150,7 +145,7 @@ export class ResponsesClient {
 
         case 'response.output_item.done': {
           if (event.item.type === 'file_search_call') {
-            if (!ragIndexId) throw new Error('how is this possible. you managed to invoke file search without ragIndexId')
+            if (!this.ragIndexId) throw new Error('how is this possible. you managed to invoke file search without ragIndexId')
 
             if (event.item.results) {
               await FileSearchResultsStore.saveResults(event.item.id, event.item.results, this.user)
@@ -164,7 +159,7 @@ export class ResponsesClient {
                   queries: event.item.queries,
                   status: event.item.status,
                   type: event.item.type,
-                  ragIndexId,
+                  ragIndexId: this.ragIndexId,
                 },
               },
               res,