Agent/agent/trace/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",  # 写入记忆
    "feedback",      # 人工反馈
]


# Step 状态
Status = Literal[
    "planned",      # 计划中（未执行）
    "in_progress",  # 执行中
    "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

    # 上下文
    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,
            "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_id: str
    trace_id: str
    step_type: StepType
    status: Status
    sequence: int  # 在 Trace 中的顺序

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

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

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

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

    # 执行指标
    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: Status = "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) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "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,
            "data": self.data,
            "summary": self.summary,
            "duration_ms": self.duration_ms,
            "tokens": self.tokens,
            "cost": self.cost,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }


# 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,
#   }
#
# feedback:
#   {
#       "target_step_id": "...",
#       "feedback_type": "positive" | "negative" | "correction",
#       "content": "..."
#   }
#
# memory_read:
#   {
#       "skills": [...],
#       "experiences": [...],
#       "skills_count": 3,
#       "experiences_count": 5
#   }
#
# memory_write:
#   {
#       "experience_id": "...",
#       "condition": "...",
#       "rule": "..."
#   }