fix: change plan parts embeddings to 256 vectprs

packages/core/src/models/lessonPlans.ts

@@ -9,9 +9,9 @@
   PrismaClientWithAccelerate,
   Subject,
 } from "@oakai/db";
 import { Prisma } from "@prisma/client";
 import { OpenAIEmbeddings } from "langchain/embeddings/openai";
 import { PrismaVectorStore } from "langchain/vectorstores/prisma";
 import yaml from "yaml";
 
 import { LLMResponseJsonSchema } from "../../../aila/src/protocol/jsonPatchProtocol";
@@ -57,7 +57,7 @@
   parts: LessonPlanPart[] | null;
 };
 
 interface FilterOptions {
   key_stage_id?: object;
   subject_id?: object;
 }
@@ -335,75 +335,4 @@
       WHERE id = ${id}`;
     return result;
   }
-
-  async search(
-    query: string,
-    keyStage: string | undefined,
-    subject: string | undefined,
-    perPage: number,
-  ): Promise<LessonPlanWithLesson[]> {
-    const filter: FilterOptions = {};
-
-    if (keyStage) {
-      const keyStageRecord = await this._rag.fetchKeyStage(keyStage);
-
-      if (keyStageRecord) {
-        filter["key_stage_id"] = {
-          equals: keyStageRecord.id,
-        };
-      }
-    }
-    if (subject) {
-      const subjectRecord = await this._rag.fetchSubject(subject);
-      if (subjectRecord) {
-        filter["subject_id"] = {
-          equals: subjectRecord.id,
-        };
-      }
-    }
-    const vectorStore = PrismaVectorStore.withModel<LessonPlan>(
-      this._prisma,
-    ).create(new OpenAIEmbeddings(), {
-      prisma: Prisma,
-      tableName: "lesson_plans" as "LessonPlan",
-      vectorColumnName: "embedding",
-      columns: {
-        id: PrismaVectorStore.IdColumn,
-        content: PrismaVectorStore.ContentColumn,
-      },
-      // @ts-expect-error TODO Bug in PrismaVectorStore which doesn't allow mapped column names
-      filter,
-    });
-    const result = await vectorStore.similaritySearch(query, perPage);
-
-    const lessonPlans: LessonPlanWithLesson[] =
-      await this._prisma.lessonPlan.findMany({
-        where: {
-          id: { in: result.map((r) => r.metadata.id) },
-        },
-        include: {
-          lesson: {
-            select: {
-              id: true,
-              slug: true,
-              title: true,
-              subjectId: true,
-              keyStageId: true,
-              isNewLesson: true,
-              newLessonContent: true,
-            },
-          },
-        },
-      });
-
-    const hydrated: LessonPlanWithLesson[] = [];
-    for (const entry of result) {
-      const lessonPlan = lessonPlans.find((ls) => ls.id === entry.metadata.id);
-      if (!lessonPlan) {
-        throw new Error("Lesson summary not found");
-      }
-      hydrated.push(lessonPlan);
-    }
-    return hydrated;
-  }
 }

packages/core/src/rag/index.ts

@@ -733,20 +733,26 @@ Thank you and happy classifying!`;
 
     const vectorStore = PrismaVectorStore.withModel<LessonPlanPart>(
       this.prisma,
-    ).create(new OpenAIEmbeddings(), {
-      prisma: Prisma,
-      tableName: "lesson_plan_parts" as "LessonPlanPart",
-      vectorColumnName: "embedding",
-      verbose: true,
-      openAIApiKey: process.env.OPENAI_API_KEY,
-      columns: {
-        id: PrismaVectorStore.IdColumn,
-        lesson_plan_id: PrismaVectorStore.IdColumn,
-        content: PrismaVectorStore.ContentColumn,
+    ).create(
+      new OpenAIEmbeddings({
+        modelName: "text-embedding-3-large",
+        dimensions: 256,
+      }),
+      {
+        prisma: Prisma,
+        tableName: "lesson_plan_parts" as "LessonPlanPart",
+        vectorColumnName: "embedding",
+        verbose: true,
+        openAIApiKey: process.env.OPENAI_API_KEY,
+        columns: {
+          id: PrismaVectorStore.IdColumn,
+          lesson_plan_id: PrismaVectorStore.IdColumn,
+          content: PrismaVectorStore.ContentColumn,
+        },
+        // @ts-expect-error TODO Bug in PrismaVectorStore which doesn't allow mapped column names
+        filter,
       },
-      // @ts-expect-error TODO Bug in PrismaVectorStore which doesn't allow mapped column names
-      filter,
-    });
+    );
 
     const similaritySearchTerm = topic ? `${title}. ${topic}` : title;
 

packages/db/prisma/migrations/20240924172451_lesson_plan_parts_embedding_vector_256/migration.sql

@@ -0,0 +1,4 @@
+-- This changes the embedding column from a 1536 vector to a 256 vector
+ALTER TABLE lesson_plan_parts ADD COLUMN embedding_temp vector(256);
+ALTER TABLE lesson_plan_parts DROP COLUMN embedding;
+ALTER TABLE lesson_plan_parts RENAME COLUMN embedding_temp TO embedding;

packages/db/prisma/schema.prisma

@@ -291,7 +291,7 @@ model LessonPlanPart {
   lessonPlan   LessonPlan                   @relation(fields: [lessonPlanId], references: [id], onDelete: Cascade)
   subject      Subject?                     @relation(fields: [subjectId], references: [id], onDelete: Cascade)
   keyStage     KeyStage?                    @relation(fields: [keyStageId], references: [id], onDelete: Cascade)
-  embedding    Unsupported("vector(1536)")?
+  embedding    Unsupported("vector(256)")?
   status       LessonPlanPartStatus         @default(PENDING)
 
   @@unique([lessonPlanId, key])