mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-12 13:50:59 +00:00
fix(provider): strip think tags for MiniMax-M3 and other OpenAI-compatible models (#2330)
* fix(provider): strip think tags for MiniMax-M3 and other OpenAI-compatible models MiniMax-M3 (and other OpenAI-compatible providers) emit chain-of-thought reasoning directly in the content field wrapped in tags, instead of using a separate reasoning_content field or the legacy CRETIRE_REASONING markers. The existing remove_think logic only handled CRETIRE_* tags, so think blocks leaked into user-visible output even when remove_think was enabled. - Add _ThinkStripState: a stateful filter that correctly handles tags split across streaming chunk boundaries. - Add _strip_think classmethod with regex patterns for both and CRETIRE_* tags. - Wire think_state into invoke_llm_stream so deltas are filtered before reaching the accumulator. - Add remove_think safety net in _StreamAccumulator so the final message from tool-call rounds also gets stripped. - Fix remove_think resolution to use defensive nested .get() so pipelines missing output.misc don't raise AttributeError. * fix(litellmchat): add missing _CLOSE_TAG class attribute on _ThinkStripState * fix(provider): handle think stripping across LiteLLM paths --------- Co-authored-by: WangCham <651122857@qq.com>
This commit is contained in:
@@ -13,6 +13,151 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
|||||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
|
|
||||||
|
|
||||||
|
class _ThinkStripState:
|
||||||
|
"""Stateful filter that drops think blocks across chunks."""
|
||||||
|
|
||||||
|
_THINK_OPEN = '<think>'
|
||||||
|
_THINK_CLOSE = '</think>'
|
||||||
|
_LEGACY_OPEN = 'CRETIRE_REASONING_BEGINk'
|
||||||
|
_LEGACY_CLOSE = 'CRETIRE_REASONING_ENDk'
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._pairs: tuple[tuple[str, str], ...] = (
|
||||||
|
(self._THINK_OPEN, self._THINK_CLOSE),
|
||||||
|
(self._LEGACY_OPEN, self._LEGACY_CLOSE),
|
||||||
|
)
|
||||||
|
self._open_tags = tuple(open_tag for open_tag, _close_tag in self._pairs)
|
||||||
|
self._buf = ''
|
||||||
|
self._close_tag: str | None = None
|
||||||
|
self._pending_initial = True
|
||||||
|
|
||||||
|
def feed(self, chunk: str) -> str:
|
||||||
|
"""Feed a streaming delta and return user-visible content."""
|
||||||
|
if not chunk:
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
text = self._buf + chunk
|
||||||
|
if self._close_tag is not None:
|
||||||
|
return self._consume_think_body(text)
|
||||||
|
|
||||||
|
return self._process_visible_text(text)
|
||||||
|
|
||||||
|
def flush(self) -> str:
|
||||||
|
"""Release buffered visible content when the stream ends."""
|
||||||
|
if self._close_tag is not None:
|
||||||
|
self._buf = ''
|
||||||
|
self._close_tag = None
|
||||||
|
return ''
|
||||||
|
|
||||||
|
pending, self._buf = self._buf, ''
|
||||||
|
self._close_tag = None
|
||||||
|
return pending
|
||||||
|
|
||||||
|
def _consume_think_body(self, text: str) -> str:
|
||||||
|
close_tag = self._close_tag
|
||||||
|
if close_tag is None:
|
||||||
|
return text
|
||||||
|
|
||||||
|
close_idx = text.find(close_tag)
|
||||||
|
if close_idx != -1:
|
||||||
|
self._close_tag = None
|
||||||
|
self._buf = ''
|
||||||
|
self._pending_initial = False
|
||||||
|
return self._process_visible_text(text[close_idx + len(close_tag) :])
|
||||||
|
|
||||||
|
self._buf = self._close_prefix(text, close_tag)
|
||||||
|
return ''
|
||||||
|
|
||||||
|
def _process_visible_text(self, text: str) -> str:
|
||||||
|
out: list[str] = []
|
||||||
|
index = 0
|
||||||
|
|
||||||
|
while index < len(text):
|
||||||
|
if self._pending_initial:
|
||||||
|
open_idx, open_tag, close_tag = self._find_next_open(text, index)
|
||||||
|
orphan_close_idx, orphan_close_tag = self._find_next_close(text, index)
|
||||||
|
|
||||||
|
if orphan_close_idx != -1 and (open_idx == -1 or orphan_close_idx < open_idx):
|
||||||
|
self._pending_initial = False
|
||||||
|
index = orphan_close_idx + len(orphan_close_tag)
|
||||||
|
continue
|
||||||
|
|
||||||
|
if open_idx == -1:
|
||||||
|
self._buf = text[index:]
|
||||||
|
return ''.join(out)
|
||||||
|
|
||||||
|
if open_idx > index:
|
||||||
|
self._pending_initial = False
|
||||||
|
out.append(text[index:open_idx])
|
||||||
|
index = open_idx
|
||||||
|
continue
|
||||||
|
|
||||||
|
open_idx, open_tag, close_tag = self._find_next_open(text, index)
|
||||||
|
if open_idx == -1:
|
||||||
|
emit_end = self._visible_emit_end(text, index)
|
||||||
|
out.append(text[index:emit_end])
|
||||||
|
if emit_end > index:
|
||||||
|
self._pending_initial = False
|
||||||
|
self._buf = text[emit_end:]
|
||||||
|
return ''.join(out)
|
||||||
|
|
||||||
|
out.append(text[index:open_idx])
|
||||||
|
if open_idx > index:
|
||||||
|
self._pending_initial = False
|
||||||
|
body_start = open_idx + len(open_tag)
|
||||||
|
close_idx = text.find(close_tag, body_start)
|
||||||
|
if close_idx == -1:
|
||||||
|
self._close_tag = close_tag
|
||||||
|
self._buf = self._close_prefix(text[body_start:], close_tag)
|
||||||
|
return ''.join(out)
|
||||||
|
|
||||||
|
self._pending_initial = False
|
||||||
|
index = close_idx + len(close_tag)
|
||||||
|
|
||||||
|
self._buf = ''
|
||||||
|
return ''.join(out)
|
||||||
|
|
||||||
|
def _find_next_open(self, text: str, start: int) -> tuple[int, str, str]:
|
||||||
|
best_idx = -1
|
||||||
|
best_open = ''
|
||||||
|
best_close = ''
|
||||||
|
for open_tag, close_tag in self._pairs:
|
||||||
|
idx = text.find(open_tag, start)
|
||||||
|
if idx != -1 and (best_idx == -1 or idx < best_idx):
|
||||||
|
best_idx = idx
|
||||||
|
best_open = open_tag
|
||||||
|
best_close = close_tag
|
||||||
|
return best_idx, best_open, best_close
|
||||||
|
|
||||||
|
def _find_next_close(self, text: str, start: int) -> tuple[int, str]:
|
||||||
|
best_idx = -1
|
||||||
|
best_close = ''
|
||||||
|
for _open_tag, close_tag in self._pairs:
|
||||||
|
idx = text.find(close_tag, start)
|
||||||
|
if idx != -1 and (best_idx == -1 or idx < best_idx):
|
||||||
|
best_idx = idx
|
||||||
|
best_close = close_tag
|
||||||
|
return best_idx, best_close
|
||||||
|
|
||||||
|
def _visible_emit_end(self, text: str, start: int) -> int:
|
||||||
|
visible = text[start:]
|
||||||
|
limit = min(len(visible), max(len(open_tag) for open_tag in self._open_tags) - 1)
|
||||||
|
for keep in range(limit, 0, -1):
|
||||||
|
suffix = visible[-keep:]
|
||||||
|
if any(open_tag.startswith(suffix) for open_tag in self._open_tags):
|
||||||
|
return len(text) - keep
|
||||||
|
return len(text)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _close_prefix(text: str, close_tag: str) -> str:
|
||||||
|
limit = min(len(text), len(close_tag) - 1)
|
||||||
|
for keep in range(limit, 0, -1):
|
||||||
|
suffix = text[-keep:]
|
||||||
|
if close_tag.startswith(suffix):
|
||||||
|
return suffix
|
||||||
|
return ''
|
||||||
|
|
||||||
|
|
||||||
class LiteLLMRequester(requester.ProviderAPIRequester):
|
class LiteLLMRequester(requester.ProviderAPIRequester):
|
||||||
"""LiteLLM unified API requester supporting chat, embedding, and rerank."""
|
"""LiteLLM unified API requester supporting chat, embedding, and rerank."""
|
||||||
|
|
||||||
@@ -237,6 +382,25 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
|
|
||||||
return req_messages
|
return req_messages
|
||||||
|
|
||||||
|
_THINK_PATTERNS: tuple[str, ...] = (
|
||||||
|
r'^\s*(?:(?!<think>).)*?</think>\s*',
|
||||||
|
r'^\s*(?:(?!CRETIRE_REASONING_BEGINk).)*?CRETIRE_REASONING_ENDk\s*',
|
||||||
|
r'<think>.*?</think>',
|
||||||
|
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:
|
def _process_thinking_content(self, content: str, reasoning_content: str | None, remove_think: bool) -> str:
|
||||||
"""Process thinking/reasoning content.
|
"""Process thinking/reasoning content.
|
||||||
|
|
||||||
@@ -248,20 +412,12 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
Returns:
|
Returns:
|
||||||
Processed content string
|
Processed content string
|
||||||
"""
|
"""
|
||||||
# Extract and handle thinking tags
|
if remove_think and content:
|
||||||
if content and 'CRETIRE_REASONING_BEGINk' in content and 'CRETIRE_REASONING_ENDk' in content:
|
content = self._strip_think(content)
|
||||||
import re
|
|
||||||
|
|
||||||
think_pattern = r'CRETIRE_REASONING_BEGINk(.*?)CRETIRE_REASONING_ENDk'
|
if reasoning_content and not remove_think:
|
||||||
|
content = f'<think>\n{reasoning_content}\n</think>\n{content or ""}'.strip()
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
# Handle separate reasoning_content field
|
|
||||||
# Currently we don't include reasoning_content in user-facing output regardless of remove_think
|
|
||||||
# because it's typically internal model reasoning, not user-visible thinking
|
|
||||||
return content or ''
|
return content or ''
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@@ -570,6 +726,7 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
chunk_idx = 0
|
chunk_idx = 0
|
||||||
role = 'assistant'
|
role = 'assistant'
|
||||||
tool_call_state: dict[int, dict[str, typing.Any]] = {}
|
tool_call_state: dict[int, dict[str, typing.Any]] = {}
|
||||||
|
think_state = _ThinkStripState() if remove_think else None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
response = await acompletion(**args)
|
response = await acompletion(**args)
|
||||||
@@ -613,6 +770,12 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
# Use reasoning_content as the displayed content
|
# Use reasoning_content as the displayed content
|
||||||
delta_content = reasoning_content
|
delta_content = reasoning_content
|
||||||
|
|
||||||
|
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)
|
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:
|
if chunk_idx == 0 and not delta_content and not tool_calls:
|
||||||
@@ -634,6 +797,15 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
yield provider_message.MessageChunk(**chunk_data)
|
yield provider_message.MessageChunk(**chunk_data)
|
||||||
chunk_idx += 1
|
chunk_idx += 1
|
||||||
|
|
||||||
|
if think_state is not None:
|
||||||
|
pending_content = think_state.flush()
|
||||||
|
if pending_content:
|
||||||
|
yield provider_message.MessageChunk(
|
||||||
|
role=role,
|
||||||
|
content=pending_content,
|
||||||
|
is_final=True,
|
||||||
|
)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self._handle_litellm_error(e)
|
self._handle_litellm_error(e)
|
||||||
|
|
||||||
|
|||||||
@@ -49,12 +49,23 @@ def _model_has_ability(model: modelmgr_requester.RuntimeLLMModel, ability: str)
|
|||||||
class _StreamAccumulator:
|
class _StreamAccumulator:
|
||||||
"""Accumulate streamed content and fragmented OpenAI-style tool calls."""
|
"""Accumulate streamed content and fragmented OpenAI-style tool calls."""
|
||||||
|
|
||||||
def __init__(self, msg_sequence: int = 0, initial_content: str | None = None):
|
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.tool_calls_map: dict[str, provider_message.ToolCall] = {}
|
||||||
self.msg_idx = 0
|
self.msg_idx = 0
|
||||||
self.accumulated_content = initial_content or ''
|
self.accumulated_content = initial_content or ''
|
||||||
self.last_role = 'assistant'
|
self.last_role = 'assistant'
|
||||||
self.msg_sequence = msg_sequence
|
self.msg_sequence = msg_sequence
|
||||||
|
self.remove_think = remove_think
|
||||||
|
self._think_state = None
|
||||||
|
if remove_think:
|
||||||
|
from ..modelmgr.requesters.litellmchat import _ThinkStripState
|
||||||
|
|
||||||
|
self._think_state = _ThinkStripState()
|
||||||
|
|
||||||
def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChunk | None:
|
def add(self, msg: provider_message.MessageChunk) -> provider_message.MessageChunk | None:
|
||||||
self.msg_idx += 1
|
self.msg_idx += 1
|
||||||
@@ -63,7 +74,10 @@ class _StreamAccumulator:
|
|||||||
self.last_role = msg.role
|
self.last_role = msg.role
|
||||||
|
|
||||||
if msg.content:
|
if msg.content:
|
||||||
self.accumulated_content += msg.content
|
content = msg.content
|
||||||
|
if self._think_state is not None:
|
||||||
|
content = self._think_state.feed(content)
|
||||||
|
self.accumulated_content += content
|
||||||
|
|
||||||
if msg.tool_calls:
|
if msg.tool_calls:
|
||||||
for tool_call in msg.tool_calls:
|
for tool_call in msg.tool_calls:
|
||||||
@@ -79,11 +93,14 @@ class _StreamAccumulator:
|
|||||||
if tool_call.function and tool_call.function.arguments:
|
if tool_call.function and tool_call.function.arguments:
|
||||||
self.tool_calls_map[tool_call.id].function.arguments += tool_call.function.arguments
|
self.tool_calls_map[tool_call.id].function.arguments += tool_call.function.arguments
|
||||||
|
|
||||||
|
if msg.is_final:
|
||||||
|
self._flush_think_state()
|
||||||
|
|
||||||
if self.msg_idx % 8 == 0 or msg.is_final:
|
if self.msg_idx % 8 == 0 or msg.is_final:
|
||||||
self.msg_sequence += 1
|
self.msg_sequence += 1
|
||||||
return provider_message.MessageChunk(
|
return provider_message.MessageChunk(
|
||||||
role=self.last_role,
|
role=self.last_role,
|
||||||
content=self.accumulated_content,
|
content=self._maybe_strip_think(self.accumulated_content),
|
||||||
tool_calls=list(self.tool_calls_map.values()) if (self.tool_calls_map and msg.is_final) else None,
|
tool_calls=list(self.tool_calls_map.values()) if (self.tool_calls_map and msg.is_final) else None,
|
||||||
is_final=msg.is_final,
|
is_final=msg.is_final,
|
||||||
msg_sequence=self.msg_sequence,
|
msg_sequence=self.msg_sequence,
|
||||||
@@ -92,13 +109,29 @@ class _StreamAccumulator:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def final_message(self) -> provider_message.MessageChunk:
|
def final_message(self) -> provider_message.MessageChunk:
|
||||||
|
self._flush_think_state()
|
||||||
return provider_message.MessageChunk(
|
return provider_message.MessageChunk(
|
||||||
role=self.last_role,
|
role=self.last_role,
|
||||||
content=self.accumulated_content,
|
content=self._maybe_strip_think(self.accumulated_content),
|
||||||
tool_calls=list(self.tool_calls_map.values()) if self.tool_calls_map else None,
|
tool_calls=list(self.tool_calls_map.values()) if self.tool_calls_map else None,
|
||||||
msg_sequence=self.msg_sequence,
|
msg_sequence=self.msg_sequence,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _maybe_strip_think(self, content: str) -> str:
|
||||||
|
if not self.remove_think or not content:
|
||||||
|
return content
|
||||||
|
|
||||||
|
from ..modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||||
|
|
||||||
|
return LiteLLMRequester._strip_think(content)
|
||||||
|
|
||||||
|
def _flush_think_state(self) -> None:
|
||||||
|
if self._think_state is None:
|
||||||
|
return
|
||||||
|
pending = self._think_state.flush()
|
||||||
|
if pending:
|
||||||
|
self.accumulated_content += pending
|
||||||
|
|
||||||
|
|
||||||
@runner.runner_class('local-agent')
|
@runner.runner_class('local-agent')
|
||||||
class LocalAgentRunner(runner.RequestRunner):
|
class LocalAgentRunner(runner.RequestRunner):
|
||||||
@@ -448,7 +481,7 @@ class LocalAgentRunner(runner.RequestRunner):
|
|||||||
except AttributeError:
|
except AttributeError:
|
||||||
is_stream = False
|
is_stream = False
|
||||||
|
|
||||||
remove_think = query.pipeline_config['output'].get('misc', '').get('remove-think')
|
remove_think = ((query.pipeline_config.get('output') or {}).get('misc') or {}).get('remove-think', False)
|
||||||
|
|
||||||
# Build ordered candidate list (primary + fallbacks)
|
# Build ordered candidate list (primary + fallbacks)
|
||||||
candidates = await self._get_model_candidates(query)
|
candidates = await self._get_model_candidates(query)
|
||||||
@@ -472,7 +505,7 @@ class LocalAgentRunner(runner.RequestRunner):
|
|||||||
final_msg = msg
|
final_msg = msg
|
||||||
else:
|
else:
|
||||||
# Streaming: invoke with fallback
|
# Streaming: invoke with fallback
|
||||||
stream_accumulator = _StreamAccumulator(msg_sequence=1)
|
stream_accumulator = _StreamAccumulator(msg_sequence=1, remove_think=remove_think)
|
||||||
|
|
||||||
stream_src, use_llm_model = await self._invoke_stream_with_fallback(
|
stream_src, use_llm_model = await self._invoke_stream_with_fallback(
|
||||||
query,
|
query,
|
||||||
@@ -576,6 +609,7 @@ class LocalAgentRunner(runner.RequestRunner):
|
|||||||
stream_accumulator = _StreamAccumulator(
|
stream_accumulator = _StreamAccumulator(
|
||||||
msg_sequence=first_end_sequence,
|
msg_sequence=first_end_sequence,
|
||||||
initial_content=first_content,
|
initial_content=first_content,
|
||||||
|
remove_think=remove_think,
|
||||||
)
|
)
|
||||||
|
|
||||||
tool_stream_src = use_llm_model.provider.invoke_llm_stream(
|
tool_stream_src = use_llm_model.provider.invoke_llm_stream(
|
||||||
|
|||||||
@@ -279,6 +279,122 @@ class TestInvokeLLMStreamUsage:
|
|||||||
|
|
||||||
assert query.variables['_stream_usage']['total_tokens'] == 12
|
assert query.variables['_stream_usage']['total_tokens'] == 12
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stream_removes_leading_think_across_chunks(self):
|
||||||
|
"""A leading think block split across chunks must be removed."""
|
||||||
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
|
|
||||||
|
mock_ap = Mock()
|
||||||
|
mock_ap.tool_mgr = Mock()
|
||||||
|
mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None)
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={})
|
||||||
|
model = MockRuntimeModel('minimax-m3', 'test-api-key')
|
||||||
|
|
||||||
|
chunks = [
|
||||||
|
self._make_chunk(content='<thi'),
|
||||||
|
self._make_chunk(content='nk>hidden'),
|
||||||
|
self._make_chunk(content=' reasoning</thi'),
|
||||||
|
self._make_chunk(content='nk>Visible answer', finish_reason='stop'),
|
||||||
|
]
|
||||||
|
|
||||||
|
async def _aiter(*args, **kwargs):
|
||||||
|
for c in chunks:
|
||||||
|
yield c
|
||||||
|
|
||||||
|
query = Mock(spec=pipeline_query.Query)
|
||||||
|
query.variables = {}
|
||||||
|
messages = [provider_message.Message(role='user', content='Hi')]
|
||||||
|
|
||||||
|
with patch.object(litellmchat, 'acompletion', new=AsyncMock(side_effect=lambda **kw: _aiter())):
|
||||||
|
collected = [
|
||||||
|
chunk
|
||||||
|
async for chunk in requester.invoke_llm_stream(
|
||||||
|
query=query,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
remove_think=True,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
assert ''.join(chunk.content or '' for chunk in collected) == 'Visible answer'
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stream_removes_initial_orphan_think_close(self):
|
||||||
|
"""Initial reasoning content without an open tag is removed until </think>."""
|
||||||
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
|
|
||||||
|
mock_ap = Mock()
|
||||||
|
mock_ap.tool_mgr = Mock()
|
||||||
|
mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None)
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={})
|
||||||
|
model = MockRuntimeModel('minimax-m3', 'test-api-key')
|
||||||
|
|
||||||
|
chunks = [
|
||||||
|
self._make_chunk(content='hidden reasoning'),
|
||||||
|
self._make_chunk(content=' still hidden</thi'),
|
||||||
|
self._make_chunk(content='nk>Visible answer', finish_reason='stop'),
|
||||||
|
]
|
||||||
|
|
||||||
|
async def _aiter(*args, **kwargs):
|
||||||
|
for c in chunks:
|
||||||
|
yield c
|
||||||
|
|
||||||
|
query = Mock(spec=pipeline_query.Query)
|
||||||
|
query.variables = {}
|
||||||
|
messages = [provider_message.Message(role='user', content='Hi')]
|
||||||
|
|
||||||
|
with patch.object(litellmchat, 'acompletion', new=AsyncMock(side_effect=lambda **kw: _aiter())):
|
||||||
|
collected = [
|
||||||
|
chunk
|
||||||
|
async for chunk in requester.invoke_llm_stream(
|
||||||
|
query=query,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
remove_think=True,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
assert ''.join(chunk.content or '' for chunk in collected) == 'Visible answer'
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stream_removes_non_leading_think_content(self):
|
||||||
|
"""A think block in the answer body is removed with its content."""
|
||||||
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
|
|
||||||
|
mock_ap = Mock()
|
||||||
|
mock_ap.tool_mgr = Mock()
|
||||||
|
mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None)
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={})
|
||||||
|
model = MockRuntimeModel('gpt-4o', 'test-api-key')
|
||||||
|
|
||||||
|
chunks = [
|
||||||
|
self._make_chunk(content='Use <think>x</think> as an XML-like example.', finish_reason='stop'),
|
||||||
|
]
|
||||||
|
|
||||||
|
async def _aiter(*args, **kwargs):
|
||||||
|
for c in chunks:
|
||||||
|
yield c
|
||||||
|
|
||||||
|
query = Mock(spec=pipeline_query.Query)
|
||||||
|
query.variables = {}
|
||||||
|
messages = [provider_message.Message(role='user', content='Hi')]
|
||||||
|
|
||||||
|
with patch.object(litellmchat, 'acompletion', new=AsyncMock(side_effect=lambda **kw: _aiter())):
|
||||||
|
collected = [
|
||||||
|
chunk
|
||||||
|
async for chunk in requester.invoke_llm_stream(
|
||||||
|
query=query,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
remove_think=True,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
assert ''.join(chunk.content or '' for chunk in collected) == 'Use as an XML-like example.'
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_stream_tool_call_delta_missing_id_and_name(self):
|
async def test_stream_tool_call_delta_missing_id_and_name(self):
|
||||||
"""LiteLLM may stream tool-call argument deltas with id/name set to None."""
|
"""LiteLLM may stream tool-call argument deltas with id/name set to None."""
|
||||||
@@ -482,6 +598,38 @@ class TestProcessThinkingContent:
|
|||||||
result = requester._process_thinking_content(content, None, remove_think=True)
|
result = requester._process_thinking_content(content, None, remove_think=True)
|
||||||
assert result == 'The answer is 42.'
|
assert result == 'The answer is 42.'
|
||||||
|
|
||||||
|
def test_remove_leading_think_tag(self):
|
||||||
|
"""Test removing a leading <think> block when remove_think=True"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
content = '<think>Let me think...</think> The answer is 42.'
|
||||||
|
result = requester._process_thinking_content(content, None, remove_think=True)
|
||||||
|
assert result == 'The answer is 42.'
|
||||||
|
|
||||||
|
def test_remove_non_leading_think_tag(self):
|
||||||
|
"""Test removing <think> and its content in the answer body"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
content = 'Use <think>example</think> in the document.'
|
||||||
|
result = requester._process_thinking_content(content, None, remove_think=True)
|
||||||
|
assert result == 'Use in the document.'
|
||||||
|
|
||||||
|
def test_remove_initial_orphan_think_close(self):
|
||||||
|
"""Test removing leading reasoning content when only </think> is visible"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
content = 'hidden reasoning</think> Visible answer.'
|
||||||
|
result = requester._process_thinking_content(content, None, remove_think=True)
|
||||||
|
assert result == 'Visible answer.'
|
||||||
|
|
||||||
|
def test_remove_multiple_think_tags(self):
|
||||||
|
"""Test removing multiple <think> blocks"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
content = '<think>hidden</think> Keep <think>example</think>.'
|
||||||
|
result = requester._process_thinking_content(content, None, remove_think=True)
|
||||||
|
assert result == 'Keep .'
|
||||||
|
|
||||||
def test_preserve_thinking_markers(self):
|
def test_preserve_thinking_markers(self):
|
||||||
"""Test preserving thinking markers when remove_think=False"""
|
"""Test preserving thinking markers when remove_think=False"""
|
||||||
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
@@ -491,6 +639,20 @@ class TestProcessThinkingContent:
|
|||||||
assert 'CRETIRE_REASONING_BEGINk' in result
|
assert 'CRETIRE_REASONING_BEGINk' in result
|
||||||
assert 'The answer is 42.' in result
|
assert 'The answer is 42.' in result
|
||||||
|
|
||||||
|
def test_preserve_reasoning_content_when_remove_think_false(self):
|
||||||
|
"""Test showing separate reasoning_content when remove_think=False"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
result = requester._process_thinking_content('The answer is 42.', 'Let me think...', remove_think=False)
|
||||||
|
assert result == '<think>\nLet me think...\n</think>\nThe answer is 42.'
|
||||||
|
|
||||||
|
def test_hide_reasoning_content_when_remove_think_true(self):
|
||||||
|
"""Test hiding separate reasoning_content when remove_think=True"""
|
||||||
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|
||||||
|
result = requester._process_thinking_content('The answer is 42.', 'Let me think...', remove_think=True)
|
||||||
|
assert result == 'The answer is 42.'
|
||||||
|
|
||||||
def test_empty_content(self):
|
def test_empty_content(self):
|
||||||
"""Test empty content"""
|
"""Test empty content"""
|
||||||
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
requester = litellmchat.LiteLLMRequester(ap=Mock(), config={})
|
||||||
|
|||||||
@@ -163,6 +163,51 @@ def test_stream_accumulator_merges_fragmented_tool_call_arguments():
|
|||||||
assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}'
|
assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}'
|
||||||
|
|
||||||
|
|
||||||
|
def test_stream_accumulator_strips_leading_think_from_tool_round_content():
|
||||||
|
accumulator = _StreamAccumulator(
|
||||||
|
msg_sequence=3,
|
||||||
|
initial_content='I will search for LangBot.',
|
||||||
|
remove_think=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert accumulator.add(provider_message.MessageChunk(role='assistant', content='<thi')) is None
|
||||||
|
assert accumulator.add(provider_message.MessageChunk(role='assistant', content='nk>hidden')) is None
|
||||||
|
emitted = accumulator.add(
|
||||||
|
provider_message.MessageChunk(
|
||||||
|
role='assistant',
|
||||||
|
content=' reasoning</think>Here is the answer.',
|
||||||
|
is_final=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert emitted is not None
|
||||||
|
assert emitted.content == 'I will search for LangBot.Here is the answer.'
|
||||||
|
assert '<think>' not in emitted.content
|
||||||
|
assert 'hidden reasoning' not in emitted.content
|
||||||
|
|
||||||
|
|
||||||
|
def test_stream_accumulator_strips_initial_orphan_think_close_from_tool_round_content():
|
||||||
|
accumulator = _StreamAccumulator(
|
||||||
|
msg_sequence=3,
|
||||||
|
initial_content='I will search for LangBot.',
|
||||||
|
remove_think=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert accumulator.add(provider_message.MessageChunk(role='assistant', content='hidden reasoning')) is None
|
||||||
|
emitted = accumulator.add(
|
||||||
|
provider_message.MessageChunk(
|
||||||
|
role='assistant',
|
||||||
|
content=' still hidden</think>Here is the answer.',
|
||||||
|
is_final=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert emitted is not None
|
||||||
|
assert emitted.content == 'I will search for LangBot.Here is the answer.'
|
||||||
|
assert '</think>' not in emitted.content
|
||||||
|
assert 'hidden reasoning' not in emitted.content
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_localagent_uses_exec_for_exact_calculation():
|
async def test_localagent_uses_exec_for_exact_calculation():
|
||||||
provider = RecordingProvider()
|
provider = RecordingProvider()
|
||||||
|
|||||||
Reference in New Issue
Block a user