improve v2 endpoint

Author: Veikkosuhonen

src/server/routes/openai.ts

@@ -63,7 +63,7 @@ const PostStreamSchemaV2 = z.object({
 openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
   const req = r as RequestWithUser
   const { options, courseId } = PostStreamSchemaV2.parse(JSON.parse(req.body.data))
-  const { userConsent } = options
+  const { userConsent, ragIndexId } = options
   const { user } = req
 
   console.log('options', options)
@@ -73,8 +73,26 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
     return
   }
 
+  const course =
+    courseId &&
+    (await ChatInstance.findOne({
+      where: { courseId },
+    }))
+
+  if (courseId && !course) {
+    res.status(404).send('Course not found')
+    return
+  }
+
   // Check if the user has usage limits for the course or model
-  const usageAllowed = courseId ? await checkCourseUsage(user, courseId) : options.model === FREE_MODEL || (await checkUsage(user, options.model))
+  let usageAllowed = false
+  if (courseId) {
+    usageAllowed = await checkCourseUsage(user, courseId)
+  } else if (options.model === FREE_MODEL) {
+    usageAllowed = true
+  } else {
+    usageAllowed = await checkUsage(user, options.model)
+  }
 
   if (!usageAllowed) {
     res.status(403).send('Usage limit reached')
@@ -123,6 +141,7 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
     return
   }
 
+  // Check context limit
   const contextLimit = getModelContextLimit(options.model)
 
   if (tokenCount > contextLimit) {
@@ -131,25 +150,40 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
     return
   }
 
+  // Check rag index
   let vectorStoreId: string | undefined = undefined
-  if (courseId) {
-    const ragIndex = await RagIndex.findOne({
-      where: { courseId },
-    })
+  let instructions: string | undefined = undefined
+
+  if (ragIndexId) {
+    const ragIndex = await RagIndex.findByPk(ragIndexId)
     if (ragIndex) {
+      if (courseId && ragIndex.courseId !== courseId) {
+        logger.error('RagIndex does not belong to the course', { ragIndexId, courseId })
+        res.status(403).send('RagIndex does not belong to the course')
+        return
+      }
+
       vectorStoreId = ragIndex.metadata.azureVectorStoreId
+      instructions = ragIndex.metadata.instructions
+
       console.log('using', ragIndex.toJSON())
+    } else {
+      logger.error('RagIndex not found', { ragIndexId })
+      res.status(404).send('RagIndex not found')
+      return
     }
   }
 
   const responsesClient = new ResponsesClient({
     model: options.model,
     courseId,
     vectorStoreId,
+    instructions,
   })
 
-  const latestMessage = options.messages[options.messages.length - 1] // Adhoc to input only the latest message
-  const events = await responsesClient.createResponse({ input: [latestMessage], prevResponseId: options.prevResponseId })
+  const latestMessage = options.messages[options.messages.length - 1]
+
+  const events = await responsesClient.createResponse({ input: [latestMessage], prevResponseId: options.prevResponseId, include: ragIndexId ? ['file_search_call.results'] : [] })
 
   if (isError(events)) {
     res.status(424)
@@ -184,12 +218,6 @@ openaiRouter.post('/stream/v2', upload.single('file'), async (r, res) => {
     courseId,
   })
 
-  const course =
-    courseId &&
-    (await ChatInstance.findOne({
-      where: { courseId },
-    }))
-
   const consentToSave = courseId && course.saveDiscussions && options.saveConsent
 
   console.log('consentToSave', options.saveConsent, user.username)

src/server/routes/rag.ts

@@ -45,7 +45,7 @@ router.post('/indices', async (req, res) => {
     return
   }
 
-  const client = getAzureOpenAIClient('curredev4omini')
+  const client = getAzureOpenAIClient()
   const vectorStore = await client.vectorStores.create({
     name,
   })
@@ -77,7 +77,7 @@ router.delete('/indices/:id', async (req, res) => {
     return
   }
 
-  const client = getAzureOpenAIClient('curredev4omini')
+  const client = getAzureOpenAIClient()
   try {
     await client.vectorStores.del(ragIndex.metadata.azureVectorStoreId)
   } catch (error) {
@@ -123,7 +123,7 @@ router.get('/indices', async (req, res) => {
   })
 
   if (includeExtras) {
-    const client = getAzureOpenAIClient('curredev4omini')
+    const client = getAzureOpenAIClient()
 
     // Add ragFileCount to each index
     const indicesWithCount = await Promise.all(
@@ -159,7 +159,7 @@ router.get('/indices/:id', async (req, res) => {
     return
   }
 
-  const client = getAzureOpenAIClient('curredev4omini')
+  const client = getAzureOpenAIClient()
   const vectorStore = await client.vectorStores.retrieve(ragIndex.metadata.azureVectorStoreId)
 
   res.json({
@@ -265,7 +265,7 @@ router.post('/indices/:id/upload', [indexUploadDirMiddleware, uploadMiddleware],
     ),
   )
 
-  const client = getAzureOpenAIClient('curredev4omini')
+  const client = getAzureOpenAIClient()
 
   const uploadDirPath = `${UPLOAD_DIR}/${id}`
 

src/server/util/azure/ResponsesAPI.ts

@@ -1,22 +1,20 @@
 import { Tiktoken } from '@dqbd/tiktoken'
 import { Response } from 'express'
-import { isError } from '../parser'
 
 import { AZURE_RESOURCE, AZURE_API_KEY } from '../config'
-import { validModels, inProduction } from '../../../config'
+import { validModels } from '../../../config'
 import logger from '../logger'
 
 import { APIError } from '../../types'
 import { AzureOpenAI } from 'openai'
 
 // import { EventStream } from '@azure/openai'
 import { Stream } from 'openai/streaming'
-import { FileSearchTool, FunctionTool, ResponseInput, ResponseInputItem, ResponseStreamEvent, ResponseTextAnnotationDeltaEvent } from 'openai/resources/responses/responses'
+import { FileSearchTool, ResponseIncludable, ResponseInput, ResponseStreamEvent } from 'openai/resources/responses/responses'
 
-import { courseAssistants } from './courseAssistants'
 import { createFileSearchTool } from './util'
 
-import type { CourseAssistant, FileCitation, ResponseStreamEventData } from '../../../shared/types'
+import type { FileCitation, ResponseStreamEventData } from '../../../shared/types'
 
 const endpoint = `https://${AZURE_RESOURCE}.openai.azure.com/`
 
@@ -34,21 +32,12 @@ export class ResponsesClient {
   model: string
   instructions: string
   tools: FileSearchTool[]
-  courseAssistant: CourseAssistant
 
-  constructor({ model, courseId, vectorStoreId }: { model: string; courseId?: string; vectorStoreId?: string }) {
+  constructor({ model, courseId, vectorStoreId, instructions }: { model: string; courseId?: string; vectorStoreId?: string; instructions?: string }) {
     const deploymentId = validModels.find((m) => m.name === model)?.deployment
 
     if (!deploymentId) throw new Error(`Invalid model: ${model}, not one of ${validModels.map((m) => m.name).join(', ')}`)
 
-    if (courseId) {
-      this.courseAssistant = courseAssistants.find((assistant) => assistant.course_id === courseId)
-
-      if (!this.courseAssistant) throw new Error(`No course assistant found for course ID: ${courseId}`)
-    } else {
-      this.courseAssistant = courseAssistants.find((assistant) => assistant.name === 'default')
-    }
-
     const fileSearchTool = courseId
       ? [
           createFileSearchTool({
@@ -58,11 +47,19 @@ export class ResponsesClient {
       : [] // needs to retrun empty array for null
 
     this.model = deploymentId
-    this.instructions = this.courseAssistant.assistant_instruction
+    this.instructions = instructions
     this.tools = fileSearchTool
   }
 
-  async createResponse({ input, prevResponseId }: { input: ResponseInput; prevResponseId?: string }): Promise<Stream<ResponseStreamEvent> | APIError> {
+  async createResponse({
+    input,
+    prevResponseId,
+    include,
+  }: {
+    input: ResponseInput
+    prevResponseId?: string
+    include?: ResponseIncludable[]
+  }): Promise<Stream<ResponseStreamEvent> | APIError> {
     try {
       return await client.responses.create({
         model: this.model,
@@ -73,6 +70,7 @@ export class ResponsesClient {
         tools: this.tools,
         tool_choice: 'auto',
         store: true,
+        include,
       })
     } catch (error: any) {
       logger.error(error)

src/server/util/azure/client.ts

@@ -1,4 +1,4 @@
-import { OpenAIClient, AzureKeyCredential, EventStream, ChatCompletions, GetEmbeddingsOptions } from '@azure/openai'
+import { OpenAIClient, AzureKeyCredential, EventStream, ChatCompletions } from '@azure/openai'
 import { Tiktoken } from '@dqbd/tiktoken'
 import { Response } from 'express'
 
@@ -12,7 +12,7 @@ const endpoint = `https://${AZURE_RESOURCE}.openai.azure.com/`
 
 const oldClient = new OpenAIClient(endpoint, new AzureKeyCredential(AZURE_API_KEY))
 
-export const getAzureOpenAIClient = (deployment: string) =>
+export const getAzureOpenAIClient = (deployment?: string) =>
   new AzureOpenAI({
     apiKey: AZURE_API_KEY,
     deployment,

src/server/util/azure/courseAssistants.ts

@@ -1,7 +1,8 @@
 export type RagIndexMetadata = {
   name: string
-  dim: number
+  dim?: number
   azureVectorStoreId: string
+  instructions?: string
 }
 
 export type RagFileAttributes = {