feat: add categoriser to fixtures in tests, fix streaming bugs (WIP)

apps/nextjs/src/app/api/chat/chatHandler.ts

@@ -1,18 +1,24 @@
 import type { Aila } from "@oakai/aila/src/core/Aila";
 import type { AilaServices } from "@oakai/aila/src/core/AilaServices";
 import type { Message } from "@oakai/aila/src/core/chat";
-import type { AilaInitializationOptions } from "@oakai/aila/src/core/types";
-import type {
-  AilaOptions,
-  AilaPublicChatOptions,
+import { MessageSchema } from "@oakai/aila/src/core/chat/types";
+import { OpenAIService } from "@oakai/aila/src/core/llm/OpenAIService";
+import {
+  AilaOptionsSchema,
+  type AilaInitializationOptions,
 } from "@oakai/aila/src/core/types";
+import type { AilaOptions } from "@oakai/aila/src/core/types";
 import { AilaAmericanisms } from "@oakai/aila/src/features/americanisms/AilaAmericanisms";
 import {
   DatadogAnalyticsAdapter,
   PosthogAnalyticsAdapter,
 } from "@oakai/aila/src/features/analytics";
+import { AilaCategorisation } from "@oakai/aila/src/features/categorisation/categorisers/AilaCategorisation";
 import { AilaRag } from "@oakai/aila/src/features/rag/AilaRag";
-import type { LooseLessonPlan } from "@oakai/aila/src/protocol/schema";
+import {
+  LessonPlanSchemaWhilstStreaming,
+  type LooseLessonPlan,
+} from "@oakai/aila/src/protocol/schema";
 import type { TracingSpan } from "@oakai/core/src/tracing/serverTracing";
 import { withTelemetry } from "@oakai/core/src/tracing/serverTracing";
 import type { PrismaClientWithAccelerate } from "@oakai/db";
@@ -24,6 +30,7 @@ import { aiLogger } from "@oakai/logger";
 import { StreamingTextResponse } from "ai";
 import type { NextRequest } from "next/server";
 import invariant from "tiny-invariant";
+import { z } from "zod";
 
 import type { Config } from "./config";
 import { handleChatException } from "./errorHandling";
@@ -40,26 +47,30 @@ export const maxDuration = 300;
 const prisma: PrismaClientWithAccelerate = globalPrisma;
 
 export async function GET() {
-  return new Response("Chat API is working", { status: 200 });
+  return new Promise((resolve) => {
+    resolve(new Response("Chat API is working", { status: 200 }));
+  });
 }
 
 async function setupChatHandler(req: NextRequest) {
   return await withTelemetry(
     "chat-setup-chat-handler",
     {},
     async (span: TracingSpan) => {
-      const json = await req.json();
+      const json = (await req.json()) as unknown;
+      const chatRequestSchema = z.object({
+        id: z.string(),
+        messages: MessageSchema.array(),
+        lessonPlan: LessonPlanSchemaWhilstStreaming,
+        options: AilaOptionsSchema,
+      });
+
       const {
         id: chatId,
         messages,
-        lessonPlan = {},
-        options: chatOptions = {},
-      }: {
-        id: string;
-        messages: Message[];
-        lessonPlan?: LooseLessonPlan;
-        options?: AilaPublicChatOptions;
-      } = json;
+        lessonPlan,
+        options: chatOptions,
+      } = chatRequestSchema.parse(json);
 
       const options: AilaOptions = {
         useRag: chatOptions.useRag ?? true,
@@ -69,6 +80,8 @@ async function setupChatHandler(req: NextRequest) {
         useModeration: true,
       };
 
+      invariant(chatId, "Chat ID is required");
+
       const llmService = getFixtureLLMService(req.headers, chatId);
       const moderationAiClient = getFixtureModerationOpenAiClient(
         req.headers,
@@ -180,9 +193,12 @@ export async function handleChatPostRequest(
               messages,
             },
             services: {
-              chatLlmService: llmService,
+              chatLlmService: () =>
+                llmService ?? new OpenAIService({ userId, chatId }),
               moderationAiClient,
               ragService: (aila: AilaServices) => new AilaRag({ aila }),
+              chatCategoriser: (aila: AilaServices) =>
+                new AilaCategorisation({ aila }),
               americanismsService: () => new AilaAmericanisms(),
               analyticsAdapters: (aila: AilaServices) => [
                 new PosthogAnalyticsAdapter(aila),
@@ -191,8 +207,8 @@ export async function handleChatPostRequest(
             },
             lessonPlan: lessonPlan ?? {},
           };
-          const result = await config.createAila(ailaOptions);
-          return result;
+          const aila = await config.createAila(ailaOptions);
+          return aila;
         },
       );
       invariant(aila, "Aila instance is required");

apps/nextjs/src/app/api/chat/config.ts

@@ -17,7 +17,7 @@ export const defaultConfig: Config = {
   prisma: globalPrisma,
   createAila: async (options) => {
     const webActionsPlugin = createWebActionsPlugin(globalPrisma);
-    return new Aila({
+    const createdAila = new Aila({
       ...options,
       plugins: [...(options.plugins || []), webActionsPlugin],
       prisma: options.prisma ?? globalPrisma,
@@ -26,5 +26,7 @@ export const defaultConfig: Config = {
         userId: undefined,
       },
     });
+    await createdAila.initialise();
+    return createdAila;
   },
 };

apps/nextjs/src/app/api/chat/fixtures/FixtureRecordLLMService.ts

@@ -8,7 +8,7 @@ import type { ZodSchema } from "zod";
 const log = aiLogger("fixtures");
 
 export class FixtureRecordLLMService implements LLMService {
-  name = "FixureRecordLLM";
+  name = "FixtureRecordLLM";
   private _openAIService: OpenAIService;
 
   constructor(
@@ -23,6 +23,11 @@ export class FixtureRecordLLMService implements LLMService {
     messages: Message[];
     temperature: number;
   }): Promise<ReadableStreamDefaultReader<string>> {
+    log.info(
+      "Creating chat completion from OpenAI",
+      this.name,
+      this.fixtureName,
+    );
     return this._openAIService.createChatCompletionStream(params);
   }
 
@@ -33,6 +38,11 @@ export class FixtureRecordLLMService implements LLMService {
     messages: Message[];
     temperature: number;
   }): Promise<ReadableStreamDefaultReader<string>> {
+    log.info(
+      "Creating chat completion object stream",
+      this.name,
+      this.fixtureName,
+    );
     const upstreamReader =
       await this._openAIService.createChatCompletionObjectStream(params);
 
@@ -76,4 +86,25 @@ export class FixtureRecordLLMService implements LLMService {
 
     return s.getReader();
   }
+
+  async generateObject<T>(params: {
+    model: string;
+    schema: ZodSchema<T>;
+    schemaName: string;
+    messages: Message[];
+    temperature: number;
+  }): Promise<T | undefined> {
+    const result = await this._openAIService.generateObject(params);
+
+    try {
+      const formattedUrl = `${process.cwd()}/tests-e2e/recordings/${this.fixtureName}.generateObject.formatted.json`;
+      const formatted = JSON.stringify(result, null, 2);
+      log.info("Writing generateObject formatted to", formattedUrl);
+      await fs.writeFile(formattedUrl, formatted);
+    } catch (e) {
+      log.error("Error writing generateObject formatted file", e);
+    }
+
+    return result;
+  }
 }

apps/nextjs/src/app/api/chat/fixtures/FixtureReplayLLMService.ts

@@ -5,9 +5,10 @@ import fs from "fs";
 const log = aiLogger("fixtures");
 
 export class FixtureReplayLLMService extends MockLLMService {
-  name = "FixureReplayLLM";
+  name = "FixtureReplayLLM";
 
   constructor(fixtureName: string) {
+    log.info("Setting up fixture", fixtureName);
     const fileUrl = `${process.cwd()}/tests-e2e/recordings/${fixtureName}.chunks.txt`;
     log.info("Loading chunks from", fileUrl);
     const fixture = fs.readFileSync(fileUrl, "utf8");

apps/nextjs/src/app/api/chat/fixtures/index.ts

@@ -18,6 +18,7 @@ export function getFixtureLLMService(headers: Headers, chatId: string) {
   const fixtureName = headers.get("x-e2e-fixture-name");
 
   if (!fixtureName) {
+    log.info("Not using fixtures");
     return undefined;
   }
 

apps/nextjs/src/app/api/chat/route.test.ts

@@ -36,30 +36,36 @@ describe("Chat API Route", () => {
     jest.spyOn(mockLLMService, "createChatCompletionObjectStream");
 
     testConfig = {
-      createAila: jest.fn().mockImplementation(async (options) => {
-        const ailaConfig: AilaInitializationOptions = {
-          options: {
-            usePersistence: false,
-            useRag: false,
-            useAnalytics: false,
-            useModeration: false,
-            useErrorReporting: false,
-            useThreatDetection: false,
+      createAila: jest
+        .fn()
+        .mockImplementation(
+          async (options: Partial<AilaInitializationOptions>) => {
+            return new Promise<Aila>((resolve) => {
+              const ailaConfig: AilaInitializationOptions = {
+                options: {
+                  usePersistence: false,
+                  useRag: false,
+                  useAnalytics: false,
+                  useModeration: false,
+                  useErrorReporting: false,
+                  useThreatDetection: false,
+                },
+                chat: {
+                  id: chatId,
+                  userId,
+                  messages: options.chat?.messages ?? [],
+                },
+                plugins: [],
+                services: {
+                  chatLlmService: () => mockLLMService,
+                  chatCategoriser: () => mockChatCategoriser,
+                },
+              };
+              resolve(new Aila(ailaConfig));
+            });
           },
-          chat: {
-            id: chatId,
-            userId,
-            messages: options.chat.messages ?? [],
-          },
-          plugins: [],
-          services: {
-            chatLlmService: mockLLMService,
-            chatCategoriser: mockChatCategoriser,
-          },
-        };
-        return new Aila(ailaConfig);
-      }),
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        ),
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-assignment
       prisma: {} as any,
     };
   }, 60000);
@@ -70,7 +76,11 @@ describe("Chat API Route", () => {
       body: JSON.stringify({
         id: "test-chat-id",
         messages: [
-          { role: "user", content: "Create a lesson about Glaciation" },
+          {
+            id: "1",
+            role: "user",
+            content: "Create a lesson about Glaciation",
+          },
         ],
         lessonPlan: {},
         options: {},