refactor: 简化 AgentScope 运行时模块与 prompt 系统

2026-03-15 17:14:15 +08:00
parent 61997f3613
commit 072c09d99d
32 changed files with 750 additions and 1863 deletions
@@ -1,32 +1,157 @@
 from __future__ import annotations

-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast

-_TYPE_MAP: dict[str, str] = {
-    "run.started": "RUN_STARTED",
-    "run.finished": "RUN_FINISHED",
-    "run.error": "RUN_ERROR",
-    "step.start": "STEP_STARTED",
-    "step.finish": "STEP_FINISHED",
-    "text.start": "TEXT_MESSAGE_START",
-    "text.delta": "TEXT_MESSAGE_CONTENT",
-    "text.end": "TEXT_MESSAGE_END",
-    "tool.start": "TOOL_CALL_START",
-    "tool.args": "TOOL_CALL_ARGS",
-    "tool.end": "TOOL_CALL_END",
-    "tool.result": "TOOL_CALL_RESULT",
-    "tool.error": "TOOL_CALL_ERROR",
-    "state.snapshot": "STATE_SNAPSHOT",
-    "messages.snapshot": "MESSAGES_SNAPSHOT",
+from ag_ui.core import (
+    BaseEvent,
+    EventType,
+    RunStartedEvent,
+    RunFinishedEvent,
+    RunErrorEvent,
+    StepStartedEvent,
+    StepFinishedEvent,
+    TextMessageStartEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    ToolCallResultEvent,
+)
+
+if TYPE_CHECKING:
+    pass
+
+_INTERNAL_TO_AGUI: dict[str, EventType] = {
+    "run.started": EventType.RUN_STARTED,
+    "run.finished": EventType.RUN_FINISHED,
+    "run.error": EventType.RUN_ERROR,
+    "step.start": EventType.STEP_STARTED,
+    "step.finish": EventType.STEP_FINISHED,
+    "text.start": EventType.TEXT_MESSAGE_START,
+    "text.delta": EventType.TEXT_MESSAGE_CONTENT,
+    "text.end": EventType.TEXT_MESSAGE_END,
+    "tool.start": EventType.TOOL_CALL_START,
+    "tool.args": EventType.TOOL_CALL_ARGS,
+    "tool.end": EventType.TOOL_CALL_END,
+    "tool.result": EventType.TOOL_CALL_RESULT,
+    "state.snapshot": EventType.STATE_SNAPSHOT,
+    "messages.snapshot": EventType.MESSAGES_SNAPSHOT,
 }


-def to_agui_wire_event(event: dict[str, Any]) -> dict[str, Any]:
-    event_type = str(event.get("type", "")).strip()
-    wire_type = _TYPE_MAP.get(event_type, event_type.upper().replace(".", "_"))
+def _convert_to_agui_type(internal_type: str) -> EventType:
+    return _INTERNAL_TO_AGUI.get(
+        internal_type, EventType(internal_type.upper().replace(".", "_"))
+    )
+
+
+def _is_agui_event(event: dict[str, Any]) -> bool:
+    event_type = event.get("type", "")
+    try:
+        EventType(event_type)
+        return True
+    except ValueError:
+        return False
+
+
+def _build_run_started(event: dict[str, Any]) -> RunStartedEvent:
+    return RunStartedEvent(
+        thread_id=event.get("threadId", ""),
+        run_id=event.get("runId", ""),
+    )
+
+
+def _build_run_finished(event: dict[str, Any]) -> RunFinishedEvent:
+    return RunFinishedEvent(
+        thread_id=event.get("threadId", ""),
+        run_id=event.get("runId", ""),
+    )
+
+
+def _build_run_error(event: dict[str, Any]) -> RunErrorEvent:
+    data = event.get("data", {})
+    return RunErrorEvent(
+        message=data.get("message", "Unknown error"),
+        code=data.get("code"),
+    )
+
+
+def _build_step_started(event: dict[str, Any]) -> StepStartedEvent:
+    data = event.get("data", {})
+    return StepStartedEvent(
+        step_name=data.get("stepName", ""),
+    )
+
+
+def _build_step_finished(event: dict[str, Any]) -> StepFinishedEvent:
+    data = event.get("data", {})
+    return StepFinishedEvent(
+        step_name=data.get("stepName", ""),
+    )
+
+
+def _build_text_start(event: dict[str, Any]) -> TextMessageStartEvent:
+    data = event.get("data", {})
+    return TextMessageStartEvent(
+        message_id=data.get("messageId", ""),
+        role=data.get("role", "assistant"),
+    )
+
+
+def _build_text_delta(event: dict[str, Any]) -> TextMessageContentEvent:
+    data = event.get("data", {})
+    return TextMessageContentEvent(
+        message_id=data.get("messageId", ""),
+        delta=data.get("delta", ""),
+    )
+
+
+def _build_text_end(event: dict[str, Any]) -> TextMessageEndEvent:
+    data = event.get("data", {})
+    return TextMessageEndEvent(
+        message_id=data.get("messageId", ""),
+    )
+
+
+def _build_tool_result(event: dict[str, Any]) -> ToolCallResultEvent:
+    data = event.get("data", {})
+    return ToolCallResultEvent(
+        message_id=data.get("messageId", ""),
+        tool_call_id=data.get("toolCallId", ""),
+        content=data.get("toolAgentOutput", ""),
+        role="tool",
+    )
+
+
+_BUILDER_MAP: dict[str, Any] = {
+    "run.started": _build_run_started,
+    "run.finished": _build_run_finished,
+    "run.error": _build_run_error,
+    "step.start": _build_step_started,
+    "step.finish": _build_step_finished,
+    "text.start": _build_text_start,
+    "text.delta": _build_text_delta,
+    "text.end": _build_text_end,
+    "tool.result": _build_tool_result,
+}
+
+
+def to_agui_wire_event(event: dict[str, Any] | BaseEvent) -> dict[str, Any]:
+    if isinstance(event, BaseEvent):
+        return event.model_dump(by_alias=True, exclude_none=True)
+
+    if _is_agui_event(event):
+        return event
+
+    internal_type = str(event.get("type", "")).strip()
+    builder = _BUILDER_MAP.get(internal_type)
+
+    if builder:
+        agui_event = builder(event)
+        return agui_event.model_dump(by_alias=True, exclude_none=True)
+
+    wire_type = _convert_to_agui_type(internal_type)

    payload: dict[str, Any] = {
-        "type": wire_type,
+        "type": wire_type.value,
    }
    thread_id = event.get("threadId")
    run_id = event.get("runId")
@@ -37,13 +162,7 @@ def to_agui_wire_event(event: dict[str, Any]) -> dict[str, Any]:

    data = event.get("data")
    if isinstance(data, dict):
-        if event_type == "tool.result":
-            for key in ("messageId", "toolCallId", "toolAgentOutput"):
-                value = data.get(key)
-                if value is not None:
-                    payload[key] = value
-            return payload
-        if event_type == "text.end":
+        if internal_type == "text.end":
            for key in ("messageId", "workerAgentOutput"):
                value = data.get(key)
                if value is not None:
@@ -57,5 +176,5 @@ def to_agui_wire_event(event: dict[str, Any]) -> dict[str, Any]:


 class AgentScopeAgUiCodec:
-    def to_wire(self, event: dict[str, Any]) -> dict[str, Any]:
+    def to_wire(self, event: dict[str, Any] | BaseEvent) -> dict[str, Any]:
        return to_agui_wire_event(event)
@@ -1,6 +1,13 @@
 from __future__ import annotations

-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
+
+from ag_ui.core import BaseEvent
+
+if TYPE_CHECKING:
+    pass
+
+T = TypeVar("T", bound=BaseEvent)


 class CodecLike(Protocol):
@@ -15,6 +22,16 @@ class BusLike(Protocol):
    async def publish(self, *, session_id: str, event: dict[str, Any]) -> str: ...


+def is_base_event(event: Any) -> bool:
+    return isinstance(event, BaseEvent)
+
+
+def to_dict(event: BaseEvent | dict[str, Any]) -> dict[str, Any]:
+    if isinstance(event, BaseEvent):
+        return event.model_dump(by_alias=True, exclude_none=True)
+    return event
+
+
 class AgentScopeEventPipeline:
    _codec: CodecLike
    _store: StoreLike
@@ -25,7 +42,13 @@ class AgentScopeEventPipeline:
        self._store = store
        self._bus = bus

-    async def emit(self, *, session_id: str, event: dict[str, Any]) -> str:
-        wire_event = self._codec.to_wire(event)
+    async def emit(
+        self,
+        *,
+        session_id: str,
+        event: "BaseEvent | dict[str, Any]",
+    ) -> str:
+        event_dict = to_dict(event)
+        wire_event = self._codec.to_wire(event_dict)
        await self._store.persist(wire_event)
        return await self._bus.publish(session_id=session_id, event=wire_event)