Log token usage of AI SDK

ref: https://sdk.vercel.ai/examples/node/streaming-structured-data/token-usage
giselles-ai · Oct 21, 2024 · bdfdbfb · bdfdbfb
1 parent a6e8332
commit bdfdbfb
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts b/app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts
@@ -23,6 +23,17 @@ type GenerateArtifactStreamParams = {
 	userPrompt: string;
 	systemPrompt?: string;
 };
+
+function recordTokenUsage({
+	promptTokens,
+	completionTokens,
+	totalTokens,
+}: TokenUsage) {
+	console.log("Prompt tokens:", promptTokens);
+	console.log("Completion tokens:", completionTokens);
+	console.log("Total tokens:", totalTokens);
+}
+
 export async function generateArtifactStream(
 	params: GenerateArtifactStreamParams,
 ) {
@@ -67,6 +78,8 @@ export async function generateArtifactStream(
 
 		const result = await object;
 
+		result.usage.then(recordTokenUsage);
+
 		stream.done();
 	})();