crewai/main.py

"""
FastAPI 主入口
提供 RESTful API 和 SSE 流式接口
"""

import asyncio
import json
from datetime import datetime
from typing import Dict, Any, Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from crew_factory import CrewFactory, run_multi_agent_task
from stream_manager import stream_manager, create_sse_generator


# ==================== 数据模型 ====================

class RunTaskRequest(BaseModel):
    """运行任务请求体"""
    user_requirement: str = Field(
        ...,
        description="用户需求描述",
        example="开发一个在线商城系统，支持用户注册、商品浏览、购物车和支付功能"
    )
    skip_confirmation: bool = Field(
        default=True,
        description="是否跳过 Coordinator 的人工确认环节"
    )


class RunTaskResponse(BaseModel):
    """运行任务响应体"""
    task_id: str
    status: str
    message: Optional[str] = None


class SSEEvent(BaseModel):
    """SSE 事件数据结构"""
    agent: str
    type: str
    content: str
    timestamp: str
    task_id: str


# ==================== 生命周期管理 ====================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时执行
    print("🚀 Multi-Agent System 启动中...")
    print("📡 SSE 流服务已就绪")
    
    # 启动后台任务：定期清理旧流
    asyncio.create_task(cleanup_streams_periodically())
    
    yield
    
    # 关闭时清理
    print("👋 正在关闭所有 SSE 流...")
    # 可以在这里添加清理逻辑


# ==================== FastAPI 应用 ====================

app = FastAPI(
    title="Multi-Agent Software Delivery System",
    description="""
    基于 CrewAI + Qwen3.5-flash 的多智能体软件交付系统
    
    ## 核心功能
    - **产品需求分析**: ProductManager Agent 分析用户需求，生成 PRD
    - **测试计划制定**: QAEngineer Agent 设计测试策略和用例
    - **技术方案设计**: SoftwareDeveloper Agent 输出架构设计和代码框架
    - **质量审核**: Coordinator Agent 审核所有产出物并生成交付报告
    
    ## 实时通信
    支持 SSE (Server-Sent Events) 协议，实时推送每个 Agent 的思考过程和任务状态
    
    ## API 端点
    - `POST /api/run_task` - 启动新任务
    - `GET /api/stream/{task_id}` - 订阅任务执行日志（SSE）
    - `GET /api/task/{task_id}/status` - 查询任务状态
    - `GET /test-ui` - 测试页面
    """,
    version="1.0.0",
    lifespan=lifespan,
)

# CORS 中间件（允许跨域）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ==================== 后台任务 ====================

async def cleanup_streams_periodically(interval: int = 300):
    """定期清理旧的流（每 5 分钟）"""
    while True:
        await asyncio.sleep(interval)
        try:
            await stream_manager.cleanup_old_streams(max_age_seconds=3600)
        except Exception as e:
            print(f"清理流失败：{e}")


# ==================== API 路由 ====================

@app.post(
    "/api/run_task",
    response_model=RunTaskResponse,
    summary="启动多智能体任务",
    description="接收用户需求，启动 CrewAI 流程，异步执行并立即返回 task_id"
)
async def run_task(request: RunTaskRequest):
    """
    启动新的多智能体任务
    
    - **user_requirement**: 用户需求描述
    - **skip_confirmation**: 是否跳过人工确认（默认 True）
    
    返回 task_id 用于后续 SSE 流订阅
    """
    try:
        # 生成 task_id 并启动任务
        task_id = await run_multi_agent_task(
            user_requirement=request.user_requirement,
            skip_confirmation=request.skip_confirmation
        )
        
        return RunTaskResponse(
            task_id=task_id,
            status="started",
            message="任务已启动，请通过 /api/stream/{task_id} 订阅执行日志"
        )
    
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"启动任务失败：{str(e)}"
        )


