diff --git a/src/llm/providers/astraflow.py b/src/llm/providers/astraflow.py index 398242a2b..2bcc6d6b2 100644 --- a/src/llm/providers/astraflow.py +++ b/src/llm/providers/astraflow.py @@ -17,7 +17,6 @@ from llm.core.interface import ( from llm.core.types import LLMInput, LLMOutput, ModelInfo, ProviderType, ToolCall from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR from llm.providers.openai_compat import to_openai_message_param -from llm.providers.reasoning import strip_reasoning ASTRAFLOW_BASE_URL = "https://api.umodelverse.ai/v1" ASTRAFLOW_CN_BASE_URL = "https://api.modelverse.cn/v1" @@ -106,7 +105,7 @@ class _AstraflowBaseProvider(LLMProvider): } return LLMOutput( - content=strip_reasoning(choice.message.content or ""), + content=choice.message.content or "", tool_calls=tool_calls, model=response.model, usage=usage, diff --git a/src/llm/providers/atlas.py b/src/llm/providers/atlas.py index 8d1dc739f..d626509cc 100644 --- a/src/llm/providers/atlas.py +++ b/src/llm/providers/atlas.py @@ -17,7 +17,6 @@ from llm.core.interface import ( from llm.core.types import LLMInput, LLMOutput, ModelInfo, ProviderType, ToolCall from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR from llm.providers.openai_compat import to_openai_message_param -from llm.providers.reasoning import strip_reasoning ATLAS_BASE_URL = "https://api.atlascloud.ai/v1" DEFAULT_ATLAS_MODEL = "deepseek-ai/deepseek-v4-pro" @@ -123,7 +122,7 @@ class AtlasProvider(LLMProvider): } return LLMOutput( - content=strip_reasoning(choice.message.content or ""), + content=choice.message.content or "", tool_calls=tool_calls, model=response.model, usage=usage, diff --git a/src/llm/providers/ollama.py b/src/llm/providers/ollama.py index 0787c250e..e0669563c 100644 --- a/src/llm/providers/ollama.py +++ b/src/llm/providers/ollama.py @@ -58,7 +58,15 @@ class OllamaProvider(LLMProvider): base_url: str | None = None, default_model: str | None = None, prefilled_reasoning: bool | None = None, + leading_reasoning: bool | None = None, ) -> None: + # Enable only for a deployment that declares balanced leading reasoning. + # Generic responses can contain literal XML and stay intact by default. + self.leading_reasoning = ( + leading_reasoning + if leading_reasoning is not None + else prefilled_reasoning_from_env("OLLAMA_LEADING_REASONING") + ) # Set for models whose chat template prefills , so replies arrive # as "...answer" without the opening tag. self.prefilled_reasoning = ( @@ -158,7 +166,9 @@ class OllamaProvider(LLMProvider): result = json.loads(response.read().decode("utf-8")) content = strip_reasoning( - result.get("message", {}).get("content", ""), prefilled=self.prefilled_reasoning + result.get("message", {}).get("content", ""), + leading=self.leading_reasoning, + prefilled=self.prefilled_reasoning ) tool_calls = None diff --git a/src/llm/providers/openai.py b/src/llm/providers/openai.py index b81d4ab08..34b21a8e5 100644 --- a/src/llm/providers/openai.py +++ b/src/llm/providers/openai.py @@ -34,7 +34,15 @@ class OpenAIProvider(LLMProvider): api_key: str | None = None, base_url: str | None = None, prefilled_reasoning: bool | None = None, + leading_reasoning: bool | None = None, ) -> None: + # Enable only for a deployment that declares balanced leading reasoning. + # Generic responses can contain literal XML and stay intact by default. + self.leading_reasoning = ( + leading_reasoning + if leading_reasoning is not None + else prefilled_reasoning_from_env("OPENAI_LEADING_REASONING") + ) # Set for self-hosted servers whose chat template prefills , so # replies arrive as "...answer" without the opening tag. self.prefilled_reasoning = ( @@ -119,7 +127,11 @@ class OpenAIProvider(LLMProvider): } return LLMOutput( - content=strip_reasoning(choice.message.content or "", prefilled=self.prefilled_reasoning), + content=strip_reasoning( + choice.message.content or "", + leading=self.leading_reasoning, + prefilled=self.prefilled_reasoning, + ), tool_calls=tool_calls, model=response.model, usage=usage, diff --git a/src/llm/providers/reasoning.py b/src/llm/providers/reasoning.py index 4a45c1fc1..70edc2ebe 100644 --- a/src/llm/providers/reasoning.py +++ b/src/llm/providers/reasoning.py @@ -12,28 +12,30 @@ _TRUTHY = {"1", "true", "yes", "on"} def prefilled_reasoning_from_env(name: str) -> bool: - """Read a ``*_PREFILLED_REASONING`` flag from the environment.""" + """Read an explicit reasoning-format flag from the environment.""" return os.environ.get(name, "").strip().lower() in _TRUTHY -def strip_reasoning(content: str, *, prefilled: bool = False) -> str: +def strip_reasoning( + content: str, *, leading: bool = False, prefilled: bool = False +) -> str: """Drop a leading reasoning block from model output. - Hosted APIs return reasoning in a separate field, but self-hosted reasoning - models (vLLM or llama.cpp without a reasoning parser, Ollama models without - a thinking template) put it in ``content``. + Some deployments put reasoning in ``content`` instead of a separate field. + This helper makes no assumptions about a hosted API or model name. - By default only a balanced block that opens the output is removed: - ``...answer``. Anything else is returned unchanged, so an - answer that merely mentions the tags (including a lone or reversed closing - tag) is never truncated, and neither is unterminated reasoning (for example - when ``max_tokens`` cut the model off before it reached an answer). + Content is preserved by default, including complete leading XML documents. + Enable ``leading=True`` only for a deployment whose declared format is + ``...answer``. Unterminated reasoning remains unchanged. ``prefilled=True`` is for deployments whose chat template prefills the opening tag in the generation prompt, so the model only ever emits ``...answer``. Only enable it when the model is known to behave this way: everything up to the first closing tag is then treated as reasoning. """ + if not leading and not prefilled: + return content + stripped = content.lstrip() if stripped.startswith(THINK_OPEN): close = stripped.find(THINK_CLOSE, len(THINK_OPEN)) diff --git a/tests/test_reasoning.py b/tests/test_reasoning.py index c8d923c17..05b115e22 100644 --- a/tests/test_reasoning.py +++ b/tests/test_reasoning.py @@ -9,6 +9,7 @@ import pytest from llm.core.interface import LLMProvider from llm.core.types import LLMInput, LLMOutput, Message, Role from llm.providers.astraflow import AstraflowProvider +from llm.providers.atlas import AtlasProvider from llm.providers.ollama import OllamaProvider from llm.providers.openai import OpenAIProvider from llm.providers.reasoning import strip_reasoning @@ -38,8 +39,8 @@ LITERAL_ANSWERS = [ ("17 * 3 = 51.51", "17 * 3 = 51.51"), ], ) -def test_strip_reasoning_default(content: str, expected: str) -> None: - assert strip_reasoning(content) == expected +def test_strip_reasoning_declared_leading(content: str, expected: str) -> None: + assert strip_reasoning(content, leading=True) == expected @pytest.mark.parametrize("content", LITERAL_ANSWERS) @@ -54,11 +55,17 @@ def test_strip_reasoning_default_preserves_literal_tags(content: str) -> None: ('build the object\n\n{"answer": 51}', '{"answer": 51}'), ("17 * 3 = 51.51", "51"), ("plain answer", "plain answer"), - ("Wrap it in and tags.", "Wrap it in and tags."), + ( + "Wrap it in and tags.", + "Wrap it in and tags.", + ), ("17 * 3 = 51.", ""), # A complete JSON answer is never reasoning plus a closing tag. ('{"closing":""}', '{"closing":""}'), - (' {"closing":"","opening":""}\n', ' {"closing":"","opening":""}\n'), + ( + ' {"closing":"","opening":""}\n', + ' {"closing":"","opening":""}\n', + ), ('["", 1]', '["", 1]'), ], ) @@ -82,18 +89,27 @@ def _openai_provider(content: str, **kwargs: Any) -> OpenAIProvider: usage=None, ) client: Any = SimpleNamespace( - chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **params: response)) + chat=SimpleNamespace( + completions=SimpleNamespace(create=lambda **params: response) + ) ) provider.client = client return provider def _ask(provider: LLMProvider, **kwargs: Any) -> LLMOutput: - return provider.generate(LLMInput(messages=[Message(role=Role.USER, content="17*3?")], **kwargs)) + return provider.generate( + LLMInput(messages=[Message(role=Role.USER, content="17*3?")], **kwargs) + ) def test_openai_provider_strips_balanced_reasoning() -> None: - assert _ask(_openai_provider("17 * 3 = 51.51")).content == "51" + assert ( + _ask( + _openai_provider("17 * 3 = 51.51", leading_reasoning=True) + ).content + == "51" + ) @pytest.mark.parametrize("content", LITERAL_ANSWERS) @@ -101,16 +117,28 @@ def test_openai_provider_keeps_literal_tags(content: str) -> None: assert _ask(_openai_provider(content)).content == content -def test_openai_provider_close_only_needs_opt_in(monkeypatch: pytest.MonkeyPatch) -> None: +def test_openai_provider_close_only_needs_opt_in( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("OPENAI_PREFILLED_REASONING", raising=False) - assert _ask(_openai_provider("17 * 3 = 51.51")).content == "17 * 3 = 51.51" - assert _ask(_openai_provider("17 * 3 = 51.51", prefilled_reasoning=True)).content == "51" + assert ( + _ask(_openai_provider("17 * 3 = 51.51")).content + == "17 * 3 = 51.51" + ) + assert ( + _ask( + _openai_provider("17 * 3 = 51.51", prefilled_reasoning=True) + ).content + == "51" + ) monkeypatch.setenv("OPENAI_PREFILLED_REASONING", "true") assert _ask(_openai_provider("17 * 3 = 51.51")).content == "51" -def test_hosted_provider_never_strips_close_only(monkeypatch: pytest.MonkeyPatch) -> None: +def test_hosted_provider_never_strips_close_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setenv("OPENAI_PREFILLED_REASONING", "true") provider = AstraflowProvider(api_key="test") message = SimpleNamespace(content='{"closing":""}', tool_calls=None) @@ -120,7 +148,9 @@ def test_hosted_provider_never_strips_close_only(monkeypatch: pytest.MonkeyPatch usage=None, ) client: Any = SimpleNamespace( - chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **params: response)) + chat=SimpleNamespace( + completions=SimpleNamespace(create=lambda **params: response) + ) ) provider.client = client @@ -132,21 +162,29 @@ def _fake_ollama(monkeypatch: pytest.MonkeyPatch, content: str) -> list[dict[str def fake_urlopen(request: urllib.request.Request, timeout: float) -> BytesIO: sent.append(json.loads(cast(bytes, request.data))) - return BytesIO(json.dumps({"message": {"content": content}, "done_reason": "stop"}).encode()) + return BytesIO( + json.dumps( + {"message": {"content": content}, "done_reason": "stop"} + ).encode() + ) monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) return sent @pytest.mark.parametrize("content", LITERAL_ANSWERS) -def test_ollama_provider_keeps_literal_tags(monkeypatch: pytest.MonkeyPatch, content: str) -> None: +def test_ollama_provider_keeps_literal_tags( + monkeypatch: pytest.MonkeyPatch, content: str +) -> None: monkeypatch.delenv("OLLAMA_PREFILLED_REASONING", raising=False) _fake_ollama(monkeypatch, content) assert _ask(OllamaProvider()).content == content -def test_ollama_provider_strips_prefilled_reasoning_when_declared(monkeypatch: pytest.MonkeyPatch) -> None: +def test_ollama_provider_strips_prefilled_reasoning_when_declared( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("OLLAMA_PREFILLED_REASONING", raising=False) _fake_ollama(monkeypatch, "17 * 3 = 51.51") assert _ask(OllamaProvider()).content == "17 * 3 = 51.51" @@ -156,13 +194,86 @@ def test_ollama_provider_strips_prefilled_reasoning_when_declared(monkeypatch: p assert _ask(OllamaProvider()).content == "51" -def test_ollama_stream_flag_still_strips_complete_message(monkeypatch: pytest.MonkeyPatch) -> None: +def test_ollama_stream_flag_still_strips_complete_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: # The provider always requests a single complete message, even when the # caller asks for streaming, so stripping sees the whole reply. If chunked # streaming is added, this needs an incremental stripper instead. sent = _fake_ollama(monkeypatch, 'plan{"closing":""}') - output = _ask(OllamaProvider(), stream=True) + output = _ask(OllamaProvider(leading_reasoning=True), stream=True) assert sent[0]["stream"] is False assert output.content == '{"closing":""}' + + +@pytest.mark.parametrize( + "content", + [ + "display this XML literally", + "display this XML literallykeep both", + ], +) +def test_default_provider_preserves_complete_leading_xml( + monkeypatch: pytest.MonkeyPatch, content: str +) -> None: + monkeypatch.delenv("OPENAI_LEADING_REASONING", raising=False) + monkeypatch.delenv("OPENAI_PREFILLED_REASONING", raising=False) + monkeypatch.delenv("OLLAMA_LEADING_REASONING", raising=False) + monkeypatch.delenv("OLLAMA_PREFILLED_REASONING", raising=False) + assert _ask(_openai_provider(content)).content == content + _fake_ollama(monkeypatch, content) + assert _ask(OllamaProvider()).content == content + + +def test_balanced_leading_reasoning_requires_explicit_format( + monkeypatch: pytest.MonkeyPatch, +) -> None: + content = "plananswer" + assert _ask(_openai_provider(content, leading_reasoning=True)).content == "answer" + monkeypatch.setenv("OPENAI_LEADING_REASONING", "true") + assert _ask(_openai_provider(content)).content == "answer" + assert _ask(_openai_provider(content, leading_reasoning=False)).content == content + _fake_ollama(monkeypatch, content) + assert _ask(OllamaProvider(leading_reasoning=True)).content == "answer" + monkeypatch.setenv("OLLAMA_LEADING_REASONING", "true") + assert _ask(OllamaProvider()).content == "answer" + assert _ask(OllamaProvider(leading_reasoning=False)).content == content + + +@pytest.mark.parametrize("provider_type", [AtlasProvider, AstraflowProvider]) +@pytest.mark.parametrize( + "content", + ["literal XML", "literal XMLwhole"], +) +def test_hosted_providers_preserve_leading_xml( + monkeypatch: pytest.MonkeyPatch, provider_type: Any, content: str +) -> None: + monkeypatch.setenv("OPENAI_LEADING_REASONING", "true") + monkeypatch.setenv("OPENAI_PREFILLED_REASONING", "true") + provider = provider_type(api_key="test") + response = SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content=content, tool_calls=None), + finish_reason="stop", + ) + ], + model="fixture", + usage=None, + ) + provider.client = SimpleNamespace( + chat=SimpleNamespace( + completions=SimpleNamespace(create=lambda **params: response) + ) + ) + assert _ask(provider).content == content + + +@pytest.mark.parametrize( + "content", + ["literal XML", "literal XMLwhole"], +) +def test_helper_defaults_preserve_leading_xml(content: str) -> None: + assert strip_reasoning(content) == content