diff --git a/src/llm/providers/astraflow.py b/src/llm/providers/astraflow.py index 2bcc6d6b2..398242a2b 100644 --- a/src/llm/providers/astraflow.py +++ b/src/llm/providers/astraflow.py @@ -17,6 +17,7 @@ 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" @@ -105,7 +106,7 @@ class _AstraflowBaseProvider(LLMProvider): } return LLMOutput( - content=choice.message.content or "", + content=strip_reasoning(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 d626509cc..8d1dc739f 100644 --- a/src/llm/providers/atlas.py +++ b/src/llm/providers/atlas.py @@ -17,6 +17,7 @@ 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" @@ -122,7 +123,7 @@ class AtlasProvider(LLMProvider): } return LLMOutput( - content=choice.message.content or "", + content=strip_reasoning(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 58505ac7b..6ee4aa8d4 100644 --- a/src/llm/providers/ollama.py +++ b/src/llm/providers/ollama.py @@ -19,6 +19,7 @@ from llm.core.types import ( ProviderType, ToolCall, ) +from llm.providers.reasoning import prefilled_reasoning_from_env, strip_reasoning class OllamaProvider(LLMProvider): @@ -28,7 +29,15 @@ class OllamaProvider(LLMProvider): self, base_url: str | None = None, default_model: str | None = None, + prefilled_reasoning: bool | None = None, ) -> None: + # Set for models whose chat template prefills , so replies arrive + # as "...answer" without the opening tag. + self.prefilled_reasoning = ( + prefilled_reasoning + if prefilled_reasoning is not None + else prefilled_reasoning_from_env("OLLAMA_PREFILLED_REASONING") + ) self.base_url = base_url or os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434") self.default_model = default_model or os.environ.get("OLLAMA_MODEL", "llama3.2") self._models = [ @@ -85,7 +94,9 @@ class OllamaProvider(LLMProvider): with urllib.request.urlopen(req, timeout=60) as response: result = json.loads(response.read().decode("utf-8")) - content = result.get("message", {}).get("content", "") + content = strip_reasoning( + result.get("message", {}).get("content", ""), prefilled=self.prefilled_reasoning + ) tool_calls = None if result.get("message", {}).get("tool_calls"): diff --git a/src/llm/providers/openai.py b/src/llm/providers/openai.py index 7d00d333c..b81d4ab08 100644 --- a/src/llm/providers/openai.py +++ b/src/llm/providers/openai.py @@ -23,12 +23,25 @@ from llm.core.types import ( ) from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR from llm.providers.openai_compat import to_openai_message_param +from llm.providers.reasoning import prefilled_reasoning_from_env, strip_reasoning class OpenAIProvider(LLMProvider): provider_type = ProviderType.OPENAI - def __init__(self, api_key: str | None = None, base_url: str | None = None) -> None: + def __init__( + self, + api_key: str | None = None, + base_url: str | None = None, + prefilled_reasoning: bool | None = None, + ) -> None: + # Set for self-hosted servers whose chat template prefills , so + # replies arrive as "...answer" without the opening tag. + self.prefilled_reasoning = ( + prefilled_reasoning + if prefilled_reasoning is not None + else prefilled_reasoning_from_env("OPENAI_PREFILLED_REASONING") + ) self.client = OpenAI( api_key=api_key or os.environ.get("OPENAI_API_KEY"), base_url=base_url, @@ -106,7 +119,7 @@ class OpenAIProvider(LLMProvider): } return LLMOutput( - content=choice.message.content or "", + content=strip_reasoning(choice.message.content or "", 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 new file mode 100644 index 000000000..4a45c1fc1 --- /dev/null +++ b/src/llm/providers/reasoning.py @@ -0,0 +1,71 @@ +"""Helpers for reasoning text that self-hosted models leave in message content.""" + +from __future__ import annotations + +import json +import os + +THINK_OPEN = "" +THINK_CLOSE = "" + +_TRUTHY = {"1", "true", "yes", "on"} + + +def prefilled_reasoning_from_env(name: str) -> bool: + """Read a ``*_PREFILLED_REASONING`` flag from the environment.""" + return os.environ.get(name, "").strip().lower() in _TRUTHY + + +def strip_reasoning(content: str, *, 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``. + + 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). + + ``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. + """ + stripped = content.lstrip() + if stripped.startswith(THINK_OPEN): + close = stripped.find(THINK_CLOSE, len(THINK_OPEN)) + if close == -1: + return content + return stripped[close + len(THINK_CLOSE) :].lstrip() + + if not prefilled: + return content + + close = content.find(THINK_CLOSE) + if close == -1 or THINK_OPEN in content[:close] or _is_json_document(content): + return content + return content[close + len(THINK_CLOSE) :].lstrip() + + +def _is_json_document(content: str) -> bool: + """Whether ``content`` is one complete JSON object or array. + + Reasoning followed by a closing tag is never valid JSON, so a structured + answer that carries the tag inside a string value is left whole. Input + nested too deeply for the parser to finish is kept whole as well: reasoning + does not open with that many brackets, so cutting it would only risk + truncating a real answer. + """ + stripped = content.strip() + if not stripped.startswith(("{", "[")): + return False + try: + json.loads(stripped) + except RecursionError: + return True + except ValueError: + return False + return True diff --git a/tests/test_reasoning.py b/tests/test_reasoning.py new file mode 100644 index 000000000..c8d923c17 --- /dev/null +++ b/tests/test_reasoning.py @@ -0,0 +1,168 @@ +import json +import urllib.request +from io import BytesIO +from types import SimpleNamespace +from typing import Any, cast + +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.ollama import OllamaProvider +from llm.providers.openai import OpenAIProvider +from llm.providers.reasoning import strip_reasoning + +# Answers that mention the tags without carrying reasoning. None of them may be +# changed, whatever the deployment's reasoning format. +LITERAL_ANSWERS = [ + "Wrap it in and tags.", + "Close it with and nothing else.", + " comes before here.", + '{"closing":"","opening":""}', + '{"note": "x"}', +] + + +@pytest.mark.parametrize( + ("content", "expected"), + [ + ("17 * 3 = 51.51", "51"), + ("\n\nplan\n\n\nanswer", "answer"), + ('build the object{"answer": 51}', '{"answer": 51}'), + ('plan{"closing":""}', '{"closing":""}'), + ("plain answer", "plain answer"), + ("", ""), + ("cut off by max_tokens", "cut off by max_tokens"), + # Close-only output is ambiguous without a declared format, so it stays. + ("17 * 3 = 51.51", "17 * 3 = 51.51"), + ], +) +def test_strip_reasoning_default(content: str, expected: str) -> None: + assert strip_reasoning(content) == expected + + +@pytest.mark.parametrize("content", LITERAL_ANSWERS) +def test_strip_reasoning_default_preserves_literal_tags(content: str) -> None: + assert strip_reasoning(content) == content + + +@pytest.mark.parametrize( + ("content", "expected"), + [ + ("17 * 3 = 51.51", "51"), + ('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."), + ("17 * 3 = 51.", ""), + # A complete JSON answer is never reasoning plus a closing tag. + ('{"closing":""}', '{"closing":""}'), + (' {"closing":"","opening":""}\n', ' {"closing":"","opening":""}\n'), + ('["", 1]', '["", 1]'), + ], +) +def test_strip_reasoning_prefilled(content: str, expected: str) -> None: + assert strip_reasoning(content, prefilled=True) == expected + + +def test_strip_reasoning_prefilled_survives_deeply_nested_json() -> None: + # Deep enough for json.loads to raise RecursionError on every supported Python. + depth = 100_000 + content = "[" * depth + '""' + "]" * depth + assert strip_reasoning(content, prefilled=True) == content + + +def _openai_provider(content: str, **kwargs: Any) -> OpenAIProvider: + provider = OpenAIProvider(api_key="test", **kwargs) + message = SimpleNamespace(content=content, tool_calls=None) + response = SimpleNamespace( + choices=[SimpleNamespace(message=message, finish_reason="stop")], + model="spark-x2.5", + usage=None, + ) + client: Any = SimpleNamespace( + 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)) + + +def test_openai_provider_strips_balanced_reasoning() -> None: + assert _ask(_openai_provider("17 * 3 = 51.51")).content == "51" + + +@pytest.mark.parametrize("content", LITERAL_ANSWERS) +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: + 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" + + 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: + monkeypatch.setenv("OPENAI_PREFILLED_REASONING", "true") + provider = AstraflowProvider(api_key="test") + message = SimpleNamespace(content='{"closing":""}', tool_calls=None) + response = SimpleNamespace( + choices=[SimpleNamespace(message=message, finish_reason="stop")], + model="deepseek-v3.2", + usage=None, + ) + client: Any = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **params: response)) + ) + provider.client = client + + assert _ask(provider).content == '{"closing":""}' + + +def _fake_ollama(monkeypatch: pytest.MonkeyPatch, content: str) -> list[dict[str, Any]]: + sent: list[dict[str, Any]] = [] + + 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()) + + 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: + 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: + monkeypatch.delenv("OLLAMA_PREFILLED_REASONING", raising=False) + _fake_ollama(monkeypatch, "17 * 3 = 51.51") + assert _ask(OllamaProvider()).content == "17 * 3 = 51.51" + assert _ask(OllamaProvider(prefilled_reasoning=True)).content == "51" + + monkeypatch.setenv("OLLAMA_PREFILLED_REASONING", "1") + assert _ask(OllamaProvider()).content == "51" + + +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) + + assert sent[0]["stream"] is False + assert output.content == '{"closing":""}'