瀏覽代碼

增加操作用户信息

xueyiming 1 周之前
父節點
當前提交
10a673d6f8

+ 1 - 1
pqai_agent/data_models/agent_test_task.py

@@ -13,7 +13,7 @@ class AgentTestTask(Base):
     create_user = Column(String(32), nullable=True, comment="创建用户")
     update_user = Column(String(32), nullable=True, comment="更新用户")
     dataset_ids = Column(Text, nullable=True, comment="数据集ids")
-    evaluate_type = Column(Integer, nullable=False, default=0, comment="数据集ids")
+    evaluate_type = Column(Integer, nullable=False, default=0, comment="评估类型(0-回复 1-唤起)")
     status = Column(Integer, nullable=True, comment="状态(0:未开始, 1:进行中, 2:已完成, 3:已取消)")
     create_time = Column(TIMESTAMP, nullable=False, server_default="CURRENT_TIMESTAMP", comment="创建时间")
     update_time = Column(TIMESTAMP, nullable=False, server_default="CURRENT_TIMESTAMP", onupdate="CURRENT_TIMESTAMP",

+ 3 - 2
pqai_agent_server/agent_task_server.py

@@ -139,7 +139,7 @@ class AgentTaskManager:
                 "list": response_data,
             }
 
-    def create_task(self, agent_id: int, task_prompt: str):
+    def create_task(self, agent_id: int, task_prompt: str, user: str):
         """创建新任务"""
         with self.session_maker() as session:
             agent_config = session.get(AgentConfiguration, agent_id)
@@ -147,7 +147,8 @@ class AgentTaskManager:
                                    status=AgentTaskStatus.NOT_STARTED.value,
                                    start_time=datetime.now(),
                                    input=task_prompt,
-                                   tools=agent_config.tools)
+                                   tools=agent_config.tools,
+                                   create_user=user)
             session.add(agent_task)
             session.commit()  # 显式提交
             task_id = agent_task.id

+ 18 - 6
pqai_agent_server/api_server.py

@@ -708,13 +708,16 @@ def create_test_task():
     agent_id = req_data.get('agentId', None)
     module_id = req_data.get('moduleId', None)
     evaluate_type = req_data.get('evaluateType', None)
+    user = req_data.get('user', None)
     if agent_id is None:
         return wrap_response(404, msg='agent id is required')
     if module_id is None:
         return wrap_response(404, msg='module id is required')
     if evaluate_type is None:
-        return wrap_response(404, msg='evaluate type id is required')
-    app.task_manager.create_task(agent_id, module_id, evaluate_type)
+        return wrap_response(404, msg='evaluate type is required')
+    if user is None:
+        return wrap_response(404, msg='user is required')
+    app.task_manager.create_task(agent_id, module_id, evaluate_type, user)
     return wrap_response(200)
 
 
@@ -726,12 +729,15 @@ def stop_test_task():
     """
     req_data = request.json
     task_id = req_data.get('taskId', None)
-    if not task_id:
+    user = req_data.get('user', None)
+    if task_id is None:
         return wrap_response(404, msg='task id is required')
+    if user is None:
+        return wrap_response(404, msg='user is required')
     task = app.task_manager.get_task(task_id)
     if task.status not in (TestTaskStatus.NOT_STARTED.value, TestTaskStatus.IN_PROGRESS.value):
         return wrap_response(404, msg='task status is invalid')
-    app.task_manager.cancel_task(task_id)
+    app.task_manager.cancel_task(task_id, user)
     return wrap_response(200)
 
 
@@ -743,12 +749,15 @@ def resume_test_task():
     """
     req_data = request.json
     task_id = req_data.get('taskId', None)
+    user = req_data.get('user', None)
     if task_id is None:
         return wrap_response(404, msg='task id is required')
+    if user is None:
+        return wrap_response(404, msg='user is required')
     task = app.task_manager.get_task(task_id)
     if task.status != TestTaskStatus.CANCELLED.value:
         return wrap_response(404, msg='task status is invalid')
-    app.task_manager.resume_task(task_id)
+    app.task_manager.resume_task(task_id, user)
     return wrap_response(200)
 
 
@@ -813,11 +822,14 @@ def create_agent_task():
     req_data = request.json
     agent_id = req_data.get('agentId', None)
     task_prompt = req_data.get('taskPrompt', None)
+    user = req_data.get('user', None)
     if agent_id is None:
         return wrap_response(404, msg='agent id is required')
     if task_prompt is None:
         return wrap_response(404, msg='task_prompt is required')
-    app.agent_task_manager.create_task(agent_id, task_prompt)
+    if user is None:
+        return wrap_response(404, msg='user is required')
+    app.agent_task_manager.create_task(agent_id, task_prompt, user)
     return wrap_response(200)
 
 @app.route("/api/getAgentTaskList", methods=["GET"])

+ 6 - 6
pqai_agent_server/task_server.py

@@ -102,11 +102,11 @@ class TaskManager:
                 "list": response_data,
             }
 
-    def create_task(self, agent_id: int, module_id: int, evaluate_type: int) -> Dict:
+    def create_task(self, agent_id: int, module_id: int, evaluate_type: int, user: str) -> Dict:
         """创建新任务"""
         with self.session_maker() as session:
             agent_test_task = AgentTestTask(agent_id=agent_id, module_id=module_id, evaluate_type=evaluate_type,
-                                            status=TestTaskStatus.CREATING.value)
+                                            status=TestTaskStatus.CREATING.value, create_user=user, update_user=user)
             session.add(agent_test_task)
             session.commit()  # 显式提交
             task_id = agent_test_task.id
@@ -249,7 +249,7 @@ class TaskManager:
                 {"status": status, "input": input, "output": output, "score": score, "update_time": datetime.now()})
             session.commit()
 
-    def cancel_task(self, task_id: int):
+    def cancel_task(self, task_id: int, user: str):
         """取消任务(带事务支持)"""
         # 设置取消事件
         if task_id in self.task_events:
@@ -261,13 +261,13 @@ class TaskManager:
         with self.session_maker() as session:
             with session.begin():
                 session.query(AgentTestTask).filter(AgentTestTask.id == task_id).update(
-                    {"status": TestTaskStatus.CANCELLED.value})
+                    {"status": TestTaskStatus.CANCELLED.value, "update_user": user, "update_time": datetime.now()})
                 session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).filter(
                     AgentTestTaskConversations.status == TestTaskConversationsStatus.PENDING.value).update(
                     {"status": TestTaskConversationsStatus.CANCELLED.value})
                 session.commit()
 
-    def resume_task(self, task_id: int) -> bool:
+    def resume_task(self, task_id: int, user: str) -> bool:
         """恢复已取消的任务"""
         task = self.get_task(task_id)
         if not task or task.status != TestTaskStatus.CANCELLED.value:
@@ -276,7 +276,7 @@ class TaskManager:
         with self.session_maker() as session:
             with session.begin():
                 session.query(AgentTestTask).filter(AgentTestTask.id == task_id).update(
-                    {"status": TestTaskStatus.NOT_STARTED.value})
+                    {"status": TestTaskStatus.NOT_STARTED.value, "update_user": user, "update_time": datetime.now()})
                 session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).filter(
                     AgentTestTaskConversations.status == TestTaskConversationsStatus.CANCELLED.value).update(
                     {"status": TestTaskConversationsStatus.PENDING.value})