Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/octop/infra/gateway/gateway.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
WebSocketChannel,
WebSocketHub,
)
from octop.infra.utils.llm_text import strip_thinking
from octop.infra.utils.locale import DEFAULT_LOCALE, Locale

if TYPE_CHECKING:
Expand Down Expand Up @@ -452,7 +453,7 @@ async def push_text_from_session(
async for chunk in self._agent_manager.stream(agent_id, request):
if chunk.get("type") in ("token", "delta"):
parts.append(str(chunk.get("content") or chunk.get("text") or ""))
outbound = "".join(parts).strip() or "(empty)"
outbound = strip_thinking("".join(parts)) or "(empty)"

if virtual_stream:
self._bump_dashboard_session(session, session_key, text)
Expand Down
4 changes: 3 additions & 1 deletion src/octop/infra/gateway/process/response_mode.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
TextContent,
)

from octop.infra.utils.llm_text import strip_thinking

ChannelResponseMode = Literal["invoke", "stream"]

DEFAULT_CHANNEL_RESPONSE_MODE: ChannelResponseMode = "invoke"
Expand Down Expand Up @@ -81,7 +83,7 @@ async def collapse_to_invoke_response(

if event.type == MessageEventType.COMPLETED:
content: list[ContentPart] = []
final_text = text_buffer.strip()
final_text = strip_thinking(text_buffer)
if final_text:
content.append(TextContent(text=final_text))
content.extend(media_buffer)
Expand Down
28 changes: 25 additions & 3 deletions src/octop/infra/utils/llm_text.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,36 @@
from typing import Any

_THINKING_RE = re.compile(
r"<think>[\s\S]*?</think>\s*",
r"<(?:think|thinking)>[\s\S]*?</(?:think|thinking)>\s*",
re.IGNORECASE,
)
_THINKING_OPEN_RE = re.compile(r"<(?:think|thinking)>", re.IGNORECASE)
_THINKING_CLOSE_RE = re.compile(r"</(?:think|thinking)>", re.IGNORECASE)


def strip_thinking(text: str) -> str:
"""Remove ``<think>...</think>`` blocks from model output."""
return _THINKING_RE.sub("", text).strip()
"""Remove tagged or malformed thinking prefixes from model output.

Some OpenAI-compatible model routers return reasoning in ordinary
``content`` and only emit the closing ``</think>`` marker. Treat a
closing marker that appears before any opening marker as the boundary
between the hidden prefix and the user-visible answer.
"""
cleaned = _THINKING_RE.sub("", text)

while True:
close_match = _THINKING_CLOSE_RE.search(cleaned)
if close_match is None:
break
open_match = _THINKING_OPEN_RE.search(cleaned)
if open_match is not None and open_match.start() < close_match.start():
break
cleaned = cleaned[close_match.end() :]

open_match = _THINKING_OPEN_RE.search(cleaned)
if open_match is not None:
cleaned = cleaned[: open_match.start()]
return cleaned.strip()


def llm_text_content(result: Any) -> str:
Expand Down
10 changes: 10 additions & 0 deletions tests/unit/api/test_chat_polish.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,16 @@ def test_strip_thinking_removes_redacted_block() -> None:
assert _strip_thinking(raw) == "Polished prompt"


def test_strip_thinking_removes_orphan_closing_prefix() -> None:
raw = "internal reasoning without an opening tag</think>\nVisible answer"
assert _strip_thinking(raw) == "Visible answer"


def test_strip_thinking_removes_unclosed_thinking_suffix() -> None:
raw = "Visible answer\n<thinking>truncated internal reasoning"
assert _strip_thinking(raw) == "Visible answer"


def test_llm_text_content_strips_thinking_from_string_message() -> None:
class Msg:
content = "<think>plan</think>Final text"
Expand Down
42 changes: 42 additions & 0 deletions tests/unit/gateway/test_gateway_push.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,48 @@ async def stream_with_token(_agent_id, _request):
assert gateway._channel_manager.push_text.await_args.args[2] == "AI reply"


@pytest.mark.asyncio
async def test_push_text_from_session_cron_agent_strips_orphan_thinking_prefix(
gateway: Gateway,
) -> None:
sk = ThreadRegistry.make_key(
agent_id="a1",
channel_type="weixin",
channel_subject_id="wx_1",
)
gateway.thread_registry._threads.insert(
thread_id="thr_weixin_cron",
agent_id="a1",
user_id=1,
channel_type="weixin",
session_key=sk,
)
gateway.thread_registry._sessions.upsert(
session_key=sk,
agent_id="a1",
user_id=1,
channel_type="weixin",
chat_type="dm",
thread_id="thr_weixin_cron",
channel_subject_id="wx_1",
channel_chat_type="dm",
channel_metadata={"channel_type": "weixin"},
channel_id="ch-1",
)

async def stream_with_thinking_leak(_agent_id, _request):
yield {"type": "token", "content": "internal reasoning"}
yield {"type": "token", "content": "</think>"}
yield {"type": "token", "content": "最终学习内容"}

gateway._agent_manager.stream = stream_with_thinking_leak

await gateway.push_text_from_session("a1", sk, "run agent", task_type="agent")

gateway._channel_manager.push_text.assert_awaited_once()
assert gateway._channel_manager.push_text.await_args.args[2] == "最终学习内容"


@pytest.mark.asyncio
async def test_push_text_from_session_dashboard_agent_pushes_ws(gateway: Gateway) -> None:
sk = ThreadRegistry.dashboard_key(agent_id="a1", user_id=1)
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/gateway/test_response_mode.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,27 @@ async def test_invoke_discards_progress_before_tool_and_emits_final_once() -> No
assert text.text == "这是最终答案。"


@pytest.mark.asyncio
async def test_invoke_strips_orphan_thinking_prefix_from_final_text() -> None:
source = _events(
MessageEvent.delta("Let me inspect another source. "),
MessageEvent.delta("This is internal reasoning."),
MessageEvent.delta("</think>"),
MessageEvent.delta("【每日指南学习】最终内容"),
MessageEvent.completed(),
)

result = [event async for event in collapse_to_invoke_response(source)]

assert [event.type for event in result] == [
MessageEventType.MESSAGE,
MessageEventType.COMPLETED,
]
text = result[0].content[0]
assert isinstance(text, TextContent)
assert text.text == "【每日指南学习】最终内容"


@pytest.mark.asyncio
async def test_invoke_preserves_tool_media_with_final_text() -> None:
attachment = FileContent(filename="report.pdf", data="cGRm")
Expand Down
Loading