Agent/agent/trace/__init__.py

"""
Trace 模块 - Context 管理 + 可视化

核心职责：
1. Trace/Step 模型定义
2. 存储接口和实现（内存/数据库）
3. RESTful API（可视化查询）
4. WebSocket 推送（实时更新）
"""

# 模型（核心，无依赖）
from agent.trace.models import Trace, Step, StepType, Status

# 存储接口（核心，无依赖）
from agent.trace.protocols import TraceStore

# 内存存储实现（核心，无依赖）
from agent.trace.memory_store import MemoryTraceStore


# API 路由（可选，需要 FastAPI）
def _get_api_router():
    """延迟导入 API Router（避免强制依赖 FastAPI）"""
    from agent.trace.api import router
    return router


def _get_ws_router():
    """延迟导入 WebSocket Router（避免强制依赖 FastAPI）"""
    from agent.trace.websocket import router
    return router


# WebSocket 广播函数（可选，需要 FastAPI）
def _get_broadcast_functions():
    """延迟导入 WebSocket 广播函数"""
    from agent.trace.websocket import (
        broadcast_step_added,
        broadcast_step_updated,
        broadcast_trace_completed,
    )
    return broadcast_step_added, broadcast_step_updated, broadcast_trace_completed


# 便捷属性（仅在访问时导入）
@property
def api_router():
    return _get_api_router()


@property
def ws_router():
    return _get_ws_router()


__all__ = [
    # 模型
    "Trace",
    "Step",
    "StepType",
    "Status",
    # 存储
    "TraceStore",
    "MemoryTraceStore",
]