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

12
lib/langfuse.ts
View File

@@ -43,12 +43,22 @@ export function setTraceInput(params: {
}
// 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;
updateActiveTrace({ output });
const activeSpan = api.trace.getActiveSpan();
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();
}
}