i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 22:32:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: add manual token usage reporting to Langfuse for Bedrock streaming

Browse Source 
Bedrock streaming responses don't auto-report token usage to OpenTelemetry.
This fix manually sets span attributes (ai.usage.promptTokens, gen_ai.usage.input_tokens)
from the AI SDK onFinish callback to ensure Langfuse captures token counts.
This commit is contained in:
dayuan.jiang
2025-12-05 00:26:02 +09:00
parent 46d2d4e078
commit 46cbc3354c
2 changed files with 17 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
app/api/chat/route.ts
View File

@@ -182,7 +182,12 @@ ${lastMessageText}
onFinish: ({ text, usage, providerMetadata }) => { onFinish: ({ text, usage, providerMetadata }) => {
console.log('[Cache] Full providerMetadata:', JSON.stringify(providerMetadata, null, 2)); console.log('[Cache] Full providerMetadata:', JSON.stringify(providerMetadata, null, 2));
console.log('[Cache] Usage:', JSON.stringify(usage, null, 2)); console.log('[Cache] Usage:', JSON.stringify(usage, null, 2));
setTraceOutput(text); // Pass usage to Langfuse (Bedrock streaming doesn't auto-report tokens to telemetry)
// AI SDK uses inputTokens/outputTokens, Langfuse expects promptTokens/completionTokens
setTraceOutput(text, {
promptTokens: usage?.inputTokens,
completionTokens: usage?.outputTokens,
});
}, },
tools: { tools: {
// Client-side tool that will be executed on the client // Client-side tool that will be executed on the client

12
lib/langfuse.ts
View File

@@ -43,12 +43,22 @@ export function setTraceInput(params: {
} }
// Update trace with output and end the span // Update trace with output and end the span
export function setTraceOutput(output: string) { export function setTraceOutput(output: string, usage?: { promptTokens?: number; completionTokens?: number }) {
if (!isLangfuseEnabled()) return; if (!isLangfuseEnabled()) return;
updateActiveTrace({ output }); updateActiveTrace({ output });
const activeSpan = api.trace.getActiveSpan(); const activeSpan = api.trace.getActiveSpan();
if (activeSpan) { if (activeSpan) {
// Manually set usage attributes since AI SDK Bedrock streaming doesn't provide them
if (usage?.promptTokens) {
activeSpan.setAttribute('ai.usage.promptTokens', usage.promptTokens);
activeSpan.setAttribute('gen_ai.usage.input_tokens', usage.promptTokens);
}
if (usage?.completionTokens) {
activeSpan.setAttribute('ai.usage.completionTokens', usage.completionTokens);
activeSpan.setAttribute('gen_ai.usage.output_tokens', usage.completionTokens);
}
activeSpan.end(); activeSpan.end();
} }
} }