generate_data_set.py

pqai_agent/agents/message_push_agent.py

@@ -134,8 +134,8 @@ class MessagePushAgent(SimpleOpenAICompatibleChatAgent):
         ])
         super().__init__(model, system_prompt, tools, generate_cfg, max_run_step)
 
-    def generate_message(self, context: Dict, dialogue_history: List[Dict]) -> str:
-        formatted_dialogue = MessagePushAgent.compose_dialogue(dialogue_history)
+    def generate_message(self, context: Dict, dialogue_history: List[Dict], timestamp_type: str='ms') -> str:
+        formatted_dialogue = MessagePushAgent.compose_dialogue(dialogue_history, timestamp_type)
         query = QUERY_PROMPT_TEMPLATE.format(**context, dialogue_history=formatted_dialogue)
         self.run(query)
         for tool_call in reversed(self.tool_call_records):
@@ -144,7 +144,7 @@ class MessagePushAgent(SimpleOpenAICompatibleChatAgent):
         return ''
 
     @staticmethod
-    def compose_dialogue(dialogue: List[Dict]) -> str:
+    def compose_dialogue(dialogue: List[Dict], timestamp_type: str='ms') -> str:
         role_map = {'user': '用户', 'assistant': '客服'}
         messages = []
         for msg in dialogue:
@@ -152,7 +152,10 @@ class MessagePushAgent(SimpleOpenAICompatibleChatAgent):
                 continue
             if msg['role'] not in role_map:
                 continue
-            format_dt = datetime.datetime.fromtimestamp(msg['timestamp'] / 1000).strftime('%Y-%m-%d %H:%M:%S')
+            if timestamp_type == 'ms':
+                format_dt = datetime.datetime.fromtimestamp(msg['timestamp'] / 1000).strftime('%Y-%m-%d %H:%M:%S')
+            else:
+                format_dt = datetime.datetime.fromtimestamp(msg['timestamp']).strftime('%Y-%m-%d %H:%M:%S')
             msg_type = msg.get('type', MessageType.TEXT).description
             messages.append('[{}][{}][{}]{}'.format(role_map[msg['role']], format_dt, msg_type, msg['content']))
         return '\n'.join(messages)