| @@ -298,34 +298,29 @@ class TraceTask: | |||
| self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") | |||
| def execute(self): | |||
| method_name, trace_info = self.preprocess() | |||
| return trace_info | |||
| return self.preprocess() | |||
| def preprocess(self): | |||
| if self.trace_type == TraceTaskName.CONVERSATION_TRACE: | |||
| return TraceTaskName.CONVERSATION_TRACE, self.conversation_trace(**self.kwargs) | |||
| if self.trace_type == TraceTaskName.WORKFLOW_TRACE: | |||
| return TraceTaskName.WORKFLOW_TRACE, self.workflow_trace(self.workflow_run, self.conversation_id) | |||
| elif self.trace_type == TraceTaskName.MESSAGE_TRACE: | |||
| return TraceTaskName.MESSAGE_TRACE, self.message_trace(self.message_id) | |||
| elif self.trace_type == TraceTaskName.MODERATION_TRACE: | |||
| return TraceTaskName.MODERATION_TRACE, self.moderation_trace(self.message_id, self.timer, **self.kwargs) | |||
| elif self.trace_type == TraceTaskName.SUGGESTED_QUESTION_TRACE: | |||
| return TraceTaskName.SUGGESTED_QUESTION_TRACE, self.suggested_question_trace( | |||
| 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.MESSAGE_TRACE: lambda: self.message_trace(self.message_id), | |||
| TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace( | |||
| self.message_id, self.timer, **self.kwargs | |||
| ) | |||
| elif self.trace_type == TraceTaskName.DATASET_RETRIEVAL_TRACE: | |||
| return TraceTaskName.DATASET_RETRIEVAL_TRACE, self.dataset_retrieval_trace( | |||
| ), | |||
| TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace( | |||
| self.message_id, self.timer, **self.kwargs | |||
| ) | |||
| elif self.trace_type == TraceTaskName.TOOL_TRACE: | |||
| return TraceTaskName.TOOL_TRACE, self.tool_trace(self.message_id, self.timer, **self.kwargs) | |||
| elif self.trace_type == TraceTaskName.GENERATE_NAME_TRACE: | |||
| return TraceTaskName.GENERATE_NAME_TRACE, self.generate_name_trace( | |||
| ), | |||
| TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace( | |||
| self.message_id, self.timer, **self.kwargs | |||
| ), | |||
| TraceTaskName.TOOL_TRACE: lambda: self.tool_trace(self.message_id, self.timer, **self.kwargs), | |||
| TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace( | |||
| self.conversation_id, self.timer, **self.kwargs | |||
| ) | |||
| else: | |||
| return '', {} | |||
| ), | |||
| } | |||
| return preprocess_map.get(self.trace_type, lambda: None)() | |||
| # process methods for different trace types | |||
| def conversation_trace(self, **kwargs): | |||