From a83318cf4b442755f066946bc2bae47596398680 Mon Sep 17 00:00:00 2001
From: Tonny <33793693+lzs0603@users.noreply.github.com>
Date: Thu, 3 Apr 2025 16:05:37 +0800
Subject: [PATCH] feat: add langfuse llm node input and output (#17368)

---
 api/core/ops/langfuse_trace/langfuse_trace.py   | 15 +++++++++++++++
 api/core/ops/langsmith_trace/langsmith_trace.py | 15 +++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/api/core/ops/langfuse_trace/langfuse_trace.py b/api/core/ops/langfuse_trace/langfuse_trace.py
index b9ba068b19..f67e270ab1 100644
--- a/api/core/ops/langfuse_trace/langfuse_trace.py
+++ b/api/core/ops/langfuse_trace/langfuse_trace.py
@@ -213,9 +213,24 @@ class LangFuseDataTrace(BaseTraceInstance):
 
             if process_data and process_data.get("model_mode") == "chat":
                 total_token = metadata.get("total_tokens", 0)
+                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:
+                        prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0)
+                        completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0)
+                except Exception:
+                    logger.error("Failed to extract usage", exc_info=True)
+
                 # add generation
                 generation_usage = GenerationUsage(
+                    input=prompt_tokens,
+                    output=completion_tokens,
                     total=total_token,
+                    unit=UnitEnum.TOKENS,
                 )
 
                 node_generation_data = LangfuseGeneration(
diff --git a/api/core/ops/langsmith_trace/langsmith_trace.py b/api/core/ops/langsmith_trace/langsmith_trace.py
index 4ffd888bdd..e3494e2f23 100644
--- a/api/core/ops/langsmith_trace/langsmith_trace.py
+++ b/api/core/ops/langsmith_trace/langsmith_trace.py
@@ -199,6 +199,7 @@ class LangSmithDataTrace(BaseTraceInstance):
             )
 
             process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
+
             if process_data and process_data.get("model_mode") == "chat":
                 run_type = LangSmithRunType.llm
                 metadata.update(
@@ -212,9 +213,23 @@ class LangSmithDataTrace(BaseTraceInstance):
             else:
                 run_type = LangSmithRunType.tool
 
+            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:
+                    prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0)
+                    completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0)
+            except Exception:
+                logger.error("Failed to extract usage", exc_info=True)
+
             node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order)
             langsmith_run = LangSmithRunModel(
                 total_tokens=node_total_tokens,
+                input_tokens=prompt_tokens,
+                output_tokens=completion_tokens,
                 name=node_type,
                 inputs=inputs,
                 run_type=run_type,