From c53b800267b666389c0856df80f60c1983ce22cf Mon Sep 17 00:00:00 2001 From: DongXiaoming <41072499+MIKAZE3@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:54:41 +0800 Subject: [PATCH] fix(provider): strip think tags for MiniMax-M3 and other OpenAI-compatible models (#2330) * fix(provider): strip think tags for MiniMax-M3 and other OpenAI-compatible models MiniMax-M3 (and other OpenAI-compatible providers) emit chain-of-thought reasoning directly in the content field wrapped in tags, instead of using a separate reasoning_content field or the legacy CRETIRE_REASONING markers. The existing remove_think logic only handled CRETIRE_* tags, so think blocks leaked into user-visible output even when remove_think was enabled. - Add _ThinkStripState: a stateful filter that correctly handles tags split across streaming chunk boundaries. - Add _strip_think classmethod with regex patterns for both and CRETIRE_* tags. - Wire think_state into invoke_llm_stream so deltas are filtered before reaching the accumulator. - Add remove_think safety net in _StreamAccumulator so the final message from tool-call rounds also gets stripped. - Fix remove_think resolution to use defensive nested .get() so pipelines missing output.misc don't raise AttributeError. * fix(litellmchat): add missing _CLOSE_TAG class attribute on _ThinkStripState * fix(provider): handle think stripping across LiteLLM paths --------- Co-authored-by: WangCham <651122857@qq.com> --- .../modelmgr/requesters/litellmchat.py | 196 ++++++++++++++++-- .../pkg/provider/runners/localagent.py | 46 +++- tests/unit_tests/provider/test_litellmchat.py | 162 +++++++++++++++ .../provider/test_localagent_sandbox_exec.py | 45 ++++ 4 files changed, 431 insertions(+), 18 deletions(-) diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index c1b5ae0b6..efad67250 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -13,6 +13,151 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message +class _ThinkStripState: + """Stateful filter that drops think blocks across chunks.""" + + _THINK_OPEN = '' + _THINK_CLOSE = '' + _LEGACY_OPEN = 'CRETIRE_REASONING_BEGINk' + _LEGACY_CLOSE = 'CRETIRE_REASONING_ENDk' + + def __init__(self) -> None: + self._pairs: tuple[tuple[str, str], ...] = ( + (self._THINK_OPEN, self._THINK_CLOSE), + (self._LEGACY_OPEN, self._LEGACY_CLOSE), + ) + self._open_tags = tuple(open_tag for open_tag, _close_tag in self._pairs) + self._buf = '' + self._close_tag: str | None = None + self._pending_initial = True + + def feed(self, chunk: str) -> str: + """Feed a streaming delta and return user-visible content.""" + if not chunk: + return chunk + + text = self._buf + chunk + if self._close_tag is not None: + return self._consume_think_body(text) + + return self._process_visible_text(text) + + def flush(self) -> str: + """Release buffered visible content when the stream ends.""" + if self._close_tag is not None: + self._buf = '' + self._close_tag = None + return '' + + pending, self._buf = self._buf, '' + self._close_tag = None + return pending + + def _consume_think_body(self, text: str) -> str: + close_tag = self._close_tag + if close_tag is None: + return text + + close_idx = text.find(close_tag) + if close_idx != -1: + self._close_tag = None + self._buf = '' + self._pending_initial = False + return self._process_visible_text(text[close_idx + len(close_tag) :]) + + self._buf = self._close_prefix(text, close_tag) + return '' + + def _process_visible_text(self, text: str) -> str: + out: list[str] = [] + index = 0 + + while index < len(text): + if self._pending_initial: + open_idx, open_tag, close_tag = self._find_next_open(text, index) + orphan_close_idx, orphan_close_tag = self._find_next_close(text, index) + + if orphan_close_idx != -1 and (open_idx == -1 or orphan_close_idx < open_idx): + self._pending_initial = False + index = orphan_close_idx + len(orphan_close_tag) + continue + + if open_idx == -1: + self._buf = text[index:] + return ''.join(out) + + if open_idx > index: + self._pending_initial = False + out.append(text[index:open_idx]) + index = open_idx + continue + + open_idx, open_tag, close_tag = self._find_next_open(text, index) + if open_idx == -1: + emit_end = self._visible_emit_end(text, index) + out.append(text[index:emit_end]) + if emit_end > index: + self._pending_initial = False + self._buf = text[emit_end:] + return ''.join(out) + + out.append(text[index:open_idx]) + if open_idx > index: + self._pending_initial = False + body_start = open_idx + len(open_tag) + close_idx = text.find(close_tag, body_start) + if close_idx == -1: + self._close_tag = close_tag + self._buf = self._close_prefix(text[body_start:], close_tag) + return ''.join(out) + + self._pending_initial = False + index = close_idx + len(close_tag) + + self._buf = '' + return ''.join(out) + + def _find_next_open(self, text: str, start: int) -> tuple[int, str, str]: + best_idx = -1 + best_open = '' + best_close = '' + for open_tag, close_tag in self._pairs: + idx = text.find(open_tag, start) + if idx != -1 and (best_idx == -1 or idx < best_idx): + best_idx = idx + best_open = open_tag + best_close = close_tag + return best_idx, best_open, best_close + + def _find_next_close(self, text: str, start: int) -> tuple[int, str]: + best_idx = -1 + best_close = '' + for _open_tag, close_tag in self._pairs: + idx = text.find(close_tag, start) + if idx != -1 and (best_idx == -1 or idx < best_idx): + best_idx = idx + best_close = close_tag + return best_idx, best_close + + def _visible_emit_end(self, text: str, start: int) -> int: + visible = text[start:] + limit = min(len(visible), max(len(open_tag) for open_tag in self._open_tags) - 1) + for keep in range(limit, 0, -1): + suffix = visible[-keep:] + if any(open_tag.startswith(suffix) for open_tag in self._open_tags): + return len(text) - keep + return len(text) + + @staticmethod + def _close_prefix(text: str, close_tag: str) -> str: + limit = min(len(text), len(close_tag) - 1) + for keep in range(limit, 0, -1): + suffix = text[-keep:] + if close_tag.startswith(suffix): + return suffix + return '' + + class LiteLLMRequester(requester.ProviderAPIRequester): """LiteLLM unified API requester supporting chat, embedding, and rerank.""" @@ -237,6 +382,25 @@ class LiteLLMRequester(requester.ProviderAPIRequester): return req_messages + _THINK_PATTERNS: tuple[str, ...] = ( + r'^\s*(?:(?!).)*?\s*', + r'^\s*(?:(?!CRETIRE_REASONING_BEGINk).)*?CRETIRE_REASONING_ENDk\s*', + r'.*?', + r'CRETIRE_REASONING_BEGINk.*?CRETIRE_REASONING_ENDk', + ) + + @classmethod + def _strip_think(cls, content: str) -> str: + """Strip chain-of-thought blocks from ``content``.""" + if not content: + return content + + import re + + for pattern in cls._THINK_PATTERNS: + content = re.sub(pattern, '', content, flags=re.DOTALL) + return content.strip() + def _process_thinking_content(self, content: str, reasoning_content: str | None, remove_think: bool) -> str: """Process thinking/reasoning content. @@ -248,20 +412,12 @@ class LiteLLMRequester(requester.ProviderAPIRequester): Returns: Processed content string """ - # Extract and handle thinking tags - if content and 'CRETIRE_REASONING_BEGINk' in content and 'CRETIRE_REASONING_ENDk' in content: - import re + if remove_think and content: + content = self._strip_think(content) - think_pattern = r'CRETIRE_REASONING_BEGINk(.*?)CRETIRE_REASONING_ENDk' + if reasoning_content and not remove_think: + content = f'\n{reasoning_content}\n\n{content or ""}'.strip() - if remove_think: - # Remove thinking tags and their content from output - content = re.sub(think_pattern, '', content, flags=re.DOTALL).strip() - # else: preserve thinking content as-is - - # Handle separate reasoning_content field - # Currently we don't include reasoning_content in user-facing output regardless of remove_think - # because it's typically internal model reasoning, not user-visible thinking return content or '' @staticmethod @@ -570,6 +726,7 @@ class LiteLLMRequester(requester.ProviderAPIRequester): chunk_idx = 0 role = 'assistant' tool_call_state: dict[int, dict[str, typing.Any]] = {} + think_state = _ThinkStripState() if remove_think else None try: response = await acompletion(**args) @@ -613,6 +770,12 @@ class LiteLLMRequester(requester.ProviderAPIRequester): # Use reasoning_content as the displayed content delta_content = reasoning_content + if think_state is not None and delta_content: + delta_content = think_state.feed(delta_content) + if not delta_content: + chunk_idx += 1 + continue + tool_calls = self._normalize_stream_tool_calls(delta.get('tool_calls'), tool_call_state) if chunk_idx == 0 and not delta_content and not tool_calls: @@ -634,6 +797,15 @@ class LiteLLMRequester(requester.ProviderAPIRequester): yield provider_message.MessageChunk(**chunk_data) chunk_idx += 1 + if think_state is not None: + pending_content = think_state.flush() + if pending_content: + yield provider_message.MessageChunk( + role=role, + content=pending_content, + is_final=True, + ) + except Exception as e: self._handle_litellm_error(e) diff --git a/src/langbot/pkg/provider/runners/localagent.py b/src/langbot/pkg/provider/runners/localagent.py index 3417c6671..6c877239c 100644 --- a/src/langbot/pkg/provider/runners/localagent.py +++ b/src/langbot/pkg/provider/runners/localagent.py @@ -49,12 +49,23 @@ def _model_has_ability(model: modelmgr_requester.RuntimeLLMModel, ability: str) class _StreamAccumulator: """Accumulate streamed content and fragmented OpenAI-style tool calls.""" - def __init__(self, msg_sequence: int = 0, initial_content: str | None = None): + def __init__( + self, + msg_sequence: int = 0, + initial_content: str | None = None, + remove_think: bool = False, + ): self.tool_calls_map: dict[str, provider_message.ToolCall] = {} self.msg_idx = 0 self.accumulated_content = initial_content or '' self.last_role = 'assistant' self.msg_sequence = msg_sequence + self.remove_think = remove_think + self._think_state = None + if remove_think: + from ..modelmgr.requesters.litellmchat import _ThinkStripState + + self._think_state = _ThinkStripState() def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChunk | None: self.msg_idx += 1 @@ -63,7 +74,10 @@ class _StreamAccumulator: self.last_role = msg.role if msg.content: - self.accumulated_content += msg.content + content = msg.content + if self._think_state is not None: + content = self._think_state.feed(content) + self.accumulated_content += content if msg.tool_calls: for tool_call in msg.tool_calls: @@ -79,11 +93,14 @@ class _StreamAccumulator: if tool_call.function and tool_call.function.arguments: self.tool_calls_map[tool_call.id].function.arguments += tool_call.function.arguments + if msg.is_final: + self._flush_think_state() + if self.msg_idx % 8 == 0 or msg.is_final: self.msg_sequence += 1 return provider_message.MessageChunk( role=self.last_role, - content=self.accumulated_content, + content=self._maybe_strip_think(self.accumulated_content), tool_calls=list(self.tool_calls_map.values()) if (self.tool_calls_map and msg.is_final) else None, is_final=msg.is_final, msg_sequence=self.msg_sequence, @@ -92,13 +109,29 @@ class _StreamAccumulator: return None def final_message(self) -> provider_message.MessageChunk: + self._flush_think_state() return provider_message.MessageChunk( role=self.last_role, - content=self.accumulated_content, + content=self._maybe_strip_think(self.accumulated_content), tool_calls=list(self.tool_calls_map.values()) if self.tool_calls_map else None, msg_sequence=self.msg_sequence, ) + def _maybe_strip_think(self, content: str) -> str: + if not self.remove_think or not content: + return content + + from ..modelmgr.requesters.litellmchat import LiteLLMRequester + + return LiteLLMRequester._strip_think(content) + + def _flush_think_state(self) -> None: + if self._think_state is None: + return + pending = self._think_state.flush() + if pending: + self.accumulated_content += pending + @runner.runner_class('local-agent') class LocalAgentRunner(runner.RequestRunner): @@ -448,7 +481,7 @@ class LocalAgentRunner(runner.RequestRunner): except AttributeError: is_stream = False - remove_think = query.pipeline_config['output'].get('misc', '').get('remove-think') + remove_think = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get('remove-think', False) # Build ordered candidate list (primary + fallbacks) candidates = await self._get_model_candidates(query) @@ -472,7 +505,7 @@ class LocalAgentRunner(runner.RequestRunner): final_msg = msg else: # Streaming: invoke with fallback - stream_accumulator = _StreamAccumulator(msg_sequence=1) + stream_accumulator = _StreamAccumulator(msg_sequence=1, remove_think=remove_think) stream_src, use_llm_model = await self._invoke_stream_with_fallback( query, @@ -576,6 +609,7 @@ class LocalAgentRunner(runner.RequestRunner): stream_accumulator = _StreamAccumulator( msg_sequence=first_end_sequence, initial_content=first_content, + remove_think=remove_think, ) tool_stream_src = use_llm_model.provider.invoke_llm_stream( diff --git a/tests/unit_tests/provider/test_litellmchat.py b/tests/unit_tests/provider/test_litellmchat.py index f7a448ab6..a22a589a7 100644 --- a/tests/unit_tests/provider/test_litellmchat.py +++ b/tests/unit_tests/provider/test_litellmchat.py @@ -279,6 +279,122 @@ class TestInvokeLLMStreamUsage: assert query.variables['_stream_usage']['total_tokens'] == 12 + @pytest.mark.asyncio + async def test_stream_removes_leading_think_across_chunks(self): + """A leading think block split across chunks must be removed.""" + import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query + import langbot_plugin.api.entities.builtin.provider.message as provider_message + + mock_ap = Mock() + mock_ap.tool_mgr = Mock() + mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None) + requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={}) + model = MockRuntimeModel('minimax-m3', 'test-api-key') + + chunks = [ + self._make_chunk(content='.""" + import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query + import langbot_plugin.api.entities.builtin.provider.message as provider_message + + mock_ap = Mock() + mock_ap.tool_mgr = Mock() + mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None) + requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={}) + model = MockRuntimeModel('minimax-m3', 'test-api-key') + + chunks = [ + self._make_chunk(content='hidden reasoning'), + self._make_chunk(content=' still hidden block when remove_think=True""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'Let me think... The answer is 42.' + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'The answer is 42.' + + def test_remove_non_leading_think_tag(self): + """Test removing and its content in the answer body""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'Use example in the document.' + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'Use in the document.' + + def test_remove_initial_orphan_think_close(self): + """Test removing leading reasoning content when only is visible""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'hidden reasoning Visible answer.' + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'Visible answer.' + + def test_remove_multiple_think_tags(self): + """Test removing multiple blocks""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'hidden Keep example.' + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'Keep .' + def test_preserve_thinking_markers(self): """Test preserving thinking markers when remove_think=False""" requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) @@ -491,6 +639,20 @@ class TestProcessThinkingContent: assert 'CRETIRE_REASONING_BEGINk' in result assert 'The answer is 42.' in result + def test_preserve_reasoning_content_when_remove_think_false(self): + """Test showing separate reasoning_content when remove_think=False""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + result = requester._process_thinking_content('The answer is 42.', 'Let me think...', remove_think=False) + assert result == '\nLet me think...\n\nThe answer is 42.' + + def test_hide_reasoning_content_when_remove_think_true(self): + """Test hiding separate reasoning_content when remove_think=True""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + result = requester._process_thinking_content('The answer is 42.', 'Let me think...', remove_think=True) + assert result == 'The answer is 42.' + def test_empty_content(self): """Test empty content""" requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) diff --git a/tests/unit_tests/provider/test_localagent_sandbox_exec.py b/tests/unit_tests/provider/test_localagent_sandbox_exec.py index 08b4c540b..9bc155343 100644 --- a/tests/unit_tests/provider/test_localagent_sandbox_exec.py +++ b/tests/unit_tests/provider/test_localagent_sandbox_exec.py @@ -163,6 +163,51 @@ def test_stream_accumulator_merges_fragmented_tool_call_arguments(): assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}' +def test_stream_accumulator_strips_leading_think_from_tool_round_content(): + accumulator = _StreamAccumulator( + msg_sequence=3, + initial_content='I will search for LangBot.', + remove_think=True, + ) + + assert accumulator.add(provider_message.MessageChunk(role='assistant', content='' not in emitted.content + assert 'hidden reasoning' not in emitted.content + + +def test_stream_accumulator_strips_initial_orphan_think_close_from_tool_round_content(): + accumulator = _StreamAccumulator( + msg_sequence=3, + initial_content='I will search for LangBot.', + remove_think=True, + ) + + assert accumulator.add(provider_message.MessageChunk(role='assistant', content='hidden reasoning')) is None + emitted = accumulator.add( + provider_message.MessageChunk( + role='assistant', + content=' still hiddenHere is the answer.', + is_final=True, + ) + ) + + assert emitted is not None + assert emitted.content == 'I will search for LangBot.Here is the answer.' + assert '' not in emitted.content + assert 'hidden reasoning' not in emitted.content + + @pytest.mark.asyncio async def test_localagent_uses_exec_for_exact_calculation(): provider = RecordingProvider()