dataflow_agent/trajectory/builder.py

"""
轨迹构建器 - 将收集的原始数据转换为标准 TRJ 格式

支持：
1. 从 State 和 Collector 构建完整轨迹
2. 自动检测 ReAct/Workflow 模式
3. 提取关键信息和统计数据
"""

from datetime import datetime
from typing import Any, Dict, List, Optional

from dataflow_agent.trajectory.models import (
    Trajectory,
    TrajectoryStep,
    TrajectoryMode,
    StepRole,
)
from dataflow_agent.trajectory.collector import TrajectoryCollector
from dataflow_agent.state import MainState
from dataflow_agent.logger import get_logger

log = get_logger(__name__)


class TrajectoryBuilder:
    """
    轨迹构建器
    
    将 TrajectoryCollector 收集的原始步骤数据和 State 对象
    转换为标准的 Trajectory 对象
    """
    
    def __init__(self):
        pass
    
    def build_from_state(self,
                        state: MainState,
                        collector: TrajectoryCollector,
                        workflow_name: str,
                        user_id: str = None,
                        session_id: str = None) -> Trajectory:
        """
        从 State 和 Collector 构建完整轨迹
        
        Args:
            state: Workflow 执行后的最终状态
            collector: 轨迹收集器
            workflow_name: Workflow 名称
            user_id: 用户 ID
            session_id: 会话 ID
            
        Returns:
            完整的 Trajectory 对象
        """
        log.info(f"[TrajectoryBuilder] 开始构建轨迹: {workflow_name}")
        
        # 1. 生成 trace_id
        trace_id = Trajectory.generate_trace_id()
        
        # 2. 获取步骤
        steps = collector.finish()
        
        # 3. 检测模式
        mode = self._detect_mode(state, steps)
        
        # 4. 提取输入
        inputs = self._extract_inputs(state, collector)
        
        # 5. 提取输出
        final_output = self._extract_final_output(state)
        
        # 6. 判断状态
        status = self._determine_status(state, steps)
        
        # 7. 计算统计信息
        total_duration_ms = self._calculate_total_duration(steps)
        total_tokens = self._calculate_total_tokens(steps)
        
        # 8. 构建 Trajectory
        trajectory = Trajectory(
            trace_id=trace_id,
            workflow_name=workflow_name,
            timestamp=datetime.now().isoformat(),
            status=status,
            mode=mode,
            user_id=user_id,
            session_id=session_id or getattr(state.request, 'session_id', None),
            inputs=inputs,
            steps=steps,
            final_output=final_output,
            total_duration_ms=total_duration_ms,
            total_tokens=total_tokens,
            metadata=collector.get_metadata()
        )
        
        # 更新统计
        trajectory.total_llm_calls = sum(len(step.llm_calls) for step in steps)
        trajectory.total_tool_calls = sum(len(step.tool_calls) for step in steps)
        
        log.info(f"[TrajectoryBuilder] 轨迹构建完成: {trace_id}, "
                f"模式={mode}, 步骤数={len(steps)}, 状态={status}")
        
        return trajectory
    
    def build_from_steps(self,
                        steps: List[TrajectoryStep],
                        workflow_name: str,
                        inputs: Dict[str, Any] = None,
                        final_output: Any = None,
                        **kwargs) -> Trajectory:
        """
        直接从步骤列表构建轨迹（不依赖 State）
        
        Args:
            steps: 步骤列表
            workflow_name: Workflow 名称
            inputs: 输入数据
            final_output: 最终输出
            **kwargs: 其他参数
            
        Returns:
            Trajectory 对象
        """
        trace_id = Trajectory.generate_trace_id()
        mode = self._detect_mode_from_steps(steps)
        status = "success" if not any(step.error for step in steps) else "failed"
        
        trajectory = Trajectory(
            trace_id=trace_id,
            workflow_name=workflow_name,
            timestamp=datetime.now().isoformat(),
            status=status,
            mode=mode,
            inputs=inputs or {},
            steps=steps,
            final_output=final_output,
            **kwargs
        )
        
        # 更新统计
        trajectory.total_llm_calls = sum(len(step.llm_calls) for step in steps)
        trajectory.total_tool_calls = sum(len(step.tool_calls) for step in steps)
        trajectory.total_duration_ms = self._calculate_total_duration(steps)
        trajectory.total_tokens = self._calculate_total_tokens(steps)
        
        return trajectory
    
    def _detect_mode(self, state: MainState, steps: List[TrajectoryStep]) -> str:
        """
        检测轨迹模式
        
        通过分析 State 和 Steps 判断是 ReAct 还是 Workflow 模式
        """
        # 检查是否有 thought 字段（ReAct 特征）
        has_thoughts = any(step.thought for step in steps)
        
        # 检查是否有 observation 字段（ReAct 特征）
        has_observations = any(step.observation for step in steps)
        
        # 检查消息历史（ReAct 通常有更多的对话轮次）
        messages = getattr(state, 'messages', [])
        has_many_messages = len(messages) > 5
        
        # 检查是否有明确的 agent 角色步骤
        has_agent_steps = any(step.role == StepRole.AGENT.value for step in steps)
        
        if has_thoughts or has_observations:
            return TrajectoryMode.REACT.value
        elif has_agent_steps and has_many_messages:
            return TrajectoryMode.HYBRID.value
        else:
            return TrajectoryMode.WORKFLOW.value
    
    def _detect_mode_from_steps(self, steps: List[TrajectoryStep]) -> str:
        """仅从步骤检测模式"""
        has_thoughts = any(step.thought for step in steps)
        has_observations = any(step.observation for step in steps)
        
        if has_thoughts or has_observations:
            return TrajectoryMode.REACT.value
        else:
            return TrajectoryMode.WORKFLOW.value
    
    def _extract_inputs(self, state: MainState, collector: TrajectoryCollector) -> Dict[str, Any]:
        """
        提取输入数据
        
        优先级：
        1. Collector 记录的初始输入
        2. State.request 中的字段
        3. State 的其他相关字段
        """
        inputs = {}
        
        # 从 collector 获取
        collector_inputs = collector.get_initial_inputs()
        if collector_inputs:
            inputs.update(collector_inputs)
        
        # 从 state.request 提取
        if hasattr(state, 'request'):
            request = state.request
            
            # 提取常见字段
            if hasattr(request, 'target') and request.target:
                inputs['query'] = request.target
            
            if hasattr(request, 'model'):
                inputs['model'] = request.model
            
            if hasattr(request, 'language'):
                inputs['language'] = request.language
            
            # 提取其他可能的输入字段
            for field in ['json_file', 'python_file_path', 'keywords', 'style']:
                if hasattr(request, field):
                    value = getattr(request, field)
                    if value:
                        inputs[field] = value
        
        return inputs
    
    def _extract_final_output(self, state: MainState) -> Any:
        """
        提取最终输出
        
        从 State 的不同字段中提取最终结果
        """
        # 尝试从 agent_results 获取
        if hasattr(state, 'agent_results') and state.agent_results:
            # 获取最后一个 agent 的结果
            last_agent_result = None
            for agent_name, result in state.agent_results.items():
                if isinstance(result, dict) and 'results' in result:
                    last_agent_result = result['results']
            
            if last_agent_result:
                return last_agent_result
        
        # 尝试从特定字段获取
        for field in ['final_output', 'execution_result', 'pipeline_structure_code', 
                     'recommendation', 'icon_prompt', 'research_summary']:
            if hasattr(state, field):
                value = getattr(state, field)
                if value:
                    return value
        
        # 如果都没有，返回整个 state 的字典表示（简化版）
        return {"status": "completed"}
    
    def _determine_status(self, state: MainState, steps: List[TrajectoryStep]) -> str:
        """
        判断执行状态
        
        Returns:
            "success" | "failed" | "partial"
        """
        # 检查是否有错误步骤
        has_errors = any(step.error for step in steps)
        
        # 检查 execution_result
        if hasattr(state, 'execution_result'):
            exec_result = state.execution_result
            if isinstance(exec_result, dict):
                if exec_result.get('success') is False:
                    return "failed"
                elif exec_result.get('success') is True:
                    return "success"
        
        # 根据错误情况判断
        if has_errors:
            # 如果所有步骤都有错误，则失败
            if all(step.error for step in steps):
                return "failed"
            else:
                return "partial"
        
        return "success"
    
    def _calculate_total_duration(self, steps: List[TrajectoryStep]) -> Optional[float]:
        """计算总耗时"""
        if not steps:
            return None
        
        total = sum(step.duration_ms for step in steps if step.duration_ms)
        return total if total > 0 else None
    
    def _calculate_total_tokens(self, steps: List[TrajectoryStep]) -> Optional[Dict[str, int]]:
        """计算总 token 使用量"""
        total_prompt = 0
        total_completion = 0
        
        for step in steps:
            for llm_call in step.llm_calls:
                if llm_call.token_usage:
                    total_prompt += llm_call.token_usage.get('prompt', 0)
                    total_completion += llm_call.token_usage.get('completion', 0)
        
        if total_prompt > 0 or total_completion > 0:
            return {
                'prompt': total_prompt,
                'completion': total_completion,
                'total': total_prompt + total_completion
            }
        
        return None


# ==================== 便捷函数 ====================

def create_builder() -> TrajectoryBuilder:
    """创建轨迹构建器"""
    return TrajectoryBuilder()