Agent/agent/execution/models.py

"""
Trace 和 Step 数据模型

Trace: 一次完整的 LLM 交互（单次调用或 Agent 任务）
Step: Trace 中的一个原子操作，形成树结构
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, Any, List, Optional, Literal
import uuid


# Step 类型
StepType = Literal[
    # 计划相关
    "goal",        # 目标/计划项（可以有子 steps）

    # LLM 输出
    "thought",     # 思考/分析（中间过程）
    "evaluation",  # 评估总结（需要 summary）
    "response",    # 最终回复

    # 工具相关（数据结构上分开以保留描述能力，可视化时候可能合并）
    "action",      # 工具调用（tool_call）
    "result",      # 工具结果（tool_result）

    # 系统相关
    "memory_read",   # 读取记忆（经验/技能）
    "memory_write",  # 写入记忆
]


# Step 状态
StepStatus = Literal[
    "planned",           # 计划中（未执行）
    "in_progress",       # 执行中
    "awaiting_approval", # 等待用户确认
    "completed",         # 已完成
    "failed",            # 失败
    "skipped",           # 跳过
]


@dataclass
class Trace:
    """
    执行轨迹 - 一次完整的 LLM 交互

    单次调用: mode="call", 只有 1 个 Step
    Agent 模式: mode="agent", 多个 Steps 形成树结构
    """
    trace_id: str
    mode: Literal["call", "agent"]

    # Prompt 标识（可选）
    prompt_name: Optional[str] = None

    # Agent 模式特有
    task: Optional[str] = None
    agent_type: Optional[str] = None

    # 状态
    status: Literal["running", "completed", "failed"] = "running"

    # 统计
    total_steps: int = 0
    total_tokens: int = 0
    total_cost: float = 0.0
    total_duration_ms: int = 0  # 总耗时（毫秒）

    # 进度追踪（head）
    last_sequence: int = 0      # 最新 step 的 sequence
    last_event_id: int = 0      # 最新事件 ID（用于 WS 续传）

    # 上下文
    uid: Optional[str] = None
    context: Dict[str, Any] = field(default_factory=dict)

    # 当前焦点 goal（用于 step 工具）
    current_goal_id: Optional[str] = None

    # 时间
    created_at: datetime = field(default_factory=datetime.now)
    completed_at: Optional[datetime] = None

    @classmethod
    def create(
        cls,
        mode: Literal["call", "agent"],
        **kwargs
    ) -> "Trace":
        """创建新的 Trace"""
        return cls(
            trace_id=str(uuid.uuid4()),
            mode=mode,
            **kwargs
        )

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "trace_id": self.trace_id,
            "mode": self.mode,
            "prompt_name": self.prompt_name,
            "task": self.task,
            "agent_type": self.agent_type,
            "status": self.status,
            "total_steps": self.total_steps,
            "total_tokens": self.total_tokens,
            "total_cost": self.total_cost,
            "total_duration_ms": self.total_duration_ms,
            "last_sequence": self.last_sequence,
            "last_event_id": self.last_event_id,
            "uid": self.uid,
            "context": self.context,
            "current_goal_id": self.current_goal_id,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
        }


@dataclass
class Step:
    """
    执行步骤 - Trace 中的一个原子操作

    Step 之间通过 parent_id 形成树结构（单父节点）

    ## 字段设计规则

    **顶层字段**（Step 类属性）：
    - 所有（或大部分）step 都有的字段
    - 需要筛选/排序/索引的字段（如 tokens, cost, duration_ms）
    - 结构化、类型明确的字段

    **data 字段**（Dict）：
    - step_type 特定的字段（不同类型有不同 schema）
    - 详细的业务数据（如 messages, content, arguments, output）
    - 可能很大的字段
    - 半结构化、动态的字段

    ## data 字段 schema（按 step_type）

    - thought/response: model, messages, content, tool_calls
    - action: tool_name, arguments
    - result: tool_name, output, error
    - memory_read: experiences_count, skills_count
    - goal: 自定义（根据具体目标）
    """
    step_id: str
    trace_id: str
    step_type: StepType
    status: StepStatus
    sequence: int  # 在 Trace 中的顺序

    # 树结构（单父节点）
    parent_id: Optional[str] = None

    # 内容
    description: str = ""  # 所有节点都有，系统自动提取

    # 类型相关数据
    data: Dict[str, Any] = field(default_factory=dict)

    # 仅 evaluation 类型需要
    summary: Optional[str] = None

    # UI 优化字段
    has_children: bool = False      # 是否有子节点
    children_count: int = 0         # 子节点数量

    # 执行指标
    duration_ms: Optional[int] = None
    tokens: Optional[int] = None
    cost: Optional[float] = None

    # 时间
    created_at: datetime = field(default_factory=datetime.now)

    @classmethod
    def create(
        cls,
        trace_id: str,
        step_type: StepType,
        sequence: int,
        status: StepStatus = "completed",
        description: str = "",
        data: Dict[str, Any] = None,
        parent_id: Optional[str] = None,
        summary: Optional[str] = None,
        duration_ms: Optional[int] = None,
        tokens: Optional[int] = None,
        cost: Optional[float] = None,
    ) -> "Step":
        """创建新的 Step"""
        return cls(
            step_id=str(uuid.uuid4()),
            trace_id=trace_id,
            step_type=step_type,
            status=status,
            sequence=sequence,
            parent_id=parent_id,
            description=description,
            data=data or {},
            summary=summary,
            duration_ms=duration_ms,
            tokens=tokens,
            cost=cost,
        )

    def to_dict(self, view: str = "full") -> Dict[str, Any]:
        """
        转换为字典

        Args:
            view: "compact" - 不返回大字段
                  "full" - 返回完整数据
        """
        result = {
            "step_id": self.step_id,
            "trace_id": self.trace_id,
            "step_type": self.step_type,
            "status": self.status,
            "sequence": self.sequence,
            "parent_id": self.parent_id,
            "description": self.description,
            "summary": self.summary,
            "has_children": self.has_children,
            "children_count": self.children_count,
            "duration_ms": self.duration_ms,
            "tokens": self.tokens,
            "cost": self.cost,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }

        # 处理 data 字段
        if view == "compact":
            # compact 模式：移除 data 中的大字段
            data_copy = self.data.copy()
            # 移除可能的大字段（如 output, content 等）
            for key in ["output", "content", "full_output", "full_content"]:
                data_copy.pop(key, None)
            result["data"] = data_copy
        else:
            # full 模式：返回完整 data
            result["data"] = self.data

        return result


# Step.data 结构说明
#
# goal:
#   {
#       "description": "探索代码库",
#   }
#
# thought:
#   {
#       "content": "需要先了解项目结构...",
#   }
#
# action:
#   {
#       "tool_name": "glob_files",
#       "arguments": {"pattern": "**/*.py"},
#   }
#
# result:
#   {
#       "tool_name": "glob_files",
#       "output": ["src/main.py", ...],
#       "title": "找到 15 个文件",
#   }
#
# evaluation:
#   {
#       "content": "分析完成...",
#   }
#   # summary 字段存储简短总结
#
# response:
#   {
#       "content": "任务已完成...",
#       "is_final": True,
#   }
#
# memory_read:
#   {
#       "skills": [...],
#       "experiences": [...],
#       "skills_count": 3,
#       "experiences_count": 5
#   }
#
# memory_write:
#   {
#       "experience_id": "...",
#       "condition": "...",
#       "rule": "..."
#   }