@app.get(
    "/api/stream/{task_id}",
    summary="订阅任务执行日志 (SSE)",
    description="建立 SSE 连接，实时接收任务执行过程中的所有事件"
)
async def stream_task_logs(task_id: str):
    """
    SSE 端点 - 实时推送任务执行日志
    
    事件类型包括:
    - **start**: 任务开始
    - **agent_start**: Agent 开始执行
    - **thought**: Agent 思考过程
    - **action**: Agent 执行动作
    - **output**: Agent 输出结果
    - **step_end**: 步骤完成
    - **end**: 任务结束
    - **error**: 发生错误
    
    数据格式:
    ```json
    {
      "agent": "ProductManager",
      "type": "thought",
      "content": "正在分析用户需求...",
      "timestamp": "2024-01-01T12:00:00",
      "task_id": "uuid"
    }
    ```
    """
    # 检查任务是否存在
    stream = await stream_manager.get_stream(task_id)
    if stream is None:
        # 任务不存在，返回错误
        return StreamingResponse(
            iter([f"data: {json.dumps({'error': 'Task not found', 'task_id': task_id})}\n\n"]),
            media_type="text/event-stream"
        )
    
    # 创建 SSE 流
    return StreamingResponse(
        create_sse_generator(task_id),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Nginx 禁用缓冲
        }
    )


@app.get(
    "/api/task/{task_id}/status",
    summary="查询任务状态",
    description="获取任务的当前状态和基本信息"
)
async def get_task_status(task_id: str):
    """查询任务状态"""
    stream = await stream_manager.get_stream(task_id)
    
    if stream is None:
        return {"task_id": task_id, "status": "not_found"}
    
    return {
        "task_id": task_id,
        "status": "closed" if stream.is_closed else "running",
        "queue_size": stream.queue.qsize() if hasattr(stream, 'queue') else 0,
    }


@app.get(
    "/api/streams",
    summary="列出所有活跃流",
    description="查看当前所有正在执行的任务"
)
async def list_active_streams():
    """列出所有活跃的 SSE 流"""
    streams = stream_manager.list_active_streams()
    return {
        "total": len(streams),
        "streams": streams
    }


