perf(scores): remove fk constraint from traceId and observationId on scores, change to projectId

packages/shared/prisma/generated/types.ts

@@ -342,7 +342,7 @@ export type Prompt = {
 export type Score = {
     id: string;
     timestamp: Generated<Timestamp>;
-    project_id: string | null;
+    project_id: string;
     name: string;
     value: number;
     source: ScoreSource;

packages/shared/prisma/migrations/20240512155021_scores_drop_fk_on_traces_and_observations/migration.sql

@@ -0,0 +1,5 @@
+-- DropForeignKey
+ALTER TABLE "scores" DROP CONSTRAINT "scores_observation_id_fkey";
+
+-- DropForeignKey
+ALTER TABLE "scores" DROP CONSTRAINT "scores_trace_id_fkey";

packages/shared/prisma/migrations/20240512155022_scores_non_null_and_add_fk_project_id/migration.sql

@@ -0,0 +1,11 @@
+/*
+  Warnings:
+
+  - Made the column `project_id` on table `scores` required. This step will fail if there are existing NULL values in that column.
+
+*/
+-- AlterTable
+ALTER TABLE "scores" ALTER COLUMN "project_id" SET NOT NULL;
+
+-- AddForeignKey
+ALTER TABLE "scores" ADD CONSTRAINT "scores_project_id_fkey" FOREIGN KEY ("project_id") REFERENCES "projects"("id") ON DELETE CASCADE ON UPDATE CASCADE;

packages/shared/prisma/migrations/20240513082203_scores_unique_id_and_projectid_instead_of_id_and_traceid/migration.sql

@@ -0,0 +1,11 @@
+/*
+  Warnings:
+
+  - A unique constraint covering the columns `[id,project_id]` on the table `scores` will be added. If there are existing duplicate values, this will fail.
+
+*/
+-- DropIndex
+DROP INDEX "scores_id_trace_id_key";
+
+-- CreateIndex
+CREATE UNIQUE INDEX "scores_id_project_id_key" ON "scores"("id", "project_id");

packages/shared/prisma/schema.prisma

@@ -113,6 +113,7 @@ model Project {
     JobExecution       JobExecution[]
     LlmApiKeys         LlmApiKeys[]
     PosthogIntegration PosthogIntegration[]
+    Score              Score[]
 
     @@map("projects")
 }
@@ -230,7 +231,6 @@ model Trace {
     createdAt  DateTime      @default(now()) @map("created_at")
     updatedAt  DateTime      @default(now()) @updatedAt @map("updated_at")
 
-    scores          Score[]
     DatasetRunItems DatasetRunItems[]
     DatasetItem     DatasetItem[]
     JobExecution    JobExecution[]
@@ -305,7 +305,6 @@ model Observation {
     outputCost          Decimal?          @map("output_cost")
     totalCost           Decimal?          @map("total_cost")
     completionStartTime DateTime?         @map("completion_start_time")
-    scores              Score[]
     project             Project           @relation(fields: [projectId], references: [id], onDelete: Cascade)
     derivedDatasetItems DatasetItem[]
     datasetRunItems     DatasetRunItems[]
@@ -392,18 +391,17 @@ enum ObservationLevel {
 model Score {
     id            String         @id @default(cuid())
     timestamp     DateTime       @default(now())
-    projectId     String?        @map("project_id")
+    projectId     String         @map("project_id")
+    project       Project        @relation(fields: [projectId], references: [id], onDelete: Cascade)
     name          String
     value         Float
     source        ScoreSource
     comment       String?
     traceId       String         @map("trace_id")
-    trace         Trace          @relation(fields: [traceId], references: [id], onDelete: Cascade)
     observationId String?        @map("observation_id")
-    observation   Observation?   @relation(fields: [observationId], references: [id], onDelete: SetNull)
     JobExecution  JobExecution[]
 
-    @@unique([id, traceId]) // used for upsert
+    @@unique([id, projectId]) // used for upserts via prisma
     @@index(timestamp)
     @@index([value])
     @@index([projectId])

web/src/__tests__/ingestion.servertest.ts

@@ -281,6 +281,7 @@ describe("/api/public/ingestion API Endpoint", () => {
       expect(dbScore?.name).toBe("score-name");
       expect(dbScore?.value).toBe(100.5);
       expect(dbScore?.observationId).toBeNull();
+      expect(dbScore?.projectId).toBe("7a88fb47-b4e2-43b8-a06c-a5ce950dc53a");
     });
   });
 

web/src/__tests__/scores.servertest.ts

@@ -60,6 +60,7 @@ describe("/api/public/scores API Endpoint", () => {
     expect(dbScore?.observationId).toBeNull();
     expect(dbScore?.comment).toBe("comment");
     expect(dbScore?.source).toBe("API");
+    expect(dbScore?.projectId).toBe("7a88fb47-b4e2-43b8-a06c-a5ce950dc53a");
   });
 
   it("should create score for a trace with int", async () => {

web/src/features/datasets/server/dataset-router.ts

@@ -147,6 +147,7 @@ export const datasetRouter = createTRPCRouter({
             JOIN scores s 
               ON s.trace_id = ri.trace_id 
               AND (ri.observation_id IS NULL OR s.observation_id = ri.observation_id) -- only include scores that are linked to the observation if observation is linked
+              AND s.project_id = ${input.projectId}
             JOIN traces t ON t.id = s.trace_id
           WHERE t.project_id = ${input.projectId}
           GROUP BY
@@ -553,13 +554,11 @@ export const datasetRouter = createTRPCRouter({
           observation: {
             select: {
               id: true,
-              scores: true,
             },
           },
           trace: {
             select: {
               id: true,
-              scores: true,
             },
           },
         },
@@ -570,6 +569,27 @@ export const datasetRouter = createTRPCRouter({
         skip: input.page * input.limit,
       });
 
+      const traceScores = await ctx.prisma.score.findMany({
+        where: {
+          projectId: ctx.session.projectId,
+          traceId: {
+            in: runItems
+              .filter((ri) => !!!ri.observationId) // only include trace scores if run is not linked to an observation
+              .map((ri) => ri.traceId),
+          },
+        },
+      });
+      const observationScores = await ctx.prisma.score.findMany({
+        where: {
+          projectId: ctx.session.projectId,
+          observationId: {
+            in: runItems
+              .filter((ri) => !!ri.observationId)
+              .map((ri) => ri.observationId) as string[],
+          },
+        },
+      });
+
       const totalRunItems = await ctx.prisma.datasetRunItems.count({
         where: {
           datasetRunId: input.datasetRunId,
@@ -582,13 +602,12 @@ export const datasetRouter = createTRPCRouter({
         },
       });
 
-      const observationIds = runItems
-        .map((ri) => ri.observation?.id)
-        .filter(Boolean) as string[];
       const observations = await ctx.prisma.observationView.findMany({
         where: {
           id: {
-            in: observationIds,
+            in: runItems
+              .map((ri) => ri.observationId)
+              .filter(Boolean) as string[],
           },
         },
         select: {
@@ -598,13 +617,10 @@ export const datasetRouter = createTRPCRouter({
         },
       });
 
-      const traceIds = runItems
-        .map((ri) => ri.trace?.id)
-        .filter(Boolean) as string[];
       const traces = await ctx.prisma.traceView.findMany({
         where: {
           id: {
-            in: traceIds,
+            in: runItems.map((ri) => ri.traceId).filter(Boolean) as string[],
           },
           projectId: ctx.session.projectId,
         },
@@ -621,10 +637,14 @@ export const datasetRouter = createTRPCRouter({
           datasetItemId: ri.datasetItemId,
           observation: observations.find((o) => o.id === ri.observationId),
           trace: traces.find((t) => t.id === ri.traceId),
-          scores:
-            // use observation scores if run is linked to an observation, otherwise use all trace scores
-            (!!ri.observationId ? ri.observation?.scores : ri.trace?.scores) ??
-            [],
+          scores: [
+            ...traceScores.filter((s) => s.traceId === ri.traceId),
+            ...observationScores.filter(
+              (s) =>
+                s.observationId === ri.observationId &&
+                s.traceId === ri.traceId,
+            ),
+          ],
         };
       });
 

web/src/features/manual-scoring/components/ManualScoreButton.tsx

@@ -89,7 +89,7 @@ export function ManualScoreButton({
 
   const handleDelete = async () => {
     if (score) {
-      await mutDeleteScore.mutateAsync(score.id);
+      await mutDeleteScore.mutateAsync({ id: score.id, projectId });
       capture("score:delete", {
         type: type,
         source: source,
@@ -130,6 +130,7 @@ export function ManualScoreButton({
         id: score.id,
         value: values.score,
         comment: values.comment,
+        projectId,
       });
       capture("score:update_form_submit", {
         type: type,
@@ -142,6 +143,7 @@ export function ManualScoreButton({
         comment: values.comment,
         traceId,
         observationId,
+        projectId,
       });
       capture("score:create_form_submit", {
         type: type,

web/src/features/public-api/server/apiScope.ts

@@ -59,7 +59,7 @@ async function isResourceInProject(resource: Resource, projectId: string) {
         (await prisma.score.count({
           where: {
             id: resource.id,
-            trace: { projectId },
+            projectId,
           },
         })) === 1;
       if (!scoreCheck)

web/src/pages/api/cron/ingestion-metrics.ts

@@ -212,18 +212,17 @@ export default async function handler(
       }>
     >`
       SELECT
-        t.project_id,
+        s.project_id,
         count(s.*)::integer count_scores
       FROM
         scores s
-        JOIN traces t ON t.id = s.trace_id
-          WHERE
-            s.timestamp < ${endTimeframe}
-            ${
-              startTimeframe
-                ? Prisma.sql`AND s.timestamp >= ${startTimeframe}`
-                : Prisma.empty
-            }
+        WHERE
+          s.timestamp < ${endTimeframe}
+          ${
+            startTimeframe
+              ? Prisma.sql`AND s.timestamp >= ${startTimeframe}`
+              : Prisma.empty
+          }
       GROUP BY
         1
     `;

web/src/pages/api/public/scores.ts

@@ -20,15 +20,6 @@ import { isPrismaException } from "@/src/utils/exceptions";
 
 const operators = ["<", ">", "<=", ">=", "!=", "="] as const;
 
-const prismaOperators: Record<(typeof operators)[number], string> = {
-  "<": "lt",
-  ">": "gt",
-  "<=": "lte",
-  ">=": "gte",
-  "!=": "not",
-  "=": "equals",
-};
-
 const ScoresGetSchema = z.object({
   ...paginationZod,
   userId: z.string().nullish(),
@@ -141,8 +132,8 @@ export default async function handler(
             s.observation_id as "observationId",
             json_build_object('userId', t.user_id) as "trace"
           FROM "scores" AS s
-          JOIN "traces" AS t ON t.id = s.trace_id
-          WHERE t.project_id = ${authCheck.scope.projectId}
+          JOIN "traces" AS t ON t.id = s.trace_id AND t.project_id = ${authCheck.scope.projectId}
+          WHERE s.project_id = ${authCheck.scope.projectId}
           ${userCondition}
           ${nameCondition}
           ${sourceCondition}
@@ -151,25 +142,23 @@ export default async function handler(
           ORDER BY t."timestamp" DESC
           LIMIT ${obj.limit} OFFSET ${skipValue}
           `);
-      const totalItems = await prisma.score.count({
-        where: {
-          name: obj.name ? obj.name : undefined,
-          source: obj.source ? obj.source : undefined,
-          timestamp: obj.fromTimestamp
-            ? { gte: new Date(obj.fromTimestamp) }
-            : undefined,
-          value:
-            obj.operator && obj.value
-              ? {
-                  [prismaOperators[obj.operator]]: obj.value,
-                }
-              : undefined,
-          trace: {
-            projectId: authCheck.scope.projectId,
-            userId: obj.userId ? obj.userId : undefined,
-          },
-        },
-      });
+
+      const totalItemsRes = await prisma.$queryRaw<{ count: bigint }[]>(
+        Prisma.sql`
+          SELECT COUNT(*) as count
+          FROM "scores" AS s
+          JOIN "traces" AS t ON t.id = s.trace_id AND t.project_id = ${authCheck.scope.projectId}
+          WHERE s.project_id = ${authCheck.scope.projectId}
+          ${userCondition}
+          ${nameCondition}
+          ${sourceCondition}
+          ${fromTimestampCondition}
+          ${valueCondition}
+        `,
+      );
+
+      const totalItems =
+        totalItemsRes[0] !== undefined ? Number(totalItemsRes[0].count) : 0;
 
       return res.status(200).json({
         data: scores,

web/src/pages/api/public/scores/[scoreId].ts

@@ -40,9 +40,7 @@ export default async function handler(
         },
         where: {
           id: scoreId,
-          trace: {
-            projectId: authCheck.scope.projectId,
-          },
+          projectId: authCheck.scope.projectId,
         },
       });
 
@@ -55,9 +53,7 @@ export default async function handler(
       await prisma.score.delete({
         where: {
           id: scoreId,
-          trace: {
-            projectId: authCheck.scope.projectId,
-          },
+          projectId: authCheck.scope.projectId,
         },
       });
 
@@ -107,9 +103,7 @@ export default async function handler(
       const score = await prisma.score.findUnique({
         where: {
           id: scoreId,
-          trace: {
-            projectId: authCheck.scope.projectId,
-          },
+          projectId: authCheck.scope.projectId,
         },
       });
 

web/src/pages/api/public/traces.ts

@@ -139,7 +139,7 @@ export default async function handler(
             COALESCE(ARRAY_AGG(DISTINCT s.id) FILTER (WHERE s.id IS NOT NULL), ARRAY[]::text[]) AS "scores"
           FROM "traces" AS t
           LEFT JOIN "observations_view" AS o ON t.id = o.trace_id AND o.project_id = ${authCheck.scope.projectId}
-          LEFT JOIN "scores" AS s ON t.id = s.trace_id
+          LEFT JOIN "scores" AS s ON t.id = s.trace_id AND s.project_id = ${authCheck.scope.projectId}
           WHERE t.project_id = ${authCheck.scope.projectId}
           ${userCondition}
           ${nameCondition}