diff --git a/src/langbot/libs/wecom_ai_bot_api/ws_client.py b/src/langbot/libs/wecom_ai_bot_api/ws_client.py index 32c11bd509..1fc5c97bdb 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -433,7 +433,9 @@ async def send_template_card(self, chat_id: str, card_payload: dict[str, Any]) - body['chatid'] = chat_id return await self._send_reply(req_id, body, cmd=CMD_SEND_MSG) - async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = False) -> bool: + async def push_stream_chunk( + self, msg_id: str, content: str, is_final: bool = False, keep_stream: bool = False + ) -> bool: """Push a streaming chunk for a given message ID. Compatible interface with WecomBotClient.push_stream_chunk. @@ -442,6 +444,9 @@ async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = Fa msg_id: The original message ID. content: The cumulative content from the pipeline. is_final: Whether this is the final chunk. + keep_stream: When True, keep the stream session alive even on + is_final so a subsequent round (e.g. tool-call loop) can + reuse the same stream_id. Returns: True if the stream session exists and chunk was sent. @@ -490,7 +495,7 @@ async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = Fa # every frame must contain the complete snapshot, not only a delta. await self.reply_stream(req_id, stream_id, next_content, finish=is_final, feedback_id=feedback_id) self._stream_last_content[msg_id] = next_content - if is_final: + if is_final and not keep_stream: self._stream_ids.pop(msg_id, None) self._stream_last_content.pop(msg_id, None) self._stream_sessions.pop(msg_id, None) diff --git a/src/langbot/pkg/pipeline/respback/respback.py b/src/langbot/pkg/pipeline/respback/respback.py index 6d8248aef7..653be29701 100644 --- a/src/langbot/pkg/pipeline/respback/respback.py +++ b/src/langbot/pkg/pipeline/respback/respback.py @@ -46,12 +46,22 @@ 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][-1] + bot_msg = query.resp_messages[-1] + # Read the keep_stream hint that the runner may have set on + # the MessageChunk's provider_specific_fields. This tells + # adapters (e.g. WeComBot WS) to keep the stream session + # alive across multi-round tool-call loops. + keep_stream = False + psf = getattr(bot_msg, 'provider_specific_fields', None) + if psf: + keep_stream = bool(psf.get('keep_stream', False)) await query.adapter.reply_message_chunk( message_source=query.message_event, - bot_message=query.resp_messages[-1], + bot_message=bot_msg, message=message_chain, quote_origin=quote_origin, is_final=is_final, + keep_stream=keep_stream, ) else: await query.adapter.reply_message( diff --git a/src/langbot/pkg/platform/sources/dingtalk.py b/src/langbot/pkg/platform/sources/dingtalk.py index 996187f08e..daf5377b2b 100644 --- a/src/langbot/pkg/platform/sources/dingtalk.py +++ b/src/langbot/pkg/platform/sources/dingtalk.py @@ -571,6 +571,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): message_id = bot_message.resp_message_id msg_seq = bot_message.msg_sequence diff --git a/src/langbot/pkg/platform/sources/discord.py b/src/langbot/pkg/platform/sources/discord.py index fae65f6a09..bfec79f910 100644 --- a/src/langbot/pkg/platform/sources/discord.py +++ b/src/langbot/pkg/platform/sources/discord.py @@ -1261,6 +1261,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): msg_id = ( bot_message.get('resp_message_id') diff --git a/src/langbot/pkg/platform/sources/http_bot.py b/src/langbot/pkg/platform/sources/http_bot.py index 16a891991d..cb2fdb3dee 100644 --- a/src/langbot/pkg/platform/sources/http_bot.py +++ b/src/langbot/pkg/platform/sources/http_bot.py @@ -468,6 +468,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ) -> dict: message_is_final = is_final and getattr(bot_message, 'tool_calls', None) is None return await self._emit_reply(message_source, message, is_final=message_is_final, stream=True) diff --git a/src/langbot/pkg/platform/sources/lark.py b/src/langbot/pkg/platform/sources/lark.py index 96e40469c2..592ec17174 100644 --- a/src/langbot/pkg/platform/sources/lark.py +++ b/src/langbot/pkg/platform/sources/lark.py @@ -2072,6 +2072,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): """ 回复消息变成更新卡片消息 diff --git a/src/langbot/pkg/platform/sources/qqofficial.py b/src/langbot/pkg/platform/sources/qqofficial.py index 35fd81f76f..320f63a198 100644 --- a/src/langbot/pkg/platform/sources/qqofficial.py +++ b/src/langbot/pkg/platform/sources/qqofficial.py @@ -532,6 +532,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): # Periodically clean up stale stream contexts await self._cleanup_stale_streams() diff --git a/src/langbot/pkg/platform/sources/telegram.py b/src/langbot/pkg/platform/sources/telegram.py index e7ad626bb6..06f37b053e 100644 --- a/src/langbot/pkg/platform/sources/telegram.py +++ b/src/langbot/pkg/platform/sources/telegram.py @@ -563,6 +563,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): message_id = bot_message.resp_message_id msg_seq = bot_message.msg_sequence diff --git a/src/langbot/pkg/platform/sources/web_page_bot_adapter.py b/src/langbot/pkg/platform/sources/web_page_bot_adapter.py index fa7f81743f..7faa0ec760 100644 --- a/src/langbot/pkg/platform/sources/web_page_bot_adapter.py +++ b/src/langbot/pkg/platform/sources/web_page_bot_adapter.py @@ -63,10 +63,11 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ) -> dict: if self._ws_adapter is not None: return await self._ws_adapter.reply_message_chunk( - message_source, bot_message, message, quote_origin, is_final + message_source, bot_message, message, quote_origin, is_final, keep_stream=keep_stream ) return {} diff --git a/src/langbot/pkg/platform/sources/websocket_adapter.py b/src/langbot/pkg/platform/sources/websocket_adapter.py index a66d4ac9ab..e20ec3dbaf 100644 --- a/src/langbot/pkg/platform/sources/websocket_adapter.py +++ b/src/langbot/pkg/platform/sources/websocket_adapter.py @@ -252,6 +252,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ) -> dict: """回复消息块 - 流式""" # 获取会话和pipeline信息 diff --git a/src/langbot/pkg/platform/sources/wecombot.py b/src/langbot/pkg/platform/sources/wecombot.py index d0febad9e5..acb3322328 100644 --- a/src/langbot/pkg/platform/sources/wecombot.py +++ b/src/langbot/pkg/platform/sources/wecombot.py @@ -410,6 +410,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): content = await self.message_converter.yiri2target(message) _ws_mode = not self.config.get('enable-webhook', False) @@ -463,7 +464,7 @@ async def reply_message_chunk( return {'stream': False, 'form': True, 'fallback': True} if _ws_mode: - success = await self.bot.push_stream_chunk(msg_id, content, is_final=is_final) + success = await self.bot.push_stream_chunk(msg_id, content, is_final=is_final, keep_stream=keep_stream) if not success and is_final: event = message_source.source_platform_object req_id = event.get('req_id', '') diff --git a/src/langbot/pkg/provider/runners/localagent.py b/src/langbot/pkg/provider/runners/localagent.py index 6c877239cb..41952d3ad7 100644 --- a/src/langbot/pkg/provider/runners/localagent.py +++ b/src/langbot/pkg/provider/runners/localagent.py @@ -481,7 +481,12 @@ async def run( except AttributeError: is_stream = False - remove_think = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get('remove-think', False) + _misc = (query.pipeline_config.get('output') or {}).get('misc') or {} + remove_think = _misc.get('remove-think', False) + keep_first_think_only = _misc.get('keep-first-think-only', False) + # When keep-first-think-only is on, the first (pre-loop) round keeps + # its CoT; subsequent tool-call rounds strip it as usual. + _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) @@ -500,23 +505,30 @@ 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, remove_think=remove_think) + 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) if chunk: + # When keep-first-think-only is active and this round + # produced tool calls, signal the adapter to keep the + # stream session alive so the next round can reuse it. + if keep_first_think_only and chunk.is_final and chunk.tool_calls: + psf = dict(chunk.provider_specific_fields or {}) + psf['keep_stream'] = True + chunk = chunk.model_copy(update={'provider_specific_fields': psf}) yield chunk initial_response_emitted = True diff --git a/src/langbot/templates/default-pipeline-config.json b/src/langbot/templates/default-pipeline-config.json index 78e2ec9584..e7230f8fae 100644 --- a/src/langbot/templates/default-pipeline-config.json +++ b/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 } } } diff --git a/src/langbot/templates/metadata/pipeline/output.yaml b/src/langbot/templates/metadata/pipeline/output.yaml index d5e0fae07a..7febc55495 100644 --- a/src/langbot/templates/metadata/pipeline/output.yaml +++ b/src/langbot/templates/metadata/pipeline/output.yaml @@ -145,4 +145,14 @@ stages: type: boolean required: true default: false + - name: keep-first-think-only + label: + en_US: Keep Only First CoT + zh_Hans: 仅保留第一条思维链 + description: + en_US: 'Only effective when "Remove CoT" is enabled. Keeps the chain-of-thought on the first LLM round but strips it from all subsequent rounds in a multi-round tool-call loop.' + zh_Hans: '仅在启用"删除思维链"时生效。多轮工具调用中，保留第一轮的思维链，后续轮次的思维链一律删除。' + type: boolean + required: true + default: false diff --git a/tests/factories/platform.py b/tests/factories/platform.py index 77b8f11f99..9ab083abc9 100644 --- a/tests/factories/platform.py +++ b/tests/factories/platform.py @@ -230,6 +230,7 @@ async def reply_message_chunk( message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, + keep_stream: bool = False, ): """Simulate streaming reply (captures for assertions).""" if self._raise_error: @@ -244,6 +245,7 @@ async def reply_message_chunk( 'message': message, 'quote_origin': quote_origin, 'is_final': is_final, + 'keep_stream': keep_stream, } )