feat: rag new schema, standalone package

packages/aila/package.json

@@ -22,6 +22,7 @@
     "@oakai/db": "*",
     "@oakai/exports": "*",
     "@oakai/logger": "*",
+    "@oakai/rag": "*",
     "@sentry/nextjs": "^8.35.0",
     "@vercel/kv": "^0.2.2",
     "ai": "^3.3.26",

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

@@ -2,6 +2,7 @@ import type { TemplateProps } from "@oakai/core/src/prompts/lesson-assistant";
 import { template } from "@oakai/core/src/prompts/lesson-assistant";
 import { prisma as globalPrisma } from "@oakai/db/client";
 import { aiLogger } from "@oakai/logger";
+import { getRelevantLessonPlans } from "@oakai/rag";
 
 import { DEFAULT_RAG_LESSON_PLANS } from "../../../constants";
 import { tryWithErrorReporting } from "../../../helpers/errorReporting";
@@ -12,6 +13,8 @@ import { compressedLessonPlanForRag } from "../../../utils/lessonPlan/compressed
 import { fetchLessonPlan } from "../../../utils/lessonPlan/fetchLessonPlan";
 import type { RagLessonPlan } from "../../../utils/rag/fetchRagContent";
 import { fetchRagContent } from "../../../utils/rag/fetchRagContent";
+import { parseKeyStage } from "../../../utils/rag/parseKeyStage";
+import { parseSubjects } from "../../../utils/rag/parseSubjects";
 import type { AilaServices } from "../../AilaServices";
 import { AilaPromptBuilder } from "../AilaPromptBuilder";
 
