update to usage schema

Author: jorgebaralt

apps/web/src/app/api/ai/journl-agent/route.ts

@@ -20,15 +20,20 @@ export async function POST(req: Request) {
         const provider = Object.keys(result.providerMetadata ?? {})[0] || "";
 
         if (result.usage && session.user?.id) {
-          await api.ai.trackLLMUsage({
-            inputTokens: result.usage.promptTokens,
-            metadata: {
-              message_count: messages.length,
-              request_type: "chat",
-            },
-            model,
-            outputTokens: result.usage.completionTokens,
-            provider,
+          await api.usage.trackAiModelUsage({
+            metrics: [
+              {
+                quantity: result.usage.promptTokens,
+                unit: "input_tokens",
+              },
+              {
+                quantity: result.usage.completionTokens,
+                unit: "output_tokens",
+              },
+            ],
+            model_id: model,
+            model_provider: provider,
+            user_id: session.user.id,
           });
         }
       },

apps/web/src/app/api/supabase/embed-document/route.ts

@@ -10,7 +10,7 @@ import removeMarkdown from "remove-markdown";
 import type { z } from "zod/v4";
 import { model } from "~/ai/providers/openai/embedding";
 import { schema } from "~/components/editor/block-schema";
-import { embedder } from "~/trpc/server";
+import { api, embedder } from "~/trpc/server";
 import { handler } from "../_lib/webhook-handler";
 
 const CHUNK_PARAMS: ChunkParams = {
@@ -102,14 +102,33 @@ export const POST = handler(zDocumentEmbeddingTask, async (payload) => {
     const mDocument =
       await MDocument.fromMarkdown(markdown).chunk(CHUNK_PARAMS);
 
-    // ! TODO: Here's where we get the usage tokens, we must process these in a different transaction
-    // ! To guarantee that the usage is tracked regardless of the success of the embedding.
-    const { embeddings, usage: _usage } = await embedMany({
+    const { embeddings, usage } = await embedMany({
       maxRetries: 5,
       model,
       values: mDocument.map((chunk) => chunk.text),
     });
 
+    if (usage) {
+      try {
+        await api.usage.trackAiModelUsage({
+          metadata: {
+            document_id: document.id,
+            model_version: model.specificationVersion,
+          },
+          metrics: [{ quantity: usage.tokens, unit: "tokens" }],
+          model_id: model.modelId,
+          model_provider: model.provider,
+          user_id: document.user_id,
+        });
+      } catch (error) {
+        console.error(
+          "Failed to track AI usage for document embedding:",
+          error,
+        );
+        // Continue with embedding even if usage tracking fails
+      }
+    }
+
     const insertions: z.infer<typeof zInsertDocumentEmbedding>[] = [];
 
     for (const [index, embedding] of embeddings.entries()) {

packages/api/src/api-router/index.ts

@@ -5,6 +5,7 @@ import { documentRouter } from "./document.js";
 import { journalRouter } from "./journal.js";
 import { notesRouter } from "./notes.js";
 import { pagesRouter } from "./pages.js";
+import { usageRouter } from "./usage.js";
 
 export const apiRouter = createTRPCRouter({
   auth: authRouter,
@@ -13,6 +14,7 @@ export const apiRouter = createTRPCRouter({
   journal: journalRouter,
   notes: notesRouter,
   pages: pagesRouter,
+  usage: usageRouter,
 });
 
 export type ApiRouter = typeof apiRouter;

packages/api/src/api-router/usage.ts

@@ -1,25 +1,23 @@
 import { eq } from "@acme/db";
 import { UsageEvent, UsageEventStatus } from "@acme/db/schema";
-import { TRPCError } from "@trpc/server";
+import { TRPCError, type TRPCRouterRecord } from "@trpc/server";
 import { z } from "zod/v4";
-import { protectedProcedure } from "../trpc.js";
+import { publicProcedure } from "../trpc.js";
 
 export const usageRouter = {
-  createUsageEvent: protectedProcedure
+  trackAiModelUsage: publicProcedure
     .input(
       z.object({
-        metadata: z.record(z.string(), z.string()),
+        metadata: z.record(z.string(), z.string()).optional(),
         metrics: z.array(z.object({ quantity: z.number(), unit: z.string() })),
         model_id: z.string(),
         model_provider: z.string(),
+        user_id: z.string(),
       }),
     )
     .mutation(async ({ ctx, input }) => {
       try {
-        return await ctx.db.insert(UsageEvent).values({
-          ...input,
-          user_id: ctx.session.user.id,
-        });
+        return await ctx.db.insert(UsageEvent).values(input);
       } catch (error) {
         console.error("Database error in usage.createUsageEvent:", error);
         throw new TRPCError({
@@ -28,7 +26,7 @@ export const usageRouter = {
         });
       }
     }),
-  processUsageEvent: protectedProcedure
+  updateUsageEventStatus: publicProcedure
     .input(
       z.object({
         id: z.string(),
@@ -51,4 +49,4 @@ export const usageRouter = {
         });
       }
     }),
-};
+} satisfies TRPCRouterRecord;