feat: 新增任务状态管理和API支持

- 添加任务状态API路由,支持获取当前用户的任务状态和待处理任务列表
- 实现任务追踪服务,使用Redis存储任务状态
- 更新回忆录处理逻辑,集成Celery任务提交和状态更新
- 增强测试用例,支持任务状态的获取和清除功能
- 优化代码结构,提升可读性和维护性
This commit is contained in:
penghanyuan
2026-01-21 23:37:00 +01:00
parent 0591e9d7c1
commit 3f899aa16c
6 changed files with 451 additions and 23 deletions

View File

@@ -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: