feat: continue migrating Aila package to support variable documents

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

@@ -60,7 +60,7 @@ async function setupChatHandler(req: NextRequest) {
       const options: AilaOptions = {
         useRag: chatOptions.useRag ?? true,
         temperature: chatOptions.temperature ?? 0.7,
-        numberOfLessonPlansInRag: chatOptions.numberOfLessonPlansInRag ?? 5,
+        numberOfRecordsInRag: chatOptions.numberOfRecordsInRag ?? 5,
         usePersistence: true,
         useModeration: true,
       };
@@ -110,10 +110,7 @@ function setTelemetryMetadata({
   span.setTag("has_lesson_plan", Object.keys(lessonPlan).length > 0);
   span.setTag("use_rag", options.useRag);
   span.setTag("temperature", options.temperature);
-  span.setTag(
-    "number_of_lesson_plans_in_rag",
-    options.numberOfLessonPlansInRag,
-  );
+  span.setTag("number_of_records_in_rag", options.numberOfRecordsInRag);
   span.setTag("use_persistence", options.usePersistence);
   span.setTag("use_moderation", options.useModeration);
 }

packages/aila/src/constants.ts

@@ -7,5 +7,5 @@ export const DEFAULT_CATEGORISE_MODEL: OpenAI.Chat.ChatModel =
   "gpt-4o-2024-08-06";
 export const DEFAULT_TEMPERATURE = 0.7;
 export const DEFAULT_MODERATION_TEMPERATURE = 0.7;
-export const DEFAULT_RAG_LESSON_PLANS = 5;
+export const DEFAULT_NUMBER_OF_RECORDS_IN_RAG = 5;
 export const BOT_USER_ID = "bot";

packages/aila/src/core/Aila.ts

@@ -5,7 +5,7 @@ import { aiLogger } from "@oakai/logger";
 import {
   DEFAULT_MODEL,
   DEFAULT_TEMPERATURE,
-  DEFAULT_RAG_LESSON_PLANS,
+  DEFAULT_NUMBER_OF_RECORDS_IN_RAG,
 } from "../constants";
 import type { AilaAmericanismsFeature } from "../features/americanisms";
 import { NullAilaAmericanisms } from "../features/americanisms/NullAilaAmericanisms";
@@ -33,7 +33,7 @@ import type { LLMService } from "./llm/LLMService";
 import { OpenAIService } from "./llm/OpenAIService";
 import type { AilaPlugin } from "./plugins/types";
 import type {
-  AilaGenerateLessonPlanOptions,
+  AilaGenerateDocumentOptions,
   AilaOptions,
   AilaOptionsWithDefaultFallbackValues,
   AilaInitializationOptions,
@@ -141,21 +141,21 @@ export class Aila implements AilaServices {
     await this.loadChatIfPersisting();
     const persistedLessonPlan = this._chat.persistedChat?.lessonPlan;
     if (persistedLessonPlan) {
-      this._document.setContent(persistedLessonPlan);
+      this._document.content = persistedLessonPlan;
     }
     await this._document.initialiseContentFromMessages(this._chat.messages);
 
     this._initialised = true;
   }
 
-  private initialiseOptions(options?: AilaOptions) {
+  private initialiseOptions(
+    options?: AilaOptions,
+  ): AilaOptionsWithDefaultFallbackValues {
     return {
       useRag: options?.useRag ?? true,
       temperature: options?.temperature ?? DEFAULT_TEMPERATURE,
-      // #TODO we should find a way to make this less specifically tied
-      // to lesson RAG
-      numberOfLessonPlansInRag:
-        options?.numberOfLessonPlansInRag ?? DEFAULT_RAG_LESSON_PLANS,
+      numberOfRecordsInRag:
+        options?.numberOfRecordsInRag ?? DEFAULT_NUMBER_OF_RECORDS_IN_RAG,
       usePersistence: options?.usePersistence ?? true,
       useAnalytics: options?.useAnalytics ?? true,
       useModeration: options?.useModeration ?? true,
@@ -177,18 +177,10 @@ export class Aila implements AilaServices {
     return this._chat;
   }
 
-  // #TODO we should refactor this to be a document
-  // and not be specifically tied to a "lesson"
-  // so that we can handle any type of generation
   public get document(): AilaDocumentService {
     return this._document;
   }
 
-  // #TODO we should not need this
-  public get lessonPlan() {
-    return this._document.content;
-  }
-
   public get snapshotStore() {
     return this._snapshotStore;
   }
@@ -259,7 +251,7 @@ export class Aila implements AilaServices {
   }
 
   // Generation methods
-  public async generateSync(opts: AilaGenerateLessonPlanOptions) {
+  public async generateSync(opts: AilaGenerateDocumentOptions) {
     this.checkInitialised();
     const stream = await this.generate(opts);
 
@@ -280,7 +272,7 @@ export class Aila implements AilaServices {
   public async generate({
     input,
     abortController,
-  }: AilaGenerateLessonPlanOptions) {
+  }: AilaGenerateDocumentOptions) {
     this.checkInitialised();
     if (this._isShutdown) {
       throw new AilaGenerationError(

packages/aila/src/core/AilaServices.ts

@@ -15,9 +15,9 @@ import type {
 import type {
   AilaPersistedChat,
   AilaRagRelevantLesson,
-  LooseLessonPlan,
 } from "../protocol/schema";
 import type { Message } from "./chat";
+import type { AilaDocumentContent } from "./document/types";
 import type { AilaPlugin } from "./plugins";
 import type { FullQuizService } from "./quiz/interfaces";
 import type { AilaOptionsWithDefaultFallbackValues } from "./types";
@@ -30,12 +30,11 @@ export interface AilaAnalyticsService {
 }
 
 export interface AilaDocumentService {
-  readonly content: LooseLessonPlan;
+  content: AilaDocumentContent;
   readonly hasInitialisedContentFromMessages: boolean;
-  setContent(content: LooseLessonPlan): void;
   extractAndApplyLlmPatches(patches: string): void;
   applyValidPatches(validPatches: ValidPatchDocument[]): void;
-  initialise(content: LooseLessonPlan): void;
+  initialise(content: AilaDocumentContent): void;
   initialiseContentFromMessages(messages: Message[]): Promise<void>;
 }
 
@@ -59,7 +58,6 @@ export interface AilaChatService {
 export interface AilaServices {
   readonly userId: string | undefined;
   readonly chatId: string;
-  readonly lessonPlan: LooseLessonPlan;
   readonly messages: Message[];
   readonly options: AilaOptionsWithDefaultFallbackValues;
   readonly analytics?: AilaAnalyticsFeature;

packages/aila/src/core/document/AilaDocument.ts

@@ -50,10 +50,6 @@ export class AilaDocument implements AilaDocumentService {
     this._content = content;
   }
 
-  public setContent(content: AilaDocumentContent) {
-    this._content = content;
-  }
-
   public get hasInitialisedContentFromMessages(): boolean {
     return this._hasInitialisedContentFromMessages;
   }

packages/aila/src/core/document/types.ts

@@ -6,6 +6,10 @@ export type AilaDummyDocumentContent = {
   keyStage: string;
   topic?: string;
   body: string;
+  basedOn?: {
+    id: string;
+    title: string;
+  };
 };
 
 export type AilaDocumentContent = LooseLessonPlan | AilaDummyDocumentContent;

packages/aila/src/core/prompt/builders/AilaLessonPromptBuilder.ts

@@ -6,7 +6,7 @@ import { prisma as globalPrisma } from "@oakai/db/client";
 import { aiLogger } from "@oakai/logger";
 import { getRelevantLessonPlans, parseSubjectsForRagSearch } from "@oakai/rag";
 
-import { DEFAULT_RAG_LESSON_PLANS } from "../../../constants";
+import { DEFAULT_NUMBER_OF_RECORDS_IN_RAG } from "../../../constants";
 import { tryWithErrorReporting } from "../../../helpers/errorReporting";
 import { LLMResponseJsonSchema } from "../../../protocol/jsonPatchProtocol";
 import type { LooseLessonPlan } from "../../../protocol/schema";
@@ -69,7 +69,8 @@ export class AilaLessonPromptBuilder extends AilaPromptBuilder {
       };
     }
 
-    const { title, subject, keyStage, topic } = this._aila?.lessonPlan ?? {};
+    const { title, subject, keyStage, topic } =
+      this._aila?.document?.content ?? {};
 
     if (!title || !subject || !keyStage) {
       log.error("Missing title, subject or keyStage, returning empty content");
@@ -119,8 +120,8 @@ export class AilaLessonPromptBuilder extends AilaPromptBuilder {
         keyStage,
         id: chatId,
         k:
-          this._aila?.options.numberOfLessonPlansInRag ??
-          DEFAULT_RAG_LESSON_PLANS,
+          this._aila?.options.numberOfRecordsInRag ??
+          DEFAULT_NUMBER_OF_RECORDS_IN_RAG,
         prisma: globalPrisma,
         chatId,
         userId,
@@ -146,7 +147,7 @@ export class AilaLessonPromptBuilder extends AilaPromptBuilder {
     relevantLessonPlans: string,
     baseLessonPlan: LooseLessonPlan | undefined,
   ): string {
-    const lessonPlan = this._aila?.lessonPlan ?? {};
+    const lessonPlan = this._aila?.document?.content ?? {};
     const args: TemplateProps = {
       lessonPlan,
       relevantLessonPlans,

packages/aila/src/core/types.ts

@@ -14,18 +14,17 @@ import type {
   AilaModerationFeature,
   AilaThreatDetectionFeature,
 } from "../features/types";
-import type { LooseLessonPlan } from "../protocol/schema";
 import type { AilaServices } from "./AilaServices";
 import type { Message } from "./chat";
 import type { AilaDocumentContent } from "./document/types";
 import type { LLMService } from "./llm/LLMService";
 import type { AilaPlugin } from "./plugins/types";
 import type { AilaPromptBuilder } from "./prompt/AilaPromptBuilder";
 
-export type AilaGenerateLessonPlanMode = "interactive" | "generate";
+export type AilaGenerateDocumentMode = "interactive" | "generate";
 
-export type AilaGenerateLessonPlanOptions = {
-  mode?: AilaGenerateLessonPlanMode;
+export type AilaGenerateDocumentOptions = {
+  mode?: AilaGenerateDocumentMode;
   input?: string;
   title?: string;
   topic?: string;
@@ -37,7 +36,7 @@ export type AilaGenerateLessonPlanOptions = {
 export type AilaPublicChatOptions = {
   useRag?: boolean;
   temperature?: number;
-  numberOfLessonPlansInRag?: number;
+  numberOfRecordsInRag?: number;
 };
 
 export type AilaOptions = AilaPublicChatOptions & {
@@ -47,7 +46,7 @@ export type AilaOptions = AilaPublicChatOptions & {
   useAnalytics?: boolean;
   useThreatDetection?: boolean;
   model?: string;
-  mode?: AilaGenerateLessonPlanMode;
+  mode?: AilaGenerateDocumentMode;
 };
 
 export type AilaOptionsWithDefaultFallbackValues = Required<AilaOptions>;

packages/aila/src/features/rag/AilaRag.ts

@@ -32,17 +32,17 @@ export class AilaRag implements AilaRagFeature {
   }
 
   public async fetchRagContent({
-    numberOfLessonPlansInRag,
+    numberOfRecordsInRag,
     lessonPlan,
   }: {
-    numberOfLessonPlansInRag?: number;
+    numberOfRecordsInRag?: number;
     lessonPlan?: LooseLessonPlan;
   }) {
     // #TODO Refactor to return an array rather than stringified JSON
     let content = "[]";
 
     const { title, keyStage, subject, topic } =
-      lessonPlan ?? this._aila?.lessonPlan ?? {};
+      lessonPlan ?? this._aila?.document?.content ?? {};
     const chatId = this._aila?.chatId ?? "anonymous"; // #TODO add proper support for CLI RAG requests without a user
     if (!title) {
       return content;
@@ -57,8 +57,8 @@ export class AilaRag implements AilaRagFeature {
               subject,
               topic,
               k:
-                numberOfLessonPlansInRag ??
-                this._aila?.options.numberOfLessonPlansInRag ??
+                numberOfRecordsInRag ??
+                this._aila?.options.numberOfRecordsInRag ??
                 5,
             })
           : [];

packages/aila/src/features/rag/index.ts

@@ -2,7 +2,7 @@ import type { LooseLessonPlan } from "../../protocol/schema";
 
 export interface AilaRagFeature {
   fetchRagContent(params: {
-    numberOfLessonPlansInRag?: number;
+    numberOfRecordsInRag?: number;
     lessonPlan?: LooseLessonPlan;
   }): Promise<string>;
 }

packages/aila/src/protocol/schema.ts

@@ -1,4 +1,3 @@
-// import dedent from "dedent";
 import dedent from "ts-dedent";
 import z from "zod";
 import { zodToJsonSchema } from "zod-to-json-schema";