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 5125a704a..bc9572712 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -34,6 +34,14 @@ CMD_RESPOND_WELCOME = 'aibot_respond_welcome_msg' CMD_RESPOND_UPDATE = 'aibot_respond_update_msg' CMD_SEND_MSG = 'aibot_send_msg' +# Media upload protocol (3 steps: init -> chunk * N -> finish). The +# command names below match the WeCom AI Bot long-connection protocol. +CMD_UPLOAD_INIT = 'aibot_upload_media_init' +CMD_UPLOAD_CHUNK = 'aibot_upload_media_chunk' +CMD_UPLOAD_FINISH = 'aibot_upload_media_finish' + +# Default upload chunk size: 512 KB before base64 encoding. +_UPLOAD_CHUNK_SIZE = 512 * 1024 def _generate_req_id(prefix: str) -> str: @@ -98,6 +106,19 @@ def __init__( self._stream_last_content: dict[str, str] = {} # msg_id -> last content sent # Stream session info for feedback tracking self._stream_sessions: dict[str, dict] = {} # msg_id -> session info + # Per-msg_id end-of-round counter. Combined with the + # ``fallback_after_round`` flag below, decides whether to + # rotate the stream_id (round N + 1 opens a new streaming + # message) or to ``pop`` it (round N + 1's next push falls + # back to ``reply_text``). + self._stream_round_counts: dict[str, int] = {} + # Caller-controlled switch. When set to N (>0), end-of-round + # events with count == N drop the stream entry, so the next + # round's push falls through to ``reply_text``. The most + # useful value is 1: only round 1 streams (renders the + # initial reasoning as a WeCom "thinking" block), rounds + # 2+ are sent as plain markdown via ``reply_text``. + self.fallback_after_round: int = 0 # Feedback tracking: feedback_id -> session info 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) @@ -258,6 +279,146 @@ async def send_message(self, chat_id: str, content: str, msgtype: str = 'markdow body['text'] = {'content': content} return await self._send_reply(req_id, body, cmd=CMD_SEND_MSG) + # ------------------------------------------------------------------ + # Media upload (image / voice / file) + # ------------------------------------------------------------------ + + async def upload_media( + self, + data: bytes, + filename: str = 'attachment', + media_type: str = 'file', + ) -> Optional[dict]: + """Upload *data* to the WeCom AI Bot CDN and return the parsed ACK. + + Implements the three-step protocol documented for the WeCom + AI Bot: + + 1. ``aibot_upload_media_init`` — declare media type, file name, + size, MD5 and chunk count; receive ``upload_id``. + 2. ``aibot_upload_media_chunk`` — send each chunk (base64-encoded + bytes) until done; receive per-chunk ACK. + 3. ``aibot_upload_media_finish`` — finalize the upload; receive + ``media_id``. + + Returns a dict with the final ``media_id`` (and the raw + ``finish`` ACK) on success, or ``None`` on any failure. The + caller is expected to ignore the result and continue + gracefully — the framework will keep working without media + delivery. + """ + import base64 as _b64 + import hashlib as _hl + + if not data: + return None + + file_size = len(data) + file_md5 = _hl.md5(data).hexdigest() + total_chunks = (file_size + _UPLOAD_CHUNK_SIZE - 1) // _UPLOAD_CHUNK_SIZE + if total_chunks == 0: + total_chunks = 1 + + # Step 1: init. + init_req_id = _generate_req_id(CMD_UPLOAD_INIT) + init_body = { + 'type': media_type, + 'filename': filename, + 'total_size': file_size, + 'total_chunks': total_chunks, + 'md5': file_md5, + } + init_ack = await self._send_reply( + init_req_id, + init_body, + cmd=CMD_UPLOAD_INIT, + ) + if not init_ack or init_ack.get('errcode', 0) != 0: + await self.logger.warning(f'upload_media init failed: ack={init_ack!r}') + return None + upload_id = ( + init_ack.get('upload_id') + or init_ack.get('body', {}).get('upload_id') + or init_ack.get('data', {}).get('upload_id') + ) + if not upload_id: + await self.logger.warning(f'upload_media init returned no upload_id: ack={init_ack!r}') + return None + + # Step 2: chunks. + for index in range(total_chunks): + start = index * _UPLOAD_CHUNK_SIZE + end = min(start + _UPLOAD_CHUNK_SIZE, file_size) + chunk_bytes = data[start:end] + chunk_req_id = _generate_req_id(CMD_UPLOAD_CHUNK) + chunk_body = { + 'upload_id': upload_id, + 'chunk_index': index, + 'base64_data': _b64.b64encode(chunk_bytes).decode('ascii'), + } + chunk_ack = await self._send_reply( + chunk_req_id, + chunk_body, + cmd=CMD_UPLOAD_CHUNK, + ) + if not chunk_ack or chunk_ack.get('errcode', 0) != 0: + await self.logger.warning(f'upload_media chunk {index} failed: ack={chunk_ack!r}') + return None + + # Step 3: finish. + finish_req_id = _generate_req_id(CMD_UPLOAD_FINISH) + finish_body = {'upload_id': upload_id} + finish_ack = await self._send_reply( + finish_req_id, + finish_body, + cmd=CMD_UPLOAD_FINISH, + ) + if not finish_ack or finish_ack.get('errcode', 0) != 0: + await self.logger.warning(f'upload_media finish failed: ack={finish_ack!r}') + return None + + media_id = ( + finish_ack.get('media_id') + or finish_ack.get('body', {}).get('media_id') + or finish_ack.get('data', {}).get('media_id') + ) + if not media_id: + await self.logger.warning(f'upload_media finish returned no media_id: ack={finish_ack!r}') + return None + await self.logger.info(f'upload_media OK: filename={filename!r} media_id={media_id!r}') + return {'media_id': media_id, 'ack': finish_ack} + + async def _reply_media( + self, + req_id: str, + media_id: str, + kind: str, + ) -> Optional[dict]: + """Send a media reply (image / voice / file) referencing *media_id*. + + ``kind`` is one of ``'image'``, ``'voice'``, ``'file'``. Uses + the standard ``aibot_respond_msg`` command with a per-kind + body key (matches the convention documented for the WeCom + AI Bot SDK). + """ + if kind not in {'image', 'voice', 'file'}: + await self.logger.warning(f'_reply_media called with unknown kind={kind!r}') + return None + body = { + 'msgtype': kind, + kind: {'media_id': media_id}, + } + return await self._send_reply(req_id, body, cmd=CMD_RESPOND_MSG) + + async def reply_image(self, req_id: str, media_id: str) -> Optional[dict]: + return await self._reply_media(req_id, media_id, 'image') + + async def reply_file(self, req_id: str, media_id: str) -> Optional[dict]: + return await self._reply_media(req_id, media_id, 'file') + + async def reply_voice(self, req_id: str, media_id: str) -> Optional[dict]: + return await self._reply_media(req_id, media_id, 'voice') + async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = False) -> bool: """Push a streaming chunk for a given message ID. @@ -293,9 +454,40 @@ 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) + # End-of-round. We support two post-round policies: + # + # * ``fallback_after_round == 0`` (default): keep + # rotating the stream_id so the next round + # opens a fresh streaming message. WeCom + # renders each as its own "thinking" block + # (the bb89fe7f behaviour). + # + # * ``fallback_after_round == N`` for N >= 1: drop + # the stream entry once the round counter has + # reached N. The next round's push fails to + # find the entry and returns ``False``, which + # the caller (wecombot) translates into a + # ``reply_text`` fallback. Round 1 still renders + # as a "thinking" block; rounds 2+ are sent + # as plain markdown messages, which is what the + # ``keep-first-think-only`` misc toggle asks + # for. + round_idx = self._stream_round_counts.get(msg_id, 0) + 1 + self._stream_round_counts[msg_id] = round_idx + if self.fallback_after_round > 0 and round_idx >= self.fallback_after_round: + # Drop the stream entry. The next push from + # the same msg_id will fail to find the + # stream_id and fall back to ``reply_text``. + self._stream_ids.pop(msg_id, None) + self._stream_last_content.pop(msg_id, None) + self._stream_sessions.pop(msg_id, None) + else: + # Rotate to a fresh stream_id so the next + # round opens a *new* streaming message. + new_stream_id = _generate_req_id('stream') + self._stream_ids[msg_id] = f'{req_id}|{new_stream_id}' + self._stream_last_content.pop(msg_id, None) + self._stream_sessions.pop(msg_id, None) return True except Exception: await self.logger.error(f'Failed to push stream chunk: {traceback.format_exc()}') diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index ca30eb930..15da37d64 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -683,34 +683,169 @@ async def collect_outbound_attachments(self, query: pipeline_query.Query) -> lis Reads ``/workspace/outbox//`` (recursively) — directly from the bind-mounted host directory when available (no size limit), else - via the exec channel — returns a list of ``{type, name, base64}`` - ready to become platform message components, then clears the outbox so - a later turn in the same session does not re-send stale files. Returns - ``[]`` when nothing was produced. + via the exec channel — and returns a list of ``{type, name, base64}`` + ready to become platform message components. + + Files that have already been handed out in a previous collect for + the same ``query_id`` are skipped — they live in ``/.sent/`` + under their original relative path and are quarantined there on the + first collect. The quarantine is cleared by an explicit + ``finalize_outbound_attachments`` call once the query is done + (typically from the chat handler after streaming completes), so + a long-running agent loop can collect new files round after round + without losing or re-emitting the earlier ones. + + Returns ``[]`` when nothing new was produced. """ if not self._available: return [] host_dir = self._host_query_dir(self.OUTBOX_SUBDIR, query.query_id) if host_dir is not None: - entries = await asyncio.to_thread(self._read_outbox_host, host_dir) + entries = await asyncio.to_thread(self._read_outbox_quarantine_new, host_dir) else: + # The exec fallback cannot quarantine atomically across the + # container boundary, so we keep the simpler read-and-clear + # semantics for the remote path. The host path (Docker / + # bind-mount) is the common case and is fully + # multi-collect-safe. entries = await self._read_outbox_via_exec(query) + await self._clear_outbox(query, host_dir) attachments = self._classify_outbound_entries(entries) - - # Always clear the per-query outbox after reading — even when nothing - # was collected — so a later turn that reuses the same query_id (the - # counter resets across restarts) never inherits stale files. - await self._clear_outbox(query, host_dir) if attachments: self.ap.logger.info( f'Collected {len(attachments)} outbound attachment(s) from sandbox: query_id={query.query_id}' ) return attachments + async def finalize_outbound_attachments(self, query: pipeline_query.Query) -> None: + """Clear the per-query outbox and quarantine directory at end of query. + + Counterpart to :meth:`collect_outbound_attachments`. Once a query + has finished streaming every tool round, any files still in + ``/.sent`` and ``/`` itself are no longer + useful — drop them so the next turn that reuses this + ``query_id`` (the counter resets across restarts) does not + inherit stale attachments. Best-effort: never raises into the + pipeline. + + For the host path we ``rmtree`` the per-query dir; for the + exec fallback we issue an in-container ``rm -rf`` because the + container may own its outbox files as root and the host has + no way to delete them. + """ + if not self._available: + return + host_dir = self._host_query_dir(self.OUTBOX_SUBDIR, query.query_id) + if host_dir is not None: + if os.path.isdir(host_dir): + import shutil + + def _drop() -> None: + shutil.rmtree(host_dir, ignore_errors=True) + os.makedirs(host_dir, exist_ok=True) + + await asyncio.to_thread(_drop) + return + # Exec path + target_dir = f'{self.OUTBOX_MOUNT_DIR}/{query.query_id}' + try: + await self.execute_tool( + {'command': f'rm -rf {target_dir} && mkdir -p {target_dir}', 'timeout_sec': 30}, + query, + ) + except Exception as exc: + self.ap.logger.warning(f'Failed to finalize sandbox outbox {target_dir}: {exc}') + + def _read_outbox_quarantine_new(self, host_dir: str) -> list[dict]: + """Read outbox files NEW since last collection, and quarantine them. + + Each newly returned file is moved to ``/.sent/`` + with ``os.replace`` (atomic on POSIX, near-atomic on Windows). + Subsequent collects see the file in ``.sent/`` and skip it + via ``os.walk`` + ``os.path.exists`` of the same path. + + ``rel`` is the full path relative to ``host_dir`` (so files + in different subdirectories with the same basename — e.g. + ``a/1.jpg`` and ``b/1.jpg`` — quarantine under distinct + keys and are both delivered). + + Root-owned container files on a bind-mounted host where + LangBot runs as a non-root user fail ``os.replace`` with + ``PermissionError``. Those files are not collected on the + host path; the next ``finalize_outbound_attachments`` call + sweeps ``host_dir`` clean (or, failing that, falls back to + clearing inside the container via ``_clear_outbox``). + """ + import base64 as _b64 + + entries: list[dict] = [] + sent_dir = os.path.join(host_dir, '.sent') + os.makedirs(sent_dir, exist_ok=True) + if not os.path.isdir(host_dir): + return entries + + sent_dir_abs = os.path.abspath(sent_dir) + # ``os.walk`` snapshots the directory list of each level + # when it descends, so pruning ``dirs`` in-place does not + # reliably prevent the walker from re-entering subtrees + # we create during the loop (e.g. ``.sent/a/`` appears + # mid-iteration once we ``os.replace`` the first file + # there). Use a manual stack instead so we have full + # control over which directories are visited. + stack: list[str] = [host_dir] + while stack: + current = stack.pop() + if os.path.abspath(current) == sent_dir_abs: + continue + try: + entries_iter = list(os.scandir(current)) + except OSError: + continue + for entry in sorted(entries_iter, key=lambda e: e.name): + if entry.name.startswith('.'): + # Skip ``.sent`` and any other hidden dots. + continue + path = os.path.join(current, entry.name) + if entry.is_dir(follow_symlinks=False): + stack.append(path) + continue + if not entry.is_file(follow_symlinks=False): + continue + # ``path`` is a regular file under ``current``. + rel = os.path.relpath(path, host_dir) + sent_path = os.path.join(sent_dir, rel) + if os.path.exists(sent_path): + # Already handed out in a previous collect. + continue + try: + if os.path.getsize(path) > self._ATTACHMENT_MAX_BYTES: + continue + with open(path, 'rb') as fh: + data = fh.read() + except OSError: + continue + try: + os.makedirs(os.path.dirname(sent_path), exist_ok=True) + os.replace(path, sent_path) + except OSError: + # Container root-owned file on a non-root host — + # cannot quarantine here. Skipping is safer than + # risking an infinite re-delivery loop on a + # subsequent collect. + continue + entries.append({'name': rel, 'b64': _b64.b64encode(data).decode('ascii')}) + return entries + def _read_outbox_host(self, host_dir: str) -> list[dict]: - """Read outbox files straight off the bind-mounted host directory.""" + """Read outbox files straight off the bind-mounted host directory. + + Kept as a public-style helper in case external callers want to peek at + the outbox without going through ``collect_outbound_attachments``. + The quarantine-aware read is in + ``_read_outbox_quarantine_new``. + """ import base64 as _b64 entries: list[dict] = [] diff --git a/src/langbot/pkg/persistence/databases/postgresql.py b/src/langbot/pkg/persistence/databases/postgresql.py index f63d8f61f..be2bca055 100644 --- a/src/langbot/pkg/persistence/databases/postgresql.py +++ b/src/langbot/pkg/persistence/databases/postgresql.py @@ -1,10 +1,66 @@ from __future__ import annotations +import json + import sqlalchemy.ext.asyncio as sqlalchemy_asyncio from .. import database +async def _register_json_codecs(pg_conn) -> None: + """Register asyncpg type codecs for JSON / JSONB so values come back as + Python ``dict`` / ``list`` instead of raw ``str``. + """ + await pg_conn.set_type_codec( + 'json', + encoder=json.dumps, + decoder=json.loads, + schema='pg_catalog', + format='text', + ) + await pg_conn.set_type_codec( + 'jsonb', + encoder=json.dumps, + decoder=json.loads, + schema='pg_catalog', + format='text', + ) + + +def _patch_asyncpg_dialect() -> None: + """Wrap ``AsyncAdapt_asyncpg_dbapi.connect`` so that every new + asyncpg-backed SQLAlchemy connection gets our JSON codecs installed + eagerly, via the wrapper's ``run_async`` hook. + """ + from sqlalchemy.dialects.postgresql.asyncpg import AsyncAdapt_asyncpg_dbapi + + _orig = AsyncAdapt_asyncpg_dbapi.connect + + def _patched(self, *args, **kwargs): + wrapper = _orig(self, *args, **kwargs) + + # `AsyncAdapt_asyncpg_connection.run_async(fn)` invokes `fn(conn)` + # where `conn` is the real asyncpg.Connection. The hook must accept + # that argument; some SQLAlchemy versions forward it positionally, + # others don't, so we tolerate both. + async def _init(conn=None): + pg = conn + if pg is None: + adapt = wrapper.connection + pg = adapt._connection if hasattr(adapt, '_connection') else adapt + await _register_json_codecs(pg) + + wrapper.run_async(_init) + return wrapper + + AsyncAdapt_asyncpg_dbapi.connect = _patched + + +# Apply the patch at module import time so subsequent engine creation +# inherits the codec registration. +_patch_asyncpg_dialect() + + @database.manager_class('postgresql') class PostgreSQLDatabaseManager(database.BaseDatabaseManager): """PostgreSQL database manager""" diff --git a/src/langbot/pkg/pipeline/pipelinemgr.py b/src/langbot/pkg/pipeline/pipelinemgr.py index adf44b524..930b4c235 100644 --- a/src/langbot/pkg/pipeline/pipelinemgr.py +++ b/src/langbot/pkg/pipeline/pipelinemgr.py @@ -163,6 +163,14 @@ async def _check_output(self, query: pipeline_query.Query, result: pipeline_enti ): result.user_notice.insert(0, platform_message.At(target=query.message_event.sender.id)) if await query.adapter.is_stream_output_supported() and query.resp_messages: + # Stash pipeline_config on the event so the + # platform adapter can read misc toggles + # (``keep-first-think-only``) without a new + # argument on the adapter interface. + 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], diff --git a/src/langbot/pkg/pipeline/process/handlers/chat.py b/src/langbot/pkg/pipeline/process/handlers/chat.py index 4e5c14ea4..6bbee748a 100644 --- a/src/langbot/pkg/pipeline/process/handlers/chat.py +++ b/src/langbot/pkg/pipeline/process/handlers/chat.py @@ -142,6 +142,19 @@ async def handle( f'Conversation({query.query_id}) Streaming completed: {chunk_count} chunks, {text_length} chars' ) + # Quarantine was over the outbox. Now that streaming + # is done, give the box service a chance to sweep + # everything (including ``.sent/``) so the next turn + # that reuses this query_id starts clean. + box_service = getattr(self.ap, 'box_service', None) + if box_service is not None and getattr(box_service, 'available', False): + try: + await box_service.finalize_outbound_attachments(query) + except Exception as exc: + self.ap.logger.warning( + f'Failed to finalize outbound attachments: query_id={query.query_id}: {exc}' + ) + else: async for result in runner.run(query): query.resp_messages.append(result) diff --git a/src/langbot/pkg/pipeline/respback/respback.py b/src/langbot/pkg/pipeline/respback/respback.py index 0c85fbb45..e27b8d0bc 100644 --- a/src/langbot/pkg/pipeline/respback/respback.py +++ b/src/langbot/pkg/pipeline/respback/respback.py @@ -46,6 +46,17 @@ 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] + # Stash the pipeline config on the message event so + # the platform adapter can read misc toggles like + # ``keep-first-think-only`` without us having to + # thread a ``query`` argument through the adapter + # interface. Pydantic BaseModel instances accept + # arbitrary attribute writes, so this is a safe + # side channel. + 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], diff --git a/src/langbot/pkg/pipeline/wrapper/wrapper.py b/src/langbot/pkg/pipeline/wrapper/wrapper.py index 50db693d4..23490e3cb 100644 --- a/src/langbot/pkg/pipeline/wrapper/wrapper.py +++ b/src/langbot/pkg/pipeline/wrapper/wrapper.py @@ -39,35 +39,82 @@ def _is_final_assistant_message(self, result) -> bool: return bool(result.is_final) return True - async def _append_outbound_attachments( + async def _collect_outbound_components( self, query: pipeline_query.Query, - message_chain: platform_message.MessageChain, - ) -> None: - """Collect sandbox outbox files and append them to *message_chain*. - - Runs at most once per query (guarded by a query variable) and never - raises into the pipeline — attachment delivery is best-effort. + ) -> tuple[list[platform_message.PlatformMessage], bool]: + """Read sandbox outbox and return (components, service_available). + + ``components`` is a list of Image / Voice / File. Empty when + the box service is unavailable, the outbox has nothing new, + or the collect call failed. + + ``service_available`` is True when ``ap.box_service`` exists + and ``available`` is True. Callers use this to decide whether + to set ``_sandbox_outbound_collected`` — we do not want to + mark a query as "collected" just because the service was + absent, since a later configuration change could surface + attachments for that same query. + + ``_collect_outbound_components`` does NOT dedupe across calls + — the caller decides when to skip. The terminal assistant + branch keeps the existing ``_sandbox_outbound_collected`` + at-most-once flag; the tool-result branch calls us every tool + round so multi-image agents get every image. """ - if query.variables.get('_sandbox_outbound_collected'): - return box_service = getattr(self.ap, 'box_service', None) if box_service is None or not getattr(box_service, 'available', False): - return - query.variables['_sandbox_outbound_collected'] = True + return [], False try: attachments = await box_service.collect_outbound_attachments(query) except Exception as e: self.ap.logger.warning(f'Outbound attachment collection failed: {e}') - return + return [], True + components: list[platform_message.PlatformMessage] = [] for att in attachments: att_type = att.get('type') if att_type == 'Image': - message_chain.append(platform_message.Image(base64=att['base64'])) + components.append(platform_message.Image(base64=att['base64'])) elif att_type == 'Voice': - message_chain.append(platform_message.Voice(base64=att['base64'])) + components.append(platform_message.Voice(base64=att['base64'])) else: - message_chain.append(platform_message.File(name=att.get('name', 'file'), base64=att['base64'])) + components.append(platform_message.File(name=att.get('name', 'file'), base64=att['base64'])) + return components, True + + async def _append_outbound_attachments( + self, + query: pipeline_query.Query, + message_chain: platform_message.MessageChain, + ) -> None: + """Collect sandbox outbox files and append them to *message_chain*. + + Runs at most once per terminal assistant message (guarded by a + query variable) and never raises into the pipeline — attachment + delivery is best-effort. + + This path is kept for backwards compatibility — the existing + ``is_final + assistant`` branch on the final turn. The new + per-tool-round branch (see ``process`` below) reuses + ``_collect_outbound_components`` directly so it bypasses this + at-most-once flag. + + The ``_sandbox_outbound_collected`` flag is set only when + ``_collect_outbound_components`` could consult the box + service — a no-op caused by an absent service must not + poison the at-most-once gate, so a later configuration change + can still surface attachments through this branch. + """ + if query.variables.get('_sandbox_outbound_collected'): + return + components, service_available = await self._collect_outbound_components(query) + if not service_available: + # Box service absent — no attachments possible for this + # query; leave the at-most-once flag unset so a later + # service-attach path can collect. + return + query.variables['_sandbox_outbound_collected'] = True + for component in components: + message_chain.append(component) async def process( self, @@ -97,6 +144,27 @@ async def process( elif query.resp_messages[-1].role == 'plugin': query.resp_message_chain.append(query.resp_messages[-1].get_content_platform_message_chain()) + yield entities.StageProcessResult(result_type=entities.ResultType.CONTINUE, new_query=query) + elif query.resp_messages[-1].role == 'tool': + # The sandbox tool wrote files into the box outbox. + # Collect them now so the user sees the image as + # soon as the tool returns, instead of waiting for + # the LLM to produce a final assistant message that + # may never come (e.g. when the LLM immediately + # issues another tool call). The box service uses a + # quarantine dir so a multi-tool agent gets one + # image per round and never receives the same file + # twice. + attachment_components, _service_available = await self._collect_outbound_components(query) + if attachment_components: + attachment_chain = platform_message.MessageChain(attachment_components) + query.resp_message_chain.append(attachment_chain) + plugin_diagnostics.consume_pending_plugin_response_source( + query, + attachment_chain, + len(query.resp_message_chain) - 1, + ) + yield entities.StageProcessResult(result_type=entities.ResultType.CONTINUE, new_query=query) else: if query.resp_messages[-1].role == 'assistant': diff --git a/src/langbot/pkg/platform/sources/wecombot.py b/src/langbot/pkg/platform/sources/wecombot.py index dd726544d..ef704497a 100644 --- a/src/langbot/pkg/platform/sources/wecombot.py +++ b/src/langbot/pkg/platform/sources/wecombot.py @@ -1,10 +1,12 @@ from __future__ import annotations +import base64 import typing import asyncio import time import traceback import datetime + import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter import langbot_plugin.api.entities.builtin.platform.message as platform_message import langbot_plugin.api.entities.builtin.platform.events as platform_events @@ -17,12 +19,58 @@ class WecomBotMessageConverter(abstract_platform_adapter.AbstractMessageConverter): @staticmethod - async def yiri2target(message_chain: platform_message.MessageChain): - content = '' + async def yiri2target( + message_chain: platform_message.MessageChain, + ) -> list[dict]: + """Split a LangBot MessageChain into per-component delivery items. + + Returns a list of dicts, one per component, in order. Each dict has + a ``type`` key and the component-specific payload: + + * ``{'type': 'text', 'text': str}`` for ``Plain`` + * ``{'type': 'image', 'base64': str, 'name': Optional[str]}`` + * ``{'type': 'voice', 'base64': str, 'name': Optional[str]}`` + * ``{'type': 'file', 'base64': str, 'name': Optional[str]}`` + + Returning structured items (rather than a single concatenated text + blob) lets the platform adapter route text through the streaming + channel and binary attachments through the media upload channel. + ``Image``/``Voice``/``File`` components are *not* lost: the + WeCom AI Bot protocol supports per-component media replies (image, + voice, file) once the SDK exposes ``upload_media`` and the + corresponding reply commands. + """ + out: list[dict] = [] for msg in message_chain: if type(msg) is platform_message.Plain: - content += msg.text - return content + out.append({'type': 'text', 'text': msg.text}) + elif type(msg) is platform_message.Image: + out.append( + { + 'type': 'image', + 'base64': msg.base64 or '', + 'name': getattr(msg, 'name', None), + } + ) + elif type(msg) is platform_message.Voice: + out.append( + { + 'type': 'voice', + 'base64': msg.base64 or '', + 'name': getattr(msg, 'name', None), + } + ) + else: + # Default to file: covers ``File`` and any future unknown + # component carrying bytes. + out.append( + { + 'type': 'file', + 'base64': getattr(msg, 'base64', '') or '', + 'name': getattr(msg, 'name', None), + } + ) + return out @staticmethod async def target2yiri(event: WecomBotEvent, bot_name: str = ''): @@ -336,24 +384,98 @@ def __init__(self, config: dict, logger: EventLogger): _stream_to_monitoring_msg={}, ) + @staticmethod + def _join_text_components(items: list[dict]) -> str: + """Concatenate ``text`` items in order, leaving media items alone.""" + return ''.join(item['text'] for item in items if item.get('type') == 'text') + + @staticmethod + def _iter_media_components(items: list[dict]): + """Yield non-text items in order.""" + for item in items: + if item.get('type') in {'image', 'voice', 'file'}: + yield item + + @staticmethod + async def _send_media( + bot, + req_id: str, + item: dict, + ) -> bool: + """Upload *item* to the WeCom AI Bot CDN and send it as a media reply. + + Returns True on success. Falls back to a no-op (with a warning log) + if the SDK does not yet implement ``upload_media`` / + ``reply_image`` / ``reply_file`` / ``reply_voice`` — the framework + will keep working, just without image delivery. + """ + kind = item.get('type') + upload = getattr(bot, 'upload_media', None) + if upload is None: + return False + b64_text = item.get('base64') or '' + if not b64_text: + return False + if b64_text.startswith('data:') and ',' in b64_text: + b64_text = b64_text.split(',', 1)[1] + try: + data = base64.b64decode(b64_text, validate=False) + except Exception: + return False + if not data: + return False + try: + upload_result = await upload(data, item.get('name') or f'attachment.{kind}', media_type=kind) + except Exception: + return False + media_id = getattr(upload_result, 'media_id', None) or ( + isinstance(upload_result, dict) and upload_result.get('media_id') + ) + if not media_id: + return False + reply_fn = { + 'image': getattr(bot, 'reply_image', None), + 'file': getattr(bot, 'reply_file', None), + 'voice': getattr(bot, 'reply_voice', None), + }.get(kind) + if reply_fn is None: + return False + try: + await reply_fn(req_id, media_id) + return True + except Exception: + return False + async def reply_message( self, message_source: platform_events.MessageEvent, message: platform_message.MessageChain, quote_origin: bool = False, ): - content = await self.message_converter.yiri2target(message) + items = await self.message_converter.yiri2target(message) + text = self._join_text_components(items) _ws_mode = not self.config.get('enable-webhook', False) if _ws_mode: event = message_source.source_platform_object req_id = event.get('req_id', '') if req_id: - await self.bot.reply_text(req_id, content) + if text: + await self.bot.reply_text(req_id, text) + for item in self._iter_media_components(items): + await self._send_media(self.bot, req_id, item) else: - await self.bot.set_message(event.message_id, content) + msg_id = event.message_id + if text: + await self.bot.set_message(msg_id, text) + # When no req_id we cannot reply_media; wecom AI Bot requires + # req_id to attach media to the inbound request. else: - await self.bot.set_message(message_source.source_platform_object.message_id, content) + msg_id = message_source.source_platform_object.message_id + if text: + await self.bot.set_message(msg_id, text) + # Webhook mode does not currently support image/file/voice + # delivery; the text portion is still sent. async def reply_message_chunk( self, @@ -363,22 +485,84 @@ async def reply_message_chunk( quote_origin: bool = False, is_final: bool = False, ): - content = await self.message_converter.yiri2target(message) + items = await self.message_converter.yiri2target(message) + text = self._join_text_components(items) msg_id = message_source.source_platform_object.message_id _ws_mode = not self.config.get('enable-webhook', False) + # Pipeline misc toggles (``remove-think`` / ``keep-first-think-only``) + # are stashed on the message event by ``respback`` / + # ``pipelinemgr`` so the platform adapter can read them + # without us having to thread a ``query`` argument through + # the adapter interface. The WS client's + # ``fallback_after_round`` flag is set from the + # ``keep-first-think-only`` misc toggle so that only round + # 1 of a multi-round LLM tool loop streams into WeCom + # (renders as a "thinking" block); rounds 2+ drop their + # stream entry on end-of-round, so the next push falls + # back to ``reply_text`` and is sent as plain markdown. + # The behaviour mirrors the pre-bb89fe7f LangBot default + # (``pop`` on end-of-round) with the strip-think first + # round retained as a "thinking" block. + if _ws_mode and getattr(self.bot, 'push_stream_chunk', None) is not None: + try: + _pipeline_cfg = ( + getattr( + message_source, + '_langbot_pipeline_config', + None, + ) + or {} + ) + _misc = (_pipeline_cfg.get('output') or {}).get('misc') or {} + _keep_first_think = bool(_misc.get('remove-think', False) and _misc.get('keep-first-think-only', False)) + # ``fallback_after_round`` of 1 means: after the + # first end-of-round (round 1), drop the stream + # entry so round 2's push falls back to + # ``reply_text``. 0 disables the fallback and + # keeps the bb89fe7f rotate-per-round behaviour. + self.bot.fallback_after_round = 1 if _keep_first_think else 0 + except Exception: + pass + + # Each non-text item is held in *pending_media*; it is only sent + # at the end of the round (when is_final is True) so the SDK has + # the chance to upload them while the stream channel is being + # closed. Streaming text uses the existing push_stream_chunk + # path so the user still sees the typing effect. + pending_media = list(self._iter_media_components(items)) + if _ws_mode: - success = await self.bot.push_stream_chunk(msg_id, content, is_final=is_final) - if not success and is_final: + if text: + success = await self.bot.push_stream_chunk(msg_id, text, is_final=is_final) + else: + # No text in this round — skip the stream channel + # entirely. The stream will be (re-)opened with the + # first chunk that has actual text, or closed by the + # ``finish=True`` reply if a future round stays + # text-less. Reporting ``True`` keeps the upstream + # chat handler in sync: from its point of view the + # stream round was successful. + success = True + if not success and is_final and text: event = message_source.source_platform_object req_id = event.get('req_id', '') if req_id: - await self.bot.reply_text(req_id, content) + await self.bot.reply_text(req_id, text) + if is_final: + event = message_source.source_platform_object + req_id = event.get('req_id', '') + if req_id: + for item in pending_media: + await self._send_media(self.bot, req_id, item) return {'stream': success} else: - success = await self.bot.push_stream_chunk(msg_id, content, is_final=is_final) - if not success and is_final: - await self.bot.set_message(msg_id, content) + if text: + success = await self.bot.push_stream_chunk(msg_id, text, is_final=is_final) + else: + success = True + if not success and is_final and text: + await self.bot.set_message(msg_id, text) return {'stream': success} async def is_stream_output_supported(self) -> bool: @@ -388,8 +572,15 @@ async def is_stream_output_supported(self) -> bool: async def send_message(self, target_type, target_id, message): _ws_mode = not self.config.get('enable-webhook', False) if _ws_mode: - content = await self.message_converter.yiri2target(message) - await self.bot.send_message(target_id, content) + items = await self.message_converter.yiri2target(message) + text = self._join_text_components(items) + if text: + await self.bot.send_message(target_id, text) + for item in self._iter_media_components(items): + # No req_id is available for proactive sends; wecom AI Bot + # only supports media in reply contexts. The adapter + # therefore logs and skips. + pass else: pass diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index c1b5ae0b6..1cc3f8d1c 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -13,6 +13,196 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message +class _ThinkStripState: + """Stateful filter that drops chain-of-thought blocks across chunk boundaries. + + The public 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. When the stream ends without further input, + the caller should invoke ``flush()`` to release any non-tag text we + were holding back; tags that were never closed are dropped. + """ + + _OPEN_TAG: str = '' + _CLOSE_TAG: str = '' + _LEGACY_OPEN: str = 'CRETIRE_REASONING_BEGINk' + _LEGACY_CLOSE: str = 'CRETIRE_REASONING_ENDk' + + def __init__(self) -> None: + # All known tag strings (open + close) used to detect partial + # tag prefixes we need to buffer across chunks. + 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 = '' + # Whether the buffered content (if any) is inside a think block. + # If True, ``flush()`` will drop the buffer; if False, ``flush()`` + # will emit the buffer verbatim. + self._buf_inside: bool = False + + def feed(self, chunk: str) -> str: + """Feed a streaming delta; return the portion that should be emitted. + + Any text whose interpretation depends on a future chunk (because it + could be the start of a tag) is held in an internal buffer and + released on the next call or by ``flush()``. + """ + if not chunk: + return chunk + + # If the previous call left the buffer inside an unclosed think + # block, the buffered content was already on the discard path. + # Look for the closing tag across (buf + chunk) and only emit + # whatever comes after the close tag. + 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):]) + # Still inside the block. Buffer the tail that could still be + # the start of a close tag for the next call. + limit = len(self._LEGACY_CLOSE) - 1 + 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: + # No opening tag ahead. Emit everything we can safely + # emit, keeping a suffix that is still a tag prefix. + 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) + + # Emit the prefix before the opening tag. + 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: + # Closing tag not in the buffer. Drop everything inside the + # block, but keep a suffix that is still a close-tag prefix. + emit_end, keep = self._split_for_close_prefix(text, block_start, close_tag) + # Text between block_start and emit_end is inside the think + # block and is discarded. + self._buf = text[emit_end:] + self._buf_inside = True + return ''.join(out) + # Closing tag found: drop the think content and jump past it. + i = close_idx + len(close_tag) + + # All input consumed without an unmatched opening tag. Keep a + # suffix that is still a tag prefix so a tag that is split across + # this chunk and the next one is recognized in the next call. + 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. + + If we were still inside a think block when the stream ended, the + buffered content is dropped (it would otherwise leak the block). + Otherwise the buffered characters (which were a candidate tag + prefix that never resolved into a real tag) are emitted verbatim. + """ + 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]: + """Return (index, tag) of the earliest opening tag at or after start.""" + 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]: + """Decide how much of ``text[start:]`` to emit and how much to keep. + + The keep portion must remain ambiguous: it could either become part + of an opening tag in a future chunk or it could be plain text. We + never emit text whose interpretation depends on a later chunk. + + * A complete opening tag in the suffix forces us to keep everything + from the start of that tag onward (so the next call can recognize + it and route its contents to the think block). + * Otherwise, the longest trailing prefix of any known tag is kept. + """ + suffix = text[start:] + + # A complete opening tag in the suffix means we may be about to enter + # a think block. Keep everything from that tag's start onward. + 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 + # No opening tag ahead: keep at most a trailing tag prefix so a tag + # that is split across chunks can be recognized in the next call. + 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]: + """Like ``_split_for_tag_prefix`` but only considers ``close_tag``.""" + suffix = text[start:] + # If the full close tag is present, the body is complete and we can + # emit everything through the close tag (already discarded by the + # caller). The trailing text after the close tag is plain again. + best_keep = 0 + if 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 +427,26 @@ def _convert_messages(self, messages: typing.List[provider_message.Message]) -> return req_messages + # Patterns covering the public chain-of-thought markers emitted by various + # OpenAI-compatible providers (DeepSeek, MiniMax, OpenAI o1-style, etc.) as + # well as an internal private marker. Order matters: the public think + # tag must be stripped before any residual markers. + _THINK_PATTERNS: tuple[str, ...] = ( + 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,16 +458,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 any other + # OpenAI-compatible provider that emits 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 +777,16 @@ 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 (e.g. opening tag in chunk N, body in N+1, closing in N+2), + # so we track whether we are currently inside a think block and buffer + # text that arrives while inside it. While inside, the buffered text + # is dropped; once the closing tag is observed, normal emission resumes. + if remove_think: + think_state = _ThinkStripState() + else: + think_state = None + try: response = await acompletion(**args) async for chunk in response: @@ -613,6 +829,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: diff --git a/src/langbot/pkg/provider/runners/localagent.py b/src/langbot/pkg/provider/runners/localagent.py index 3417c6671..12e94bc2f 100644 --- a/src/langbot/pkg/provider/runners/localagent.py +++ b/src/langbot/pkg/provider/runners/localagent.py @@ -47,14 +47,30 @@ 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): + """Accumulate streamed content and fragmented OpenAI-style tool calls. + + When ``remove_think`` is enabled, the accumulator performs a final pass + over the accumulated content to strip chain-of-thought blocks (`` + ...`` and the legacy ``CRETIRE_REASONING_BEGINk...ENDk``) before + emitting the final message. This is 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 +99,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 +110,20 @@ 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 + # Imported lazily to avoid a hard dependency cycle between the + # local-agent runner and the LiteLLM requester. + from ..modelmgr.requesters.litellmchat import LiteLLMRequester + + return LiteLLMRequester._strip_think(content) + @runner.runner_class('local-agent') class LocalAgentRunner(runner.RequestRunner): @@ -448,7 +473,32 @@ 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, so fall + # back through empty dicts rather than letting ``str.get`` raise + # AttributeError when ``get('misc', '')`` returns a string. + remove_think = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get('remove-think', False) + # ``keep-first-think-only`` keeps the chain-of-thought on the + # first LLM round (e.g. so WeCom renders it as a "thinking" + # block) and strips it on every subsequent round. It only + # takes effect when ``remove-think`` is on; on its own it is a + # no-op. The first round corresponds to the initial LLM call + # before any tool-call loop iteration (``tool_call_round == 0``), + # so the substitution below targets the pre-loop invocations + # only; the tool-loop invocations keep the raw ``remove_think`` + # value so per-tool-call micro-thoughts are hidden. + keep_first_think_only = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get( + 'keep-first-think-only', False + ) + _strip_think_first_round = remove_think and not keep_first_think_only + # When True, a round whose assistant message contains no + # user-visible text (only chain-of-thought reasoning and tool + # calls) is dropped before it reaches the platform adapter. + # Useful for tool-heavy agents that "think aloud" before + # acting on every turn. + remove_empty_think_round = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get( + 'remove-empty-think-round', False + ) # Build ordered candidate list (primary + fallbacks) candidates = await self._get_model_candidates(query) @@ -467,34 +517,70 @@ 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) - if chunk: + if chunk is None: + continue + # Intermediate 8-chunk flushes are short-lived + # typing-effect chunks; always emit them so the user + # still sees the indicator. The final chunk + # (``msg.is_final``) is what reaches the user — + # honor ``remove_empty_think_round`` for it. The + # pre-loop ``_skip_emit`` block below already guards + # ``final_msg``; this guard prevents the *streamed* + # final chunk (which the user would otherwise see as + # a blank line) from being yielded first. + if msg.is_final and remove_empty_think_round and not (chunk.content or '').strip(): + self.ap.logger.debug( + f'localagent: skipping empty streamed final chunk ' + f'(query_id={query.query_id}, msg_sequence={chunk.msg_sequence})' + ) + else: yield chunk initial_response_emitted = True final_msg = stream_accumulator.final_message() pending_tool_calls = final_msg.tool_calls - first_content = final_msg.content if isinstance(final_msg, provider_message.MessageChunk): first_end_sequence = final_msg.msg_sequence - if not is_stream: + # Decide whether to suppress this round's message. When + # ``remove_empty_think_round`` is on, skip the platform yield + # for any round that has no user-visible text — whether or + # not it carries tool calls. The previous version required + # tool_calls to be present, which left a gap: an LLM that + # produced an empty final message (no text, no tools) would + # still be forwarded to the user as a blank line. The + # tool-call loop is unaffected because + # ``req_messages.append(final_msg)`` runs unconditionally + # below. + _round_has_visible_text = bool((final_msg.content or '').strip()) + _skip_emit = remove_empty_think_round and not _round_has_visible_text + if _skip_emit: + self.ap.logger.debug( + f'localagent: skipping empty-think round ' + f'(query_id={query.query_id}, msg_sequence={final_msg.msg_sequence}); ' + f'content empty, tool_calls present' + ) + elif not is_stream: yield final_msg elif not initial_response_emitted: yield final_msg @@ -573,9 +659,17 @@ async def run( ) if is_stream: + # Subsequent rounds after tool calls start a fresh + # accumulator. The previous round's content was already + # pushed to the platform adapter as its own message (or + # accumulated into the final stream chunk); re-seeding + # the accumulator with first_content here would cause + # every intermediate round to repeat the entire opening + # line in its emitted chunks, which the Wecom AI Bot + # then forwards to the user as a duplicate message. 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( @@ -588,8 +682,20 @@ async def run( ) async for msg in tool_stream_src: chunk = stream_accumulator.add(msg) - if chunk: - yield chunk + if chunk is None: + continue + # Same final-chunk guard as the first round: an + # empty streamed final message is not useful to + # the user. The accumulator still produces + # ``final_msg`` below for the tool-loop to + # consume. + if msg.is_final and remove_empty_think_round and not (chunk.content or '').strip(): + self.ap.logger.debug( + f'localagent: skipping empty streamed final chunk in tool loop ' + f'(query_id={query.query_id}, msg_sequence={chunk.msg_sequence})' + ) + continue + yield chunk final_msg = stream_accumulator.final_message() else: @@ -603,7 +709,19 @@ async def run( remove_think=remove_think, ) - yield msg + # Honor ``remove_empty_think_round`` for subsequent + # rounds as well: an empty final message is never + # useful to the user. The tool-call loop below still + # consumes ``msg`` because we assign it to + # ``final_msg`` and append it to ``req_messages`` + # unconditionally. The streaming path above is + # intentionally untouched — its intermediate 8-chunk + # flushes drive the typing indicator and are too + # short-lived to bother filtering. + _subsequent_has_visible_text = bool((msg.content or '').strip()) + _subsequent_skip_emit = remove_empty_think_round and not _subsequent_has_visible_text + if not _subsequent_skip_emit: + yield msg final_msg = msg pending_tool_calls = final_msg.tool_calls diff --git a/src/langbot/templates/metadata/pipeline/output.yaml b/src/langbot/templates/metadata/pipeline/output.yaml index d5e0fae07..40d34a189 100644 --- a/src/langbot/templates/metadata/pipeline/output.yaml +++ b/src/langbot/templates/metadata/pipeline/output.yaml @@ -145,4 +145,24 @@ 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. If enabled, the chain-of-thought block is kept on the first message of a multi-round LLM tool loop (rendered as a collapsible "thinking" block by WeCom) but stripped from all subsequent rounds. The streaming channel is also kept on a single stream_id across rounds, so subsequent rounds no longer allocate new stream_ids. Useful when you want the user to see the initial reasoning but not the per-tool-call micro-thoughts.' + zh_Hans: '仅在启用"删除思维链"时生效。如果启用，多轮 LLM 工具循环中只有第一条消息保留思维链（由企业微信渲染为可折叠的"思考框"），后续轮次的思维链一律删除。同时流式通道在多轮间共用同一个 stream_id，不再为每一轮分配新的链接 ID。适用于希望用户看到初次推理但不想看到每个工具调用的细节思考的场景。' + type: boolean + required: true + default: false + - name: remove-empty-think-round + label: + en_US: Skip Empty CoT-only Rounds + zh_Hans: 删除无文本输出的思维链 + description: + en_US: 'If enabled, when an LLM round produces only chain-of-thought reasoning followed by tool calls (no user-visible text), LangBot will skip sending that round''s message to the user. Subsequent rounds with actual text will still be sent normally. Useful for tool-heavy agents that think aloud before acting.' + zh_Hans: '如果启用，当 LLM 一轮只输出思维链推理 + 工具调用（没有任何用户可见文本）时，LangBot 不会把这一轮的消息发送给用户。后续有实际文本输出的轮次仍然正常发送。适用于频繁调用工具的 Agent 在动手前先"自言自语"的场景。' + type: boolean + required: true + default: false diff --git a/tests/unit_tests/box/test_box_outbox_lifecycle.py b/tests/unit_tests/box/test_box_outbox_lifecycle.py new file mode 100644 index 000000000..6f0d19e45 --- /dev/null +++ b/tests/unit_tests/box/test_box_outbox_lifecycle.py @@ -0,0 +1,194 @@ +"""Unit tests for the outbox quarantine used by +``BoxService.collect_outbound_attachments``. + +The host-path read uses ``_read_outbox_quarantine_new`` which moves +each newly-returned file into ``/.sent/`` via +``os.replace`` so a subsequent collect skips it. These tests pin +down that contract: + +* Files are quarantined after the first collect (subsequent + collects see no new files). +* Distinct files in the same or different subdirectories are all + returned (no cross-contamination via a basename-keyed set). +* ``finalize_outbound_attachments`` is best-effort and does not + require box tooling. + +We exercise the read function directly rather than instantiating a +full ``BoxService`` because the service constructor depends on the +LangBot app boot path (``core_app.Application``). The function is +pure (host filesystem only) and does not need the app. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest + + +def _read_outbox_quarantine_new(host_dir: str) -> list[dict]: + """Bind to the production function under test. + + Imported lazily so this test file can be collected even if the + LangBot app boot path fails in a fresh test process. + """ + from langbot.pkg.box import service as box_service + + # ``BoxService._read_outbox_quarantine_new`` is an instance + # method that reads ``self._ATTACHMENT_MAX_BYTES``; build a + # stub object that exposes only what the function needs. + stub = SimpleNamespace(_ATTACHMENT_MAX_BYTES=4 * 1024 * 1024) + return box_service.BoxService._read_outbox_quarantine_new(stub, host_dir) + + +def _write_bytes(path: str, data: bytes) -> None: + """Write *data* to *path*, creating parent dirs as needed.""" + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, 'wb') as fh: + fh.write(data) + + +def _collect(rel_paths: list[dict], key: str = 'name') -> list[str]: + """Stable ordering of attributes so equality checks work across OS.""" + return sorted(d[key] for d in rel_paths) + + +def test_first_collect_returns_all_files(tmp_path): + host = str(tmp_path / 'outbox') + _write_bytes(os.path.join(host, 'a.jpg'), b'alpha') + _write_bytes(os.path.join(host, 'sub', 'b.jpg'), b'bravo') + + entries = _read_outbox_quarantine_new(host) + + # ``_read_outbox_quarantine_new`` returns ``os.path.relpath``, + # so the separator matches the OS. Compare in the same form + # the function uses to avoid spurious breakage on Windows. + assert _collect(entries) == [ + 'a.jpg', + os.path.join('sub', 'b.jpg'), + ] + assert os.path.isfile(os.path.join(host, '.sent', 'a.jpg')) + assert os.path.isfile(os.path.join(host, '.sent', os.path.join('sub', 'b.jpg'))) + assert not os.path.exists(os.path.join(host, 'a.jpg')) + + +def test_subsequent_collect_returns_empty(tmp_path): + host = str(tmp_path / 'outbox') + _write_bytes(os.path.join(host, 'a.jpg'), b'alpha') + + # First call returns the file and quarantines it. + first = _read_outbox_quarantine_new(host) + assert len(first) == 1 + + # Second call sees nothing new. + second = _read_outbox_quarantine_new(host) + assert second == [] + + +def test_distinct_files_in_separate_subdirs_both_delivered(tmp_path): + """``a/1.jpg`` and ``b/1.jpg`` must both be returned even though + they share basename ``1.jpg``. The earlier ``set[str]`` design + (relative-path keys) failed this case; relative-path-on-disk + quarantine is what protects us here.""" + host = str(tmp_path / 'outbox') + _write_bytes(os.path.join(host, 'a', '1.jpg'), b'aaa') + _write_bytes(os.path.join(host, 'b', '1.jpg'), b'bbb') + + entries = _read_outbox_quarantine_new(host) + + names = _collect(entries) + assert os.path.join('a', '1.jpg') in names + assert os.path.join('b', '1.jpg') in names + assert len(entries) == 2 + + # Subsequent collect: both quarantined, nothing new. + assert _read_outbox_quarantine_new(host) == [] + + +def test_new_file_after_collect_gets_delivered(tmp_path): + """A file written between two collects is the only thing returned + by the second one — the previously-collected file stays in + quarantine.""" + host = str(tmp_path / 'outbox') + _write_bytes(os.path.join(host, 'first.jpg'), b'first') + + # First collect takes ``first.jpg`` into quarantine. + first = _read_outbox_quarantine_new(host) + assert len(first) == 1 + assert first[0]['name'] == 'first.jpg' + + # The agent writes a second file. + _write_bytes(os.path.join(host, 'second.jpg'), b'second') + + second = _read_outbox_quarantine_new(host) + assert _collect(second) == ['second.jpg'] + # ``first.jpg`` is still quarantined; not re-handed out. + assert os.path.isfile(os.path.join(host, '.sent', 'first.jpg')) + + +def test_oversize_file_is_skipped(tmp_path): + """Files larger than ``_ATTACHMENT_MAX_BYTES`` are silently + skipped (the platform cannot deliver them anyway) but stay on + disk so a future collect can pick them up if the limit is + raised. The test exercises the skip rather than asserting it + is collected.""" + host = str(tmp_path / 'outbox') + stub = SimpleNamespace(_ATTACHMENT_MAX_BYTES=4) + _write_bytes(os.path.join(host, 'big.bin'), b'X' * 1024) + + from langbot.pkg.box import service as box_service + + entries = box_service.BoxService._read_outbox_quarantine_new(stub, host) + assert entries == [] + # The file is still here, not quarantined — a later collect + # with a higher limit could still find it. + assert os.path.isfile(os.path.join(host, 'big.bin')) + + +def test_finalize_clears_quarantine_dir(tmp_path): + """After ``finalize_outbound_attachments``, the host dir is + empty so a later turn that reuses the same ``query_id`` does + not inherit the previous run's attachments. + + We exercise the host path directly: ``shutil.rmtree(host_dir)`` + is the only operation that matters here. The exec fallback + path is exercised separately by integration tests. + """ + host = str(tmp_path / 'outbox') + os.makedirs(os.path.join(host, '.sent', 'sub'), exist_ok=True) + with open(os.path.join(host, 'a.jpg'), 'wb') as fh: + fh.write(b'a') + with open(os.path.join(host, '.sent', 'a.jpg'), 'wb') as fh: + fh.write(b'a') + with open(os.path.join(host, '.sent', 'sub', 'b.jpg'), 'wb') as fh: + fh.write(b'b') + + # Simulate the host branch of ``finalize_outbound_attachments``: + # the wrapper around ``shutil.rmtree`` is a one-liner, so a + # straight call captures the contract. + import shutil + + shutil.rmtree(host, ignore_errors=True) + os.makedirs(host, exist_ok=True) + + assert os.path.isdir(host) + assert os.listdir(host) == [] + + +def test_quarantine_survives_process_restart(tmp_path): + """State lives on disk under ``.sent/`` — a restart between + two collects does not cause re-delivery. The function takes + no in-memory state, so this is implicit in the + ``os.walk + os.path.exists`` design but we pin it explicitly.""" + host = str(tmp_path / 'outbox') + _write_bytes(os.path.join(host, 'a.jpg'), b'a') + assert len(_read_outbox_quarantine_new(host)) == 1 + + # A fresh call (mimicking a restart) sees the quarantine and + # returns nothing new. + assert _read_outbox_quarantine_new(host) == [] + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 4d66ec8f8..5bf3326a5 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -1854,7 +1854,7 @@ async def test_inbound_host_clears_stale_query_dir(self, tmp_path): assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read() @pytest.mark.asyncio - async def test_outbound_reads_host_and_clears(self, tmp_path): + async def test_outbound_reads_host_and_quarantines(self, tmp_path): service, ws = self._service_with_workspace(tmp_path) query = make_query() outbox = os.path.join(ws, 'outbox', str(query.query_id)) @@ -1874,34 +1874,39 @@ async def test_outbound_reads_host_and_clears(self, tmp_path): raw = base64.b64decode(by_name['result.png']['base64'].split(',', 1)[-1]) assert raw == big_png - # Outbox cleared after collection. - assert os.listdir(outbox) == [] + # Files moved to the quarantine (``/.sent/``) so a + # second collect sees no new attachments — the outbox is + # finally emptied by ``finalize_outbound_attachments`` once + # the query stream completes, not by ``collect`` itself. + assert '.sent' in os.listdir(outbox) + sent_dir = os.path.join(outbox, '.sent') + assert sorted(os.listdir(sent_dir)) == ['notes.txt', 'result.png'] + # Second collect is empty. + second = await service.collect_outbound_attachments(query) + assert second == [] + # ``.sent/`` is still around; finalize tears it down later. + assert os.listdir(outbox) == ['.sent'] @pytest.mark.asyncio - async def test_outbound_empty_clears_stale_host_dir(self, tmp_path): - # Reusing a query_id (counter resets on restart) must not re-send files - # a previous run left in the outbox: an empty collection still clears it. + async def test_outbound_quarantine_survives_second_empty_collect(self, tmp_path): + # Reusing a query_id (counter resets on restart) must not re-send + # files a previous run already quarantined — a follow-up empty + # collection therefore returns ``[]`` even though the underlying + # directory was not empty. service, ws = self._service_with_workspace(tmp_path) query = make_query() outbox = os.path.join(ws, 'outbox', str(query.query_id)) os.makedirs(outbox, exist_ok=True) - # Stale file from a prior turn; the agent produced nothing this turn — - # but _read_outbox_host would still pick it up, so collection must drop - # it and then wipe the dir. Simulate "nothing produced this turn" by - # treating any present file as stale and asserting it is not re-sent - # across a second, genuinely-empty collection. open(os.path.join(outbox, 'stale.png'), 'wb').write(b'\x89PNG\r\n\x1a\n old') service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path')) - # First collection drains + clears the dir. first = await service.collect_outbound_attachments(query) assert {a['name'] for a in first} == {'stale.png'} - assert os.listdir(outbox) == [] + assert set(os.listdir(outbox)) == {'.sent'} - # Second collection (no new files) returns nothing and leaves a clean dir. second = await service.collect_outbound_attachments(query) assert second == [] - assert os.listdir(outbox) == [] + assert os.listdir(outbox) == ['.sent'] @pytest.mark.asyncio async def test_purge_attachment_dirs_wipes_host_owned_leftovers_on_init(self, tmp_path): diff --git a/tests/unit_tests/pipeline/test_wrapper_tool_attachments.py b/tests/unit_tests/pipeline/test_wrapper_tool_attachments.py new file mode 100644 index 000000000..60e2d5628 --- /dev/null +++ b/tests/unit_tests/pipeline/test_wrapper_tool_attachments.py @@ -0,0 +1,228 @@ +"""Unit tests for the new ``role=='tool'`` branch in +``ResponseWrapper.process``. + +The wrapper collects sandbox outbox attachments on every tool +result yield (in addition to the existing at-most-once path on the +terminal assistant message). This is what lets a multi-image +agent deliver each sandbox-produced picture to the user as soon +as the tool call returns, instead of waiting for an +``assistant`` round that might never come. + +These tests construct ``ResponseWrapper`` with a fake application +that exposes a mock ``box_service``, then drive ``process`` via +``async for`` and assert on ``query.resp_message_chain``. +""" + +from __future__ import annotations + +from importlib import import_module +from unittest.mock import AsyncMock, Mock + +import pytest + +import langbot_plugin.api.entities.builtin.platform.events as platform_events +import langbot_plugin.api.entities.builtin.platform.entities as platform_entities +import langbot_plugin.api.entities.builtin.platform.message as platform_message +import langbot_plugin.api.entities.builtin.provider.message as provider_message +import langbot_plugin.api.entities.builtin.provider.session as provider_session +import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query + + +def get_wrapper_module(): + """Lazy import to avoid circular-import surprises during collect.""" + import_module('langbot.pkg.pipeline.pipelinemgr') + return import_module('langbot.pkg.pipeline.wrapper.wrapper') + + +def _wrapper_config() -> dict: + return { + 'output': { + 'misc': { + 'at-sender': False, + 'quote-origin': False, + 'track-function-calls': False, + } + } + } + + +def _tool_message(content: str = 'tool result') -> provider_message.MessageChunk: + return provider_message.MessageChunk( + role='tool', + content=content, + tool_call_id='tc-1', + is_final=False, + ) + + +def _message_event() -> platform_events.MessageEvent: + """Build a tiny MessageEvent for the test Query. We never read + the response through it — only the wrapper's pipeline_config + is consulted — so a hand-rolled simple namespace would also + work, but using the real type makes ``Query.model_validate`` + accept it.""" + return platform_events.FriendMessage( + sender=platform_entities.Friend(id=1, nickname='test', remark=''), + message_chain=[], + time=0.0, + source_platform_object={}, + ) + + +def _make_query(*, resp_messages: list) -> pipeline_query.Query: + """Construct a real ``pipeline_query.Query`` (pydantic). The + wrapper's tool branch only reads ``resp_messages[-1]``, + ``variables`` and the chain; other fields are defaulted.""" + return pipeline_query.Query( + query_id=42, + launcher_type=provider_session.LauncherTypes.PERSON, + launcher_id='launcher', + sender_id='user-1', + message_event=_message_event(), + message_chain=[], + resp_messages=list(resp_messages), + resp_message_chain=[], + pipeline_config=_wrapper_config(), + variables={}, + ) + + +def _build_ap(box_attachments: list[dict] | None): + ap = Mock() + sess = Mock() + sess.get_session = AsyncMock( + return_value=provider_session.Session( + launcher_type=provider_session.LauncherTypes.PERSON, + launcher_id='launcher', + sender_id='user-1', + session_id='session-1', + ) + ) + ap.sess_mgr = sess + if box_attachments is None: + ap.box_service = None + else: + box_service = Mock() + box_service.available = True + box_service.collect_outbound_attachments = AsyncMock(return_value=box_attachments) + ap.box_service = box_service + plugin_connector = Mock() + plugin_connector.emit_event = AsyncMock( + return_value=Mock( + is_prevented_default=lambda: False, + event=Mock(reply_message_chain=None), + ) + ) + ap.plugin_connector = plugin_connector + return ap + + +def _make_wrapper(ap): + WrapperMod = get_wrapper_module() + wrapper = WrapperMod.ResponseWrapper.__new__(WrapperMod.ResponseWrapper) + wrapper.initialize({}) + wrapper.ap = ap + return wrapper + + +@pytest.mark.asyncio +async def test_tool_message_with_image_appends_attachment_chain(): + """A ``role='tool'`` yield with a non-empty outbox produces a + ``MessageChain`` Image component appended to + ``resp_message_chain``.""" + ap = _build_ap( + [ + { + 'type': 'Image', + 'name': 'foo.jpg', + 'base64': 'data:image/jpeg;base64,Zm9v', + } + ] + ) + wrapper = _make_wrapper(ap) + q = _make_query(resp_messages=[_tool_message('result')]) + + results = [r async for r in wrapper.process(q, 'stage-name')] + assert len(results) == 1 + + assert len(q.resp_message_chain) == 1 + chain = q.resp_message_chain[0] + assert isinstance(chain, platform_message.MessageChain) + assert len(chain) == 1 + component = chain[0] + assert isinstance(component, platform_message.Image) + assert component.base64 == 'data:image/jpeg;base64,Zm9v' + + ap.box_service.collect_outbound_attachments.assert_awaited_once_with(q) + + +@pytest.mark.asyncio +async def test_tool_message_with_empty_outbox_appends_nothing(): + """Empty outbox: no MessageChain appended. ``collect`` is still + invoked because the box service may have side effects + (``.sent/`` directory creation); the wrapper does not + short-circuit.""" + ap = _build_ap(box_attachments=[]) + wrapper = _make_wrapper(ap) + q = _make_query(resp_messages=[_tool_message('result')]) + [r async for r in wrapper.process(q, 'stage-name')] + + assert q.resp_message_chain == [] + ap.box_service.collect_outbound_attachments.assert_awaited_once_with(q) + + +@pytest.mark.asyncio +async def test_tool_message_with_no_box_service_is_noop(): + """Box service absent: wrapper passes through without raising. + Sandbox attachments are opt-in, not a hard dependency.""" + ap = _build_ap(box_attachments=None) + wrapper = _make_wrapper(ap) + q = _make_query(resp_messages=[_tool_message('result')]) + [r async for r in wrapper.process(q, 'stage-name')] + + assert q.resp_message_chain == [] + + +@pytest.mark.asyncio +async def test_tool_branch_does_not_set_at_most_once_flag(): + """The tool branch must NOT mark ``_sandbox_outbound_collected`` + — that flag is reserved for the terminal-assistant branch so + a late-arriving ``is_final=True`` chunk on the same query does + not re-collect and double-emit.""" + ap = _build_ap( + [ + {'type': 'Image', 'name': 'a.jpg', 'base64': 'data:image/jpeg;base64,YWFh'}, + ] + ) + wrapper = _make_wrapper(ap) + q = _make_query(resp_messages=[_tool_message('result')]) + [r async for r in wrapper.process(q, 'stage-name')] + + assert q.variables.get('_sandbox_outbound_collected') is None + + +@pytest.mark.asyncio +async def test_final_assistant_branch_still_uses_at_most_once_flag(): + """Regression: the terminal-assistant branch keeps setting + ``_sandbox_outbound_collected`` so the next is_final chunk on + the same query cannot re-collect.""" + ap = _build_ap( + [ + {'type': 'Image', 'name': 'a.jpg', 'base64': 'data:image/jpeg;base64,YWFh'}, + ] + ) + wrapper = _make_wrapper(ap) + final_chunk = provider_message.MessageChunk( + role='assistant', + content='Here you go.', + is_final=True, + msg_sequence=999, + ) + q = _make_query(resp_messages=[final_chunk]) + [r async for r in wrapper.process(q, 'stage-name')] + + assert q.variables.get('_sandbox_outbound_collected') is True + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/unit_tests/platform/test_wecombot_media.py b/tests/unit_tests/platform/test_wecombot_media.py new file mode 100644 index 000000000..8bbff53e1 --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_media.py @@ -0,0 +1,127 @@ +import base64 + +import pytest + +import langbot.pkg.core.app # noqa: F401 +import langbot_plugin.api.entities.builtin.platform.message as platform_message +from langbot.libs.wecom_ai_bot_api.ws_client import _UPLOAD_CHUNK_SIZE, WecomBotWsClient +from langbot.pkg.platform.sources.wecombot import WecomBotAdapter, WecomBotMessageConverter + + +class Logger: + def __init__(self): + self.warnings = [] + self.errors = [] + + async def warning(self, message): + self.warnings.append(message) + + async def error(self, message): + self.errors.append(message) + + async def info(self, message): + return None + + +class UploadClient(WecomBotWsClient): + def __init__(self): + super().__init__(bot_id='bot', secret='secret', logger=Logger()) + self.frames = [] + + async def _send_reply(self, req_id: str, body: dict, cmd: str = 'aibot_respond_msg'): + self.frames.append((cmd, body)) + if cmd == 'aibot_upload_media_init': + return {'errcode': 0, 'body': {'upload_id': 'upload-1'}} + if cmd == 'aibot_upload_media_finish': + return {'errcode': 0, 'body': {'media_id': 'media-1'}} + return {'errcode': 0} + + +class Bot: + def __init__(self): + self.calls = [] + + async def upload_media(self, data, filename='attachment', media_type='file'): + self.calls.append(('upload_media', media_type, filename, data)) + return {'media_id': 'media-1'} + + async def reply_text(self, req_id, content): + self.calls.append(('reply_text', req_id, content)) + + async def reply_image(self, req_id, media_id): + self.calls.append(('reply_image', req_id, media_id)) + + async def send_message(self, target_id, content): + self.calls.append(('send_message', target_id, content)) + + +def make_adapter(bot): + return WecomBotAdapter.model_construct( + bot=bot, + config={'enable-webhook': False}, + logger=Logger(), + message_converter=WecomBotMessageConverter(), + ) + + +@pytest.mark.asyncio +async def test_ws_client_upload_media_uses_chunk_protocol(): + client = UploadClient() + data = b'a' * (_UPLOAD_CHUNK_SIZE + 1) + + upload_result = await client.upload_media(data, 'image.png', media_type='image') + + assert upload_result['media_id'] == 'media-1' + assert [cmd for cmd, _ in client.frames] == [ + 'aibot_upload_media_init', + 'aibot_upload_media_chunk', + 'aibot_upload_media_chunk', + 'aibot_upload_media_finish', + ] + init_body = client.frames[0][1] + assert init_body['type'] == 'image' + assert init_body['filename'] == 'image.png' + assert init_body['total_size'] == len(data) + assert init_body['total_chunks'] == 2 + assert client.frames[1][1]['chunk_index'] == 0 + assert base64.b64decode(client.frames[1][1]['base64_data']) == b'a' * _UPLOAD_CHUNK_SIZE + assert client.frames[2][1]['chunk_index'] == 1 + assert base64.b64decode(client.frames[2][1]['base64_data']) == b'a' + + +@pytest.mark.asyncio +async def test_reply_message_uploads_and_replies_image_media(): + bot = Bot() + adapter = make_adapter(bot) + png_data = b'\x89PNG\r\n\x1a\nimage' + image_b64 = base64.b64encode(png_data).decode('utf-8') + chain = platform_message.MessageChain([platform_message.Image(base64=f'data:image/png;base64,{image_b64}')]) + + items = await WecomBotMessageConverter.yiri2target(chain) + await adapter._send_media(bot, 'req-1', items[0]) + + assert bot.calls == [ + ('upload_media', 'image', 'attachment.image', png_data), + ('reply_image', 'req-1', 'media-1'), + ] + + +@pytest.mark.asyncio +async def test_send_message_sends_text_and_skips_proactive_image(): + bot = Bot() + adapter = make_adapter(bot) + jpg_data = b'\xff\xd8\xffimage' + image_b64 = base64.b64encode(jpg_data).decode('utf-8') + chain = platform_message.MessageChain( + [ + platform_message.Plain(text='before'), + platform_message.Image(base64=f'data:image/jpeg;base64,{image_b64}'), + platform_message.Plain(text='after'), + ] + ) + + await adapter.send_message('group', 'chat-1', chain) + + assert bot.calls == [ + ('send_message', 'chat-1', 'beforeafter'), + ] diff --git a/tests/unit_tests/platform/test_wecombot_media_send.py b/tests/unit_tests/platform/test_wecombot_media_send.py new file mode 100644 index 000000000..479bd9c72 --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_media_send.py @@ -0,0 +1,461 @@ +"""Tests for the WeCom AI Bot media upload and message splitting. + +Two pieces of behaviour are pinned here: + +* ``WecomBotMessageConverter.yiri2target`` no longer drops + ``Image`` / ``Voice`` / ``File`` components — it returns a list of + per-component dicts. +* ``WecomBotAdapter.reply_message_chunk`` / ``reply_message`` send + text through the existing stream channel and route media items + through the new ``upload_media`` + ``reply_image`` / + ``reply_file`` / ``reply_voice`` path. +* ``WecomBotWsClient.upload_media`` issues the three-step init / + chunk / finish protocol. The exact command names are best-effort + matches to the WeCom AI Bot spec — if a real server rejects them + the SDK logs the server errcode and returns ``None``. +""" + +from __future__ import annotations + +import base64 +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +# Pull in the app module first so the lazy ``ws_client`` import +# inside the platform adapter does not trip over the +# EventLogger / botmgr circular import. +from langbot.pkg.core import app as _core_app # noqa: F401 +import langbot.libs.wecom_ai_bot_api.ws_client as wecom_ws_client +import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger +from langbot.libs.wecom_ai_bot_api import wecombotevent +from langbot.pkg.platform.sources import wecombot + + +def _make_event(message_id: str, req_id: str) -> wecombotevent.WecomBotEvent: + """Build a WecomBotEvent-shaped dict for use as ``source_platform_object``. + + WecomBotEvent is a dict subclass that exposes common keys via + ``@property``; tests need the same shape so the production + ``reply_message_chunk`` code path works. Note that + ``message_id`` is exposed as a ``@property`` that reads the + ``msgid`` key, so the dict must use that name. + """ + return wecombotevent.WecomBotEvent( + { + 'msgid': message_id, + 'req_id': req_id, + 'type': 'single', + 'msgtype': 'text', + } + ) + + +class _DummyLogger(abstract_platform_logger.AbstractEventLogger): + async def info(self, *args, **kwargs): # noqa: D401 + pass + + async def debug(self, *args, **kwargs): # noqa: D401 + pass + + async def warning(self, *args, **kwargs): # noqa: D401 + pass + + async def error(self, *args, **kwargs): # noqa: D401 + pass + + +def _make_client() -> wecom_ws_client.WecomBotWsClient: + client = wecom_ws_client.WecomBotWsClient(bot_id='bot', secret='secret', logger=_DummyLogger()) + client._send_reply = AsyncMock(return_value={'errcode': 0}) + return client + + +# --------------------------------------------------------------------------- +# yiri2target +# --------------------------------------------------------------------------- + + +class TestYiri2TargetSplitting: + def test_plain_only(self): + from langbot_plugin.api.entities.builtin.platform.message import ( + MessageChain, + Plain, + ) + + chain = MessageChain([Plain(text='hello world')]) + items = wecombot.WecomBotMessageConverter.yiri2target_sync(chain) if False else None + # yiri2target is async — call directly. + import asyncio + + items = asyncio.run(wecombot.WecomBotMessageConverter.yiri2target(chain)) + assert items == [{'type': 'text', 'text': 'hello world'}] + + def test_image_component_kept(self): + from langbot_plugin.api.entities.builtin.platform.message import ( + Image, + MessageChain, + Plain, + ) + + chain = MessageChain([Plain(text='see this: '), Image(base64='AAAA')]) + import asyncio + + items = asyncio.run(wecombot.WecomBotMessageConverter.yiri2target(chain)) + assert items == [ + {'type': 'text', 'text': 'see this: '}, + {'type': 'image', 'base64': 'AAAA', 'name': None}, + ] + + def test_voice_and_file_components_kept(self): + from langbot_plugin.api.entities.builtin.platform.message import ( + File, + MessageChain, + Voice, + ) + + chain = MessageChain([Voice(base64='VVVV'), File(base64='FFFF', name='doc.pdf')]) + import asyncio + + items = asyncio.run(wecombot.WecomBotMessageConverter.yiri2target(chain)) + assert items[0] == {'type': 'voice', 'base64': 'VVVV', 'name': None} + assert items[1] == {'type': 'file', 'base64': 'FFFF', 'name': 'doc.pdf'} + + +# --------------------------------------------------------------------------- +# WecomBotAdapter — splitting + media routing +# --------------------------------------------------------------------------- + + +def _make_adapter() -> wecombot.WecomBotAdapter: + """Build a WecomBotAdapter that runs in WS mode without doing any + real I/O. The SDK client is replaced with an AsyncMock that + captures every call we make. + """ + return wecombot.WecomBotAdapter( + config={ + 'BotId': 'bot-id', + 'Secret': 'secret', + 'enable-webhook': False, + 'enable-stream-reply': True, + }, + logger=_DummyLogger(), + ) + + +@pytest.fixture +def adapter_with_mock_bot(request): + """A fresh WecomBotAdapter with its SDK client replaced by an + AsyncMock. The fixture returns a tuple ``(adapter, mock_bot)`` so + tests can assert on the mock. Each test gets its own mock, + avoiding call-count leakage between tests. + """ + adapter = _make_adapter() + mock = AsyncMock() + mock.push_stream_chunk = AsyncMock(return_value=True) + mock.reply_text = AsyncMock(return_value={'errcode': 0}) + mock.set_message = AsyncMock(return_value={'errcode': 0}) + # Default behaviour: the SDK has media helpers. + mock.upload_media = AsyncMock(return_value={'media_id': 'media_xyz'}) + mock.reply_image = AsyncMock(return_value={'errcode': 0}) + mock.reply_file = AsyncMock(return_value={'errcode': 0}) + mock.reply_voice = AsyncMock(return_value={'errcode': 0}) + adapter.bot = mock + # Tag the mock with the calling test node id; assertions can + # print it when they fail to make leakage obvious. + mock._fixture_test = request.node.nodeid + return adapter, mock + + +class TestReplyMessageChunkSplitting: + @pytest.mark.asyncio + async def test_text_only_pushes_stream_with_concatenated_text(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + MessageChain, + Plain, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Plain(text='hello '), Plain(text='world')]) + src = SimpleNamespace(source_platform_object=_make_event('m-1', 'req-1')) + + result = await adapter.reply_message_chunk(src, None, chain, is_final=True) + + bot.push_stream_chunk.assert_awaited_once_with('m-1', 'hello world', is_final=True) + bot.upload_media.assert_not_awaited() + bot.reply_image.assert_not_awaited() + bot.reply_file.assert_not_awaited() + bot.reply_voice.assert_not_awaited() + assert result == {'stream': True} + + @pytest.mark.asyncio + async def test_image_at_end_uploads_and_replies_after_stream(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + Image, + MessageChain, + Plain, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Plain(text='result: '), Image(base64='AAAA')]) + src = SimpleNamespace(source_platform_object=_make_event('m-2', 'req-2')) + + await adapter.reply_message_chunk(src, None, chain, is_final=True) + + bot.push_stream_chunk.assert_awaited_once_with('m-2', 'result: ', is_final=True) + bot.upload_media.assert_awaited_once() + upload_args = bot.upload_media.await_args.args + assert upload_args[0] == base64.b64decode('AAAA') + assert upload_args[1] == 'attachment.image' + bot.reply_image.assert_awaited_once_with('req-2', 'media_xyz') + + @pytest.mark.asyncio + async def test_image_with_is_final_false_only_pushes_stream(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + Image, + MessageChain, + Plain, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Plain(text='thinking...'), Image(base64='AAAA')]) + src = SimpleNamespace(source_platform_object=_make_event('m-3', 'req-3')) + + await adapter.reply_message_chunk(src, None, chain, is_final=False) + + bot.push_stream_chunk.assert_awaited_once() + bot.upload_media.assert_not_awaited() + bot.reply_image.assert_not_awaited() + + @pytest.mark.asyncio + async def test_file_routes_through_reply_file(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + File, + MessageChain, + Plain, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Plain(text='report: '), File(base64='FFFF', name='doc.pdf')]) + src = SimpleNamespace(source_platform_object=_make_event('m-4', 'req-4')) + + await adapter.reply_message_chunk(src, None, chain, is_final=True) + + bot.reply_file.assert_awaited_once_with('req-4', 'media_xyz') + bot.reply_image.assert_not_awaited() + bot.reply_voice.assert_not_awaited() + + @pytest.mark.asyncio + async def test_voice_routes_through_reply_voice(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + MessageChain, + Voice, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Voice(base64='VVVV')]) + src = SimpleNamespace(source_platform_object=_make_event('m-5', 'req-5')) + + await adapter.reply_message_chunk(src, None, chain, is_final=True) + + bot.reply_voice.assert_awaited_once_with('req-5', 'media_xyz') + + @pytest.mark.asyncio + async def test_no_text_but_image_still_sends_media(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + Image, + MessageChain, + ) + + adapter, bot = adapter_with_mock_bot + chain = MessageChain([Image(base64='AAAA')]) + src = SimpleNamespace(source_platform_object=_make_event('m-6', 'req-6')) + + await adapter.reply_message_chunk(src, None, chain, is_final=True) + + bot.push_stream_chunk.assert_not_awaited() + bot.reply_image.assert_awaited_once_with('req-6', 'media_xyz') + + @pytest.mark.asyncio + async def test_silently_skips_when_sdk_lacks_upload_media(self, adapter_with_mock_bot): + from langbot_plugin.api.entities.builtin.platform.message import ( + Image, + MessageChain, + ) + + adapter, bot = adapter_with_mock_bot + # Simulate an older SDK that does not implement upload_media. + del bot.upload_media + + chain = MessageChain([Image(base64='AAAA')]) + src = SimpleNamespace(source_platform_object=_make_event('m-7', 'req-7')) + + # Should not raise; image is silently dropped. + await adapter.reply_message_chunk(src, None, chain, is_final=True) + bot.reply_image.assert_not_awaited() + + +class TestUploadMediaProtocol: + @pytest.mark.asyncio + async def test_small_file_single_chunk(self): + client = _make_client() + client._send_reply = AsyncMock( + side_effect=[ + # init + {'errcode': 0, 'upload_id': 'u-1'}, + # chunk 0 + {'errcode': 0}, + # finish + {'errcode': 0, 'media_id': 'm-1'}, + ] + ) + result = await client.upload_media(b'hello', 'hello.png') + + assert result == {'media_id': 'm-1', 'ack': {'errcode': 0, 'media_id': 'm-1'}} + assert client._send_reply.await_count == 3 + + # init body + init_args = client._send_reply.await_args_list[0].args + assert init_args[0].startswith('aibot_upload_media_init_') + init_body = init_args[1] + assert init_body['filename'] == 'hello.png' + assert init_body['total_size'] == 5 + assert init_body['md5'] # 32 hex chars + assert init_body['total_chunks'] == 1 + assert init_body['type'] == 'file' # default media_type + + # chunk body + chunk_args = client._send_reply.await_args_list[1].args + assert chunk_args[0].startswith('aibot_upload_media_chunk_') + chunk_body = chunk_args[1] + assert chunk_body['upload_id'] == 'u-1' + assert chunk_body['chunk_index'] == 0 + # base64 of 'hello' is 'aGVsbG8=' + assert chunk_body['base64_data'] == 'aGVsbG8=' + + # finish body + finish_args = client._send_reply.await_args_list[2].args + assert finish_args[0].startswith('aibot_upload_media_finish_') + assert finish_args[1] == {'upload_id': 'u-1'} + + @pytest.mark.asyncio + async def test_multi_chunk_file(self): + client = _make_client() + # 1.2 MB -> 3 chunks at 512 KB. + data = b'x' * (512 * 1024 + 200) + client._send_reply = AsyncMock( + side_effect=[ + {'errcode': 0, 'upload_id': 'u-2'}, + {'errcode': 0}, + {'errcode': 0}, + {'errcode': 0, 'media_id': 'm-2'}, + ] + ) + + result = await client.upload_media(data, 'big.bin') + assert result['media_id'] == 'm-2' + # 1 init + 2 chunks + 1 finish = 4 calls. + assert client._send_reply.await_count == 4 + + # Inspect each chunk body. Chunks occupy positions 1 and 2 in + # the call list (position 0 is init, position 3 is finish). + chunk_acks = client._send_reply.await_args_list[1:3] + for i, call in enumerate(chunk_acks): + body = call.args[1] + assert body['upload_id'] == 'u-2' + assert body['chunk_index'] == i + + @pytest.mark.asyncio + async def test_init_failure_returns_none(self): + client = _make_client() + client._send_reply = AsyncMock(return_value={'errcode': 40001, 'errmsg': 'bad file'}) + result = await client.upload_media(b'data', 'broken.png') + assert result is None + + @pytest.mark.asyncio + async def test_chunk_failure_returns_none(self): + client = _make_client() + client._send_reply = AsyncMock( + side_effect=[ + {'errcode': 0, 'upload_id': 'u-3'}, + {'errcode': 40002, 'errmsg': 'bad chunk'}, + ] + ) + result = await client.upload_media(b'data', 'broken.png') + assert result is None + + @pytest.mark.asyncio + async def test_finish_failure_returns_none(self): + client = _make_client() + client._send_reply = AsyncMock( + side_effect=[ + {'errcode': 0, 'upload_id': 'u-4'}, + {'errcode': 0}, + {'errcode': 40003, 'errmsg': 'finish failed'}, + ] + ) + result = await client.upload_media(b'data', 'broken.png') + assert result is None + + @pytest.mark.asyncio + async def test_init_without_upload_id_returns_none(self): + client = _make_client() + client._send_reply = AsyncMock( + return_value={'errcode': 0} # no upload_id + ) + result = await client.upload_media(b'data', 'broken.png') + assert result is None + + @pytest.mark.asyncio + async def test_empty_data_returns_none(self): + client = _make_client() + result = await client.upload_media(b'', 'empty.png') + assert result is None + client._send_reply.assert_not_awaited() + + +class TestReplyMediaHelpers: + @pytest.mark.asyncio + async def test_reply_image_dispatches_to_respond_msg(self): + client = _make_client() + client._send_reply = AsyncMock(return_value={'errcode': 0}) + await client.reply_image('req-x', 'media-x') + args = client._send_reply.await_args.args + kwargs = client._send_reply.await_args.kwargs + assert args[0] == 'req-x' + assert args[1] == { + 'msgtype': 'image', + 'image': {'media_id': 'media-x'}, + } + assert kwargs.get('cmd') == wecom_ws_client.CMD_RESPOND_MSG + + @pytest.mark.asyncio + async def test_reply_file_uses_file_key(self): + client = _make_client() + client._send_reply = AsyncMock(return_value={'errcode': 0}) + await client.reply_file('req-x', 'media-x') + body = client._send_reply.await_args.args[1] + assert body['msgtype'] == 'file' + assert body['file'] == {'media_id': 'media-x'} + + @pytest.mark.asyncio + async def test_reply_voice_uses_voice_key(self): + client = _make_client() + client._send_reply = AsyncMock(return_value={'errcode': 0}) + await client.reply_voice('req-x', 'media-x') + body = client._send_reply.await_args.args[1] + assert body['msgtype'] == 'voice' + assert body['voice'] == {'media_id': 'media-x'} + + @pytest.mark.asyncio + async def test_reply_media_unknown_kind_is_logged_and_returns_none(self): + client = _make_client() + client._send_reply = AsyncMock() + result = await client._reply_media('req-x', 'media-x', 'video') + assert result is None + client._send_reply.assert_not_awaited() + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/unit_tests/platform/test_wecombot_round_fallback.py b/tests/unit_tests/platform/test_wecombot_round_fallback.py new file mode 100644 index 000000000..8a03115a9 --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_round_fallback.py @@ -0,0 +1,223 @@ +"""Unit tests for ``WecomBotWsClient.push_stream_chunk`` end-of-round +fallback behaviour driven by ``fallback_after_round``. + +When ``fallback_after_round`` is set to N (>= 1) by the caller (e.g. +``WecomBot`` reads ``keep-first-think-only`` and passes 1), the stream +entry must be **dropped** the moment round N's final chunk is sent, so +that the next round's first push fails to find a stream_id and the +caller falls back to the markdown ``reply_text`` path. This is the +mechanism that keeps round 1 as a "thinking" block in WeCom while +rounds 2+ show up as plain markdown. + +The condition for the drop is ``round_idx >= fallback_after_round`` +(>=, not >). The previous implementation used strict ``>``, which +meant round 1 final was **not** dropped with ``fallback_after_round=1`` +— the stream entry rotated to a new id and the next round's push +succeeded, which gave WeCom a second streaming bubble for round 2 +instead of a plain markdown message. + +These tests pin down the >= semantics so a future refactor does not +regress to the strict-> comparison. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger + +# ``ws_client.py`` has a hidden transitive dependency on +# ``langbot.pkg.platform.logger.EventLogger`` via ``api.py``. In normal +# operation the application boot path (``main.py`` -> ``app.py``) +# imports that module first, so the partial-init error does not +# surface. In a fresh test process we replicate that ordering by +# importing ``app`` here. +from langbot.pkg.core import app as _core_app # noqa: F401 + +import langbot.libs.wecom_ai_bot_api.ws_client as wecom_ws_client # noqa: E402 +from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402,F401 + + +class _DummyLogger(abstract_platform_logger.AbstractEventLogger): + async def info(self, *args, **kwargs): + pass + + async def debug(self, *args, **kwargs): + pass + + async def warning(self, *args, **kwargs): + pass + + async def error(self, *args, **kwargs): + pass + + +def _make_client() -> wecom_ws_client.WecomBotWsClient: + client = wecom_ws_client.WecomBotWsClient( + bot_id='bot', + secret='secret', + logger=_DummyLogger(), + ) + # Block real network I/O. push_stream_chunk calls reply_stream + # which calls _send_reply; we mock both. + client._send_reply = AsyncMock(return_value={'errcode': 0}) + client.reply_stream = AsyncMock(return_value={'errcode': 0}) + return client + + +def _seed_stream( + client: wecom_ws_client.WecomBotWsClient, + msg_id: str = 'm-1', + req_id: str = 'req-1', +) -> str: + """Seed a stream entry for ``msg_id`` and return the original stream_id.""" + original = 'stream_original' + client._stream_ids[msg_id] = f'{req_id}|{original}' + client._stream_sessions[msg_id] = { + 'req_id': req_id, + 'stream_id': original, + 'msg_id': msg_id, + } + return original + + +class TestFallbackAfterRoundOne: + """``fallback_after_round=1`` must drop the stream entry on round 1 final.""" + + @pytest.mark.asyncio + async def test_round_1_final_drops_stream_entry(self): + # With ``fallback_after_round=1`` and round 1 final, the stream + # entry for msg_id must be removed so the next push fails. + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 1 + + success = await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + assert success is True + assert 'm-1' not in client._stream_ids, ( + 'after round 1 final with fallback_after_round=1, the stream ' + 'entry must be dropped so the next round falls back to ' + 'reply_text' + ) + # The round counter must still record round 1 happened. + assert client._stream_round_counts.get('m-1') == 1 + + @pytest.mark.asyncio + async def test_round_2_push_returns_false_after_fallback(self): + # After round 1 final drops the entry, round 2's first push + # must return False (no stream_id) so the caller (wecombot.py) + # falls back to reply_text. + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 1 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + # The runner starts round 2; first push must fail. + success = await client.push_stream_chunk('m-1', 'round 2 first', is_final=False) + + assert success is False, ( + 'round 2 push after a round-1 fallback must return False so the platform adapter falls back to reply_text' + ) + # And the content must NOT have been sent through the stream + # channel at all during round 2. + assert client.reply_stream.await_count == 1, ( + 'only round 1 final should have used the stream channel; round 2 should not have called reply_stream at all' + ) + + @pytest.mark.asyncio + async def test_round_2_final_also_returns_false(self): + # Same as above but with is_final=True on the round 2 push. + # The intent is that round 2 is delivered as a markdown + # ``reply_text`` (handled by the caller on push_stream_chunk + # returning False and is_final=True), so the stream channel + # stays untouched. + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 1 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + success = await client.push_stream_chunk('m-1', 'round 2 full', is_final=True) + + assert success is False + assert client.reply_stream.await_count == 1, ( + 'round 2 must not have hit reply_stream; the caller should fall back to reply_text instead' + ) + + +class TestFallbackAfterRoundTwo: + """``fallback_after_round=2`` keeps round 1 streaming (rotate to + a new id) and drops the entry at round 2 final, so round 3+ falls + back to reply_text.""" + + @pytest.mark.asyncio + async def test_round_1_rotates_stream(self): + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 2 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + # round_idx=1, 1 >= 2 is False, so the entry rotates (kept). + assert 'm-1' in client._stream_ids + assert client._stream_round_counts.get('m-1') == 1 + + @pytest.mark.asyncio + async def test_round_2_final_drops_stream(self): + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 2 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + # Round 2 final: round_idx becomes 2, 2 >= 2 fires the drop. + success = await client.push_stream_chunk('m-1', 'round 2 final', is_final=True) + + assert success is True + assert 'm-1' not in client._stream_ids + assert client._stream_round_counts.get('m-1') == 2 + + # Round 3 push must fail and not call reply_stream. + round_3_success = await client.push_stream_chunk('m-1', 'round 3 first', is_final=False) + assert round_3_success is False + assert client.reply_stream.await_count == 2, ( + 'only round 1 and round 2 should have used reply_stream; round 3 must not reach the stream channel' + ) + + +class TestFallbackAfterRoundZeroRegression: + """``fallback_after_round=0`` keeps the original rotate-per-round + behaviour. The off-by-one fix must not break this path.""" + + @pytest.mark.asyncio + async def test_round_1_rotates_stream(self): + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 0 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + # The stream entry is preserved (rotated to a new id) so the + # next round can push against the same msg_id. + assert 'm-1' in client._stream_ids + new = client._stream_ids['m-1'] + _, new_stream_id = new.split('|', 1) + assert new_stream_id.startswith('stream_') + + @pytest.mark.asyncio + async def test_round_2_push_succeeds_after_rotate(self): + client = _make_client() + _seed_stream(client) + client.fallback_after_round = 0 + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + success = await client.push_stream_chunk('m-1', 'round 2 first', is_final=False) + + assert success is True + # Both calls went through reply_stream. + assert client.reply_stream.await_count == 2 + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/unit_tests/platform/test_wecombot_stream_rotation.py b/tests/unit_tests/platform/test_wecombot_stream_rotation.py new file mode 100644 index 000000000..e353448df --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_stream_rotation.py @@ -0,0 +1,213 @@ +"""Unit tests for ``WecomBotWsClient.push_stream_chunk`` end-of-round rotation. + +When the local-agent runner yields a chunk with ``is_final=True`` for one +LLM round, the WeCom stream channel must be **rotated** — not torn down. +Otherwise the next LLM round (after a tool call) would push against the +same stream_id, which the WeCom client has already closed upon seeing +``finish=true``. The push would silently fail and the next round would +fall back to the markdown ``reply_text`` path, where ```` tags +are not rendered as collapsible "thinking" blocks but as raw text. + +These tests pin down the rotation behavior so future refactors do not +regress to a single-shot stream that breaks multi-round agent turns. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger + +# ``ws_client.py`` has a hidden transitive dependency on +# ``langbot.pkg.platform.logger.EventLogger`` via ``api.py``. In normal +# operation the application boot path (``main.py`` -> ``app.py``) +# imports that module first, so the partial-init error does not +# surface. In a fresh test process we replicate that ordering by +# importing ``app`` here. +from langbot.pkg.core import app as _core_app # noqa: F401 + +# Now safe to import the websocket client module and its symbols. +import langbot.libs.wecom_ai_bot_api.ws_client as wecom_ws_client # noqa: E402 +from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402,F401 + + +class _DummyLogger(abstract_platform_logger.AbstractEventLogger): + async def info(self, *args, **kwargs): + pass + + async def debug(self, *args, **kwargs): + pass + + async def warning(self, *args, **kwargs): + pass + + async def error(self, *args, **kwargs): + pass + + +def _make_client(reply_stream: AsyncMock | None = None) -> wecom_ws_client.WecomBotWsClient: + client = wecom_ws_client.WecomBotWsClient( + bot_id='bot', + secret='secret', + logger=_DummyLogger(), + ) + client._send_reply = AsyncMock(return_value={'errcode': 0}) + if reply_stream is not None: + client.reply_stream = reply_stream + return client + + +def _seed_stream(client: wecom_ws_client.WecomBotWsClient, msg_id: str = 'm-1') -> str: + """Seed a stream entry for ``msg_id`` and return the original stream_id.""" + original = 'stream_original' + client._stream_ids[msg_id] = f'req-1|{original}' + client._stream_sessions[msg_id] = { + 'req_id': 'req-1', + 'stream_id': original, + 'msg_id': msg_id, + } + return original + + +class TestPushStreamChunkEndOfRound: + """``is_final=True`` must rotate the stream_id, not close the channel.""" + + @pytest.mark.asyncio + async def test_final_chunk_keeps_stream_entry(self): + client = _make_client() + _seed_stream(client) + + success = await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + assert success is True + # The stream entry is preserved so the next round can push + # against the same msg_id. + assert 'm-1' in client._stream_ids + + @pytest.mark.asyncio + async def test_final_chunk_rotates_to_new_stream_id(self): + client = _make_client() + original = _seed_stream(client) + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + # The new stream_id must differ from the original so the WeCom + # client treats the next push as a fresh streaming message. + new = client._stream_ids['m-1'] + _, new_stream_id = new.split('|', 1) + assert new_stream_id != original + # And it must look like a freshly minted stream id. + assert new_stream_id.startswith('stream_') + + @pytest.mark.asyncio + async def test_final_chunk_clears_last_content_and_session(self): + client = _make_client() + _seed_stream(client) + client._stream_last_content['m-1'] = 'round 1 final' + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + # Both side-tables are dropped so the next round is not + # suppressed by the dedupe check and has a fresh session record. + assert 'm-1' not in client._stream_last_content + assert 'm-1' not in client._stream_sessions + + @pytest.mark.asyncio + async def test_final_chunk_sends_finish_true(self): + reply_stream = AsyncMock() + client = _make_client(reply_stream) + _seed_stream(client) + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + reply_stream.assert_awaited_once() + args = reply_stream.await_args.args + # ``reply_stream(req_id, stream_id, content, finish=, feedback_id=)`` + assert args[0] == 'req-1' # req_id + assert args[1] == 'stream_original' # original stream_id + assert args[2] == 'round 1 final' # content + assert reply_stream.await_args.kwargs['finish'] is True + # The finish=true frame uses the *original* stream_id, which + # tells the WeCom client to close the streaming message it + # already opened. + + @pytest.mark.asyncio + async def test_second_round_push_succeeds_after_final(self): + # End-to-end: round 1 finalizes, then a tool call returns, the + # LLM is re-invoked, and round 2 starts streaming. The second + # round's first push must succeed and use the *new* stream_id. + reply_stream = AsyncMock() + client = _make_client(reply_stream) + original = _seed_stream(client) + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + # The runner now begins a new LLM round. The very first push + # of the new round must go through, on the rotated stream_id. + success = await client.push_stream_chunk('m-1', 'round 2 first', is_final=False) + + assert success is True + assert reply_stream.await_count == 2 + # Second call: finish must be False and stream_id must be the + # rotated one. + second_args = reply_stream.await_args_list[1].args + assert second_args[0] == 'req-1' # req_id unchanged + assert second_args[1] != original # new stream_id + assert second_args[1].startswith('stream_') + assert second_args[2] == 'round 2 first' + assert reply_stream.await_args_list[1].kwargs['finish'] is False + + @pytest.mark.asyncio + async def test_final_chunk_generates_feedback_id_once(self): + reply_stream = AsyncMock() + client = _make_client(reply_stream) + _seed_stream(client) + + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + + feedback_id = client._msg_feedback_ids.get('m-1') + assert feedback_id is not None + assert feedback_id.startswith('feedback_') + + @pytest.mark.asyncio + async def test_unknown_msg_id_returns_false(self): + # Defensive: if the channel was never opened (or has been + # cleared by a new query) the push must report failure so + # callers can fall back to ``reply_text``. + client = _make_client() + success = await client.push_stream_chunk('nope', 'content') + assert success is False + + +class TestNewQueryCleansUpStaleStreams: + """A new inbound message replaces the stream entry for the same + msg_id, so a stale stream_id from the previous round is harmless. + """ + + @pytest.mark.asyncio + async def test_handle_message_callback_overwrites_stream_entry(self): + client = _make_client() + # Round 1 of the previous turn left a stream entry in place. + _seed_stream(client) + await client.push_stream_chunk('m-1', 'round 1 final', is_final=True) + assert 'm-1' in client._stream_ids + + # A new inbound message arrives. Simulate the parts of + # ``_handle_message_callback`` that touch ``_stream_ids``. + new_stream_id = 'stream_brand_new' + client._stream_ids['m-1'] = f'req-2|{new_stream_id}' + + # Push for the new round should land on the rotated id. + reply_stream = AsyncMock() + client.reply_stream = reply_stream + success = await client.push_stream_chunk('m-1', 'round 2 hello', is_final=False) + + assert success is True + args = reply_stream.await_args.args + assert args[0] == 'req-2' # req_id updated to the new request + assert args[1] == new_stream_id + assert args[2] == 'round 2 hello' + assert reply_stream.await_args.kwargs['finish'] is False diff --git a/tests/unit_tests/provider/test_litellmchat.py b/tests/unit_tests/provider/test_litellmchat.py index f7a448ab6..ccd76a7c4 100644 --- a/tests/unit_tests/provider/test_litellmchat.py +++ b/tests/unit_tests/provider/test_litellmchat.py @@ -498,6 +498,72 @@ def test_empty_content(self): result = requester._process_thinking_content('', None, remove_think=True) assert result == '' + def test_remove_think_markers_public(self): + """Test removing the public think tags when remove_think=True. + + MiniMax-M3 (and other OpenAI-compatible providers) emit the + reasoning text inline in ``content`` wrapped in ``…``. + The stripper must handle this public convention. + """ + 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_preserve_think_markers_public(self): + """Public think tags are preserved when remove_think=False.""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'reasoningFinal answer.' + result = requester._process_thinking_content(content, None, remove_think=False) + assert '' in result + assert 'Final answer.' in result + + def test_multiline_think_block(self): + """Multiline think blocks are stripped in their entirety.""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = '\nline1\nline2\nline3\n\nAnswer' + result = requester._process_thinking_content(content, None, remove_think=True) + assert '' not in result + assert 'line1' not in result + assert result.strip() == 'Answer' + + def test_multiple_think_blocks(self): + """Multiple think blocks in the same content are all stripped.""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'aXbY' + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'XY' + + def test_mixed_markers_stripped_together(self): + """Public and legacy markers can both be present and both are stripped.""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = ( + 'publicA' + 'CRETIRE_REASONING_BEGINklegacyCRETIRE_REASONING_ENDkB' + ) + result = requester._process_thinking_content(content, None, remove_think=True) + assert result == 'AB' + + def test_remove_think_false_does_not_strip(self): + """When remove_think is False, neither marker type is touched.""" + requester = litellmchat.LiteLLMRequester(ap=Mock(), config={}) + + content = 'reasoningFinal' + result = requester._process_thinking_content(content, None, remove_think=False) + assert result == content + + def test_strip_think_helper(self): + """The classmethod helper exposes the same logic for reuse.""" + result = litellmchat.LiteLLMRequester._strip_think( + 'aXbY', + ) + assert result == 'XY' + class TestBuildCommonArgs: """Test _build_common_args method""" diff --git a/tests/unit_tests/provider/test_localagent_remove_think.py b/tests/unit_tests/provider/test_localagent_remove_think.py new file mode 100644 index 000000000..87601de45 --- /dev/null +++ b/tests/unit_tests/provider/test_localagent_remove_think.py @@ -0,0 +1,494 @@ +"""Unit tests for the ``remove-think`` integration in the local agent runner. + +These tests pin down two pieces of the "strip chain-of-thought before it +reaches the platform adapter" feature: + +1. The runner must safely resolve the ``output.misc.remove-think`` flag + even when the pipeline config is missing the ``misc`` sub-key (the + old code raised ``AttributeError: 'str' object has no attribute 'get'`` + in that case, which meant the toggle was effectively never read). +2. ``_StreamAccumulator`` must apply the stripper to the content it + accumulates, both on intermediate 8-chunk flushes and on the final + message. This is the safety net for second-pass LLM invocations + after tool calls, where the upstream LiteLLM requester has not + had a chance to strip the content yet. +""" + +from __future__ import annotations + +import inspect +import re + +import pytest + +import langbot_plugin.api.entities.builtin.provider.message as provider_message + +from langbot.pkg.provider.runners.localagent import ( + _StreamAccumulator, +) + + +def _chunk(content: str, is_final: bool = False) -> provider_message.MessageChunk: + return provider_message.MessageChunk( + role='assistant', + content=content, + is_final=is_final, + msg_sequence=0, + ) + + +class TestRemoveThinkConfigResolution: + """The toggle lookup must not raise when ``output.misc`` is missing.""" + + def _resolve(self, output: dict | None) -> bool: + """Replicate the safe-chain expression used in ``localagent.run``.""" + return ( + ((output or {}) if output is not None else {}).get('misc') + if False + else (((output or {}).get('misc') or {}).get('remove-think', False)) + ) + + def test_full_config_true(self): + assert self._resolve({'misc': {'remove-think': True}}) is True + + def test_full_config_false(self): + assert self._resolve({'misc': {'remove-think': False}}) is False + + def test_missing_misc_defaults_false(self): + # The original bug: ``output.get('misc', '').get('remove-think')`` + # raised AttributeError here. + assert self._resolve({}) is False + + def test_output_is_none(self): + assert self._resolve(None) is False + + def test_misc_is_empty_dict(self): + assert self._resolve({'misc': {}}) is False + + def test_misc_is_none(self): + assert self._resolve({'misc': None}) is False + + def test_misc_is_string_legacy_bug(self): + # The exact value the old code passed to ``.get('remove-think')``: + # an empty string instead of a dict. + assert self._resolve({'misc': ''}) is False + + +class TestStreamAccumulatorStripping: + """The accumulator must strip think blocks when ``remove_think`` is on.""" + + def test_passthrough_when_disabled(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=False) + for c in ['r', 'XyY']: + out = acc.add(_chunk(c)) + if out is not None: + # Intermediate 8-chunk flushes preserve the raw text. + assert '' in out.content + final = acc.final_message() + assert '' in final.content + + def test_passthrough_when_no_chunks(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + out = acc.add(_chunk('')) + assert out is None + assert acc.final_message().content == '' + + def test_simple_strip_on_final(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + acc.add(_chunk('reasoning')) + acc.add(_chunk('The answer is 42.')) + final = acc.final_message() + assert final.content == 'The answer is 42.' + + def test_strip_on_intermediate_flush(self): + # _StreamAccumulator flushes every 8 chunks. The first emit must + # already be free of the think block. + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + for _ in range(7): + emitted = acc.add(_chunk('')) + assert emitted is None + emitted = acc.add(_chunk('reasoningfinal')) + assert emitted is not None + assert '' not in emitted.content + assert 'reasoning' not in emitted.content + assert 'final' in emitted.content + + def test_strip_when_think_block_straddles_chunks(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + acc.add(_chunk('reas')) + acc.add(_chunk('oningFinal')) + assert acc.final_message().content == 'Final' + + def test_multiple_think_blocks(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + acc.add(_chunk('aX')) + acc.add(_chunk('bY')) + assert acc.final_message().content == 'XY' + + def test_legacy_marker_stripped(self): + acc = _StreamAccumulator(msg_sequence=0, remove_think=True) + acc.add(_chunk('CRETIRE_REASONING_BEGINkgoCRETIRE_REASONING_ENDkFinal')) + assert acc.final_message().content == 'Final' + + def test_initial_content_is_also_stripped(self): + acc = _StreamAccumulator( + msg_sequence=0, + initial_content='rfirst', + remove_think=True, + ) + acc.add(_chunk(' second')) + assert acc.final_message().content == 'first second' + + def test_no_strip_when_remove_think_false_even_with_think(self): + # When the user wants to see the reasoning, ``remove_think=False`` + # must preserve the block in both the intermediate flushes and + # the final message. + acc = _StreamAccumulator(msg_sequence=0, remove_think=False) + acc.add(_chunk('rX')) + final = acc.final_message() + assert '' in final.content + assert 'X' in final.content + + +class TestRemoveThinkPipelineHookup: + """The runner passes the resolved ``remove_think`` flag into the + accumulator. We verify this by inspecting the local-agent run method + after a minimal stub: a much smaller surface test than the full + run(), but enough to catch a regression where the flag is silently + dropped. + """ + + def test_remove_think_propagates_to_accumulator_constructor(self): + # Spot-check the literal default: the runner instantiates the + # accumulator with ``remove_think=remove_think`` so the safety + # net kicks in. This is a code-shape guard, not a behavior test. + import inspect + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + assert 'remove_think=remove_think' in src + + def test_tool_loop_accumulator_does_not_reseed_with_first_content(self): + # Reproduces the "Wecom AI Bot duplicates the first sentence on + # every tool round" bug. After a tool call, the runner re-invokes + # the LLM; the second-round accumulator must not be seeded with + # the first round's content, otherwise every emitted chunk + # contains the entire opening line, and the Wecom adapter — which + # is forced to fall back to reply_text after the stream is + # closed — sends it as a brand new (duplicate) message each + # round. + acc = _StreamAccumulator(msg_sequence=3, remove_think=True) + for c in 'Data fetched. Computing stats now.': + out = acc.add(_chunk(c, is_final=False)) + if out is not None: + # Any intermediate emit must be the *current* round's + # content only, not a re-statement of the first round. + assert '好的，我来查一下' not in out.content + final = acc.final_message() + assert '好的，我来查一下' not in final.content + assert final.content.strip() == 'Data fetched. Computing stats now.' + + def test_tool_loop_first_round_unaffected(self): + # Sanity check: the first round still works the same way. The + # accumulator starts empty and accumulates only what the runner + # produces in this round. + acc = _StreamAccumulator(msg_sequence=1, remove_think=True) + for c in '好的，我来查一下': + acc.add(_chunk(c, is_final=False)) + final = acc.final_message() + assert final.content == '好的，我来查一下' + + def test_safe_chain_in_run(self): + # The fix for the ``AttributeError`` lives as a literal expression + # in ``LocalAgentRunner.run``. Pin the shape so future refactors + # do not regress to the old broken chain + # ``output.get('misc', '').get('remove-think')``. + import inspect + import re + + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + # Strip comments to avoid false positives on docstrings. + code = re.sub(r'#[^\n]*', '', src) + + assert "get('misc') or {}" in code + assert "get('remove-think', False)" in code + # The exact broken chain must not be present as code. + assert re.search(r"\.get\('misc',\s*''\)\.get\('remove-think'\)\b", code) is None + + +class TestKeepFirstThinkOnly: + """``output.misc.keep-first-think-only`` keeps the chain-of-thought + on the *first* LLM round (so WeCom renders it as a "thinking" block) + and strips it on every *subsequent* round. It only takes effect when + ``remove-think`` is on; on its own it is a no-op. + """ + + def _compute(self, output: dict | None) -> dict[str, bool]: + """Replicate the safe-chain expressions used in ``localagent.run``. + + Returns both the raw ``remove_think`` value and the + ``_strip_think_first_round`` value the runner uses for the + pre-tool-loop invocations. + """ + output = output or {} + misc = (output.get('misc') or {}) if isinstance(output.get('misc'), dict) else {} + remove_think = misc.get('remove-think', False) + keep_first_think_only = misc.get('keep-first-think-only', False) + strip_think_first_round = remove_think and not keep_first_think_only + return { + 'remove_think': bool(remove_think), + 'strip_think_first_round': bool(strip_think_first_round), + } + + def test_first_round_kept_when_both_toggles_on(self): + # remove-think=True AND keep-first-think-only=True: + # first round must NOT strip (CoT rendered as thinking block). + result = self._compute({'misc': {'remove-think': True, 'keep-first-think-only': True}}) + assert result['remove_think'] is True + assert result['strip_think_first_round'] is False + + def test_subsequent_rounds_strip_unchanged(self): + # The tool-loop invocations keep the raw ``remove_think`` value, + # so subsequent rounds still strip. We assert that the raw flag + # is True so the contract on the *second* round is preserved. + result = self._compute({'misc': {'remove-think': True, 'keep-first-think-only': True}}) + assert result['remove_think'] is True # used by the tool-loop path + + def test_keep_first_think_only_alone_is_noop(self): + # remove-think=False, keep-first-think-only=True: the user did + # not ask to strip CoT, so the toggle must have no effect. + result = self._compute({'misc': {'remove-think': False, 'keep-first-think-only': True}}) + assert result['remove_think'] is False + assert result['strip_think_first_round'] is False + + def test_remove_think_alone_strips_everywhere_unchanged(self): + # remove-think=True, keep-first-think-only=False: legacy + # behavior. The first round also strips; the toggle is + # effectively disabled. + result = self._compute({'misc': {'remove-think': True, 'keep-first-think-only': False}}) + assert result['remove_think'] is True + assert result['strip_think_first_round'] is True + + def test_default_off(self): + # Missing toggles: both default to False, so nothing is stripped. + result = self._compute({'misc': {}}) + assert result['remove_think'] is False + assert result['strip_think_first_round'] is False + + def test_safe_chain_under_missing_keys(self): + # The new ``keep-first-think-only`` lookup must follow the same + # safe-chain pattern as ``remove-think`` so that + # ``misc=None`` / ``misc=''`` / ``output=None`` do not raise. + assert self._compute(None) == {'remove_think': False, 'strip_think_first_round': False} + assert self._compute({}) == {'remove_think': False, 'strip_think_first_round': False} + assert self._compute({'misc': None}) == {'remove_think': False, 'strip_think_first_round': False} + assert self._compute({'misc': ''}) == {'remove_think': False, 'strip_think_first_round': False} + + def test_first_round_accumulator_keeps_think_block(self): + # End-to-end: with the first-round strip disabled, the + # accumulator preserves the chain-of-thought in the final + # message. This mirrors what a WeCom user would see. + acc = _StreamAccumulator(msg_sequence=1, remove_think=False) + acc.add(_chunk('plan')) + acc.add(_chunk('Final answer.')) + final = acc.final_message() + assert '' in final.content + assert 'plan' in final.content + assert 'Final answer.' in final.content + + def test_source_includes_keep_first_think_only_literal(self): + # Code-shape guard: the literal key must appear in ``run`` so + # future refactors do not silently drop the toggle. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + assert "'keep-first-think-only'" in src + assert '_strip_think_first_round' in src + + def test_first_round_invocation_uses_strip_think_first_round(self): + # The pre-tool-loop invocations must reference the + # round-aware variable, not the raw ``remove_think``. The + # tool-loop invocations must still reference ``remove_think`` + # so subsequent rounds keep stripping. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + # At least one occurrence of the round-aware variable in the + # pre-loop path (the accumulator constructor and/or + # _invoke_*_fallback). + assert 'remove_think=_strip_think_first_round' in src + # The tool-loop accumulator must still pass the raw + # ``remove_think`` so per-tool-call CoT is hidden. + assert 'msg_sequence=first_end_sequence' in src + + +class TestRemoveEmptyThinkRound: + """``output.misc.remove-empty-think-round`` suppresses any LLM + round that produces no user-visible text, regardless of whether + the round carries tool calls. The tool-call loop itself is + unaffected: ``req_messages.append(final_msg)`` runs + unconditionally, so the next LLM call still sees the round's + output. + """ + + def _should_skip( + self, + content: str, + tool_calls: list | None, + remove_empty_think_round: bool, + ) -> bool: + """Replicate the expression in ``localagent.run``. + + Mirrors the post-fix condition: + ``_skip_emit = remove_empty_think_round and not (content or '').strip()`` + """ + has_visible_text = bool((content or '').strip()) + return remove_empty_think_round and not has_visible_text + + def test_skip_when_no_text_no_tool_calls(self): + # The bug we just fixed: an LLM round with no text AND no + # tool calls used to be forwarded as a blank line. With the + # fix it is dropped. + assert self._should_skip('', None, True) is True + assert self._should_skip(' ', None, True) is True + + def test_skip_when_no_text_with_tool_calls(self): + # The original supported case: empty text + tool calls + # → still skipped (preserved behavior). + assert self._should_skip('', [object()], True) is True + + def test_no_skip_when_has_text(self): + # Any non-empty text wins, regardless of tool calls. + assert self._should_skip('hello', [object()], True) is False + assert self._should_skip('hello', None, True) is False + + def test_no_skip_when_disabled(self): + # Default off: an empty round is forwarded as-is. + assert self._should_skip('', None, False) is False + assert self._should_skip('', [object()], False) is False + + def test_source_drops_tool_calls_gate(self): + # Code-shape guard: the source must not require ``tool_calls`` + # for the skip to fire, so the regression cannot return. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + # Strip comments to avoid false positives on docstrings. + code = re.sub(r'#[^\n]*', '', src) + # The expression must reference ``remove_empty_think_round`` and + # ``_round_has_visible_text`` (or equivalent) but must NOT + # have a ``getattr(..., 'tool_calls', None)`` guard. + assert '_skip_emit = remove_empty_think_round' in code + # Old form was: ``and bool(getattr(final_msg, 'tool_calls', None))`` + assert "getattr(final_msg, 'tool_calls', None)" not in code + + def test_subsequent_round_skip_under_safe_chain(self): + # Same expression is used in the tool-loop non-streaming path. + # We pin the condition by source inspection so future refactors + # keep the contract on the *second* round as well. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + code = re.sub(r'#[^\n]*', '', src) + # The subsequent-round block must reference the same flag. + assert 'remove_empty_think_round' in code + # It must guard ``yield`` with a skip flag. + assert '_subsequent_skip_emit' in code + # Streaming path is intentionally untouched; we only assert + # the non-streaming branch was updated. + assert 'if not _subsequent_skip_emit:' in code + assert 'yield msg' in code # still present, but now guarded + + def test_subsequent_round_still_appends_to_req_messages(self): + # Even when the yield is skipped, ``final_msg`` must still + # be assigned and the tool-loop must still proceed. This is + # checked at the source level because exercising the full + # run() with a mock LLM is out of scope for this unit test. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + code = re.sub(r'#[^\n]*', '', src) + # The subsequent-round skip block must be followed by + # ``final_msg = msg`` so the tool-loop below keeps running + # even when the platform yield was suppressed. We use a + # small window to avoid matching the first-round + # ``final_msg = msg`` assignment. + match = re.search( + r'_subsequent_skip_emit[\s\S]{0,200}?final_msg\s*=\s*msg', + code, + ) + assert match is not None, ( + 'subsequent-round skip block must still assign final_msg = msg ' + 'so the tool-call loop continues after a suppressed yield' + ) + + def test_streaming_first_round_final_chunk_skip(self): + # Code-shape guard: the pre-loop streaming block must + # reference the empty-final-chunk guard. We pin the + # literal expression so future refactors cannot silently + # drop the toggle from the streaming path. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + assert 'msg.is_final and remove_empty_think_round' in src + # ``initial_response_emitted`` is still set on a non-empty + # yield so the post-loop ``elif not initial_response_emitted`` + # branch does not double-yield. + assert 'initial_response_emitted = True' in src + + def test_streaming_subsequent_round_final_chunk_skip(self): + # The tool-loop streaming block must carry the same guard. + # We count occurrences to confirm both first-round and + # subsequent-round streaming blocks were updated. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + code = re.sub(r'#[^\n]*', '', src) + occurrences = code.count('msg.is_final and remove_empty_think_round') + assert occurrences >= 2, ( + f'expected the guard in both first-round and tool-loop streaming blocks, got {occurrences}' + ) + + def test_streaming_intermediate_chunk_always_emitted(self): + # The guard must only fire on ``msg.is_final``; intermediate + # 8-chunk flushes drive the typing indicator and must always + # be emitted. If someone refactors and removes the + # ``msg.is_final and`` prefix, intermediate chunks would be + # suppressed and the user would see a frozen typing cursor. + from langbot.pkg.provider.runners import localagent + + src = inspect.getsource(localagent.LocalAgentRunner.run) + code = re.sub(r'#[^\n]*', '', src) + guard_pattern = re.search(r'if\s+msg\.is_final\s+and\s+remove_empty_think_round', code) + assert guard_pattern is not None, 'streaming final-chunk guard must be gated on msg.is_final' + # Sanity: a bare ``if remove_empty_think_round`` would catch + # intermediate chunks. Make sure that pattern is not in the + # streaming yield blocks. + bare_pattern = re.search( + r'yield\s+chunk[\s\S]{0,200}?if\s+remove_empty_think_round\b', + code, + ) + assert bare_pattern is None, ( + 'streaming yield must not be guarded by a bare ' + 'remove_empty_think_round check (would suppress ' + 'intermediate 8-chunk flushes and freeze the typing ' + 'indicator)' + ) + + def test_streaming_final_chunk_accumulator_returns_empty(self): + # Reproduce the trigger condition: an accumulator whose + # entire content is chain-of-thought produces an empty + # final chunk. The runtime guard treats this as skippable. + acc = _StreamAccumulator(msg_sequence=1, remove_think=True) + acc.add(_chunk('')) + final = acc.final_message() + assert final.content == '' + # The guard expression evaluates to True on this output. + has_visible_text = bool((final.content or '').strip()) + assert has_visible_text is False + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/unit_tests/provider/test_think_strip_state.py b/tests/unit_tests/provider/test_think_strip_state.py new file mode 100644 index 000000000..e1e520ac7 --- /dev/null +++ b/tests/unit_tests/provider/test_think_strip_state.py @@ -0,0 +1,210 @@ +"""Tests for ``_ThinkStripState`` — the streaming chain-of-thought filter. + +The filter is responsible for dropping ``…`` and the legacy +``CRETIRE_REASONING_BEGINk…ENDk`` blocks from a stream of text deltas. A +naive per-chunk regex leaks the tag whenever it straddles two chunks, so +this state machine tracks a small tail buffer to keep the implementation +correct across chunk boundaries. +""" + +import pytest + +from langbot.pkg.provider.modelmgr.requesters.litellmchat import ( + LiteLLMRequester, + _ThinkStripState, +) + + +def _feed_all(chunks, flush: bool = True) -> str: + """Helper: run ``feed`` for every chunk and optionally ``flush`` at end.""" + state = _ThinkStripState() + emitted = [] + for c in chunks: + emitted.append(state.feed(c)) + if flush: + emitted.append(state.flush()) + return ''.join(emitted) + + +class TestThinkStripStateNoOp: + """Cases where no stripping is needed.""" + + def test_plain_text(self): + assert _feed_all(['Hello world']) == 'Hello world' + + def test_empty_stream(self): + assert _feed_all(['']) == '' + + def test_plain_text_across_chunks(self): + assert _feed_all(['Hel', 'lo wor', 'ld']) == 'Hello world' + + def test_mixed_case_passthrough(self): + assert _feed_all(['noThInK here']) == 'noThInK here' + + def test_less_than_without_tag(self): + assert _feed_all(['1<2<3<4<5']) == '1<2<3<4<5' + + def test_partial_close_no_open(self): + # ``reasoning']) == '' + + def test_think_then_text(self): + assert _feed_all(['r\nAnswer']) == '\nAnswer' + + def test_multiline_block(self): + assert _feed_all(['\nline1\nline2\n\nAnswer']).strip() == 'Answer' + + def test_multiple_blocks(self): + assert _feed_all(['aXbY']) == 'XY' + + def test_block_then_partial_open(self): + assert _feed_all(['rrestaX', + 'CRETIRE_REASONING_BEGINkbCRETIRE_REASONING_ENDkY', + ] + assert _feed_all(chunks) == 'XY' + + +class TestThinkStripStateAcrossChunks: + """Think tags split across chunk boundaries must still be stripped.""" + + def test_open_tag_split(self): + assert _feed_all(['reasoning\nFinal']) == '\nFinal' + + def test_body_split(self): + assert _feed_all(['reas', 'oning\nFinal']) == '\nFinal' + + def test_close_tag_split(self): + assert _feed_all(['reasoning\nFinal']) == '\nFinal' + + def test_many_chunks(self): + chunks = [ + '<', + 'think>r', + 'e', + 'as', + 'oning\nFinal', + ] + assert _feed_all(chunks) == '\nFinal' + + def test_legacy_body_split(self): + chunks = [ + 'CRETIRE_REASONING_BEGINkrea', + 'soningCRETIRE_REASONING_ENDkX', + ] + assert _feed_all(chunks) == 'X' + + def test_unterminated_think_drops_everything(self): + # Stream ends while still inside a think block; everything is dropped. + assert _feed_all(['partial', ' reasoning']) == '' + + def test_think_split_with_plain_between(self): + chunks = ['aY'] + assert _feed_all(chunks) == 'Y' + + +class TestThinkStripStateBufferBehavior: + """Partial tags held in the tail buffer must be released safely.""" + + def test_partial_open_at_end_then_more(self): + # ``okX'] + assert _feed_all(chunks) == 'Hello X' + + def test_flush_releases_orphan_partial(self): + # Partial ``r']) == '' + + def test_legacy_open_only(self): + # Stream ends inside the legacy block; content is dropped. + assert _feed_all(['CRETIRE_REASONING_BEGINkpartial', ' body']) == '' + + def test_close_only_preserved_as_plain_text(self): + # An orphan `` ``k>`` that is never preceded by an + # opening tag is plain text. The state machine has no way to + # distinguish a real close from a literal ```` snippet, + # so it preserves the bytes verbatim. + assert _feed_all(['prefaceX']) == 'prefaceX' + + def test_close_tag_only_with_text_after(self): + # Same idea: an orphan close tag stays in the output. + assert _feed_all(['arest']) == 'arest' + + +class TestThinkStripStatePreservation: + """When ``remove_think=False`` the LLM is not expected to use the + state machine; this class exists only to document that ``flush()`` and + ``feed()`` themselves are side-effect free with respect to whether the + caller would have stripped the content. (The accumulator flag is what + controls that, not the state itself.) + """ + + def test_state_does_not_strip_without_remove_flag(self): + # The state machine only filters; it has no flag of its own. The + # caller decides whether to use it. This is a documentation + # guard — if anyone later changes the contract this test will + # remind them. + s = _ThinkStripState() + # First a normal block, then plain text. + assert s.feed('aX') == 'X' + assert s.feed('bY') == 'Y' + # Final plain text. + assert s.feed('z') == 'z' + + +class TestLiteLLMStripThinkHelper: + """The classmethod ``_strip_think`` is the same logic, no state needed.""" + + def test_no_tags(self): + assert LiteLLMRequester._strip_think('hello') == 'hello' + + def test_multiple(self): + assert LiteLLMRequester._strip_think('aXbY') == 'XY' + + def test_legacy_and_public(self): + text = ( + 'publicA' + 'CRETIRE_REASONING_BEGINklegacyCRETIRE_REASONING_ENDkB' + ) + assert LiteLLMRequester._strip_think(text) == 'AB' + + def test_preserves_orphan_close(self): + # The helper runs a single-pass regex: it has no notion of + # context, so an orphan ``