diff --git a/modules/chat.py b/modules/chat.py index 02ae46e4..daecd50b 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -1028,6 +1028,13 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess thinking_prefix = start_tag break + # When tools are active, buffer streaming output during potential tool + # call generation to prevent raw markup from leaking into the display. + _check_tool_markers = bool(state.get('tools')) + if _check_tool_markers: + from modules.tool_parsing import streaming_tool_buffer_check + _tool_names = [t['function']['name'] for t in state['tools'] if 'function' in t and 'name' in t['function']] + # Generate reply = None for j, reply in enumerate(generate_reply(prompt, state, stopping_strings=stopping_strings, is_chat=True, for_ui=for_ui)): @@ -1077,6 +1084,10 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess }) if is_stream: + if _check_tool_markers: + if streaming_tool_buffer_check(output['internal'][-1][1], _tool_names): + continue + yield output if _continue: diff --git a/modules/tool_parsing.py b/modules/tool_parsing.py index 460188d3..418503ad 100644 --- a/modules/tool_parsing.py +++ b/modules/tool_parsing.py @@ -9,6 +9,55 @@ def get_tool_call_id() -> str: return "call_" + "".join(b).lower() +# Known opening markers for tool calls across model formats. +# Used during streaming to buffer output that might be tool call markup, +# preventing raw markup from leaking into displayed/streamed content. +TOOL_CALL_OPENING_MARKERS = [ + '', + '', + '', + '<|tool_call_begin|>', + '<|tool_calls_section_begin|>', + '<|tool▁call▁begin|>', + '<|tool▁calls▁begin|>', + '[TOOL_CALLS]', + 'to=functions.', + '<|channel|>commentary', +] + +def streaming_tool_buffer_check(text, tool_names=None): + ''' + Check whether streaming output should be withheld because it may + contain tool-call markup. + ''' + # Full marker found → buffer permanently + for marker in TOOL_CALL_OPENING_MARKERS: + if marker in text: + return True + + # Bare function-name style (e.g. Devstral): "get_weather{...}" + # Only match tool name followed by '{' to avoid false positives on + # common words that happen to be tool names (e.g. "get", "search"). + if tool_names: + for name in tool_names: + if name + '{' in text or name + ' {' in text: + return True + # Partial: text ends with tool name (or prefix of it) but '{' hasn't arrived yet + if text.endswith(name): + return True + for prefix_len in range(min(len(name) - 1, len(text)), 0, -1): + if text.endswith(name[:prefix_len]): + return True + + # Tail might be a partial marker forming across tokens + for marker in TOOL_CALL_OPENING_MARKERS: + for prefix_len in range(min(len(marker) - 1, len(text)), 0, -1): + if text.endswith(marker[:prefix_len]): + return True + + return False + + def check_and_sanitize_tool_call_candidate(candidate_dict: dict, tool_names: list[str]): # check if property 'function' exists and is a dictionary, otherwise adapt dict if 'function' not in candidate_dict and 'name' in candidate_dict and isinstance(candidate_dict['name'], str):