vent_agent_2.0/runtime/events.py

"""执行事件结构：事件字段独立于 HTTP / SSE，编码方式由调用方提供。"""

from __future__ import annotations

from typing import Any, Callable, Generic, Iterable, TypeVar

ExecutionEvent = dict[str, Any]
EventT = TypeVar("EventT")

# 事件类型全集（与 1.5 对齐）
AGENT_START = "agent_start"
THINKING = "thinking"
THINKING_TOKEN = "thinking_token"
EXECUTING = "executing"
TOOL_CALL = "tool_call"
TOOL_RESULT = "tool_result"
UPDATED_TODO_LIST = "updated_todo_list"
TOKEN = "token"
PROGRESS = "progress"
INTERRUPT = "interrupt"
CITATIONS = "citations"
AGENT_DONE = "agent_done"
AGENT_ERROR = "agent_error"
ERROR = "error"
DONE = "done"

ALL_EVENT_TYPES: tuple[str, ...] = (
    AGENT_START,
    THINKING,
    THINKING_TOKEN,
    EXECUTING,
    TOOL_CALL,
    TOOL_RESULT,
    UPDATED_TODO_LIST,
    TOKEN,
    PROGRESS,
    INTERRUPT,
    CITATIONS,
    AGENT_DONE,
    AGENT_ERROR,
    ERROR,
    DONE,
)

# 错误码 → 前端提示文案。前端会用 event.code 查这张表兜底。
ERROR_MESSAGES: dict[str, str] = {
    "quota_exceeded": "今日额度已用尽",
    "too_many_tasks": "并发任务已达上限，请稍后重试",
    "task_timeout": "任务执行超时",
    "recursion_limit": "任务步骤过多已中止，请拆分为更小的任务",
    "model_error": "模型调用失败",
    "internal_error": "服务内部错误",
}


