Agent/agent/trace/api.py

"""
Trace RESTful API

提供 Trace、GoalTree、Message 的查询接口
"""

from typing import List, Optional, Dict, Any
from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel

from .protocols import TraceStore


router = APIRouter(prefix="/api/traces", tags=["traces"])


# ===== Response 模型 =====


class TraceListResponse(BaseModel):
    """Trace 列表响应"""
    traces: List[Dict[str, Any]]


class TraceDetailResponse(BaseModel):
    """Trace 详情响应（包含 GoalTree 和 Sub-Traces 元数据）"""
    trace: Dict[str, Any]
    goal_tree: Optional[Dict[str, Any]] = None
    sub_traces: Dict[str, Dict[str, Any]] = {}


class MessagesResponse(BaseModel):
    """Messages 响应"""
    messages: List[Dict[str, Any]]


# ===== 全局 TraceStore（由 api_server.py 注入）=====


_trace_store: Optional[TraceStore] = None


def set_trace_store(store: TraceStore):
    """设置 TraceStore 实例"""
    global _trace_store
    _trace_store = store


def get_trace_store() -> TraceStore:
    """获取 TraceStore 实例"""
    if _trace_store is None:
        raise RuntimeError("TraceStore not initialized")
    return _trace_store


# ===== 路由 =====


@router.get("", response_model=TraceListResponse)
async def list_traces(
    mode: Optional[str] = None,
    agent_type: Optional[str] = None,
    uid: Optional[str] = None,
    status: Optional[str] = None,
    limit: int = Query(20, le=100)
):
    """
    列出 Traces

    Args:
        mode: 模式过滤（call/agent）
        agent_type: Agent 类型过滤
        uid: 用户 ID 过滤
        status: 状态过滤（running/completed/failed）
        limit: 最大返回数量
    """
    store = get_trace_store()
    traces = await store.list_traces(
        mode=mode,
        agent_type=agent_type,
        uid=uid,
        status=status,
        limit=limit
    )
    return TraceListResponse(
        traces=[t.to_dict() for t in traces]
    )


@router.get("/{trace_id}", response_model=TraceDetailResponse)
async def get_trace(trace_id: str):
    """
    获取 Trace 详情

    返回 Trace 元数据、GoalTree、Sub-Traces 元数据（不含 Sub-Trace 内 GoalTree）

    Args:
        trace_id: Trace ID
    """
    store = get_trace_store()

    # 获取 Trace
    trace = await store.get_trace(trace_id)
    if not trace:
        raise HTTPException(status_code=404, detail="Trace not found")

    # 获取 GoalTree
    goal_tree = await store.get_goal_tree(trace_id)

    # 获取所有 Sub-Traces（通过 parent_trace_id 查询）
    sub_traces = {}
    all_traces = await store.list_traces(limit=1000)  # 获取所有 traces
    for t in all_traces:
        if t.parent_trace_id == trace_id:
            sub_traces[t.trace_id] = t.to_dict()

    return TraceDetailResponse(
        trace=trace.to_dict(),
        goal_tree=goal_tree.to_dict() if goal_tree else None,
        sub_traces=sub_traces
    )


@router.get("/{trace_id}/messages", response_model=MessagesResponse)
async def get_messages(
    trace_id: str,
    goal_id: Optional[str] = Query(None, description="过滤指定 Goal 的消息。使用 '_init' 查询初始阶段（goal_id=None）的消息")
):
    """
    获取 Messages

    Args:
        trace_id: Trace ID
        goal_id: 可选，过滤指定 Goal 的消息
                - 不指定: 返回所有消息
                - "_init" 或 "null": 返回初始阶段（goal_id=None）的消息
                - 其他值: 返回指定 Goal 的消息
    """
    store = get_trace_store()

    # 验证 Trace 存在
    trace = await store.get_trace(trace_id)
    if not trace:
        raise HTTPException(status_code=404, detail="Trace not found")

    # 获取 Messages
    if goal_id is None:
        # 没有指定 goal_id，返回所有消息
        messages = await store.get_trace_messages(trace_id)
    elif goal_id in ("_init", "null"):
        # 特殊值：查询初始阶段的消息（goal_id=None）
        all_messages = await store.get_trace_messages(trace_id)
        messages = [m for m in all_messages if m.goal_id is None]
    else:
        # 查询指定 Goal 的消息
        messages = await store.get_messages_by_goal(trace_id, goal_id)

    return MessagesResponse(
        messages=[m.to_dict() for m in messages]
    )