@app.get(
    "/health",
    summary="健康检查",
    description="检查服务是否正常运行"
)
async def health_check():
    """健康检查端点"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "service": "Multi-Agent Software Delivery System"
    }


# ==================== 测试页面 ====================

@app.get(
    "/test-ui",
    response_class=HTMLResponse,
    summary="测试 UI 页面",
    description="一个简单的 HTML 页面，用于测试 SSE 流功能"
)
async def test_ui():
    """返回测试页面"""
    return HTMLResponse(content=get_test_page_html())


def get_test_page_html() -> str:
    """生成测试页面 HTML"""
    return """<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>多智能体系统 - 测试 UI</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        
        body {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        
        .container {
            max-width: 1200px;
            margin: 0 auto;
        }
        
        h1 {
            color: white;
            text-align: center;
            margin-bottom: 30px;
            font-size: 2.5em;
            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        }
        
        .card {
            background: white;
            border-radius: 10px;
            padding: 25px;
            margin-bottom: 20px;
            box-shadow: 0 10px 40px rgba(0,0,0,0.2);
        }
        
        .input-group {
            margin-bottom: 20px;
        }
        
        label {
            display: block;
            margin-bottom: 8px;
            font-weight: 600;
            color: #333;
        }
        
        textarea {
            width: 100%;
            min-height: 120px;
            padding: 12px;
            border: 2px solid #e0e0e0;
            border-radius: 6px;
            font-size: 14px;
            resize: vertical;
            transition: border-color 0.3s;
        }
        
        textarea:focus {
            outline: none;
            border-color: #667eea;
        }
        
        .checkbox-group {
            display: flex;
            align-items: center;
            gap: 8px;
        }
        
        button {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            border: none;
            padding: 12px 30px;
            border-radius: 6px;
            font-size: 16px;
            font-weight: 600;
            cursor: pointer;
            transition: transform 0.2s, box-shadow 0.2s;
        }
        
        button:hover {
            transform: translateY(-2px);
            box-shadow: 0 5px 20px rgba(102, 126, 234, 0.4);
        }
        
        button:disabled {
            opacity: 0.6;
            cursor: not-allowed;
            transform: none;
        }
        
        .status-bar {
            background: #f5f5f5;
            padding: 15px;
            border-radius: 6px;
            margin-bottom: 20px;
            display: flex;
            justify-content: space-between;
            align-items: center;
        }
        
        .status-indicator {
            display: inline-block;
            width: 12px;
            height: 12px;
            border-radius: 50%;
            margin-right: 8px;
        }
        
        .status-idle { background: #ccc; }
        .status-running { background: #4caf50; animation: pulse 1s infinite; }
        .status-error { background: #f44336; }
        
        @keyframes pulse {
            0%, 100% { opacity: 1; }
            50% { opacity: 0.5; }
        }
        
        .event-log {
            background: #1e1e1e;
            color: #d4d4d4;
            border-radius: 6px;
            padding: 15px;
            max-height: 600px;
            overflow-y: auto;
            font-family: 'Consolas', 'Monaco', monospace;
            font-size: 13px;
            line-height: 1.6;
        }
        
        .event-item {
            padding: 8px 12px;
            margin-bottom: 8px;
            border-left: 3px solid;
            background: rgba(255,255,255,0.05);
            border-radius: 0 4px 4px 0;
        }
        
        .event-start { border-color: #2196f3; }
        .event-thought { border-color: #ff9800; }
        .event-output { border-color: #4caf50; }
        .event-end { border-color: #9c27b0; }
        .event-error { border-color: #f44336; background: rgba(244, 67, 54, 0.1); }
        
        .event-header {
            display: flex;
            justify-content: space-between;
            margin-bottom: 5px;
            font-size: 12px;
            color: #aaa;
        }
        
        .event-agent {
            font-weight: 600;
            color: #667eea;
        }
        
        .event-type {
            padding: 2px 8px;
            border-radius: 3px;
            background: rgba(102, 126, 234, 0.2);
            color: #667eea;
        }
        
        .event-content {
            color: #d4d4d4;
            white-space: pre-wrap;
            word-break: break-word;
        }
        
        .clear-btn {
            background: #f44336;
            padding: 8px 16px;
            font-size: 14px;
        }
        
        .example-requirements {
            margin-top: 15px;
        }
        
        .example-btn {
            background: #f5f5f5;
            color: #333;
            padding: 8px 16px;
            font-size: 13px;
            margin-right: 10px;
            margin-bottom: 10px;
            border: 1px solid #ddd;
        }
        
        .example-btn:hover {
            background: #e0e0e0;
            transform: none;
            box-shadow: none;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>🤖 多智能体软件交付系统</h1>
        
        <!-- 控制面板 -->
        <div class="card">
            <div class="input-group">
                <label for="requirement">用户需求描述：</label>
                <textarea 
                    id="requirement" 
                    placeholder="请输入您的需求，例如：开发一个在线商城系统，支持用户注册、商品浏览、购物车和支付功能..."
                ></textarea>
                
                <div class="example-requirements">
                    <strong>示例需求：</strong><br>
                    <button class="example-btn" onclick="useExample('mall')">在线商城</button>
                    <button class="example-btn" onclick="useExample('blog')">博客系统</button>
                    <button class="example-btn" onclick="useExample('task')">任务管理</button>
                </div>
            </div>
            
            <div class="input-group">
                <div class="checkbox-group">
                    <input type="checkbox" id="skipConfirmation" checked>
                    <label for="skipConfirmation" style="margin: 0;">跳过人工确认环节</label>
                </div>
            </div>
            
            <button id="startBtn" onclick="startTask()">🚀 启动任务</button>
        </div>
        
        <!-- 状态栏 -->
        <div class="card">
            <div class="status-bar">
                <div>
                    <span class="status-indicator status-idle" id="statusIndicator"></span>
                    <span id="statusText">等待中...</span>
                </div>
                <div>
                    <strong>Task ID:</strong> <span id="taskIdDisplay">-</span>
                </div>
                <button class="clear-btn" onclick="clearLog()">清空日志</button>
            </div>
        </div>
        
        <!-- 事件日志 -->
        <div class="card">
            <h2 style="margin-bottom: 15px;">📊 实时事件日志</h2>
            <div class="event-log" id="eventLog">
                <div style="color: #666; text-align: center; padding: 40px;">
                    暂无事件，点击上方"启动任务"按钮开始
                </div>
            </div>
        </div>
    </div>
    
    <script>
        let currentTaskId = null;
        let eventSource = null;
        
        // 示例需求
        const examples = {
            mall: '开发一个在线商城系统，支持用户注册登录、商品分类浏览、搜索功能、购物车管理、订单创建和支付宝/微信支付集成。需要包含后台管理系统用于商品上架和订单处理。',
            blog: '构建一个个人博客系统，支持 Markdown 编辑器、文章分类和标签、评论功能、SEO 优化。还需要有访客统计、文章阅读量统计，以及响应式设计适配移动端。',
            task: '创建一个团队协作任务管理工具，类似简化版 Jira。支持创建项目、分解任务、分配负责人、设置优先级和截止日期、进度跟踪。需要有看板视图和甘特图展示。'
        };
        
        function useExample(type) {
            document.getElementById('requirement').value = examples[type];
        }
        
        function updateStatus(status, text) {
            const indicator = document.getElementById('statusIndicator');
            indicator.className = `status-indicator status-${status}`;
            document.getElementById('statusText').textContent = text;
        }
        
        function addEventToLog(event) {
            const log = document.getElementById('eventLog');
            
            // 清除初始提示
            if (log.querySelector('[style*="text-align: center"]')) {
                log.innerHTML = '';
            }
            
            const eventDiv = document.createElement('div');
            eventDiv.className = `event-item event-${event.type}`;
            
            const time = new Date(event.timestamp).toLocaleTimeString();
            
            eventDiv.innerHTML = `
                <div class="event-header">
                    <span class="event-agent">${escapeHtml(event.agent)}</span>
                    <span>
                        <span class="event-type">${event.type}</span>
                        <span style="margin-left: 10px;">${time}</span>
                    </span>
                </div>
                <div class="event-content">${escapeHtml(event.content)}</div>
            `;
            
            log.appendChild(eventDiv);
            log.scrollTop = log.scrollHeight;
        }
        
        function escapeHtml(text) {
            const div = document.createElement('div');
            div.textContent = text;
            return div.innerHTML;
        }
        
        async function startTask() {
            const requirement = document.getElementById('requirement').value.trim();
            if (!requirement) {
                alert('请输入用户需求描述');
                return;
            }
            
            const skipConfirmation = document.getElementById('skipConfirmation').checked;
            
            try {
                updateStatus('running', '任务启动中...');
                document.getElementById('startBtn').disabled = true;
                
                const response = await fetch('/api/run_task', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                    },
                    body: JSON.stringify({
                        user_requirement: requirement,
                        skip_confirmation: skipConfirmation
                    })
                });
                
                if (!response.ok) {
                    throw new Error(`HTTP ${response.status}`);
                }
                
                const data = await response.json();
                currentTaskId = data.task_id;
                document.getElementById('taskIdDisplay').textContent = currentTaskId;
                
                // 连接 SSE
                connectSSE(currentTaskId);
                
            } catch (error) {
                console.error('启动任务失败:', error);
                updateStatus('error', '启动失败：' + error.message);
                document.getElementById('startBtn').disabled = false;
            }
        }
        
        function connectSSE(taskId) {
            // 关闭旧的连接
            if (eventSource) {
                eventSource.close();
            }
            
            const eventSourceUrl = `/api/stream/${taskId}`;
            eventSource = new EventSource(eventSourceUrl);
            
            eventSource.onopen = () => {
                console.log('SSE 连接已建立');
                updateStatus('running', '任务执行中...');
            };
            
            eventSource.onmessage = (event) => {
                try {
                    const data = JSON.parse(event.data);
                    
                    // 忽略连接结束事件
                    if (data.type === 'connection_end') {
                        return;
                    }
                    
                    addEventToLog(data);
                    
                    // 如果是结束或错误事件，更新状态
                    if (data.type === 'end') {
                        updateStatus('idle', '任务已完成');
                        document.getElementById('startBtn').disabled = false;
                        eventSource.close();
                    } else if (data.type === 'error') {
                        updateStatus('error', '发生错误');
                        document.getElementById('startBtn').disabled = false;
                    }
                    
                } catch (e) {
                    console.error('解析事件数据失败:', e);
                }
            };
            
            eventSource.onerror = (error) => {
                console.error('SSE 错误:', error);
                // 不立即显示错误，等待服务端发送的 error 事件
            };
        }
        
        function clearLog() {
            document.getElementById('eventLog').innerHTML = `
                <div style="color: #666; text-align: center; padding: 40px;">
                    暂无事件，点击上方"启动任务"按钮开始
                </div>
            `;
            currentTaskId = null;
            document.getElementById('taskIdDisplay').textContent = '-';
            if (eventSource) {
                eventSource.close();
                eventSource = null;
            }
            updateStatus('idle', '等待中...');
            document.getElementById('startBtn').disabled = false;
        }
    </script>
</body>
</html>"""


# ==================== 主程序入口 ====================

if __name__ == "__main__":
    import uvicorn
    
    # 加载环境变量
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass
    
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )