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

@@ -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)
 

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)
 

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)