From 2270e41ec8eba16dd155762a0c9d0b2e3151c45e Mon Sep 17 00:00:00 2001 From: Davide Delbianco Date: Tue, 1 Jul 2025 14:44:20 +0200 Subject: [PATCH] fix: Prefer process_data to track usage, because ParameterExtractorNode can override the name of the output variable "usage" --- api/core/ops/langfuse_trace/langfuse_trace.py | 8 ++++---- api/core/ops/langsmith_trace/langsmith_trace.py | 8 ++++---- api/core/ops/opik_trace/opik_trace.py | 12 ++++++++---- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/api/core/ops/langfuse_trace/langfuse_trace.py b/api/core/ops/langfuse_trace/langfuse_trace.py index d37342012d..155fc06f2f 100644 --- a/api/core/ops/langfuse_trace/langfuse_trace.py +++ b/api/core/ops/langfuse_trace/langfuse_trace.py @@ -180,12 +180,12 @@ class LangFuseDataTrace(BaseTraceInstance): prompt_tokens = 0 completion_tokens = 0 try: - if outputs.get("usage"): - prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) - completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) - else: + if process_data.get("usage"): prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0) completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0) + else: + prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) except Exception: logger.error("Failed to extract usage", exc_info=True) diff --git a/api/core/ops/langsmith_trace/langsmith_trace.py b/api/core/ops/langsmith_trace/langsmith_trace.py index 8a392940db..e025220208 100644 --- a/api/core/ops/langsmith_trace/langsmith_trace.py +++ b/api/core/ops/langsmith_trace/langsmith_trace.py @@ -206,12 +206,12 @@ class LangSmithDataTrace(BaseTraceInstance): prompt_tokens = 0 completion_tokens = 0 try: - if outputs.get("usage"): - prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) - completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) - else: + if process_data.get("usage"): prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0) completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0) + else: + prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) except Exception: logger.error("Failed to extract usage", exc_info=True) diff --git a/api/core/ops/opik_trace/opik_trace.py b/api/core/ops/opik_trace/opik_trace.py index f4d2760ba5..b7e49c2650 100644 --- a/api/core/ops/opik_trace/opik_trace.py +++ b/api/core/ops/opik_trace/opik_trace.py @@ -222,10 +222,14 @@ class OpikDataTrace(BaseTraceInstance): ) try: - if outputs.get("usage"): - total_tokens = outputs["usage"].get("total_tokens", 0) - prompt_tokens = outputs["usage"].get("prompt_tokens", 0) - completion_tokens = outputs["usage"].get("completion_tokens", 0) + if process_data.get("usage"): + total_tokens = process_data.get("usage", {}).get("total_tokens", 0) + prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0) + else: + total_tokens = outputs.get("usage", {}).get("total_tokens", 0) + prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) except Exception: logger.error("Failed to extract usage", exc_info=True)