multi-agent-bidding/app/orchestrator.py

"""
规划Agent (Orchestrator) - 意图理解、任务拆分、竞标管理、结果验证
"""

import time
import json
from typing import Dict, List, Optional, Any
from .models import (
    Task, TaskType, TaskStatus, TaskGraph, Bid, Execution, 
    Attempt, ValidationResult, UserRequest, AgentProfile
)
from .task_board import TaskBoard
from .bid_evaluator import BidEvaluator
from .execution_monitor import ExecutionMonitorManager
from .llm_client import LLMClient, default_client


class Orchestrator:
    """规划Agent - 核心调度引擎"""
    
    def __init__(
        self,
        task_board: TaskBoard,
        llm_client: LLMClient = None
    ):
        self.task_board = task_board
        self.llm = llm_client or default_client
        
        self.bid_evaluator = BidEvaluator()
        self.execution_manager = ExecutionMonitorManager()
        
        # 配置
        self.bidding_window = 30       # 竞标窗口时间(秒)
        self.max_retry_per_agent = 1   # 每个Agent最大重试次数
        self.max_total_attempts = 5    # 总尝试次数上限
        self.max_total_time = 1800     # 最大总执行时间(秒)
    
    # === 意图理解 ===
    
    def understand_intent(self, user_request: str) -> Dict:
        """
        意图理解
        
        Args:
            user_request: 用户原始请求
        
        Returns:
            意图分析结果
        """
        try:
            intent = self.llm.analyze_intent(user_request)
            return intent
        except Exception as e:
            # 简化处理
            return {
                "intent": user_request,
                "keywords": [],
                "need_clarification": False,
                "questions": []
            }
    
    # === 任务拆分 ===
    
    def split_tasks(self, intent: Dict, user_request: str) -> TaskGraph:
        """
        任务拆分
        
        Args:
            intent: 意图分析结果
            user_request: 用户原始请求
        
        Returns:
            TaskGraph
        """
        try:
            task_list = self.llm.split_tasks(intent, user_request)
            
            graph = TaskGraph()
            
            for t in task_list:
                task_type = TaskType.PARALLEL if t.get('type') == 'parallel' else TaskType.SERIAL
                task = Task(
                    type=task_type,
                    description=t.get('description', ''),
                    input_data=t.get('input_schema', {}),
                    output_schema=t.get('output_schema', {}),
                    dependencies=t.get('dependencies', [])
                )
                graph.add_task(task)
            
            return graph
        
        except Exception as e:
            # 简化：创建单个任务
            graph = TaskGraph()
            task = Task(
                type=TaskType.INDEPENDENT,
                description=user_request
            )
            graph.add_task(task)
            return graph
    
    # === 任务发布 ===
    
    def publish_task(self, task: Task) -> List[AgentProfile]:
        """
        发布任务
        
        Args:
            task: 任务
        
        Returns:
            能处理该任务的Agent列表
        """
        # 更新状态
        task.status = TaskStatus.PUBLISHED
        self.task_board.add_task(task)
        
        # 找能处理的Agent
        capable_agents = self.task_board.find_capable_agents(task)
        
        return capable_agents
    
    # === 竞标收集 ===
    
    def collect_bids(
        self,
        task: Task,
        agents: List[AgentProfile],
        window_seconds: int = None
    ) -> List[Bid]:
        """
        收集竞标
        
        Args:
            task: 任务
            agents: Agent列表
            window_seconds: 竞标窗口时间
        
        Returns:
            竞标列表
        """
        window = window_seconds or self.bidding_window
        task.status = TaskStatus.BIDDING
        
        bids = []
        
        # 让每个Agent生成竞标
        for agent in agents:
            try:
                bid_data = self.llm.generate_bid(task.to_dict(), agent.to_dict())
                
                bid = Bid(
                    task_id=task.id,
                    agent_id=agent.id,
                    capability_match=float(bid_data.get('capability_match', 0.5)),
                    estimated_time=int(bid_data.get('estimated_time', 60)),
                    confidence=float(bid_data.get('confidence', 0.5)),
                    approach=bid_data.get('approach', ''),
                    prerequisites=bid_data.get('prerequisites', []),
                    alternative_approaches=bid_data.get('alternative_approaches', [])
                )
                
                bids.append(bid)
                self.task_board.add_bid(bid)
            
            except Exception as e:
                # Agent无法竞标，跳过
                pass
        
        return bids
    
    # === 竞标评估 ===
    
    def evaluate_bids(self, task: Task, bids: List[Bid]) -> tuple:
        """
        评估竞标
        
        Args:
            task: 任务
            bids: 竞标列表
        
        Returns:
            (selected_agent, selected_bid, backup_agents)
        """
        agents_dict = {
            a.id: a for a in self.task_board.list_agents()
        }
        
        best_bid, best_agent, scores = self.bid_evaluator.select_best_bid(bids, agents_dict)
        
        if best_agent:
            backup_agents = self.bid_evaluator.get_backup_agents(
                bids, agents_dict, best_agent.id
            )
            
            return (best_agent, best_bid, backup_agents)
        
        return (None, None, [])
    
    # === 任务执行 ===
    
    def assign_and_execute(
        self,
        task: Task,
        agent: AgentProfile,
        bid: Bid,
        backup_agents: List[AgentProfile] = None
    ) -> Execution:
        """
        分配并执行任务
        
        Args:
            task: 任务
            agent: 执行Agent
            bid: 竞标书
            backup_agents: 备选Agent
        
        Returns:
            Execution对象
        """
        task.status = TaskStatus.ASSIGNED
        
        # 定义执行函数
        def execute_func(t, a, b):
            try:
                result = self.llm.execute_task(t.to_dict(), b.approach)
                return result.get('result') or result
            except Exception as e:
                raise e
        
        # 创建监控器
        monitor = self.execution_manager.create_monitor(
            task=task,
            agent=agent,
            bid=bid,
            callbacks={
                'complete': lambda exec: self._on_execution_complete(exec, task),
                'timeout': lambda exec: self._on_execution_timeout(exec, task, backup_agents),
                'error': lambda exec, e: self._on_execution_error(exec, task, e)
            }
        )
        
        # 添加备选
        if backup_agents:
            for backup in backup_agents:
                monitor.add_backup_agent(backup)
        
        # 启动执行
        task.status = TaskStatus.EXECUTING
        execution = monitor.start(execute_func)
        
        self.task_board.add_execution(execution)
        
        return execution
    
    def _on_execution_complete(self, execution: Execution, task: Task):
        """执行完成回调"""
        task.status = TaskStatus.VALIDATING
        
        # 验证结果
        validation = self.validate_result(task, execution.result)
        
        if validation.passed:
            task.status = TaskStatus.COMPLETED
            self.task_board.update_task_status(task.id, task.status)
            
            # 更新Agent统计
            duration = execution.end_time - execution.start_time
            self.task_board.update_agent_stats(
                execution.agent_id,
                success=True,
                duration=duration,
                quality_score=validation.score
            )
        else:
            # 验证失败，尝试重试
            self._handle_validation_failure(task, execution, validation)
    
    def _on_execution_timeout(self, execution: Execution, task: Task, backup_agents: List[AgentProfile]):
        """执行超时回调"""
        # 记录失败
        attempt = Attempt(
            task_id=task.id,
            agent_id=execution.agent_id,
            execution=execution,
            success=False,
            failure_type='timeout',
            duration=execution.end_time - execution.start_time
        )
        self.task_board.add_attempt(attempt)
        
        # 尝试备选Agent
        if backup_agents and self.task_board.get_attempt_count(task.id) < self.max_total_attempts:
            # 找一个备选Agent重新执行
            for backup in backup_agents:
                attempts = [a for a in self.task_board.get_attempts(task.id) if a.agent_id == backup.id]
                if len(attempts) < self.max_retry_per_agent:
                    # 创建新的竞标
                    bid = Bid(
                        task_id=task.id,
                        agent_id=backup.id,
                        capability_match=0.5,
                        estimated_time=60,
                        confidence=0.5,
                        approach="备选执行"
                    )
                    self.assign_and_execute(task, backup, bid, [])
                    return
        
        # 无备选或已达上限，任务失败
        task.status = TaskStatus.FAILED
        self.task_board.update_task_status(task.id, task.status)
    
    def _on_execution_error(self, execution: Execution, task: Task, error: Exception):
        """执行错误回调"""
        # 记录失败
        attempt = Attempt(
            task_id=task.id,
            agent_id=execution.agent_id,
            execution=execution,
            success=False,
            failure_type='error',
            duration=time.time() - execution.start_time,
            feedback=str(error)
        )
        self.task_board.add_attempt(attempt)
        
        task.status = TaskStatus.FAILED
        self.task_board.update_task_status(task.id, task.status)
    
    def _handle_validation_failure(self, task: Task, execution: Execution, validation: ValidationResult):
        """处理验证失败"""
        # 记录失败
        attempt = Attempt(
            task_id=task.id,
            agent_id=execution.agent_id,
            execution=execution,
            validation=validation,
            success=False,
            failure_type='validation',
            duration=execution.end_time - execution.start_time,
            feedback=json.dumps(validation.issues)
        )
        self.task_board.add_attempt(attempt)
        
        # 检查重试次数
        agent_attempts = [a for a in self.task_board.get_attempts(task.id) if a.agent_id == execution.agent_id]
        
        if len(agent_attempts) < self.max_retry_per_agent:
            # 给予重试机会
            task.status = TaskStatus.RETRY
            self.task_board.update_task_status(task.id, task.status)
            
            # 重新执行（带上反馈）
            # 简化处理，直接失败
            task.status = TaskStatus.FAILED
            self.task_board.update_task_status(task.id, task.status)
        else:
            # 已达重试上限，任务失败
            task.status = TaskStatus.FAILED
            self.task_board.update_task_status(task.id, task.status)
    
    # === 结果验证 ===
    
    def validate_result(self, task: Task, result: Any) -> ValidationResult:
        """
        验证结果
        
        Args:
            task: 任务
            result: 执行结果
        
        Returns:
            ValidationResult
        """
        try:
            validation_data = self.llm.validate_result(task.to_dict(), result)
            
            validation = ValidationResult(
                passed=validation_data.get('passed', False),
                issues=validation_data.get('issues', []),
                score=float(validation_data.get('score', 0.5))
            )
            
            return validation
        
        except Exception as e:
            # 简化：假设通过
            return ValidationResult(passed=True, score=0.5)
    
    # === 主流程 ===
    
    def process_request(self, user_request: str) -> UserRequest:
        """
        处理用户请求（完整流程）
        
        Args:
            user_request: 用户原始请求
        
        Returns:
            UserRequest对象
        """
        # 1. 创建请求记录
        request = self.task_board.create_request(user_request)
        
        try:
            # 2. 意图理解
            intent = self.understand_intent(user_request)
            
            # 3. 如果需要澄清
            if intent.get('need_clarification'):
                request.status = 'clarification_needed'
                request.error = json.dumps(intent.get('questions', []))
                self.task_board.update_request(request)
                return request
            
            # 4. 任务拆分
            task_graph = self.split_tasks(intent, user_request)
            request.task_graph = task_graph
            request.status = 'processing'
            self.task_board.update_request(request)
            
            # 5. 按执行顺序处理任务
            execution_order = task_graph.get_execution_order()
            completed = set()
            results = {}
            
            for layer_tasks in execution_order:
                # 并行处理当前层的任务
                layer_results = {}
                
                for task in layer_tasks:
                    # 发布任务
                    capable_agents = self.publish_task(task)
                    
                    if not capable_agents:
                        task.status = TaskStatus.FAILED
                        continue
                    
                    # 收集竞标
                    bids = self.collect_bids(task, capable_agents)
                    
                    if not bids:
                        task.status = TaskStatus.FAILED
                        continue
                    
                    # 评估竞标
                    agent, bid, backup = self.evaluate_bids(task, bids)
                    
                    if not agent:
                        task.status = TaskStatus.FAILED
                        continue
                    
                    # 执行任务
                    execution = self.assign_and_execute(task, agent, bid, backup)
                    
                    # 等待执行完成（简化：直接等待）
                    # 实际应该用异步等待或轮询
                    import time
                    max_wait = task.timeout + 10
                    waited = 0
                    while waited < max_wait:
                        exec_record = self.task_board.get_execution(execution.id)
                        if exec_record and exec_record.status in ['completed', 'failed', 'timeout']:
                            break
                        time.sleep(1)
                        waited += 1
                    
                    # 获取结果
                    if task.status == TaskStatus.COMPLETED:
                        layer_results[task.id] = execution.result
                        completed.add(task.id)
                    else:
                        # 任务失败，整个请求失败
                        request.status = 'failed'
                        request.error = f"任务 {task.id} 执行失败"
                        self.task_board.update_request(request)
                        return request
                
                results.update(layer_results)
            
            # 6. 整合结果
            request.final_result = results
            request.status = 'completed'
            self.task_board.update_request(request)
            
            return request
        
        except Exception as e:
            request.status = 'failed'
            request.error = str(e)
            self.task_board.update_request(request)
            return request
    
    def get_request_status(self, request_id: str) -> Dict:
        """获取请求状态"""
        request = self.task_board.get_request(request_id)
        if request:
            return request.to_dict()
        return None