diff --git a/src/llm/providers/reasoning.py b/src/llm/providers/reasoning.py index 4f2add4ef..cbb8b1270 100644 --- a/src/llm/providers/reasoning.py +++ b/src/llm/providers/reasoning.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os THINK_OPEN = "" @@ -44,6 +45,22 @@ def strip_reasoning(content: str, *, prefilled: bool = False) -> str: return content close = content.find(THINK_CLOSE) - if close == -1 or THINK_OPEN in content[: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. + """ + stripped = content.strip() + if not stripped.startswith(("{", "[")): + return False + try: + json.loads(stripped) + except ValueError: + return False + return True diff --git a/tests/test_reasoning.py b/tests/test_reasoning.py index 5ac45563d..b2ade433b 100644 --- a/tests/test_reasoning.py +++ b/tests/test_reasoning.py @@ -2,10 +2,12 @@ import json import urllib.request from io import BytesIO from types import SimpleNamespace +from typing import Any, cast import pytest -from llm.core.types import LLMInput, Message, Role +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 @@ -36,12 +38,12 @@ LITERAL_ANSWERS = [ ("17 * 3 = 51.51", "17 * 3 = 51.51"), ], ) -def test_strip_reasoning_default(content, expected): +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): +def test_strip_reasoning_default_preserves_literal_tags(content: str) -> None: assert strip_reasoning(content) == content @@ -54,13 +56,17 @@ def test_strip_reasoning_default_preserves_literal_tags(content): ("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, expected): +def test_strip_reasoning_prefilled(content: str, expected: str) -> None: assert strip_reasoning(content, prefilled=True) == expected -def _openai_provider(content, **kwargs): +def _openai_provider(content: str, **kwargs: Any) -> OpenAIProvider: provider = OpenAIProvider(api_key="test", **kwargs) message = SimpleNamespace(content=content, tool_calls=None) response = SimpleNamespace( @@ -68,26 +74,27 @@ def _openai_provider(content, **kwargs): model="spark-x2.5", usage=None, ) - provider.client = SimpleNamespace( + client: Any = SimpleNamespace( chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **params: response)) ) + provider.client = client return provider -def _ask(provider, **kwargs): +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(): +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): +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): +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" @@ -96,7 +103,7 @@ def test_openai_provider_close_only_needs_opt_in(monkeypatch): assert _ask(_openai_provider("17 * 3 = 51.51")).content == "51" -def test_hosted_provider_never_strips_close_only(monkeypatch): +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) @@ -105,18 +112,19 @@ def test_hosted_provider_never_strips_close_only(monkeypatch): model="deepseek-v3.2", usage=None, ) - provider.client = SimpleNamespace( + client: Any = SimpleNamespace( chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **params: response)) ) + provider.client = client assert _ask(provider).content == '{"closing":""}' -def _fake_ollama(monkeypatch, content): - sent = [] +def _fake_ollama(monkeypatch: pytest.MonkeyPatch, content: str) -> list[dict[str, Any]]: + sent: list[dict[str, Any]] = [] - def fake_urlopen(request, timeout): - sent.append(json.loads(request.data)) + 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) @@ -124,14 +132,14 @@ def _fake_ollama(monkeypatch, content): @pytest.mark.parametrize("content", LITERAL_ANSWERS) -def test_ollama_provider_keeps_literal_tags(monkeypatch, content): +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): +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" @@ -141,7 +149,7 @@ def test_ollama_provider_strips_prefilled_reasoning_when_declared(monkeypatch): assert _ask(OllamaProvider()).content == "51" -def test_ollama_stream_flag_still_strips_complete_message(monkeypatch): +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.