mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-09-04 02:01:48 +03:00
refactor: decouple webui runtime state via events
This commit is contained in:
chengyongru
Xubin Ren
+150
-82
@@ -1,8 +1,4 @@
|
||||
"""Session turn helpers for WebUI-capable WebSocket sessions.
|
||||
|
||||
AgentLoop uses these without importing a concrete channel plugin; only
|
||||
``channel == "websocket"`` messages are affected.
|
||||
"""
|
||||
"""Session turn helpers for WebUI-capable WebSocket sessions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -14,8 +10,19 @@ from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.bus import progress as bus_progress
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.bus.runtime_events import (
|
||||
GoalStateChanged,
|
||||
RuntimeEvent,
|
||||
RuntimeEventBus,
|
||||
RuntimeEventContext,
|
||||
RuntimeModelChanged,
|
||||
SessionTurnStarted,
|
||||
TurnCompleted,
|
||||
TurnRunStatusChanged,
|
||||
)
|
||||
from nanobot.providers.base import LLMProvider
|
||||
from nanobot.session.goal_state import goal_state_ws_blob
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
@@ -178,6 +185,14 @@ def websocket_turn_wall_started_at(chat_id: str) -> float | None:
|
||||
return _WEBSOCKET_TURN_WALL_STARTED_AT.get(chat_id)
|
||||
|
||||
|
||||
def build_bus_progress_callback(
|
||||
bus: MessageBus,
|
||||
msg: InboundMessage,
|
||||
) -> Callable[..., Awaitable[None]]:
|
||||
"""Compatibility wrapper for the generic bus progress callback."""
|
||||
return bus_progress.build_bus_progress_callback(bus, msg)
|
||||
|
||||
|
||||
async def publish_turn_run_status(
|
||||
bus: MessageBus,
|
||||
msg: InboundMessage,
|
||||
@@ -212,91 +227,110 @@ async def publish_turn_run_status(
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_bus_progress_callback(
|
||||
bus: MessageBus,
|
||||
msg: InboundMessage,
|
||||
) -> Callable[..., Awaitable[None]]:
|
||||
"""Return the bus progress callback for agent runtime events."""
|
||||
|
||||
async def _publish_progress(
|
||||
content: str,
|
||||
*,
|
||||
tool_hint: bool = False,
|
||||
tool_events: list[dict[str, Any]] | None = None,
|
||||
file_edit_events: list[dict[str, Any]] | None = None,
|
||||
reasoning: bool = False,
|
||||
reasoning_end: bool = False,
|
||||
) -> None:
|
||||
meta = dict(msg.metadata or {})
|
||||
meta["_progress"] = True
|
||||
meta["_tool_hint"] = tool_hint
|
||||
if reasoning:
|
||||
meta["_reasoning_delta"] = True
|
||||
if reasoning_end:
|
||||
meta["_reasoning_end"] = True
|
||||
if tool_events:
|
||||
meta["_tool_events"] = tool_events
|
||||
if file_edit_events:
|
||||
meta["_file_edit_events"] = file_edit_events
|
||||
await bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel=msg.channel,
|
||||
chat_id=msg.chat_id,
|
||||
content=content,
|
||||
metadata=meta,
|
||||
)
|
||||
)
|
||||
|
||||
if msg.channel == "websocket":
|
||||
async def _websocket_progress(
|
||||
content: str,
|
||||
*,
|
||||
tool_hint: bool = False,
|
||||
tool_events: list[dict[str, Any]] | None = None,
|
||||
file_edit_events: list[dict[str, Any]] | None = None,
|
||||
reasoning: bool = False,
|
||||
reasoning_end: bool = False,
|
||||
) -> None:
|
||||
await _publish_progress(
|
||||
content,
|
||||
tool_hint=tool_hint,
|
||||
tool_events=tool_events,
|
||||
file_edit_events=file_edit_events,
|
||||
reasoning=reasoning,
|
||||
reasoning_end=reasoning_end,
|
||||
)
|
||||
|
||||
return _websocket_progress
|
||||
|
||||
async def _bus_progress(
|
||||
content: str,
|
||||
*,
|
||||
tool_hint: bool = False,
|
||||
tool_events: list[dict[str, Any]] | None = None,
|
||||
reasoning: bool = False,
|
||||
reasoning_end: bool = False,
|
||||
) -> None:
|
||||
await _publish_progress(
|
||||
content,
|
||||
tool_hint=tool_hint,
|
||||
tool_events=tool_events,
|
||||
reasoning=reasoning,
|
||||
reasoning_end=reasoning_end,
|
||||
)
|
||||
|
||||
return _bus_progress
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebuiTurnCoordinator:
|
||||
"""Own the WebUI/WebSocket wire details that hang off AgentLoop turns."""
|
||||
"""Translate generic runtime events into WebUI/WebSocket wire messages."""
|
||||
|
||||
bus: MessageBus
|
||||
sessions: SessionManager
|
||||
schedule_background: Callable[[Awaitable[None]], None]
|
||||
_title_contexts: dict[str, LLMRuntime] = field(default_factory=dict)
|
||||
|
||||
def subscribe(self, runtime_events: RuntimeEventBus) -> Callable[[], None]:
|
||||
"""Subscribe this coordinator to runtime events."""
|
||||
return runtime_events.subscribe(self.handle_runtime_event)
|
||||
|
||||
async def handle_runtime_event(self, event: RuntimeEvent) -> None:
|
||||
if isinstance(event, SessionTurnStarted):
|
||||
self._handle_session_turn_started(event)
|
||||
return
|
||||
if isinstance(event, TurnRunStatusChanged):
|
||||
await self._handle_run_status_changed(event)
|
||||
return
|
||||
if isinstance(event, TurnCompleted):
|
||||
await self._handle_turn_completed_event(event)
|
||||
return
|
||||
if isinstance(event, GoalStateChanged):
|
||||
await self._handle_goal_state_changed(event)
|
||||
return
|
||||
if isinstance(event, RuntimeModelChanged):
|
||||
await self._handle_runtime_model_changed(event)
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _ctx_msg(ctx: RuntimeEventContext) -> InboundMessage:
|
||||
return InboundMessage(
|
||||
channel=ctx.channel,
|
||||
sender_id="runtime",
|
||||
chat_id=ctx.chat_id,
|
||||
content="",
|
||||
metadata=dict(ctx.metadata or {}),
|
||||
session_key_override=ctx.session_key,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_websocket_event(ctx: RuntimeEventContext) -> bool:
|
||||
return ctx.channel == "websocket"
|
||||
|
||||
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
session = self.sessions.get_or_create(event.context.session_key)
|
||||
mark_webui_session(session, event.context.metadata)
|
||||
|
||||
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
await publish_turn_run_status(
|
||||
self.bus,
|
||||
self._ctx_msg(event.context),
|
||||
event.status,
|
||||
started_at=event.started_at,
|
||||
)
|
||||
|
||||
async def _handle_turn_completed_event(self, event: TurnCompleted) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
msg = self._ctx_msg(event.context)
|
||||
await self.handle_turn_end(
|
||||
msg,
|
||||
session_key=event.context.session_key,
|
||||
latency_ms=event.latency_ms,
|
||||
)
|
||||
self._schedule_title_update_from_event(event)
|
||||
|
||||
async def _handle_goal_state_changed(self, event: GoalStateChanged) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
cid = str(event.context.chat_id or "").strip()
|
||||
if not cid:
|
||||
return
|
||||
await self.bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel=event.context.channel,
|
||||
chat_id=cid,
|
||||
content="",
|
||||
metadata={
|
||||
"_goal_state_sync": True,
|
||||
"goal_state": goal_state_ws_blob(event.session_metadata),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
async def _handle_runtime_model_changed(self, event: RuntimeModelChanged) -> None:
|
||||
await self.bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel="websocket",
|
||||
chat_id="*",
|
||||
content="",
|
||||
metadata={
|
||||
"_runtime_model_updated": True,
|
||||
"model": event.model,
|
||||
"model_preset": event.model_preset,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def capture_title_context(
|
||||
self,
|
||||
session_key: str,
|
||||
@@ -370,3 +404,37 @@ class WebuiTurnCoordinator:
|
||||
))
|
||||
|
||||
self.schedule_background(_generate_title_and_notify())
|
||||
|
||||
def _schedule_title_update_from_event(self, event: TurnCompleted) -> None:
|
||||
title_context = event.runtime
|
||||
if (
|
||||
event.context.metadata.get("webui") is not True
|
||||
or title_context is None
|
||||
or not isinstance(title_context, LLMRuntime)
|
||||
):
|
||||
return
|
||||
|
||||
async def _generate_title_and_notify(
|
||||
title_llm: LLMRuntime = title_context,
|
||||
) -> None:
|
||||
generated = await maybe_generate_webui_title_after_turn(
|
||||
channel=event.context.channel,
|
||||
metadata=event.context.metadata,
|
||||
sessions=self.sessions,
|
||||
session_key=event.context.session_key,
|
||||
provider=title_llm.provider,
|
||||
model=title_llm.model,
|
||||
)
|
||||
if generated:
|
||||
await self.bus.publish_outbound(OutboundMessage(
|
||||
channel=event.context.channel,
|
||||
chat_id=event.context.chat_id,
|
||||
content="",
|
||||
metadata={
|
||||
**event.context.metadata,
|
||||
"_session_updated": True,
|
||||
"_session_update_scope": "metadata",
|
||||
},
|
||||
))
|
||||
|
||||
self.schedule_background(_generate_title_and_notify())
|
||||
|
||||
Reference in New Issue
Block a user