diff --git a/app/agent/callback/__init__.py b/app/agent/callback/__init__.py index 8669792a0..4cf840d97 100644 --- a/app/agent/callback/__init__.py +++ b/app/agent/callback/__init__.py @@ -1,6 +1,7 @@ import asyncio import re import threading +import time from typing import Any, Optional, Tuple from app.agent.policy.sanitizer import sanitize_for_host @@ -84,6 +85,44 @@ class StreamingHandler: self._live_tool_stats: dict[str, dict[str, Any]] = {} # 本轮已写入缓冲区的工具展示行,供 Telegram 富文本渲染时做区分样式 self._tool_summaries: set[str] = set() + # 当前模型思考状态,仅用于实时展示,不写入最终助手正文。 + self._thinking_started_at: Optional[float] = None + self._thinking_active = False + self._thinking_status_text = "" + + def thinking_started(self) -> None: + """开始展示模型思考状态,并记录本轮思考起始时间。""" + with self._lock: + if self._thinking_active: + return + self._thinking_started_at = time.monotonic() + self._thinking_active = True + self._thinking_status_text = self._format_thinking_status(0) + if self._can_stream(): + self._buffer = self._thinking_status_text + + def thinking_finished(self) -> None: + """结束模型思考状态,移除尚未被正文替换的状态占位文本。""" + with self._lock: + if not self._thinking_active: + return + if self._buffer == self._thinking_status_text: + self._buffer = "" + self._thinking_active = False + self._thinking_status_text = "" + + def _format_thinking_status(self, elapsed_seconds: int) -> str: + """生成消息渠道可直接展示的思考状态文本。""" + return f"🤔 思考中 · 已用时 {max(0, elapsed_seconds)} 秒" + + def _refresh_thinking_status(self) -> None: + """更新可编辑消息中的思考计时文本,保持正文尚未生成时也有反馈。""" + with self._lock: + if not self._thinking_active or not self._can_stream() or self._buffer != self._thinking_status_text: + return + elapsed = int(max(0.0, time.monotonic() - (self._thinking_started_at or time.monotonic()))) + self._thinking_status_text = self._format_thinking_status(elapsed) + self._buffer = self._thinking_status_text def set_dispatch_policy(self, allow_dispatch_without_context: bool = False) -> None: """ @@ -99,6 +138,11 @@ class StreamingHandler: """ with self._lock: emitted = token or "" + if emitted and self._thinking_active: + if self._buffer == self._thinking_status_text: + self._buffer = "" + self._thinking_active = False + self._thinking_status_text = "" if self._pending_tool_stats: if self._live_tool_summary: @@ -300,6 +344,9 @@ class StreamingHandler: self._tool_summaries = set() self._live_tool_summary = None self._live_tool_stats = {} + self._thinking_started_at = None + self._thinking_active = False + self._thinking_status_text = "" # 检查渠道是否支持消息编辑,不支持则仅收集 token 到 buffer,不实时推送 if not self._can_stream(): @@ -334,6 +381,7 @@ class StreamingHandler: return False, "" self._streaming_enabled = False + self.thinking_finished() # 取消定时任务 await self._cancel_flush_task() @@ -366,6 +414,9 @@ class StreamingHandler: self._tool_summaries = set() self._live_tool_summary = None self._live_tool_stats = {} + self._thinking_started_at = None + self._thinking_active = False + self._thinking_status_text = "" if all_sent: # 所有内容已通过流式发送,清空缓冲区 self._buffer = "" @@ -724,6 +775,7 @@ class StreamingHandler: while self._streaming_enabled: await asyncio.sleep(self.FLUSH_INTERVAL) if self._streaming_enabled: + self._refresh_thinking_status() await self._flush() except asyncio.CancelledError: pass diff --git a/app/agent/orchestrator.py b/app/agent/orchestrator.py index cf1779860..fa2d8998f 100644 --- a/app/agent/orchestrator.py +++ b/app/agent/orchestrator.py @@ -2250,7 +2250,13 @@ class MoviePilotAgent: return attachments @staticmethod - async def _stream_agent_tokens(agent, messages: dict, config: dict, on_token: Callable[[str], None]): + async def _stream_agent_tokens( + agent, + messages: dict, + config: dict, + on_token: Callable[[str], None], + stream_handler: Any = None, + ): """ 流式运行智能体,过滤工具调用token和思考内容,将模型生成的内容通过回调输出。 :param agent: LangGraph Agent 实例 @@ -2259,6 +2265,11 @@ class MoviePilotAgent: :param on_token: 收到有效 token 时的回调 """ stripper = _ThinkTagStripper() + thinking_started = stream_handler + if thinking_started is not None: + start_thinking = getattr(thinking_started, "thinking_started", None) + if callable(start_thinking): + start_thinking() async for chunk in agent.astream( messages, @@ -2290,6 +2301,9 @@ class MoviePilotAgent: # content 可能是字符串或内容块列表,过滤掉思考类型的块 content = LLMHelper.extract_text_content(token.content) if content: + finish_thinking = getattr(thinking_started, "thinking_finished", None) + if callable(finish_thinking): + finish_thinking() stripper.process(content, on_token) stripper.flush(on_token) @@ -2505,6 +2519,7 @@ class MoviePilotAgent: messages={"messages": input_messages}, config=agent_config, on_token=self._handle_stream_text, + stream_handler=self.stream_handler, ) # 输出流式过程中可能残留的工具调用统计信息 diff --git a/app/agent/web.py b/app/agent/web.py index f83a4d4de..f7dfcf036 100644 --- a/app/agent/web.py +++ b/app/agent/web.py @@ -1,5 +1,6 @@ """WebAgent 运行时类型适配,不包含 HTTP 或 SSE 编码。""" +import time import uuid from threading import Lock from typing import Any, Awaitable, Callable, Optional @@ -57,6 +58,18 @@ class _WebAgentStreamingHandlerMixin: except Exception as error: logger.debug(f"Web工具生命周期回调失败: {error}") + def thinking_started(self) -> None: + """发布 Web SSE 的模型思考开始事件,由前端负责持续计时。""" + super().thinking_started() # type: ignore[misc] + self._publish_tool_event({"type": "thinking", "status": "running", "started_at": int(time.time() * 1000)}) + + def thinking_finished(self) -> None: + """发布 Web SSE 的模型思考结束事件,并清理基础流式状态。""" + was_active = bool(getattr(self, "_thinking_active", False)) + super().thinking_finished() # type: ignore[misc] + if was_active: + self._publish_tool_event({"type": "thinking", "status": "done"}) + def tool_call_started( self, tool_name: str, @@ -186,6 +199,9 @@ class _WebAgentStreamingHandlerMixin: self._msg_start_offset = 0 self._pending_tool_stats = {} self._live_tool_summary = None + self._thinking_started_at = None + self._thinking_active = False + self._thinking_status_text = "" async def stop_streaming(self) -> tuple[bool, str]: """停止 Web SSE 流式状态,保留缓冲区给 Agent 收口逻辑去重。""" diff --git a/app/application/messaging/agent.py b/app/application/messaging/agent.py index 846478234..a637f7b3d 100644 --- a/app/application/messaging/agent.py +++ b/app/application/messaging/agent.py @@ -189,6 +189,9 @@ def build_web_agent_display_message( "content": normalized_content, "createdAt": int(datetime.now().timestamp() * 1000), "status": status, + "thinking": False, + "thinking_started_at": None, + "thinking_elapsed_ms": 0, "tools": [], "segments": ([{"type": "text", "content": normalized_content}] if normalized_content else []), "attachments": attachments or [], @@ -458,6 +461,14 @@ def apply_web_agent_display_event(event: dict[str, Any], assistant_message: dict else build_legacy_web_agent_segments(assistant_message["content"], assistant_message["tools"]) ) assistant_message["status"] = target_message.get("status") or "done" + elif event_type == "thinking": + is_thinking = str(event.get("status") or "") == "running" + assistant_message["thinking"] = is_thinking + if is_thinking: + assistant_message["thinking_started_at"] = event.get("started_at") + assistant_message["thinking_elapsed_ms"] = 0 + else: + assistant_message["thinking_elapsed_ms"] = event.get("elapsed_ms") or assistant_message.get("thinking_elapsed_ms") or 0 elif event_type == "error": assistant_message["status"] = "error" if not assistant_message["content"]: diff --git a/tests/test_agent_streaming_lifecycle.py b/tests/test_agent_streaming_lifecycle.py index 6ddce47cb..87d253566 100644 --- a/tests/test_agent_streaming_lifecycle.py +++ b/tests/test_agent_streaming_lifecycle.py @@ -8,6 +8,21 @@ from app.agent.callback import StreamingHandler from app.schemas.types import NotificationChannel +def test_thinking_status_is_replaced_by_first_answer_token() -> None: + """思考占位消息应在首个正文 token 到达时被移除,避免污染最终回复。""" + handler = StreamingHandler() + handler._can_stream = lambda: True + + handler.thinking_started() + assert handler._thinking_active is True + assert "思考中" in handler._buffer + + handler.emit("最终答案") + + assert handler._thinking_active is False + assert handler._buffer == "最终答案" + + @pytest.mark.asyncio async def test_repeated_streaming_start_retains_previous_flush_owner() -> None: """重复启动必须先等待旧 flush owner 结束,再发布新一轮上下文。"""