Fix think tag streaming in agent output

This commit is contained in:
jxxghp
2026-06-21 19:58:33 +08:00
parent 78ddd6093f
commit 1f97870fa9
2 changed files with 31 additions and 4 deletions

View File

@@ -192,10 +192,11 @@ class _ThinkTagStripper:
self.buffer = self.buffer[-i:]
partial_match = True
break
if not partial_match:
on_output(self.buffer)
emitted = True
self.buffer = ""
if partial_match:
break
on_output(self.buffer)
emitted = True
self.buffer = ""
else:
end_idx = self.buffer.find("</think>")
if end_idx != -1:

View File

@@ -7,6 +7,7 @@ import langchain.agents as langchain_agents
if not hasattr(langchain_agents, "create_agent"):
langchain_agents.create_agent = lambda *args, **kwargs: None
from app.agent import _ThinkTagStripper
from app.agent.callback import StreamingHandler
from app.agent.middleware.subagents import is_subagent_stream_metadata
from app.agent.tools.base import MoviePilotTool
@@ -17,6 +18,31 @@ from app.schemas.message import MessageResponse
from app.schemas.types import MessageChannel, NotificationType
def test_think_tag_stripper_waits_for_partial_open_tag():
"""普通文本后出现不完整 think 开始标签时不应进入死循环。"""
stripper = _ThinkTagStripper()
outputs = []
emitted = stripper.process("你好<", outputs.append)
emitted_next = stripper.process("世界", outputs.append)
assert emitted is True
assert emitted_next is True
assert outputs == ["你好", "<世界"]
def test_think_tag_stripper_hides_split_think_tag_content():
"""think 标签被拆分到多个 token 时应继续隐藏思考内容。"""
stripper = _ThinkTagStripper()
outputs = []
stripper.process("回答前<", outputs.append)
stripper.process("thi", outputs.append)
stripper.process("nk>隐藏内容</think>回答后", outputs.append)
assert outputs == ["回答前", "回答后"]
class DummyTool(MoviePilotTool):
"""用于流式输出测试的固定结果工具。"""