|
@@ -271,6 +271,7 @@ class TraceTask:
|
|
|
message_id: Optional[str] = None,
|
|
|
workflow_run: Optional[WorkflowRun] = None,
|
|
|
conversation_id: Optional[str] = None,
|
|
|
+ user_id: Optional[str] = None,
|
|
|
timer: Optional[Any] = None,
|
|
|
**kwargs
|
|
|
):
|
|
@@ -278,6 +279,7 @@ class TraceTask:
|
|
|
self.message_id = message_id
|
|
|
self.workflow_run = workflow_run
|
|
|
self.conversation_id = conversation_id
|
|
|
+ self.user_id = user_id
|
|
|
self.timer = timer
|
|
|
self.kwargs = kwargs
|
|
|
self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
|
|
@@ -290,7 +292,9 @@ class TraceTask:
|
|
|
def preprocess(self):
|
|
|
preprocess_map = {
|
|
|
TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs),
|
|
|
- TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace(self.workflow_run, self.conversation_id),
|
|
|
+ TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace(
|
|
|
+ self.workflow_run, self.conversation_id, self.user_id
|
|
|
+ ),
|
|
|
TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(self.message_id),
|
|
|
TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace(
|
|
|
self.message_id, self.timer, **self.kwargs
|
|
@@ -313,7 +317,7 @@ class TraceTask:
|
|
|
def conversation_trace(self, **kwargs):
|
|
|
return kwargs
|
|
|
|
|
|
- def workflow_trace(self, workflow_run: WorkflowRun, conversation_id):
|
|
|
+ def workflow_trace(self, workflow_run: WorkflowRun, conversation_id, user_id):
|
|
|
workflow_id = workflow_run.workflow_id
|
|
|
tenant_id = workflow_run.tenant_id
|
|
|
workflow_run_id = workflow_run.id
|
|
@@ -358,6 +362,7 @@ class TraceTask:
|
|
|
"total_tokens": total_tokens,
|
|
|
"file_list": file_list,
|
|
|
"triggered_form": workflow_run.triggered_from,
|
|
|
+ "user_id": user_id,
|
|
|
}
|
|
|
|
|
|
workflow_trace_info = WorkflowTraceInfo(
|
|
@@ -654,10 +659,11 @@ trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100))
|
|
|
|
|
|
|
|
|
class TraceQueueManager:
|
|
|
- def __init__(self, app_id=None):
|
|
|
+ def __init__(self, app_id=None, user_id=None):
|
|
|
global trace_manager_timer
|
|
|
|
|
|
self.app_id = app_id
|
|
|
+ self.user_id = user_id
|
|
|
self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)
|
|
|
self.flask_app = current_app._get_current_object()
|
|
|
if trace_manager_timer is None:
|