From 3f0da88ff7ae1b157ec2ada876ce9bdf5fed427b Mon Sep 17 00:00:00 2001
From: Joe <79627742+ZhouhaoJiang@users.noreply.github.com>
Date: Fri, 5 Jul 2024 18:37:26 +0800
Subject: [PATCH] fix: update workflow trace query (#6010)

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

diff --git a/api/core/ops/langfuse_trace/langfuse_trace.py b/api/core/ops/langfuse_trace/langfuse_trace.py
index 0441475921..cb86396420 100644
--- a/api/core/ops/langfuse_trace/langfuse_trace.py
+++ b/api/core/ops/langfuse_trace/langfuse_trace.py
@@ -107,7 +107,20 @@ class LangFuseDataTrace(BaseTraceInstance):
 
         # through workflow_run_id get all_nodes_execution
         workflow_nodes_executions = (
-            db.session.query(WorkflowNodeExecution)
+            db.session.query(
+                WorkflowNodeExecution.id,
+                WorkflowNodeExecution.tenant_id,
+                WorkflowNodeExecution.app_id,
+                WorkflowNodeExecution.title,
+                WorkflowNodeExecution.node_type,
+                WorkflowNodeExecution.status,
+                WorkflowNodeExecution.inputs,
+                WorkflowNodeExecution.outputs,
+                WorkflowNodeExecution.created_at,
+                WorkflowNodeExecution.elapsed_time,
+                WorkflowNodeExecution.process_data,
+                WorkflowNodeExecution.execution_metadata,
+            )
             .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id)
             .all()
         )
diff --git a/api/core/ops/langsmith_trace/langsmith_trace.py b/api/core/ops/langsmith_trace/langsmith_trace.py
index 93d74cc9e3..0ce91db335 100644
--- a/api/core/ops/langsmith_trace/langsmith_trace.py
+++ b/api/core/ops/langsmith_trace/langsmith_trace.py
@@ -100,7 +100,20 @@ class LangSmithDataTrace(BaseTraceInstance):
 
         # through workflow_run_id get all_nodes_execution
         workflow_nodes_executions = (
-            db.session.query(WorkflowNodeExecution)
+            db.session.query(
+                WorkflowNodeExecution.id,
+                WorkflowNodeExecution.tenant_id,
+                WorkflowNodeExecution.app_id,
+                WorkflowNodeExecution.title,
+                WorkflowNodeExecution.node_type,
+                WorkflowNodeExecution.status,
+                WorkflowNodeExecution.inputs,
+                WorkflowNodeExecution.outputs,
+                WorkflowNodeExecution.created_at,
+                WorkflowNodeExecution.elapsed_time,
+                WorkflowNodeExecution.process_data,
+                WorkflowNodeExecution.execution_metadata,
+            )
             .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id)
             .all()
         )