crewai/main.py

"""
SDLC Agent Demo - FastAPI 主服务（异步版本）
多智能体端到端软件交付协同系统
"""

import json
import uuid
import asyncio
from typing import Dict, Optional, AsyncGenerator
from datetime import datetime
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, RedirectResponse, JSONResponse
from pydantic import BaseModel, Field
import uvicorn

from crews.sdlc_crew import SDLCCrew
from models.qwen_config import get_qwen_config

# ========== FastAPI 应用初始化 ==========
app = FastAPI(
    title="SDLC Agent Demo",
    description="多智能体端到端软件交付协同系统 - 基于 CrewAI + Qwen3.5-flash",
    version="1.0.0"
)

# 启用 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")


# ========== 数据模型 ==========
class StartRequest(BaseModel):
    """启动请求模型"""
    requirement: str = Field(..., description="用户需求描述", min_length=10)


# ========== 任务管理（内存存储） ==========
class TaskManager:
    """任务管理器 - 负责任务状态持久化"""
    
    def __init__(self):
        self.tasks: Dict[str, Dict] = {}
        self.task_queues: Dict[str, asyncio.Queue] = {}
    
    def create_task(self, requirement: str) -> str:
        """创建新任务"""
        task_id = str(uuid.uuid4())
        self.tasks[task_id] = {
            "task_id": task_id,
            "requirement": requirement,
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat()
        }
        # 创建异步队列用于 SSE 推送
        self.task_queues[task_id] = asyncio.Queue()
        return task_id
    
    def update_task_status(self, task_id: str, status: str):
        """更新任务状态"""
        if task_id in self.tasks:
            self.tasks[task_id]["status"] = status
            self.tasks[task_id]["updated_at"] = datetime.now().isoformat()
    
    async def send_event(self, task_id: str, event: dict):
        """发送事件到队列"""
        if task_id in self.task_queues:
            await self.task_queues[task_id].put(event)
    
    async def get_event(self, task_id: str, timeout: float = 60.0) -> Optional[dict]:
        """从队列获取事件"""
        if task_id in self.task_queues:
            try:
                return await asyncio.wait_for(
                    self.task_queues[task_id].get(), 
                    timeout=timeout
                )
            except asyncio.TimeoutError:
                return None
        return None
    
    def get_task(self, task_id: str) -> Optional[Dict]:
        """获取任务信息"""
        return self.tasks.get(task_id)


# 全局任务管理器
task_manager = TaskManager()


# ========== API 端点 ==========
@app.post("/api/v1/sdlc/start", response_model=Dict[str, str])
async def start_sdlc_process(request: StartRequest):
    """
    启动 SDLC 流程（异步执行）
    """
    # 验证配置
    try:
        get_qwen_config()
    except ValueError as e:
        raise HTTPException(status_code=500, detail=str(e))
    
    # 创建任务
    task_id = task_manager.create_task(request.requirement)
    
    # 异步执行 SDLC 流程
    asyncio.create_task(execute_sdlc_flow(task_id, request.requirement))
    
    return {
        "task_id": task_id,
        "status": "processing"
    }


