Agent/agent/models/trace.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


StepType = Literal[
    "llm_call",      # LLM 调用
    "tool_call",     # 工具调用
    "tool_result",   # 工具结果
    "conclusion",    # 中间/最终结论
    "feedback",      # 人工反馈
    "memory_read",   # 读取记忆（经验/技能）
    "memory_write",  # 写入记忆
]


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

    单次调用: mode="call", 只有 1 个 Step
    Agent 模式: mode="agent", 多个 Steps 形成 DAG
    """
    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)

    # 时间
    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,
            "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_ids 形成 DAG 结构
    """
    step_id: str
    trace_id: str
    step_type: StepType
    sequence: int  # 在 Trace 中的顺序

    # DAG 结构（支持多父节点）
    parent_ids: List[str] = field(default_factory=list)

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

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

    @classmethod
    def create(
        cls,
        trace_id: str,
        step_type: StepType,
        sequence: int,
        data: Dict[str, Any] = None,
        parent_ids: List[str] = None,
    ) -> "Step":
        """创建新的 Step"""
        return cls(
            step_id=str(uuid.uuid4()),
            trace_id=trace_id,
            step_type=step_type,
            sequence=sequence,
            parent_ids=parent_ids or [],
            data=data or {},
        )

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "step_id": self.step_id,
            "trace_id": self.trace_id,
            "step_type": self.step_type,
            "sequence": self.sequence,
            "parent_ids": self.parent_ids,
            "data": self.data,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }


# Step.data 结构说明
#
# llm_call:
#   {
#       "messages": [...],
#       "response": "...",
#       "model": "gpt-4o",
#       "prompt_tokens": 100,
#       "completion_tokens": 50,
#       "cost": 0.01,
#       "tool_calls": [...]  # 如果有
#   }
#
# tool_call:
#   {
#       "tool_name": "search_blocks",
#       "arguments": {...},
#       "llm_step_id": "..."  # 哪个 LLM 调用触发的
#   }
#
# tool_result:
#   {
#       "tool_call_step_id": "...",
#       "result": "...",
#       "duration_ms": 123
#   }
#
# conclusion:
#   {
#       "content": "...",
#       "is_final": True/False
#   }
#
# 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": "..."
#   }