feat: 新增任务状态管理和API支持
- 添加任务状态API路由,支持获取当前用户的任务状态和待处理任务列表 - 实现任务追踪服务,使用Redis存储任务状态 - 更新回忆录处理逻辑,集成Celery任务提交和状态更新 - 增强测试用例,支持任务状态的获取和清除功能 - 优化代码结构,提升可读性和维护性
This commit is contained in:
@@ -117,6 +117,25 @@ class ConversationTester:
|
||||
return {"message": "获取失败"}
|
||||
return resp.json()
|
||||
|
||||
async def get_tasks_status(self):
|
||||
"""获取任务状态"""
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
resp = await client.get(
|
||||
f"{BASE_URL}/api/tasks/status",
|
||||
headers={"Authorization": f"Bearer {self.token}"}
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
return {"total": 0, "all_completed": True, "tasks": []}
|
||||
return resp.json()
|
||||
|
||||
async def clear_tasks(self):
|
||||
"""清除任务记录"""
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
await client.delete(
|
||||
f"{BASE_URL}/api/tasks/clear",
|
||||
headers={"Authorization": f"Bearer {self.token}"}
|
||||
)
|
||||
|
||||
async def run_conversation(self):
|
||||
"""运行多轮对话"""
|
||||
print(f"\n🔗 连接 WebSocket: {self.conversation_id}")
|
||||
@@ -187,18 +206,75 @@ class ConversationTester:
|
||||
except asyncio.TimeoutError:
|
||||
print("⏰ 等待结束确认超时(但后台处理可能仍在进行)")
|
||||
|
||||
async def wait_for_processing(self, max_wait_seconds: int = 300, check_interval: int = 3):
|
||||
"""
|
||||
等待后台处理完成
|
||||
通过查询 Celery 任务状态来判断处理是否完成
|
||||
|
||||
Args:
|
||||
max_wait_seconds: 最大等待时间(秒),默认 5 分钟
|
||||
check_interval: 检查间隔(秒)
|
||||
|
||||
Returns:
|
||||
是否在超时前完成
|
||||
"""
|
||||
print(f"\n⏳ 等待后台任务完成(最多 {max_wait_seconds} 秒)...")
|
||||
print(" 提示: 通过 Celery 任务状态 API 追踪任务进度")
|
||||
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
|
||||
while True:
|
||||
elapsed = asyncio.get_event_loop().time() - start_time
|
||||
|
||||
if elapsed >= max_wait_seconds:
|
||||
print(f"\n⚠️ 已等待 {max_wait_seconds} 秒,超时退出")
|
||||
return False
|
||||
|
||||
# 检查任务状态
|
||||
tasks_status = await self.get_tasks_status()
|
||||
total = tasks_status.get("total", 0)
|
||||
pending = tasks_status.get("pending", 0)
|
||||
running = tasks_status.get("running", 0)
|
||||
success = tasks_status.get("success", 0)
|
||||
failure = tasks_status.get("failure", 0)
|
||||
all_completed = tasks_status.get("all_completed", False)
|
||||
|
||||
# 同时检查章节内容
|
||||
chapters = await self.get_chapters()
|
||||
chapter_count = len(chapters)
|
||||
total_content_length = sum(len(ch.get('content', '')) for ch in chapters)
|
||||
|
||||
status_str = f"📊 总:{total} 等待:{pending} 运行:{running} 成功:{success} 失败:{failure}"
|
||||
content_str = f"📚 章节:{chapter_count} 内容:{total_content_length}字符"
|
||||
print(f" [{int(elapsed):3d}s] {status_str} | {content_str}")
|
||||
|
||||
# 判断是否完成:
|
||||
# 1. 有任务且全部完成
|
||||
# 2. 或者没有任务但有章节内容(兼容旧逻辑)
|
||||
if total > 0 and all_completed:
|
||||
print(f"\n✅ 所有任务已完成!共 {total} 个任务,等待 {int(elapsed)} 秒")
|
||||
return True
|
||||
|
||||
# 如果没有任务记录,等待一会儿任务提交
|
||||
if total == 0 and elapsed < 15:
|
||||
await asyncio.sleep(check_interval)
|
||||
continue
|
||||
|
||||
# 如果长时间没有任务但有内容,也认为完成
|
||||
if total == 0 and chapter_count > 0 and elapsed > 30:
|
||||
print(f"\n✅ 无待处理任务,已有 {chapter_count} 个章节。等待 {int(elapsed)} 秒")
|
||||
return True
|
||||
|
||||
await asyncio.sleep(check_interval)
|
||||
|
||||
async def check_results(self):
|
||||
"""检查回忆录生成结果"""
|
||||
print(f"\n{'='*60}")
|
||||
print("📊 检查结果")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# 等待后台处理完成(后台任务有5秒debounce + 多次LLM调用)
|
||||
# 每条消息可能触发2-3次LLM调用,7条消息 = 约20次调用,每次约3秒
|
||||
print("\n⏳ 等待后台处理完成...")
|
||||
for i in range(6): # 60秒,每10秒检查一次
|
||||
print(f" 等待中... {(i+1)*10}秒")
|
||||
await asyncio.sleep(10)
|
||||
# 等待后台处理完成(使用智能轮询)
|
||||
await self.wait_for_processing(max_wait_seconds=180, check_interval=5)
|
||||
|
||||
# 获取回忆录状态
|
||||
print("\n📋 回忆录状态:")
|
||||
@@ -223,7 +299,8 @@ class ConversationTester:
|
||||
if chapters:
|
||||
for ch in chapters:
|
||||
is_new = "🆕" if ch.get("is_new") else ""
|
||||
print(f" {is_new} [{ch.get('category', 'N/A')}] {ch.get('title', 'N/A')}")
|
||||
content_len = len(ch.get('content', ''))
|
||||
print(f" {is_new} [{ch.get('category', 'N/A')}] {ch.get('title', 'N/A')} ({content_len} 字符)")
|
||||
else:
|
||||
print(" (暂无章节)")
|
||||
|
||||
@@ -248,7 +325,7 @@ class ConversationTester:
|
||||
content = ch.get('content', '')
|
||||
|
||||
print(f"\n{'─'*60}")
|
||||
print(f"【{title}】")
|
||||
print(f"【{title}】({category})")
|
||||
print(f"{'─'*60}")
|
||||
if content:
|
||||
print(content)
|
||||
@@ -270,15 +347,19 @@ async def main():
|
||||
# 1. 注册/登录
|
||||
await tester.register_or_login()
|
||||
|
||||
# 2. 查看初始状态
|
||||
# 2. 清除旧的任务记录
|
||||
await tester.clear_tasks()
|
||||
print("\n🧹 已清除旧的任务记录")
|
||||
|
||||
# 3. 查看初始状态
|
||||
print("\n📋 初始回忆录状态:")
|
||||
state = await tester.get_memoir_state()
|
||||
print(f" 当前阶段: {state.get('current_stage', 'N/A')}")
|
||||
|
||||
# 3. 运行多轮对话
|
||||
# 4. 运行多轮对话
|
||||
await tester.run_conversation()
|
||||
|
||||
# 4. 检查结果
|
||||
# 5. 检查结果
|
||||
await tester.check_results()
|
||||
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user