Skip to content

Commit 47ae1c4

Browse files
jorgebaraltjorgebaralt
committed
Clean up - create ai trpc procedure
1 parent 0f31b27 commit 47ae1c4

File tree

6 files changed

+169
-148
lines changed

6 files changed

+169
-148
lines changed

‎apps/web/src/app/api/chat/route.ts‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
import { db, insertLLMTokenUsage } from "@acme/db";
21
import { journlAgent } from "~/ai/mastra/agents/journl-agent";
32
import { getSession } from "~/auth/server";
3+
import { api } from "~/trpc/server";
44

55
export const maxDuration = 30; // Allow streaming responses up to 30 seconds
66

@@ -20,18 +20,15 @@ export async function POST(req: Request) {
2020
const provider = Object.keys(result.providerMetadata ?? {})[0] || "";
2121

2222
if (result.usage && session.user?.id) {
23-
await db.transaction(async (tx) => {
24-
await insertLLMTokenUsage(tx, {
25-
inputTokens: result.usage.promptTokens,
26-
metadata: {
27-
message_count: messages.length,
28-
request_type: "chat",
29-
}, // Update this based on your agent configuration
30-
model,
31-
outputTokens: result.usage.completionTokens,
32-
provider, // fallback since provider not available in result
33-
userId: session.user.id,
34-
});
23+
await api.ai.trackLLMUsage({
24+
inputTokens: result.usage.promptTokens,
25+
metadata: {
26+
message_count: messages.length,
27+
request_type: "chat",
28+
},
29+
model,
30+
outputTokens: result.usage.completionTokens,
31+
provider,
3532
});
3633
}
3734
},

‎apps/web/src/app/api/webhooks/journal_entry/route.ts‎

Lines changed: 10 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
1-
import { db, insertEmbeddingTokenUsage } from "@acme/db";
2-
import { JournalEmbedding, zJournalEntry } from "@acme/db/schema";
3-
import { openai } from "@ai-sdk/openai";
4-
import { embed } from "ai";
1+
import { zJournalEntry } from "@acme/db/schema";
52
import { NextResponse } from "next/server";
3+
import { api } from "~/trpc/server";
64
import { handler } from "../_lib/webhook-handler";
75

86
/**
@@ -17,41 +15,14 @@ export const POST = handler(zJournalEntry, async (payload) => {
1715
);
1816
}
1917

20-
const { embedding, usage } = await embed({
21-
maxRetries: 5,
22-
model: openai.embedding("text-embedding-3-small"),
23-
value: payload.record.content,
24-
});
25-
26-
await db.transaction(async (tx) => {
27-
// Store token usage with automatic pricing calculation
28-
await insertEmbeddingTokenUsage(tx, {
29-
metadata: {
30-
content_id: payload.record.id,
31-
content_type: "journal_entry",
32-
},
33-
model: "text-embedding-3-small",
34-
provider: "openai",
35-
tokenCount: usage.tokens,
36-
userId: payload.record.user_id,
37-
});
38-
39-
await tx
40-
.insert(JournalEmbedding)
41-
.values({
42-
chunk_text: payload.record.content,
43-
date: payload.record.date,
44-
embedding,
45-
journal_entry_id: payload.record.id,
46-
user_id: payload.record.user_id,
47-
})
48-
.onConflictDoUpdate({
49-
set: {
50-
chunk_text: payload.record.content,
51-
embedding,
52-
},
53-
target: [JournalEmbedding.journal_entry_id],
54-
});
18+
// Call the usage procedure to create the embedding and track usage
19+
await api.ai.createJournalEmbedding({
20+
content: payload.record.content,
21+
date: payload.record.date,
22+
journalEntryId: payload.record.id,
23+
model: "text-embedding-3-small",
24+
provider: "openai",
25+
userId: payload.record.user_id,
5526
});
5627

5728
console.debug(

‎packages/api/src/root.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { aiRouter } from "./router/ai.js";
12
import { authRouter } from "./router/auth.js";
23
import { blocksRouter } from "./router/blocks.js";
34
import { documentRouter } from "./router/document.js";
@@ -7,6 +8,7 @@ import { pagesRouter } from "./router/pages.js";
78
import { createTRPCRouter } from "./trpc.js";
89

910
export const appRouter = createTRPCRouter({
11+
ai: aiRouter,
1012
auth: authRouter,
1113
blocks: blocksRouter,
1214
document: documentRouter,

‎packages/api/src/router/ai.ts‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
import { sql } from "@acme/db";
2+
import {
3+
EmbeddingTokenUsage,
4+
JournalEmbedding,
5+
LLMTokenUsage,
6+
ModelPrices,
7+
} from "@acme/db/schema";
8+
import { openai } from "@ai-sdk/openai";
9+
import type { TRPCRouterRecord } from "@trpc/server";
10+
import { TRPCError } from "@trpc/server";
11+
import { embed } from "ai";
12+
import { z } from "zod/v4";
13+
import { protectedProcedure, publicProcedure } from "../trpc.js";
14+
15+
export const aiRouter = {
16+
createJournalEmbedding: publicProcedure
17+
.input(
18+
z.object({
19+
content: z.string(),
20+
date: z.string(),
21+
journalEntryId: z.string().uuid(),
22+
model: z.string(),
23+
provider: z.string(),
24+
userId: z.string().uuid(),
25+
}),
26+
)
27+
.mutation(async ({ ctx, input }) => {
28+
try {
29+
const { embedding, usage } = await embed({
30+
maxRetries: 5,
31+
model: openai.embedding("text-embedding-3-small"),
32+
value: input.content,
33+
});
34+
35+
return await ctx.db.transaction(async (tx) => {
36+
// Store token usage with automatic pricing calculation
37+
await tx.insert(EmbeddingTokenUsage).values({
38+
metadata: {
39+
content_id: input.journalEntryId,
40+
content_type: "journal_entry",
41+
},
42+
model_price_id: sql`(
43+
SELECT id FROM ${ModelPrices}
44+
WHERE provider = ${input.provider}
45+
AND model = ${input.model}
46+
AND model_type = 'embedding'
47+
AND is_active = 'true'
48+
LIMIT 1
49+
)`,
50+
token_count: usage.tokens,
51+
total_cost: sql`(
52+
SELECT (${usage.tokens} * embedding_price_per_1m_tokens / 1000000.0)::decimal(10,6)
53+
FROM ${ModelPrices}
54+
WHERE provider = ${input.provider}
55+
AND model = ${input.model}
56+
AND model_type = 'embedding'
57+
AND is_active = 'true'
58+
LIMIT 1
59+
)`,
60+
user_id: input.userId,
61+
});
62+
63+
// Store the embedding
64+
const [journalEmbedding] = await tx
65+
.insert(JournalEmbedding)
66+
.values({
67+
chunk_text: input.content,
68+
date: input.date,
69+
embedding,
70+
journal_entry_id: input.journalEntryId,
71+
user_id: input.userId,
72+
})
73+
.onConflictDoUpdate({
74+
set: {
75+
chunk_text: input.content,
76+
embedding,
77+
},
78+
target: [JournalEmbedding.journal_entry_id],
79+
})
80+
.returning();
81+
82+
return journalEmbedding;
83+
});
84+
} catch (error) {
85+
console.error("Error creating journal embedding:", error);
86+
throw new TRPCError({
87+
code: "INTERNAL_SERVER_ERROR",
88+
message: "Failed to create journal embedding",
89+
});
90+
}
91+
}),
92+
// Procedure for tracking LLM token usage
93+
trackLLMUsage: protectedProcedure
94+
.input(
95+
z.object({
96+
inputTokens: z.number(),
97+
metadata: z.record(z.string(), z.unknown()).optional(),
98+
model: z.string(),
99+
outputTokens: z.number(),
100+
provider: z.string(),
101+
}),
102+
)
103+
.mutation(async ({ ctx, input }) => {
104+
try {
105+
return await ctx.db.transaction(async (tx) => {
106+
const totalTokens = input.inputTokens + input.outputTokens;
107+
108+
return await tx
109+
.insert(LLMTokenUsage)
110+
.values({
111+
input_tokens: input.inputTokens,
112+
metadata: input.metadata,
113+
model_price_id: sql`(
114+
SELECT id FROM ${ModelPrices}
115+
WHERE provider = ${input.provider}
116+
AND model = ${input.model}
117+
AND model_type = 'llm'
118+
AND is_active = 'true'
119+
LIMIT 1
120+
)`,
121+
output_tokens: input.outputTokens,
122+
total_cost: sql`(
123+
SELECT (
124+
(${input.inputTokens} * input_price_per_1m_tokens / 1000000.0) +
125+
(${input.outputTokens} * output_price_per_1m_tokens / 1000000.0)
126+
)::decimal(10,6)
127+
FROM ${ModelPrices}
128+
WHERE provider = ${input.provider}
129+
AND model = ${input.model}
130+
AND model_type = 'llm'
131+
AND is_active = 'true'
132+
LIMIT 1
133+
)`,
134+
total_tokens: totalTokens,
135+
user_id: ctx.session.user.id,
136+
})
137+
.returning();
138+
});
139+
} catch (error) {
140+
console.error("Error tracking LLM usage:", error);
141+
throw new TRPCError({
142+
code: "INTERNAL_SERVER_ERROR",
143+
message: "Failed to track LLM usage",
144+
});
145+
}
146+
}),
147+
} satisfies TRPCRouterRecord;

‎packages/db/src/index.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
export { alias } from "drizzle-orm/pg-core";
22
export * from "drizzle-orm/sql";
33
export * from "./client.js";
4-
export * from "./lib/token-usage.js";
54
export * from "./schema.js";

‎packages/db/src/lib/token-usage.ts‎

Lines changed: 0 additions & 95 deletions
This file was deleted.

0 commit comments

Comments
 (0)