mirror of
https://github.com/affaan-m/ECC.git
synced 2026-10-14 12:15:29 +02:00
280 lines
9.8 KiB
Python
280 lines
9.8 KiB
Python
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.atlas import AtlasProvider
|
|
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 <think> and </think> tags.",
|
|
"Close it with </think> and nothing else.",
|
|
"</think> comes before <think> here.",
|
|
'{"closing":"</think>","opening":"<think>"}',
|
|
'{"note": "<think>x</think>"}',
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("content", "expected"),
|
|
[
|
|
("<think>17 * 3 = 51.</think>51", "51"),
|
|
("\n<think>\nplan\n</think>\n\nanswer", "answer"),
|
|
('<think>build the object</think>{"answer": 51}', '{"answer": 51}'),
|
|
('<think>plan</think>{"closing":"</think>"}', '{"closing":"</think>"}'),
|
|
("plain answer", "plain answer"),
|
|
("", ""),
|
|
("<think>cut off by max_tokens", "<think>cut off by max_tokens"),
|
|
# Close-only output is ambiguous without a declared format, so it stays.
|
|
("17 * 3 = 51.</think>51", "17 * 3 = 51.</think>51"),
|
|
],
|
|
)
|
|
def test_strip_reasoning_declared_leading(content: str, expected: str) -> None:
|
|
assert strip_reasoning(content, leading=True) == 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.</think>51", "51"),
|
|
('build the object\n</think>\n{"answer": 51}', '{"answer": 51}'),
|
|
("<think>17 * 3 = 51.</think>51", "51"),
|
|
("plain answer", "plain answer"),
|
|
(
|
|
"Wrap it in <think> and </think> tags.",
|
|
"Wrap it in <think> and </think> tags.",
|
|
),
|
|
("17 * 3 = 51.</think>", ""),
|
|
# A complete JSON answer is never reasoning plus a closing tag.
|
|
('{"closing":"</think>"}', '{"closing":"</think>"}'),
|
|
(
|
|
' {"closing":"</think>","opening":"<think>"}\n',
|
|
' {"closing":"</think>","opening":"<think>"}\n',
|
|
),
|
|
('["</think>", 1]', '["</think>", 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 + '"</think>"' + "]" * 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("<think>17 * 3 = 51.</think>51", leading_reasoning=True)
|
|
).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.</think>51")).content
|
|
== "17 * 3 = 51.</think>51"
|
|
)
|
|
assert (
|
|
_ask(
|
|
_openai_provider("17 * 3 = 51.</think>51", prefilled_reasoning=True)
|
|
).content
|
|
== "51"
|
|
)
|
|
|
|
monkeypatch.setenv("OPENAI_PREFILLED_REASONING", "true")
|
|
assert _ask(_openai_provider("17 * 3 = 51.</think>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":"</think>"}', 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":"</think>"}'
|
|
|
|
|
|
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.</think>51")
|
|
assert _ask(OllamaProvider()).content == "17 * 3 = 51.</think>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, '<think>plan</think>{"closing":"</think>"}')
|
|
|
|
output = _ask(OllamaProvider(leading_reasoning=True), stream=True)
|
|
|
|
assert sent[0]["stream"] is False
|
|
assert output.content == '{"closing":"</think>"}'
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"content",
|
|
[
|
|
"<think>display this XML literally</think>",
|
|
"<think>display this XML literally</think><answer>keep both</answer>",
|
|
],
|
|
)
|
|
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 = "<think>plan</think>answer"
|
|
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",
|
|
["<think>literal XML</think>", "<think>literal XML</think><answer>whole</answer>"],
|
|
)
|
|
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",
|
|
["<think>literal XML</think>", "<think>literal XML</think><answer>whole</answer>"],
|
|
)
|
|
def test_helper_defaults_preserve_leading_xml(content: str) -> None:
|
|
assert strip_reasoning(content) == content
|