UI: Fix tool buffer check truncating visible text at end of generation

This commit is contained in:
oobabooga 2026-04-04 09:57:07 -07:00
parent 16af11f868
commit e0ad4e60df
2 changed files with 9 additions and 2 deletions

View file

@ -1183,7 +1183,7 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess
# visible text from before buffering started so raw markup doesn't flash
# in the UI. The internal text is left intact so the caller can still
# parse tool calls from it.
if is_stream and _check_tool_markers and streaming_tool_buffer_check(output['internal'][-1][1], markers=_streaming_markers, tool_names=_tool_names, check_bare_names=_check_bare_names):
if is_stream and _check_tool_markers and streaming_tool_buffer_check(output['internal'][-1][1], markers=_streaming_markers, tool_names=_tool_names, check_bare_names=_check_bare_names, partial_match=False):
output['visible'][-1][1] = _last_visible_before_tool_buffer or ''
yield output

View file

@ -31,7 +31,7 @@ TOOL_CALL_OPENING_MARKERS = [
]
def streaming_tool_buffer_check(text, markers=None, tool_names=None, check_bare_names=False):
def streaming_tool_buffer_check(text, markers=None, tool_names=None, check_bare_names=False, partial_match=True):
'''
Check whether streaming output should be withheld because it may
contain tool-call markup.
@ -43,6 +43,10 @@ def streaming_tool_buffer_check(text, markers=None, tool_names=None, check_bare_
tool_names: List of tool function names.
check_bare_names: Whether to do partial-prefix matching on tool
names (for models with unknown template format).
partial_match: Whether to check partial prefixes of markers/names.
Set to False for end-of-generation checks where a
partial prefix is just normal text, not an incomplete
tool call.
'''
# Strip thinking blocks so tool-call syntax inside <think> doesn't
# trigger false positives.
@ -60,6 +64,9 @@ def streaming_tool_buffer_check(text, markers=None, tool_names=None, check_bare_
if name + '{' in text or name + ' {' in text:
return True
if not partial_match:
return False
# Partial-prefix matching: only for template-specific markers.
for marker in (markers if markers is not None else TOOL_CALL_OPENING_MARKERS):
for prefix_len in range(min(len(marker) - 1, len(text)), 0, -1):