async def execute_sdlc_flow(task_id: str, requirement: str):
    """
    异步执行 SDLC 流程（使用 asyncio.to_thread 运行同步生成器）
    
    Args:
        task_id: 任务 ID
        requirement: 用户需求
    """
    task_manager.update_task_status(task_id, "processing")
    
    try:
        # 先发送一个任务启动事件
        await task_manager.send_event(task_id, {
            "event": "task_started",
            "data": {
                "status": "starting",
                "message": "SDLC 流程已启动",
                "timestamp": datetime.now().isoformat()
            }
        })
        
        # 在线程池中执行同步生成器
        crew = SDLCCrew()
        for event in crew.execute(requirement):
            # 发送事件到队列
            await task_manager.send_event(task_id, event)
            
            # 如果是最终结果或错误，更新状态
            event_type = event.get("event", "")
            if event_type == "final_result":
                task_manager.update_task_status(task_id, "completed")
            elif event_type == "error":
                task_manager.update_task_status(task_id, "failed")
    
    except Exception as e:
        task_manager.update_task_status(task_id, "failed")
        await task_manager.send_event(task_id, {
            "event": "error",
            "data": {
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
        })


@app.get("/api/v1/sdlc/stream/{task_id}")
async def stream_task_progress(task_id: str):
    """
    SSE流式输出任务进度
    
    直接在异步函数中使用 async for 和 yield，
    FastAPI 会自动将其转换为 StreamingResponse
    """
    # 验证任务存在
    task = task_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    # 直接使用 async for 和 yield
    while True:
        event = await task_manager.get_event(task_id, timeout=120.0)
        
        if event is None:
            # 超时，检查任务状态
            task_data = task_manager.get_task(task_id)
            if task_data and task_data["status"] in ["completed", "failed"]:
                yield f"event: end\ndata: {json.dumps({'status': task_data['status']}, ensure_ascii=False)}\n\n"
                break
            continue
        
        # 格式化 SSE事件
        event_type = event.get("event", "message")
        event_data = event.get("data", {})
        
        yield f"event: {event_type}\ndata: {json.dumps(event_data, ensure_ascii=False)}\n\n"
        
        # 如果是结束事件，断开连接
        if event_type in ["final_result", "error"]:
            break


@app.get("/api/v1/sdlc/status/{task_id}")
def get_task_status(task_id: str):
    """
    获取任务状态（非流式）
    """
    task = task_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    return {
        "task_id": task["task_id"],
        "status": task["status"],
        "created_at": task["created_at"],
        "updated_at": task["updated_at"]
    }


@app.get("/api/v1/sdlc/result/{task_id}")
def get_task_result(task_id: str):
    """
    获取任务完整结果
    """
    task = task_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    if task["status"] != "completed":
        raise HTTPException(
            status_code=400, 
            detail=f"Task not completed yet. Status: {task['status']}"
        )
    
    return task


@app.get("/api/v1/sdlc/download/{task_id}")
def download_result(task_id: str):
    """
    打包下载任务结果（ZIP 文件）
    """
    import zipfile
    import io
    from fastapi.responses import StreamingResponse
    
    task = task_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    if task["status"] != "completed":
        raise HTTPException(
            status_code=400, 
            detail=f"Task not completed yet. Status: {task['status']}"
        )
    
    # 创建 ZIP 文件
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        # 1. SRS 文档
        srs_content = ""
        # 注意：异步模式下需要从 events 中获取，这里简化处理
        zip_file.writestr("01_SRS_需求规格说明书.md", "需求文档内容")
        
        # 2. 测试用例
        zip_file.writestr("02_Test_测试用例.md", "测试用例内容")
        
        # 3. 代码实现
        zip_file.writestr("03_Code_代码实现.md", "代码实现内容")
        
        # 4. 项目摘要
        summary = f"""# SDLC 项目交付摘要

## 项目信息
- 任务 ID: {task['task_id']}
- 创建时间：{task['created_at']}
- 完成时间：{task['updated_at']}
- 原始需求：{task['requirement']}

## 生成说明
本项目由 SDLC Agent Demo 自动生成
基于 CrewAI + Qwen3.5-flash + FastAPI
"""
        zip_file.writestr("README_项目摘要.md", summary)
    
    # 准备下载
    zip_buffer.seek(0)
    
    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={
            "Content-Disposition": f"attachment; filename=SDLC_Result_{task_id[:8]}.zip"
        }
    )


@app.get("/")
def root():
    """根路径重定向到测试页面"""
    return RedirectResponse(url="/static/index.html")


@app.get("/health")
def health_check():
    """健康检查端点"""
    return {"status": "healthy", "version": "1.0.0"}


# ========== 主程序入口 ==========
if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=False,
        log_level="info"
    )