class EventBuilder(Generic[EventT]):
    """构造统一的执行事件。

    一次请求一个实例，绑定 ``session_id`` 与 ``trace_id``——这两个字段
    每帧都要带，放在实例里避免每个方法都传一遍。
    """

    def __init__(self, session_id: str, trace_id: str = "", *, encode: Callable[[ExecutionEvent], EventT]) -> None:
        self.encode = encode
        self.session_id = session_id
        self.trace_id = trace_id

    # -------- 底层 --------

    def event(self, event_type: str, **fields: Any) -> EventT:
        """构造一个事件。``type`` 与 ``session_id`` 由本方法统一注入。"""
        payload: dict[str, Any] = {"type": event_type, "session_id": self.session_id}
        if self.trace_id:
            payload["trace_id"] = self.trace_id
        payload.update(fields)
        return self.encode(payload)

    # -------- 生命周期 --------

    def agent_start(self, agent: str, cn_agent: str = "") -> EventT:
        return self.event(AGENT_START, agent=agent, cn_agent=cn_agent or agent)

    def agent_done(
        self, agent: str, cn_agent: str = "", duration_ms: int | None = None, **extra: Any
    ) -> EventT:
        fields: dict[str, Any] = {"agent": agent, "cn_agent": cn_agent or agent}
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        fields.update(extra)
        return self.event(AGENT_DONE, **fields)

    def agent_error(
        self, agent: str, message: str, cn_agent: str = "", duration_ms: int | None = None
    ) -> EventT:
        fields: dict[str, Any] = {
            "agent": agent,
            "cn_agent": cn_agent or agent,
            "message": message,
        }
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        return self.event(AGENT_ERROR, **fields)

    def done(self, duration_ms: int | None = None, **extra: Any) -> EventT:
        fields: dict[str, Any] = {}
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        fields.update(extra)
        return self.event(DONE, **fields)

    def error(
        self, message: str, code: str = "internal_error", retry_after: int | None = None
    ) -> EventT:
        fields: dict[str, Any] = {"message": message, "code": code}
        if retry_after is not None:
            fields["retry_after"] = retry_after
        return self.event(ERROR, **fields)

    # -------- 过程 --------

    def thinking(self, message: str, cn_agent: str = "") -> EventT:
        fields: dict[str, Any] = {"message": message}
        if cn_agent:
            fields["cn_agent"] = cn_agent
        return self.event(THINKING, **fields)

    def thinking_token(self, content: str) -> EventT:
        return self.event(THINKING_TOKEN, content=content)

    def executing(
        self,
        message: str,
        cn_agent: str = "",
        tools: Iterable[str] | None = None,
        cn_tools: Iterable[str] | None = None,
    ) -> EventT:
        fields: dict[str, Any] = {"message": message}
        if cn_agent:
            fields["cn_agent"] = cn_agent
        tool_list = list(tools or [])
        if tool_list:
            fields["tools"] = tool_list
            fields["cn_tools"] = list(cn_tools or []) or tool_list
        return self.event(EXECUTING, **fields)

    def tool_call(
        self,
        tool: str,
        args: Any = None,
        cn_tool: str = "",
        agent_id: str = "",
        tool_call_id: str = "",
    ) -> EventT:
        """工具调用卡。

        ``tool_call_id`` 是 **2.0 的附加字段**（1.5 没有）：同一轮里同一个工具
        可能被调多次，只按 ``tool`` 名索引卡片会导致结果挂到错误的卡上。
        附加字段不破坏兼容——旧前端忽略未知字段即可。
        """
        fields: dict[str, Any] = {"tool": tool, "cn_tool": cn_tool or tool, "args": args}
        if tool_call_id:
            fields["tool_call_id"] = tool_call_id
        if agent_id:
            fields["agent_id"] = agent_id
        return self.event(TOOL_CALL, **fields)

    def tool_result(
        self,
        tool: str,
        content: Any = None,
        cn_tool: str = "",
        agent_id: str = "",
        error: bool = False,
        tool_call_id: str = "",
    ) -> EventT:
        """工具结果卡。``tool_call_id`` 与对应的 tool_call 帧一致。"""
        fields: dict[str, Any] = {
            "tool": tool,
            "cn_tool": cn_tool or tool,
            "content": content,
            "error": error,
        }
        if tool_call_id:
            fields["tool_call_id"] = tool_call_id
        if agent_id:
            fields["agent_id"] = agent_id
        return self.event(TOOL_RESULT, **fields)

    def token(self, content: str) -> EventT:
        return self.event(TOKEN, content=content)

    def todo_list(self, todos: list[dict[str, Any]]) -> EventT:
        return self.event(UPDATED_TODO_LIST, todos=todos)

    def progress(self, message: str, **extra: Any) -> EventT:
        fields: dict[str, Any] = {"message": message}
        fields.update(extra)
        return self.event(PROGRESS, **fields)

    def citations(self, items: list[dict[str, Any]], model_only: bool = False) -> EventT:
        return self.event(CITATIONS, items=items, model_only=model_only)

    # -------- 中断 --------

    def interrupt_tool_approval(
        self,
        tool: str,
        args: Any = None,
        reason: str = "",
        side_effect: str = "",
        **extra: Any,
    ) -> EventT:
        """工具审批卡（对应前端 ``showToolApprovalCard``）。"""
        payload: dict[str, Any] = {
            "tool": tool,
            "args": args,
            "reason": reason,
            "side_effect": side_effect,
        }
        payload.update(extra)
        return self.event(INTERRUPT, tool_approval=payload)

    def interrupt_ask_user(
        self, questions: list[dict[str, Any]], message: str = ""
    ) -> EventT:
        """提问卡（对应前端 ``showAskUserCard``）。

        每个问题形如 ``{"question": str, "type": "text"|"multiple_choice",
        "choices": [...], "required": bool}``。
        """
        return self.event(
            INTERRUPT,
            ask_user={"questions": questions},
            message=message or "智能体暂停，等待输入",
        )

    def interrupt_raw(self, message: str = "", **extra: Any) -> EventT:
        """无法归类的中断：只给一条提示，前端显示 toast。"""
        return self.event(INTERRUPT, message=message, **extra)