From 821c4033276239cf6802f40b284ed407e9de19bf Mon Sep 17 00:00:00 2001 From: Jason Jensen Date: Tue, 25 Feb 2025 17:04:14 +0000 Subject: [PATCH] Get usage from usage field in message response body --- examples/filters/langfuse_filter_pipeline.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/filters/langfuse_filter_pipeline.py b/examples/filters/langfuse_filter_pipeline.py index 8c05f46..7789fdd 100644 --- a/examples/filters/langfuse_filter_pipeline.py +++ b/examples/filters/langfuse_filter_pipeline.py @@ -239,10 +239,10 @@ class Pipeline: # Extract usage if available usage = None if assistant_message_obj: - info = assistant_message_obj.get("info", {}) - if isinstance(info, dict): - input_tokens = info.get("prompt_eval_count") or info.get("prompt_tokens") - output_tokens = info.get("eval_count") or info.get("completion_tokens") + message_usage = assistant_message_obj.get("usage", {}) + if isinstance(message_usage, dict): + input_tokens = message_usage.get("prompt_eval_count") or message_usage.get("prompt_tokens") + output_tokens = message_usage.get("eval_count") or message_usage.get("completion_tokens") if input_tokens is not None and output_tokens is not None: usage = { "input": input_tokens,