@@ -64,36 +67,75 @@ export class AilaLessonPromptBuilder extends AilaPromptBuilder {
 
     const { title, subject, keyStage, topic } = this._aila?.lessonPlan ?? {};
 
-    let relevantLessonPlans: RagLessonPlan[] = [];
-    await tryWithErrorReporting(async () => {
-      relevantLessonPlans = await fetchRagContent({
-        title: title ?? "unknown",
-        subject,
-        topic,
-        keyStage,
-        id: chatId,
-        k:
-          this._aila?.options.numberOfLessonPlansInRag ??
-          DEFAULT_RAG_LESSON_PLANS,
-        prisma: globalPrisma,
-        chatId,
-        userId,
+    const NEW_RAG_ENABLED = true;
+
+    if (NEW_RAG_ENABLED) {
+      if (!title || !subject || !keyStage) {
+        log.error(
+          "Missing title, subject or keyStage, returning empty content",
+        );
+        return {
+          ragLessonPlans: [],
+          stringifiedRelevantLessonPlans: noRelevantLessonPlans,
+        };
+      }
+
+      const keyStageSlugs = keyStage ? [parseKeyStage(keyStage)] : null;
+      const subjectSlugs = subject ? parseSubjects(subject) : null;
+
+      const relevantLessonPlans = await getRelevantLessonPlans({
+        title,
+        keyStageSlugs,
+        subjectSlugs,
       });
-    }, "Did not fetch RAG content. Continuing");
+      const stringifiedRelevantLessonPlans = JSON.stringify(
+        relevantLessonPlans,
+        null,
+        2,
+      );
 
-    log.info("Fetched relevant lesson plans", relevantLessonPlans.length);
-    const stringifiedRelevantLessonPlans = JSON.stringify(
-      relevantLessonPlans,
-      null,
-      2,
-    );
+      return {
+        ragLessonPlans: relevantLessonPlans.map((l) => ({
+          ...l.lessonPlan,
+          id: l.ragLessonPlanId,
+        })),
+        stringifiedRelevantLessonPlans,
+      };
+    } else {
+      let relevantLessonPlans: RagLessonPlan[] = [];
+      await tryWithErrorReporting(async () => {
+        relevantLessonPlans = await fetchRagContent({
+          title: title ?? "unknown",
+          subject,
+          topic,
+          keyStage,
+          id: chatId,
+          k:
+            this._aila?.options.numberOfLessonPlansInRag ??
+            DEFAULT_RAG_LESSON_PLANS,
+          prisma: globalPrisma,
+          chatId,
+          userId,
+        });
+      }, "Did not fetch RAG content. Continuing");
 
-    log.info("Got RAG content, length:", stringifiedRelevantLessonPlans.length);
+      log.info("Fetched relevant lesson plans", relevantLessonPlans.length);
+      const stringifiedRelevantLessonPlans = JSON.stringify(
+        relevantLessonPlans,
+        null,
+        2,
+      );
 
-    return {
-      ragLessonPlans: relevantLessonPlans,
-      stringifiedRelevantLessonPlans,
-    };
+      log.info(
+        "Got RAG content, length:",
+        stringifiedRelevantLessonPlans.length,
+      );
+
+      return {
+        ragLessonPlans: relevantLessonPlans,
+        stringifiedRelevantLessonPlans,
+      };
+    }
   }
 
   private systemPrompt(

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

@@ -12,9 +12,9 @@ import { minifyLessonPlanForRelevantLessons } from "../../utils/lessonPlan/minif
 const log = aiLogger("aila:rag");
 
 export class AilaRag implements AilaRagFeature {
-  private _aila: AilaServices;
-  private _rag: RAG;
-  private _prisma: PrismaClientWithAccelerate;
+  private readonly _aila: AilaServices;
+  private readonly _rag: RAG;
+  private readonly _prisma: PrismaClientWithAccelerate;
 
   constructor({
     aila,

packages/aila/src/protocol/schema.ts

@@ -588,8 +588,7 @@ export const LessonPlanJsonSchema = zodToJsonSchema(
 );
 
 const AilaRagRelevantLessonSchema = z.object({
-  // @todo add this after next ingest
-  // oakLessonId: z.number(),
+  oakLessonId: z.number().nullish(),
   lessonPlanId: z.string(),
   title: z.string(),
 });

packages/aila/src/utils/rag/parseKeyStage.ts

@@ -0,0 +1,20 @@
+const keyStageMap: Record<string, string> = {
+  1: "ks1",
+  2: "ks2",
+  3: "ks3",
+  4: "ks4",
+  5: "ks5",
+  keystage1: "ks1",
+  keystage2: "ks2",
+  keystage3: "ks3",
+  keystage4: "ks4",
+  keystage5: "ks5",
+  eyfs: "early-years-foundation-stage",
+};
+
+export function parseKeyStage(maybeKeyStage: string): string {
+  maybeKeyStage = maybeKeyStage.toLowerCase().replace(/[^a-z0-9]/g, "");
+  const keyStageSlug = keyStageMap[maybeKeyStage];
+
+  return keyStageSlug ?? maybeKeyStage;
+}

packages/aila/src/utils/rag/parseSubjects.ts

@@ -0,0 +1,17 @@
+const subjectMap: Record<string, string[]> = {
+  science: ["biology", "chemistry", "physics", "science", "combined-science"],
+  biology: ["biology", "science", "combined-science"],
+  chemistry: ["chemistry", "science", "combined-science"],
+  physics: ["physics", "science", "combined-science"],
+  "combined-science": [
+    "combined-science",
+    "science",
+    "biology",
+    "chemistry",
+    "physics",
+  ],
+};
+
+export function parseSubjects(subject: string): string[] {
+  return subjectMap[subject] ?? [subject];
+}

packages/core/src/prompts/lesson-assistant/parts/interactingWithTheUser.ts

@@ -49,6 +49,8 @@ These Oak lessons might be relevant:
 1. Introduction to the Periodic Table 
 2. Chemical Reactions and Equations
 3. The Structure of the Atom
+4. The Mole Concept
+5. Acids, Bases and Salts
 \n
 To base your lesson on one of these existing Oak lessons, type the lesson number. Tap **Continue** to start from scratch.
 END OF EXAMPLE RESPONSE`,

packages/core/src/rag/index.ts

@@ -21,6 +21,7 @@ import { DEFAULT_CATEGORISE_MODEL } from "../../../aila/src/constants";
 import type { OpenAICompletionWithLoggingOptions } from "../../../aila/src/lib/openai/OpenAICompletionWithLogging";
 import { OpenAICompletionWithLogging } from "../../../aila/src/lib/openai/OpenAICompletionWithLogging";
 import type { JsonValue } from "../models/prompts";
+import { shortenKeyStage } from "../utils/shortenKeyStage";
 import { slugify } from "../utils/slugify";
 import { keyStages, subjects } from "../utils/subjects";
 import { CategoriseKeyStageAndSubjectResponse } from "./categorisation";
@@ -552,19 +553,19 @@ Thank you and happy classifying!`;
     if (!keyStage) {
       return null;
     }
-    let cachedKeyStage: KeyStage | null;
-    try {
-      cachedKeyStage = await kv.get<KeyStage>(`keyStage:${keyStage}`);
-      if (cachedKeyStage) {
-        return cachedKeyStage;
-      }
-    } catch (e) {
-      log.error(
-        "Error parsing cached keyStage. Continuing without cached value",
-        e,
-      );
-      await kv.del(`keyStage:${keyStage}`);
-    }
+    // let cachedKeyStage: KeyStage | null;
+    // try {
+    //   cachedKeyStage = await kv.get<KeyStage>(`keyStage:${keyStage}`);
+    //   if (cachedKeyStage) {
+    //     return cachedKeyStage;
+    //   }
+    // } catch (e) {
+    //   log.error(
+    //     "Error parsing cached keyStage. Continuing without cached value",
+    //     e,
+    //   );
+    //   await kv.del(`keyStage:${keyStage}`);
+    // }
 
     let foundKeyStage: KeyStage | null = null;
     foundKeyStage = await this.prisma.keyStage.findFirst({
@@ -575,6 +576,12 @@ Thank you and happy classifying!`;
           { slug: slugify(keyStage) },
           { title: { equals: keyStage.toLowerCase(), mode: "insensitive" } },
           { slug: { equals: keyStage.toLowerCase(), mode: "insensitive" } },
+          {
+            slug: {
+              equals: shortenKeyStage(slugify(keyStage)),
+              mode: "insensitive",
+            },
+          },
         ],
       },
       cacheStrategy: { ttl: 60 * 5, swr: 60 * 2 },
@@ -692,6 +699,8 @@ Thank you and happy classifying!`;
       };
     }
 
+    log.info("Filter:", filter);
+
     const vectorStore = PrismaVectorStore.withModel<LessonPlanPart>(
       this.prisma,
     ).create(
@@ -723,6 +732,8 @@ Thank you and happy classifying!`;
         similaritySearchTerm,
         k * 5, // search for more records than we need
       );
+
+      log.info("Initial search result", result);
     } catch (e) {
       if (e instanceof TypeError && e.message.includes("join([])")) {
         log.warn("Caught TypeError with join([]), returning empty array");

packages/core/src/utils/shortenKeyStage.ts

@@ -0,0 +1,10 @@
+export const shortenKeyStage = (keyStage: string) => {
+  const keyStageMap: Record<string, string> = {
+    "key-stage-1": "KS1",
+    "key-stage-2": "KS2",
+    "key-stage-3": "KS3",
+    "key-stage-4": "KS4",
+  };
+
+  return keyStageMap[keyStage] || keyStage;
+};

packages/db/package.json

@@ -23,6 +23,8 @@
     "db-migrate": "pnpm with-env prisma migrate dev",
     "db-migrate:dev": "pnpm with-env prisma migrate dev",
     "db-migrate:status": "pnpm with-env prisma migrate status",
+    "db-migrate:status:stg": "DB_ENV=stg doppler run --config stg -- prisma migrate status",
+    "db-migrate:status:prd": "DB_ENV=prd doppler run --config prd -- prisma migrate status",
     "db-migrate-resolve-applied:prd": "doppler run --config prd -- prisma migrate resolve --applied",
     "db-migrate-resolve-applied:stg": "doppler run --config stg -- prisma migrate resolve --applied",
     "db-migrate-resolve-rolled-back:prd": "doppler run --config prd -- prisma migrate resolve --rolled-back",

packages/db/prisma/additions/README.md

@@ -0,0 +1,3 @@
+# Schema additions
+
+This directory houses SQL additions to the Prisma schema where the Prisma schema is not expressive enough to capture the desired schema.

packages/db/prisma/additions/rag_lesson_plans_ivfflat_index.sql

@@ -0,0 +1,4 @@
+CREATE INDEX IF NOT EXISTS idx_rag_lesson_plan_parts_embedding_ann
+ON rag.rag_lesson_plan_parts
+USING ivfflat (embedding vector_cosine_ops)
+WITH (lists = 100);

packages/db/prisma/additions/rag_lesson_plans_unique_slug_index.sql

@@ -0,0 +1,3 @@
+CREATE UNIQUE INDEX IF NOT EXISTS idx_rag_lesson_plans_unique_published_oak_lesson_slug
+ON rag.rag_lesson_plans (oak_lesson_slug)
+WHERE is_published = TRUE;

packages/db/prisma/schema.prisma

@@ -1117,16 +1117,21 @@ model IngestError {
 }
 
 model RagLessonPlan {
-  id                String              @id @default(cuid())
-  oakLessonId       Int                 @map("oak_lesson_id")
-  ingestLessonId    String?             @map("ingest_lesson_id")
-  lessonPlan        Json                @map("lesson_plan") @db.JsonB
-  subjectSlug       String              @map("subject_slug")
-  keyStageSlug      String              @map("key_stage_slug")
-  createdAt         DateTime            @default(now()) @map("created_at")
-  updatedAt         DateTime            @updatedAt @map("updated_at")
-  ragLessonPlanPart RagLessonPlanPart[]
-
+  id                 String              @id @default(cuid())
+  oakLessonId        Int?                @map("oak_lesson_id")
+  oakLessonSlug      String              @map("oak_lesson_slug")
+  ingestLessonId     String?             @map("ingest_lesson_id")
+  lessonPlan         Json                @map("lesson_plan") @db.JsonB
+  subjectSlug        String              @map("subject_slug")
+  keyStageSlug       String              @map("key_stage_slug")
+  isPublished        Boolean             @default(false) @map("is_published")
+  createdAt          DateTime            @default(now()) @map("created_at")
+  updatedAt          DateTime            @updatedAt @map("updated_at")
+  ragLessonPlanParts RagLessonPlanPart[]
+  // The following index is not supported by prisma so is applied manually in ./additions/rag_lesson_plans_unique_slug_index.sql
+  //  @@index([oakLessonSlug], name: "unique_published_oak_lesson_slug", dbIndex: false) @db.PartialIndex("is_published = TRUE")
+
+  @@index([isPublished, keyStageSlug, subjectSlug], name: "idx_rag_lesson_plans_published_key_stage_subject")
   @@map("rag_lesson_plans")
   @@schema("rag")
 }

packages/ingest/package.json

@@ -23,6 +23,7 @@
     "@oakai/core": "*",
     "@oakai/db": "*",
     "@oakai/logger": "*",
+    "@paralleldrive/cuid2": "^2.2.2",
     "csv-parser": "^3.0.0",
     "graphql-request": "^6.1.0",
     "webvtt-parser": "^2.2.0",

packages/ingest/src/_data/.gitignore

@@ -1,3 +1,4 @@
 *.jsonl
 *.json
-*.csv
+*.csv
+*

packages/ingest/src/config/ingestConfig.ts

@@ -20,5 +20,7 @@ export const IngestConfigSchema = z.object({
       filePath: z.string(),
     }),
   ]),
+  title: z.string().optional(),
+  description: z.string().optional(),
 });
 export type IngestConfig = z.infer<typeof IngestConfigSchema>;

packages/ingest/src/db-helpers/step.ts

@@ -6,6 +6,7 @@ export const STEP = [
   "lesson_plan_generation",
   "chunking",
   "embedding",
+  "publishing",
 ] as const;
 
 const STEP_STATUS = ["started", "completed", "failed"] as const;

packages/ingest/src/generate-lesson-plans/parseBatchLessonPlan.ts

@@ -2,6 +2,7 @@ import { CompletedLessonPlanSchema } from "@oakai/aila/src/protocol/schema";
 
 import { IngestError } from "../IngestError";
 import { CompletionBatchResponseSchema } from "../zod-schema/zodSchema";
+import { parseKeyStage } from "./parseKeyStage";
 
 export function parseBatchLessonPlan(line: unknown) {
   let result;
@@ -37,6 +38,8 @@ export function parseBatchLessonPlan(line: unknown) {
     lessonPlan = CompletedLessonPlanSchema.parse(
       JSON.parse(maybeLessonPlanString),
     );
+
+    lessonPlan.keyStage = parseKeyStage(lessonPlan.keyStage);
   } catch (cause) {
     throw new IngestError("Failed to parse lesson plan", {
       cause,