mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-31 12:24:39 +08:00
feat(agent): expose model thinking progress
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
# 输出流式过程中可能残留的工具调用统计信息
|
||||
|
||||
@@ -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 收口逻辑去重。"""
|
||||
|
||||
@@ -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"]:
|
||||
|
||||
@@ -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 结束,再发布新一轮上下文。"""
|
||||
|
||||
Reference in New Issue
Block a user