feat(pipeline): add keep-first-think-only toggle for multi-round CoT

src/langbot/libs/wecom_ai_bot_api/ws_client.py

@@ -102,6 +102,11 @@ def __init__(
         self._feedback_sessions: dict[str, dict] = {}  # feedback_id -> {msg_id, user_id, chat_id, stream_id, req_id}
         # msg_id -> feedback_id (for associating feedback with message)
         self._msg_feedback_ids: dict[str, str] = {}  # msg_id -> feedback_id
+        # Round counter per msg_id for keep-first-think-only: when
+        # fallback_after_round >= 1, the stream entry is NOT cleared on
+        # end-of-round so the next round can reuse the same stream_id.
+        self._stream_rounds: dict[str, int] = {}  # msg_id -> round count
+        self.fallback_after_round: int = 0  # 0 = disabled
 
     # ── Public API ──────────────────────────────────────────────────
 
@@ -293,9 +298,20 @@ async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = Fa
             await self.reply_stream(req_id, stream_id, content, finish=is_final, feedback_id=feedback_id)
             self._stream_last_content[msg_id] = content
             if is_final:
-                self._stream_ids.pop(msg_id, None)
-                self._stream_last_content.pop(msg_id, None)
-                self._stream_sessions.pop(msg_id, None)
+                # When fallback_after_round is set, keep the stream entry
+                # alive across rounds so the next LLM round can reuse
+                # the same stream_id. Only clear when the round counter
+                # has crossed the threshold (final round).
+                round_idx = self._stream_rounds.get(msg_id, 0) + 1
+                self._stream_rounds[msg_id] = round_idx
+                if self.fallback_after_round > 0 and round_idx < self.fallback_after_round:
+                    # Keep stream alive for the next round
+                    pass
+                else:
+                    self._stream_ids.pop(msg_id, None)
+                    self._stream_last_content.pop(msg_id, None)
+                    self._stream_sessions.pop(msg_id, None)
+                    self._stream_rounds.pop(msg_id, None)
             return True
         except Exception:
             await self.logger.error(f'Failed to push stream chunk: {traceback.format_exc()}')

src/langbot/pkg/pipeline/respback/respback.py

@@ -46,6 +46,14 @@ async def process(self, query: pipeline_query.Query, stage_inst_name: str) -> en
         try:
             if await query.adapter.is_stream_output_supported() and has_chunks:
                 is_final = [msg.is_final for msg in query.resp_messages][0]
+                self.ap.logger.info(f'respback: calling reply_message_chunk, is_final={is_final}')
+                # Stash pipeline config on the message event so the
+                # platform adapter can read misc toggles like
+                # keep-first-think-only without threading query through.
+                try:
+                    query.message_event._langbot_pipeline_config = query.pipeline_config
+                except Exception:
+                    pass
                 await query.adapter.reply_message_chunk(
                     message_source=query.message_event,
                     bot_message=query.resp_messages[-1],

src/langbot/pkg/platform/sources/wecombot.py

@@ -368,6 +368,15 @@ async def reply_message_chunk(
         _ws_mode = not self.config.get('enable-webhook', False)
 
         if _ws_mode:
+            # Read keep-first-think-only from pipeline config stashed by respback
+            try:
+                _pipeline_cfg = getattr(message_source, '_langbot_pipeline_config', None) or {}
+                _misc = (_pipeline_cfg.get('output') or {}).get('misc') or {}
+                _keep_first = bool(_misc.get('remove-think', False) and _misc.get('keep-first-think-only', False))
+                self.bot.fallback_after_round = 1 if _keep_first else 0
+            except Exception:
+                pass
+
             success = await self.bot.push_stream_chunk(msg_id, content, is_final=is_final)
             if not success and is_final:
                 event = message_source.source_platform_object

src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py

@@ -13,6 +13,136 @@
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 
 
+class _ThinkStripState:
+    """Stateful filter that drops ```` think blocks across chunk boundaries.
+
+    The think tag (and the legacy CRETIRE_* marker) can straddle two or
+    more streaming chunks.  A naive per-chunk regex either fails to match
+    (when split mid-tag) or leaks the tag into user-visible output.  This
+    state machine tracks the longest suffix of the consumed input that is
+    also a prefix of any known tag, and keeps that suffix in an internal
+    buffer for the next call.
+    """
+
+    _OPEN_TAG: str = '<think>'
+    _CLOSE_TAG: str = '</think>'
+    _LEGACY_OPEN: str = 'CRETIRE_REASONING_BEGINk'
+    _LEGACY_CLOSE: str = 'CRETIRE_REASONING_ENDk'
+
+    def __init__(self) -> None:
+        self._tags: tuple[str, ...] = (
+            self._OPEN_TAG,
+            self._CLOSE_TAG,
+            self._LEGACY_OPEN,
+            self._LEGACY_CLOSE,
+        )
+        self._max_tag_len: int = max(len(t) for t in self._tags)
+        self._buf: str = ''
+        self._buf_inside: bool = False
+
+    def feed(self, chunk: str) -> str:
+        """Feed a streaming delta; return the portion that should be emitted."""
+        if not chunk:
+            return chunk
+
+        if self._buf_inside:
+            text = self._buf + chunk
+            for close_tag in (self._CLOSE_TAG, self._LEGACY_CLOSE):
+                idx = text.find(close_tag)
+                if idx != -1:
+                    self._buf = ''
+                    self._buf_inside = False
+                    return self._process(text[idx + len(close_tag) :])
+            keep = 0
+            for tag in (self._CLOSE_TAG, self._LEGACY_CLOSE):
+                _, k = self._split_for_close_prefix(text, 0, tag)
+                if k > keep:
+                    keep = k
+            self._buf = text[len(text) - keep :]
+            return ''
+
+        return self._process(chunk)
+
+    def _process(self, chunk: str) -> str:
+        text = self._buf + chunk
+
+        out: list[str] = []
+        i = 0
+        n = len(text)
+        while i < n:
+            open_idx, open_tag = self._find_open(text, i)
+            if open_idx == -1:
+                emit_end, _keep = self._split_for_tag_prefix(text, i)
+                out.append(text[i:emit_end])
+                self._buf = text[emit_end:]
+                self._buf_inside = False
+                return ''.join(out)
+
+            if open_idx > i:
+                out.append(text[i:open_idx])
+            block_start = open_idx + len(open_tag)
+            close_tag = self._LEGACY_CLOSE if open_tag == self._LEGACY_OPEN else self._CLOSE_TAG
+            close_idx = text.find(close_tag, block_start)
+            if close_idx == -1:
+                emit_end, _keep = self._split_for_close_prefix(text, block_start, close_tag)
+                self._buf = text[emit_end:]
+                self._buf_inside = True
+                return ''.join(out)
+            i = close_idx + len(close_tag)
+
+        emit_end, _keep = self._split_for_tag_prefix(text, i)
+        out.append(text[i:emit_end])
+        self._buf = text[emit_end:]
+        self._buf_inside = False
+        return ''.join(out)
+
+    def flush(self) -> str:
+        """Release any text still held in the internal buffer."""
+        pending, self._buf = self._buf, ''
+        inside, self._buf_inside = self._buf_inside, False
+        if inside:
+            return ''
+        return pending
+
+    def _find_open(self, text: str, start: int) -> tuple[int, str]:
+        best_idx = -1
+        best_tag = ''
+        for tag in (self._OPEN_TAG, self._LEGACY_OPEN):
+            idx = text.find(tag, start)
+            if idx != -1 and (best_idx == -1 or idx < best_idx):
+                best_idx = idx
+                best_tag = tag
+        return best_idx, best_tag
+
+    def _split_for_tag_prefix(self, text: str, start: int) -> tuple[int, int]:
+        suffix = text[start:]
+        best_keep = 0
+        for tag in (self._OPEN_TAG, self._LEGACY_OPEN):
+            idx = suffix.find(tag)
+            if idx != -1:
+                candidate = len(suffix) - idx
+                if candidate > best_keep:
+                    best_keep = candidate
+        if best_keep == 0:
+            limit = min(len(suffix), self._max_tag_len - 1)
+            for k in range(limit, 0, -1):
+                tail = suffix[-k:]
+                if any(tag.startswith(tail) for tag in self._tags):
+                    best_keep = k
+                    break
+        return len(text) - best_keep, best_keep
+
+    def _split_for_close_prefix(self, text: str, start: int, close_tag: str) -> tuple[int, int]:
+        suffix = text[start:]
+        best_keep = 0
+        limit = min(len(suffix), len(close_tag) - 1)
+        for k in range(limit, 0, -1):
+            if close_tag.startswith(suffix[-k:]):
+                best_keep = k
+                break
+        return len(text) - best_keep, best_keep
+
+
 class LiteLLMRequester(requester.ProviderAPIRequester):
     """LiteLLM unified API requester supporting chat, embedding, and rerank."""
 
@@ -237,6 +367,25 @@ def _convert_messages(self, messages: typing.List[provider_message.Message]) ->
 
         return req_messages
 
+    # Patterns covering chain-of-thought markers emitted by various
+    # OpenAI-compatible providers (DeepSeek, MiniMax-M3, etc.) as well
+    # as an internal private marker.
+    _THINK_PATTERNS: tuple[str, ...] = (
+        r'</think>',
+        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,16 +397,12 @@ def _process_thinking_content(self, content: str, reasoning_content: str | None,
         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
-
-            think_pattern = r'CRETIRE_REASONING_BEGINk(.*?)CRETIRE_REASONING_ENDk'
-
-            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
+        # Strip chain-of-thought blocks when requested. The think pattern is
+        # the public convention used by DeepSeek, MiniMax-M3 and other
+        # OpenAI-compatible providers that emit raw reasoning text in the
+        # content field. The CRETIRE_* pattern is an internal marker.
+        if remove_think and content:
+            content = self._strip_think(content)
 
         # Handle separate reasoning_content field
         # Currently we don't include reasoning_content in user-facing output regardless of remove_think
@@ -571,6 +716,13 @@ async def invoke_llm_stream(
         role = 'assistant'
         tool_call_state: dict[int, dict[str, typing.Any]] = {}
 
+        # Stream-level state for `` stripping. `` tags can straddle
+        # chunk boundaries, so use a stateful filter.
+        if remove_think:
+            think_state = _ThinkStripState()
+        else:
+            think_state = None
+
         try:
             response = await acompletion(**args)
             async for chunk in response:
@@ -613,6 +765,14 @@ async def invoke_llm_stream(
                         # Use reasoning_content as the displayed content
                         delta_content = reasoning_content
 
+                # Strip `` tags from the delta when remove_think is set.
+                # Think blocks can straddle chunks, so use a stateful filter.
+                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:

src/langbot/pkg/provider/runners/localagent.py

@@ -47,14 +47,28 @@ def _model_has_ability(model: modelmgr_requester.RuntimeLLMModel, ability: str)
 
 
 class _StreamAccumulator:
-    """Accumulate streamed content and fragmented OpenAI-style tool calls."""
+    """Accumulate streamed content and fragmented OpenAI-style tool calls.
 
-    def __init__(self, msg_sequence: int = 0, initial_content: str | None = None):
+    When ``remove_think`` is enabled, the accumulator performs a final pass
+    over the accumulated content to strip chain-of-thought blocks as a
+    safety net. The upstream LLM requester is expected to have already
+    stripped its own streaming chunks, but the local agent runner can
+    re-invoke the LLM after tool calls and that second pass may produce
+    content that has not been filtered yet.
+    """
+
+    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
 
     def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChunk | None:
         self.msg_idx += 1
@@ -83,7 +97,7 @@ def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChu
             self.msg_sequence += 1
             return provider_message.MessageChunk(
                 role=self.last_role,
-                content=self.accumulated_content,
+                content=self._maybe_strip(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,
@@ -94,11 +108,18 @@ def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChu
     def final_message(self) -> provider_message.MessageChunk:
         return provider_message.MessageChunk(
             role=self.last_role,
-            content=self.accumulated_content,
+            content=self._maybe_strip(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(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)
+
 
 @runner.runner_class('local-agent')
 class LocalAgentRunner(runner.RequestRunner):
@@ -448,7 +469,14 @@ async def run(
         except AttributeError:
             is_stream = False
 
-        remove_think = query.pipeline_config['output'].get('misc', '').get('remove-think')
+        # Safely resolve the "remove think blocks" toggle. ``output.misc``
+        # is not guaranteed to exist on every pipeline configuration.
+        _misc = (query.pipeline_config.get('output') or {}).get('misc') or {}
+        remove_think = _misc.get('remove-think', False)
+        # ``keep-first-think-only``: only effective when ``remove-think`` is on.
+        # Round 1 (pre-loop) keeps CoT; rounds 2+ (tool-loop) strip it.
+        keep_first_think_only = _misc.get('keep-first-think-only', False)
+        _strip_think_first_round = remove_think and not keep_first_think_only
 
         # Build ordered candidate list (primary + fallbacks)
         candidates = await self._get_model_candidates(query)
@@ -467,19 +495,19 @@ async def run(
                 candidates,
                 req_messages,
                 query.use_funcs,
-                remove_think,
+                _strip_think_first_round,
             )
             final_msg = msg
         else:
             # Streaming: invoke with fallback
-            stream_accumulator = _StreamAccumulator(msg_sequence=1)
+            stream_accumulator = _StreamAccumulator(msg_sequence=1, remove_think=_strip_think_first_round)
 
             stream_src, use_llm_model = await self._invoke_stream_with_fallback(
                 query,
                 candidates,
                 req_messages,
                 query.use_funcs,
-                remove_think,
+                _strip_think_first_round,
             )
             async for msg in stream_src:
                 chunk = stream_accumulator.add(msg)
@@ -576,6 +604,7 @@ async def run(
                 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(

src/langbot/templates/default-pipeline-config.json

@@ -107,7 +107,8 @@
             "at-sender": true,
             "quote-origin": true,
             "track-function-calls": false,
-            "remove-think": false
+            "remove-think": false,
+            "keep-first-think-only": false
         }
     }
 }