diff --git a/strix/llm/llm.py b/strix/llm/llm.py
index 5e6a01f..c3f07ff 100644
--- a/strix/llm/llm.py
+++ b/strix/llm/llm.py
@@ -1,4 +1,5 @@
import asyncio
+import re
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any
@@ -25,6 +26,19 @@ from strix.utils.resource_paths import get_strix_resource_path
litellm.drop_params = True
litellm.modify_params = True
+_THINKING_BLOCK_RE = re.compile(r"]*>.*?", re.DOTALL)
+_THINKING_BLOCK_OR_OPEN_RE = re.compile(r"]*>.*?(?:|\Z)", re.DOTALL)
+
+
+def _find_end_tag_outside_thinking(content: str, end_tag: str) -> int:
+ thinking_spans = [(m.start(), m.end()) for m in _THINKING_BLOCK_OR_OPEN_RE.finditer(content)]
+ start = 0
+ while (idx := content.find(end_tag, start)) != -1:
+ if not any(s <= idx < e for s, e in thinking_spans):
+ return idx
+ start = idx + 1
+ return -1
+
class LLMRequestFailedError(Exception):
def __init__(self, message: str, details: str | None = None):
@@ -197,9 +211,10 @@ class LLM:
delta = self._get_chunk_content(chunk)
if delta:
accumulated += delta
- if "" in accumulated or "" in accumulated:
- end_tag = "" if "" in accumulated else ""
- pos = accumulated.find(end_tag)
+ check_content = _THINKING_BLOCK_OR_OPEN_RE.sub("", accumulated)
+ if "" in check_content or "" in check_content:
+ end_tag = "" if "" in check_content else ""
+ pos = _find_end_tag_outside_thinking(accumulated, end_tag)
accumulated = accumulated[: pos + len(end_tag)]
yield LLMResponse(content=accumulated)
done_streaming = 1
@@ -209,8 +224,10 @@ class LLM:
if chunks:
self._update_usage_stats(stream_chunk_builder(chunks))
+ accumulated = _THINKING_BLOCK_RE.sub("", accumulated)
accumulated = normalize_tool_format(accumulated)
accumulated = fix_incomplete_tool_call(_truncate_to_first_function(accumulated))
+
yield LLMResponse(
content=accumulated,
tool_invocations=parse_tool_invocations(accumulated),
diff --git a/strix/llm/utils.py b/strix/llm/utils.py
index 9771854..0f56bed 100644
--- a/strix/llm/utils.py
+++ b/strix/llm/utils.py
@@ -20,7 +20,12 @@ def normalize_tool_format(content: str) -> str:
→
→
"""
- if "