From 8abcd8dbcbe8bd946bb4b5e56847808133956e1e Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 15:23:36 +0000 Subject: [PATCH 1/9] feat: only trust the lesson plan from the server side --- apps/nextjs/src/app/api/chat/chatHandler.ts | 104 +++++++++++++++--- .../components/AppComponents/Chat/AiSdk.tsx | 2 - package.json | 3 +- .../aila/src/core/document/AilaHomework.ts | 0 .../aila/src/core/document/AilaLessonPlan.ts | 0 5 files changed, 90 insertions(+), 19 deletions(-) create mode 100644 packages/aila/src/core/document/AilaHomework.ts create mode 100644 packages/aila/src/core/document/AilaLessonPlan.ts diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index eaf82f1a8..77c17bc63 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -50,12 +50,10 @@ async function setupChatHandler(req: NextRequest) { const { id: chatId, messages, - lessonPlan = {}, options: chatOptions = {}, }: { id: string; messages: Message[]; - lessonPlan?: LooseLessonPlan; options?: AilaPublicChatOptions; } = json; @@ -85,7 +83,6 @@ async function setupChatHandler(req: NextRequest) { return { chatId, messages, - lessonPlan, options, llmService, moderationAiClient, @@ -95,13 +92,19 @@ async function setupChatHandler(req: NextRequest) { ); } -function setTelemetryMetadata( - span: TracingSpan, - id: string, - messages: Message[], - lessonPlan: LooseLessonPlan, - options: AilaOptions, -) { +function setTelemetryMetadata({ + span, + id, + messages, + lessonPlan, + options, +}: { + span: TracingSpan; + id: string; + messages: Message[]; + lessonPlan: LooseLessonPlan; + options: AilaOptions; +}) { span.setTag("chat_id", id); span.setTag("messages.count", messages.length); span.setTag("has_lesson_plan", Object.keys(lessonPlan).length > 0); @@ -167,6 +170,69 @@ async function generateChatStream( ); } +function hasLessonPlan(obj: unknown): obj is { lessonPlan: unknown } { + return obj !== null && typeof obj === "object" && "lessonPlan" in obj; +} + +function isValidLessonPlan(lessonPlan: unknown): boolean { + return lessonPlan !== null && typeof lessonPlan === "object"; +} + +function parseLessonPlanFromOutput(output: unknown): LooseLessonPlan { + if (!output) return {}; + + try { + const parsedOutput = + typeof output === "string" ? JSON.parse(output) : output; + + if ( + hasLessonPlan(parsedOutput) && + isValidLessonPlan(parsedOutput.lessonPlan) + ) { + return parsedOutput.lessonPlan as LooseLessonPlan; + } + } catch (error) { + log.error("Error parsing output to extract lesson plan", error); + } + + return {}; +} + +async function loadLessonPlanFromDatabase( + chatId: string, + userId: string, +): Promise { + try { + const chat = await prisma.appSession.findUnique({ + where: { id: chatId }, + select: { + id: true, + userId: true, + output: true, + }, + }); + + if (!chat) { + log.info(`No existing chat found for id: ${chatId}`); + return {}; + } + + if (chat.userId !== userId) { + log.error( + `User ${userId} attempted to access chat ${chatId} which belongs to ${chat.userId}`, + ); + throw new Error("Unauthorized access to chat"); + } + + const lessonPlan = parseLessonPlanFromOutput(chat.output); + log.info(`Loaded lesson plan for chat ${chatId}`); + return lessonPlan; + } catch (error) { + log.error(`Error loading lesson plan for chat ${chatId}`, error); + return {}; + } +} + export async function handleChatPostRequest( req: NextRequest, config: Config, @@ -175,22 +241,29 @@ export async function handleChatPostRequest( const { chatId, messages, - lessonPlan, options, llmService, moderationAiClient, threatDetectors, } = await setupChatHandler(req); - setTelemetryMetadata(span, chatId, messages, lessonPlan, options); - let userId: string | undefined; let aila: Aila | undefined; try { userId = await fetchAndCheckUser(chatId); - span.setTag("user_id", userId); + + const dbLessonPlan = await loadLessonPlanFromDatabase(chatId, userId); + + setTelemetryMetadata({ + span, + id: chatId, + messages, + lessonPlan: dbLessonPlan, + options, + }); + aila = await withTelemetry( "chat-create-aila", { chat_id: chatId, user_id: userId }, @@ -213,8 +286,7 @@ export async function handleChatPostRequest( ], threatDetectors: () => threatDetectors, }, - - lessonPlan: lessonPlan ?? {}, + lessonPlan: dbLessonPlan ?? {}, }; const result = await config.createAila(ailaOptions); return result; diff --git a/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx b/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx index 9b5e6b369..b99318393 100644 --- a/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx +++ b/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx @@ -69,8 +69,6 @@ export function AiSdk({ id }: Readonly) { id, body: { id, - // NOTE: this lesson plan is used by the chat endpoint - lessonPlan, options: { useRag: true, temperature: 0.7, diff --git a/package.json b/package.json index ad5cae0cd..03ba6b35c 100644 --- a/package.json +++ b/package.json @@ -24,7 +24,8 @@ "doppler:pull:stg": "doppler secrets download --config stg --no-file --format env > .env", "doppler:run:stg": "doppler run -c stg --silent", "format": "prettier --write \"**/*.{ts,tsx,md}\"", - "lint": "turbo lint", + "lint": "turbo lint --pass-with-no-tests", + "lint:file": "eslint", "lint:fix": "pnpm lint -- --fix", "lint:debug": "pnpm lint -- --debug", "prompts": "turbo prompts", diff --git a/packages/aila/src/core/document/AilaHomework.ts b/packages/aila/src/core/document/AilaHomework.ts new file mode 100644 index 000000000..e69de29bb diff --git a/packages/aila/src/core/document/AilaLessonPlan.ts b/packages/aila/src/core/document/AilaLessonPlan.ts new file mode 100644 index 000000000..e69de29bb From 3edac7807ad3c3107cc2801aeb9480091107cd64 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 15:27:23 +0000 Subject: [PATCH 2/9] Revert lint change --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 03ba6b35c..b5cff5ee7 100644 --- a/package.json +++ b/package.json @@ -24,7 +24,7 @@ "doppler:pull:stg": "doppler secrets download --config stg --no-file --format env > .env", "doppler:run:stg": "doppler run -c stg --silent", "format": "prettier --write \"**/*.{ts,tsx,md}\"", - "lint": "turbo lint --pass-with-no-tests", + "lint": "turbo lint", "lint:file": "eslint", "lint:fix": "pnpm lint -- --fix", "lint:debug": "pnpm lint -- --debug", From 5945091ae00644bf222aeb0bccecb3cf8d2e0221 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 16:10:06 +0000 Subject: [PATCH 3/9] feat: load messages from the db --- apps/nextjs/src/app/api/chat/chatHandler.ts | 204 +++++++++++++++----- 1 file changed, 155 insertions(+), 49 deletions(-) diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index 77c17bc63..15285097a 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -12,6 +12,7 @@ import { PosthogAnalyticsAdapter, } from "@oakai/aila/src/features/analytics"; import { AilaRag } from "@oakai/aila/src/features/rag/AilaRag"; +import type { AilaThreatDetector } from "@oakai/aila/src/features/threatDetection"; import { HeliconeThreatDetector } from "@oakai/aila/src/features/threatDetection/detectors/helicone/HeliconeThreatDetector"; import { LakeraThreatDetector } from "@oakai/aila/src/features/threatDetection/detectors/lakera/LakeraThreatDetector"; import type { LooseLessonPlan } from "@oakai/aila/src/protocol/schema"; @@ -178,30 +179,59 @@ function isValidLessonPlan(lessonPlan: unknown): boolean { return lessonPlan !== null && typeof lessonPlan === "object"; } -function parseLessonPlanFromOutput(output: unknown): LooseLessonPlan { - if (!output) return {}; +function hasMessages(obj: unknown): obj is { messages: unknown } { + return obj !== null && typeof obj === "object" && "messages" in obj; +} + +function isValidMessages(messages: unknown): boolean { + return Array.isArray(messages); +} + +function verifyChatOwnership( + chat: { userId: string }, + requestUserId: string, + chatId: string, +): void { + if (chat.userId !== requestUserId) { + log.error( + `User ${requestUserId} attempted to access chat ${chatId} which belongs to ${chat.userId}`, + ); + throw new Error("Unauthorized access to chat"); + } +} + +function parseChatOutput( + output: unknown, + chatId: string, +): { messages: Message[]; lessonPlan: LooseLessonPlan } { + let messages: Message[] = []; + let lessonPlan: LooseLessonPlan = {}; try { const parsedOutput = typeof output === "string" ? JSON.parse(output) : output; + if (hasMessages(parsedOutput) && isValidMessages(parsedOutput.messages)) { + messages = parsedOutput.messages as Message[]; + } + if ( hasLessonPlan(parsedOutput) && isValidLessonPlan(parsedOutput.lessonPlan) ) { - return parsedOutput.lessonPlan as LooseLessonPlan; + lessonPlan = parsedOutput.lessonPlan as LooseLessonPlan; } } catch (error) { - log.error("Error parsing output to extract lesson plan", error); + log.error(`Error parsing output for chat ${chatId}`, error); } - return {}; + return { messages, lessonPlan }; } -async function loadLessonPlanFromDatabase( +async function loadChatDataFromDatabase( chatId: string, userId: string, -): Promise { +): Promise<{ messages: Message[]; lessonPlan: LooseLessonPlan }> { try { const chat = await prisma.appSession.findUnique({ where: { id: chatId }, @@ -214,25 +244,111 @@ async function loadLessonPlanFromDatabase( if (!chat) { log.info(`No existing chat found for id: ${chatId}`); - return {}; + return { messages: [], lessonPlan: {} }; } - if (chat.userId !== userId) { - log.error( - `User ${userId} attempted to access chat ${chatId} which belongs to ${chat.userId}`, - ); - throw new Error("Unauthorized access to chat"); - } + verifyChatOwnership(chat, userId, chatId); + + const { messages, lessonPlan } = parseChatOutput(chat.output, chatId); - const lessonPlan = parseLessonPlanFromOutput(chat.output); - log.info(`Loaded lesson plan for chat ${chatId}`); - return lessonPlan; + log.info( + `Loaded ${messages.length} messages and lesson plan for chat ${chatId}`, + ); + return { messages, lessonPlan }; } catch (error) { - log.error(`Error loading lesson plan for chat ${chatId}`, error); - return {}; + log.error(`Error loading chat data for chat ${chatId}`, error); + return { messages: [], lessonPlan: {} }; } } +// Extract the latest user message from frontend messages +function extractLatestUserMessage(frontendMessages: Message[]): Message | null { + if (!frontendMessages || frontendMessages.length === 0) { + return null; + } + + // Find the last user message + for (let i = frontendMessages.length - 1; i >= 0; i--) { + const message = frontendMessages[i]; + if (message && message.role === "user") { + return message; + } + } + + return null; +} + +function prepareMessages( + dbMessages: Message[], + frontendMessages: Message[], + chatId: string, +): Message[] { + const latestUserMessage = extractLatestUserMessage(frontendMessages); + + let messages = [...dbMessages]; + if ( + latestUserMessage && + !messages.some((m) => m.id === latestUserMessage.id) + ) { + messages.push(latestUserMessage); + log.info(`Appended new user message to history for chat ${chatId}`); + } + + return messages; +} + +// Helper function to create Aila instance +async function createAilaInstance({ + config, + options, + chatId, + userId, + messages, + lessonPlan, + llmService, + moderationAiClient, + threatDetectors, +}: { + config: Config; + options: AilaOptions; + chatId: string; + userId: string | undefined; + messages: Message[]; + lessonPlan: LooseLessonPlan; + llmService: ReturnType; + moderationAiClient: ReturnType; + threatDetectors: AilaThreatDetector[]; +}): Promise { + return await withTelemetry( + "chat-create-aila", + { chat_id: chatId, user_id: userId }, + async (): Promise => { + const ailaOptions: Partial = { + options, + chat: { + id: chatId, + userId, + messages, + }, + services: { + chatLlmService: llmService, + moderationAiClient, + ragService: (aila: AilaServices) => new AilaRag({ aila }), + americanismsService: () => new AilaAmericanisms(), + analyticsAdapters: (aila: AilaServices) => [ + new PosthogAnalyticsAdapter(aila), + new DatadogAnalyticsAdapter(aila), + ], + threatDetectors: () => threatDetectors, + }, + lessonPlan: lessonPlan ?? {}, + }; + const result = await config.createAila(ailaOptions); + return result; + }, + ); +} + export async function handleChatPostRequest( req: NextRequest, config: Config, @@ -240,7 +356,7 @@ export async function handleChatPostRequest( return await withTelemetry("chat-api", {}, async (span: TracingSpan) => { const { chatId, - messages, + messages: frontendMessages, options, llmService, moderationAiClient, @@ -254,7 +370,12 @@ export async function handleChatPostRequest( userId = await fetchAndCheckUser(chatId); span.setTag("user_id", userId); - const dbLessonPlan = await loadLessonPlanFromDatabase(chatId, userId); + // Load both message history and lesson plan from database + const { messages: dbMessages, lessonPlan: dbLessonPlan } = + await loadChatDataFromDatabase(chatId, userId); + + // Prepare messages by combining database messages with the latest user message + const messages = prepareMessages(dbMessages, frontendMessages, chatId); setTelemetryMetadata({ span, @@ -264,34 +385,19 @@ export async function handleChatPostRequest( options, }); - aila = await withTelemetry( - "chat-create-aila", - { chat_id: chatId, user_id: userId }, - async (): Promise => { - const ailaOptions: Partial = { - options, - chat: { - id: chatId, - userId, - messages, - }, - services: { - chatLlmService: llmService, - moderationAiClient, - ragService: (aila: AilaServices) => new AilaRag({ aila }), - americanismsService: () => new AilaAmericanisms(), - analyticsAdapters: (aila: AilaServices) => [ - new PosthogAnalyticsAdapter(aila), - new DatadogAnalyticsAdapter(aila), - ], - threatDetectors: () => threatDetectors, - }, - lessonPlan: dbLessonPlan ?? {}, - }; - const result = await config.createAila(ailaOptions); - return result; - }, - ); + // Create Aila instance + aila = await createAilaInstance({ + config, + options, + chatId, + userId, + messages, + lessonPlan: dbLessonPlan, + llmService, + moderationAiClient, + threatDetectors, + }); + invariant(aila, "Aila instance is required"); const abortController = handleConnectionAborted(req); From 6a0f82a4db380b6b0b27389488ad95aad440c625 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 16:10:19 +0000 Subject: [PATCH 4/9] Delete unused files --- packages/aila/src/core/document/AilaHomework.ts | 0 packages/aila/src/core/document/AilaLessonPlan.ts | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 packages/aila/src/core/document/AilaHomework.ts delete mode 100644 packages/aila/src/core/document/AilaLessonPlan.ts diff --git a/packages/aila/src/core/document/AilaHomework.ts b/packages/aila/src/core/document/AilaHomework.ts deleted file mode 100644 index e69de29bb..000000000 diff --git a/packages/aila/src/core/document/AilaLessonPlan.ts b/packages/aila/src/core/document/AilaLessonPlan.ts deleted file mode 100644 index e69de29bb..000000000 From a5725e2e143a96a02ea7402e9545c838f1eff81f Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 16:14:16 +0000 Subject: [PATCH 5/9] remove comments --- apps/nextjs/src/app/api/chat/chatHandler.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index 15285097a..3ee37ad6c 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -261,13 +261,11 @@ async function loadChatDataFromDatabase( } } -// Extract the latest user message from frontend messages function extractLatestUserMessage(frontendMessages: Message[]): Message | null { if (!frontendMessages || frontendMessages.length === 0) { return null; } - // Find the last user message for (let i = frontendMessages.length - 1; i >= 0; i--) { const message = frontendMessages[i]; if (message && message.role === "user") { From 2c093778a48afde2e20cf24adcc3aac19b937181 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 16:31:51 +0000 Subject: [PATCH 6/9] Remove comments --- apps/nextjs/src/app/api/chat/chatHandler.ts | 4 ---- 1 file changed, 4 deletions(-) diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index 3ee37ad6c..5751c66a5 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -295,7 +295,6 @@ function prepareMessages( return messages; } -// Helper function to create Aila instance async function createAilaInstance({ config, options, @@ -368,11 +367,9 @@ export async function handleChatPostRequest( userId = await fetchAndCheckUser(chatId); span.setTag("user_id", userId); - // Load both message history and lesson plan from database const { messages: dbMessages, lessonPlan: dbLessonPlan } = await loadChatDataFromDatabase(chatId, userId); - // Prepare messages by combining database messages with the latest user message const messages = prepareMessages(dbMessages, frontendMessages, chatId); setTelemetryMetadata({ @@ -383,7 +380,6 @@ export async function handleChatPostRequest( options, }); - // Create Aila instance aila = await createAilaInstance({ config, options, From 8adbd36736ac69e808c89c8bf01ead1ef2418c6b Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Fri, 28 Feb 2025 18:36:41 +0000 Subject: [PATCH 7/9] Touch --- apps/nextjs/src/app/api/chat/chatHandler.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index 5751c66a5..d71c39785 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -123,7 +123,7 @@ function handleConnectionAborted(req: NextRequest) { const abortController = new AbortController(); req.signal.addEventListener("abort", () => { - log.info("Client has disconnected"); + log.info("Connection aborted: client has disconnected"); abortController.abort(); }); return abortController; From 77caad87e103688277a0900ea63155343ba70913 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Mon, 3 Mar 2025 11:08:58 +0000 Subject: [PATCH 8/9] Sentry errors --- apps/nextjs/src/app/api/chat/chatHandler.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/apps/nextjs/src/app/api/chat/chatHandler.ts b/apps/nextjs/src/app/api/chat/chatHandler.ts index 1b744e093..5b0d51f5d 100644 --- a/apps/nextjs/src/app/api/chat/chatHandler.ts +++ b/apps/nextjs/src/app/api/chat/chatHandler.ts @@ -224,7 +224,6 @@ function parseChatOutput( } } catch (error) { log.error(`Error parsing output for chat ${chatId}`, error); - // Report to Sentry captureException(error, { extra: { chatId, output }, tags: { context: "parseChatOutput" }, @@ -263,12 +262,11 @@ async function loadChatDataFromDatabase( return { messages, lessonPlan }; } catch (error) { log.error(`Error loading chat data for chat ${chatId}`, error); - // Report to Sentry captureException(error, { extra: { chatId, userId }, tags: { context: "loadChatDataFromDatabase" }, }); - throw error; // Re-throw after logging to Sentry + throw error; } } From eee426454283a37331185ea129f991894b667383 Mon Sep 17 00:00:00 2001 From: Stef Lewandowski Date: Mon, 3 Mar 2025 11:10:35 +0000 Subject: [PATCH 9/9] Unused variables --- apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx b/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx index b99318393..e2ef36b7c 100644 --- a/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx +++ b/apps/nextjs/src/components/AppComponents/Chat/AiSdk.tsx @@ -1,4 +1,3 @@ -import type React from "react"; import { useEffect, useState } from "react"; import { toast } from "react-hot-toast"; @@ -46,7 +45,6 @@ export function AiSdk({ id }: Readonly) { const [hasFinished, setHasFinished] = useState(true); const initialMessages = useChatStore((state) => state.initialMessages); - const lessonPlan = useLessonPlanStore((state) => state.lessonPlan); const streamingFinished = useChatStore((state) => state.streamingFinished); const scrollToBottom = useChatStore((state) => state.scrollToBottom); const messageStarted = useLessonPlanStore((state) => state.messageStarted);