feat: add cost column to generations table (langfuse#725)

Author: maxdeichmann

src/components/table/use-cases/generations.tsx

@@ -30,6 +30,7 @@ import { type LangfuseColumnDef } from "@/src/components/table/types";
 import { type ObservationLevel } from "@prisma/client";
 import { cn } from "@/src/utils/tailwind";
 import { LevelColors } from "@/src/components/level-colors";
+import { usdFormatter } from "@/src/utils/numbers";
 
 export type GenerationsTableRow = {
   id: string;
@@ -200,6 +201,18 @@ export default function GenerationsTable({ projectId }: GenerationsTableProps) {
       },
       enableHiding: true,
     },
+    {
+      accessorKey: "cost",
+      header: "Cost",
+      cell: ({ row }) => {
+        const value: number | undefined = row.getValue("cost");
+        console.log(value);
+        return value !== undefined ? (
+          <span>{usdFormatter(value)}</span>
+        ) : undefined;
+      },
+      enableHiding: true,
+    },
     {
       accessorKey: "level",
       header: "Level",
@@ -293,29 +306,32 @@ export default function GenerationsTable({ projectId }: GenerationsTableProps) {
     );
 
   const rows: GenerationsTableRow[] = generations.isSuccess
-    ? generations.data.map((generation) => ({
-        id: generation.id,
-        traceId: generation.traceId,
-        traceName: generation.traceName,
-        startTime: generation.startTime.toLocaleString(),
-        endTime: generation.endTime?.toLocaleString() ?? undefined,
-        latency: generation.latency === null ? undefined : generation.latency,
-        name: generation.name ?? undefined,
-        version: generation.version ?? "",
-        model: generation.model ?? "",
-        input: generation.input,
-        output: generation.output,
-        level: generation.level,
-        metadata: generation.metadata
-          ? JSON.stringify(generation.metadata)
-          : undefined,
-        statusMessage: generation.statusMessage ?? undefined,
-        usage: {
-          promptTokens: generation.promptTokens,
-          completionTokens: generation.completionTokens,
-          totalTokens: generation.totalTokens,
-        },
-      }))
+    ? generations.data.map((generation) => {
+        return {
+          id: generation.id,
+          traceId: generation.traceId,
+          traceName: generation.traceName,
+          startTime: generation.startTime.toLocaleString(),
+          endTime: generation.endTime?.toLocaleString() ?? undefined,
+          latency: generation.latency === null ? undefined : generation.latency,
+          cost: generation.cost,
+          name: generation.name ?? undefined,
+          version: generation.version ?? "",
+          model: generation.model ?? "",
+          input: generation.input,
+          output: generation.output,
+          level: generation.level,
+          metadata: generation.metadata
+            ? JSON.stringify(generation.metadata)
+            : undefined,
+          statusMessage: generation.statusMessage ?? undefined,
+          usage: {
+            promptTokens: generation.promptTokens,
+            completionTokens: generation.completionTokens,
+            totalTokens: generation.totalTokens,
+          },
+        };
+      })
     : [];
 
   return (

src/pages/api/public/ingestion.ts

@@ -206,7 +206,6 @@ const handleSingleEvent = async (
 
   await persistEventMiddleware(
     prisma,
-
     apiScope.projectId,
     req,
     cleanedEvent,

src/server/api/routers/generations.ts

@@ -16,6 +16,9 @@ import {
   type ObservationOptions,
   observationsTableCols,
 } from "@/src/server/api/definitions/observationsTable";
+import { calculateTokenCost } from "@/src/features/ingest/lib/usage";
+import Decimal from "decimal.js";
+import { usdFormatter } from "@/src/utils/numbers";
 
 const exportFileFormats = ["CSV", "JSON", "OPENAI-JSONL"] as const;
 export type ExportFileFormats = (typeof exportFileFormats)[number];
@@ -118,7 +121,25 @@ export const generationsRouter = createTRPCRouter({
         `,
       );
 
-      return generations;
+      const pricings = await ctx.prisma.pricing.findMany();
+
+      return generations.map(({ input, output, ...rest }) => {
+        return {
+          ...rest,
+          input,
+          output,
+          cost: rest.model
+            ? calculateTokenCost(pricings, {
+                model: rest.model,
+                totalTokens: new Decimal(rest.totalTokens),
+                promptTokens: new Decimal(rest.promptTokens),
+                completionTokens: new Decimal(rest.completionTokens),
+                input: input,
+                output: output,
+              })
+            : undefined,
+        };
+      });
     }),
 
   export: protectedProjectProcedure
@@ -175,6 +196,28 @@ export const generationsRouter = createTRPCRouter({
         `,
       );
 
+      const pricings = await ctx.prisma.pricing.findMany();
+
+      const enrichedGenerations = generations.map(
+        ({ input, output, ...rest }) => {
+          return {
+            ...rest,
+            input,
+            output,
+            cost: rest.model
+              ? calculateTokenCost(pricings, {
+                  model: rest.model,
+                  totalTokens: new Decimal(rest.totalTokens),
+                  promptTokens: new Decimal(rest.promptTokens),
+                  completionTokens: new Decimal(rest.completionTokens),
+                  input: input,
+                  output: output,
+                })
+              : undefined,
+          };
+        },
+      );
+
       // create csv
       switch (input.fileFormat) {
         case "CSV":
@@ -185,19 +228,23 @@ export const generationsRouter = createTRPCRouter({
               "model",
               "startTime",
               "endTime",
+              "cost",
               "prompt",
               "completion",
               "metadata",
             ],
           ]
             .concat(
-              generations.map((generation) =>
+              enrichedGenerations.map((generation) =>
                 [
                   generation.traceId,
                   generation.name ?? "",
                   generation.model ?? "",
                   generation.startTime.toISOString(),
                   generation.endTime?.toISOString() ?? "",
+                  generation.cost
+                    ? usdFormatter(generation.cost.toNumber())
+                    : "",
                   JSON.stringify(generation.input),
                   JSON.stringify(generation.output),
                   JSON.stringify(generation.metadata),
@@ -210,7 +257,7 @@ export const generationsRouter = createTRPCRouter({
             .map((row) => row.join(","))
             .join("\n");
         case "JSON":
-          return JSON.stringify(generations);
+          return JSON.stringify(enrichedGenerations);
         case "OPENAI-JSONL":
           const inputSchemaOpenAI = z.array(
             z.object({
@@ -223,7 +270,7 @@ export const generationsRouter = createTRPCRouter({
           });
 
           return (
-            generations
+            enrichedGenerations
               .map((generation) => ({
                 parsedInput: inputSchemaOpenAI.safeParse(generation.input),
                 parsedOutput: outputSchema.safeParse(generation.output),