feat: aila instantiation includes option to specify a schema and categorisation approach

.vscode/settings.json

@@ -37,6 +37,8 @@
     "catchall",
     "categorisation",
     "Categorised",
+    "Categorisе",
+    "Categorizе",
     "centered",
     "cloudinary",
     "clsx",

apps/nextjs/src/app/actions.ts

@@ -1,7 +1,7 @@
 "use server";
 
 import type { AilaPersistedChat } from "@oakai/aila/src/protocol/schema";
-import { chatSchema } from "@oakai/aila/src/protocol/schema";
+import { AilaPersistedChatSchema } from "@oakai/aila/src/protocol/schema";
 import type { Prisma } from "@oakai/db";
 import { prisma } from "@oakai/db";
 import * as Sentry from "@sentry/nextjs";
@@ -18,7 +18,7 @@ function parseChatAndReportError({
   if (typeof sessionOutput !== "object") {
     throw new Error("sessionOutput is not an object");
   }
-  const parseResult = chatSchema.safeParse({
+  const parseResult = AilaPersistedChatSchema.safeParse({
     ...sessionOutput,
     userId,
     id,

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

@@ -1,6 +1,8 @@
 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 { LessonPlanCategorisationPlugin } from "@oakai/aila/src/core/document/plugins/LessonPlanCategorisationPlugin";
+import { LessonPlanSchema } from "@oakai/aila/src/core/document/schemas/lessonPlan";
 import type {
   AilaOptions,
   AilaPublicChatOptions,
@@ -11,6 +13,7 @@ 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 { AilaThreatDetector } from "@oakai/aila/src/features/threatDetection";
 import { HeliconeThreatDetector } from "@oakai/aila/src/features/threatDetection/detectors/helicone/HeliconeThreatDetector";
@@ -328,7 +331,7 @@ async function createAilaInstance({
           chatLlmService: llmService,
           moderationAiClient,
           ragService: (aila: AilaServices) => new AilaRag({ aila }),
-          americanismsService: () => new AilaAmericanisms(),
+          americanismsService: () => new AilaAmericanisms<LooseLessonPlan>(),
           analyticsAdapters: (aila: AilaServices) => [
             new PosthogAnalyticsAdapter(aila),
             new DatadogAnalyticsAdapter(aila),
@@ -337,6 +340,11 @@ async function createAilaInstance({
         },
         document: {
           content: lessonPlan ?? {},
+          schema: LessonPlanSchema,
+          categorisationPlugin: (aila: AilaServices) =>
+            new LessonPlanCategorisationPlugin(
+              new AilaCategorisation({ aila }),
+            ),
         },
       };
       const result = await config.createAila(ailaOptions);

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

@@ -1,4 +1,6 @@
 import { Aila } from "@oakai/aila/src/core/Aila";
+import { LessonPlanCategorisationPlugin } from "@oakai/aila/src/core/document/plugins/LessonPlanCategorisationPlugin";
+import { LessonPlanSchema } from "@oakai/aila/src/core/document/schemas/lessonPlan";
 import { MockLLMService } from "@oakai/aila/src/core/llm/MockLLMService";
 import type { AilaInitializationOptions } from "@oakai/aila/src/core/types";
 import { MockCategoriser } from "@oakai/aila/src/features/categorisation/categorisers/MockCategoriser";
@@ -54,10 +56,15 @@ describe("Chat API Route", () => {
                 userId,
                 messages: options?.chat?.messages ?? [],
               },
+              document: {
+                content: {},
+                schema: LessonPlanSchema,
+                categorisationPlugin: () =>
+                  new LessonPlanCategorisationPlugin(mockChatCategoriser),
+              },
               plugins: [],
               services: {
                 chatLlmService: mockLLMService,
-                chatCategoriser: mockChatCategoriser,
               },
             };
             const ailaInstance = new Aila(ailaConfig);

packages/aila/src/core/Aila.liveWithOpenAI.test.ts

@@ -1,6 +1,8 @@
 import { MockCategoriser } from "../features/categorisation/categorisers/MockCategoriser";
 import { Aila } from "./Aila";
 import { checkLastMessage, expectPatch, expectText } from "./Aila.testHelpers";
+import { LessonPlanCategorisationPlugin } from "./document/plugins/LessonPlanCategorisationPlugin";
+import { LessonPlanSchema } from "./document/schemas/lessonPlan";
 import type { AilaInitializationOptions } from "./types";
 
 const runInCI = process.env.CI === "true";
@@ -13,7 +15,21 @@ const runManually = process.env.RUN_LLM_TESTS === "true";
 
     beforeEach(() => {
       ailaInstance = new Aila({
-        document: { content: {} },
+        document: {
+          content: {},
+          schema: LessonPlanSchema,
+          categorisationPlugin: () =>
+            new LessonPlanCategorisationPlugin(
+              new MockCategoriser({
+                mockedContent: {
+                  keyStage: "specialist",
+                  subject: "design-technology",
+                  title: "Motorcycle Maintenance",
+                  topic: "Basics and Advanced Techniques",
+                },
+              }),
+            ),
+        },
         chat: { id: "test-chat", userId: "test-user" },
         options: {
           usePersistence: false,
@@ -22,16 +38,6 @@ const runManually = process.env.RUN_LLM_TESTS === "true";
           useModeration: false,
         },
         plugins: [],
-        services: {
-          chatCategoriser: new MockCategoriser({
-            mockedContent: {
-              keyStage: "specialist",
-              subject: "design-technology",
-              title: "Motorcycle Maintenance",
-              topic: "Basics and Advanced Techniques",
-            },
-          }),
-        },
       });
     });
 
@@ -70,7 +76,21 @@ const runManually = process.env.RUN_LLM_TESTS === "true";
 
     beforeEach(() => {
       const options: AilaInitializationOptions = {
-        document: { content: {} },
+        document: {
+          content: {},
+          schema: LessonPlanSchema,
+          categorisationPlugin: () =>
+            new LessonPlanCategorisationPlugin(
+              new MockCategoriser({
+                mockedContent: {
+                  keyStage: "key-stage-3",
+                  subject: "geography",
+                  title: "Glaciation",
+                  topic: "The Formation of Glacial Landscapes",
+                },
+              }),
+            ),
+        },
         chat: { id: "test-chat", userId: "test-user" },
         options: {
           usePersistence: false,
@@ -81,16 +101,6 @@ const runManually = process.env.RUN_LLM_TESTS === "true";
           useErrorReporting: false,
         },
         plugins: [],
-        services: {
-          chatCategoriser: new MockCategoriser({
-            mockedContent: {
-              keyStage: "key-stage-3",
-              subject: "geography",
-              title: "Glaciation",
-              topic: "The Formation of Glacial Landscapes",
-            },
-          }),
-        },
       };
       ailaInstance = new Aila(options);
     });