strip thinking blocks in place so streaming detection uses clean content

Previously the streaming loop computed a separate check_content for the
end-tag presence check but still resolved end_tag and find() against the
unstripped accumulated, so a <thinking> block containing </function>
while the real output used </invoke> (or vice-versa) caused the wrong
tag to be picked and truncation to land inside the thinking block.

The post-loop strip also ran after _truncate_to_first_function, which
counts every <function= / <invoke name= occurrence — so a thinking
block mentioning a tool call ("I should call <function=read_file>...")
became the "first" function and the real tool call got sliced off.

Move the strip into the streaming loop and apply it to accumulated
itself; the post-loop pipeline now operates on already-clean content
and the redundant trailing re.sub is dropped.
This commit is contained in:
0xallam 2026-05-03 15:55:40 -07:00
parent 806b735a1b
commit a3115bf686

View file

@ -26,6 +26,8 @@ from strix.utils.resource_paths import get_strix_resource_path
litellm.drop_params = True
litellm.modify_params = True
_THINKING_BLOCK_RE = re.compile(r"<thinking[^>]*>.*?</thinking>", re.DOTALL)
class LLMRequestFailedError(Exception):
def __init__(self, message: str, details: str | None = None):
@ -188,11 +190,8 @@ class LLM:
continue
delta = self._get_chunk_content(chunk)
if delta:
accumulated += delta
check_content = re.sub(
r"<thinking[^>]*>.*?</thinking>", "", accumulated, flags=re.DOTALL
)
if "</function>" in check_content or "</invoke>" in check_content:
accumulated = _THINKING_BLOCK_RE.sub("", accumulated + delta)
if "</function>" in accumulated or "</invoke>" in accumulated:
end_tag = "</function>" if "</function>" in accumulated else "</invoke>"
pos = accumulated.find(end_tag)
accumulated = accumulated[: pos + len(end_tag)]
@ -207,8 +206,6 @@ class LLM:
accumulated = normalize_tool_format(accumulated)
accumulated = fix_incomplete_tool_call(_truncate_to_first_function(accumulated))
accumulated = re.sub(r"<thinking[^>]*>.*?</thinking>", "", accumulated, flags=re.DOTALL)
yield LLMResponse(
content=accumulated,
tool_invocations=parse_tool_invocations(accumulated),