diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index a756fefcb..1fed0a1a9 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -282,6 +282,51 @@ def _cleanup_kv_cache(cache_info: HFAloraCacheInfo) -> None: _CHAT_TEMPLATE_THINKING_VARS: tuple[str, ...] = ("think", "thinking", "enable_thinking") +_THINK_OPEN_TAG: str = "" +_THINK_CLOSE_TAG: str = "" + + +def _split_think_tags(text: str) -> tuple[str | None, str]: + r"""Split raw HF output into (thinking, answer) on the closing tag. + + A string-level fallback for `transformers.PreTrainedTokenizerBase.parse_response()` + (schema-driven, token-decode-then-parse), used because no tokenizer available to + this backend today declares a `response_schema`. Even that upstream mechanism is + text-level, not token-level. + + A token-position check against the raw generated sequence (scanning for + Granite's `` token id rather than re-scanning decoded text) would be + strictly more conservative than this string match — it would not misfire on + a mention the model spells out as separate pieces rather than emitting the + single vocab token. It would not fully eliminate false positives either: a + model is free to emit that same single token for a literal in-prose mention, + so a genuine close and a canonical-token literal mention still decode + identically. Not adopted here because HF's streaming path + (`TextIteratorStreamer`) only exposes decoded text, not token ids, so a + token check could not apply uniformly to both the streaming and + non-streaming paths — see #1604 for a scoped follow-up. + + Splits on alone (not a ... pair) because some chat + templates (e.g. granite-4.2 with enable_thinking) bake the opening tag into + the prompt, so it never appears in the model's own output. A fixed pattern + match for Granite's convention, not a general reasoning-parser; other + delimiters pass through unchanged. See #1604 for generalizing this. + + Deliberately strips leading/trailing whitespace from both `thinking` and + `answer`: the whitespace immediately around the tags is delimiter framing, + not meaningful content, so the user-visible completion loses that framing + whitespace as part of this split, not incidentally. This framing comes + from the *generation-prompt* path (chat_template.jinja:178-182, e.g. + `<|im_start|>assistant\n\n`), not the replay-reconstruction path + (chat_template.jinja:83-84) — this function only ever sees freshly + generated text, never a replayed history turn. + """ + if _THINK_CLOSE_TAG not in text: + return None, text + reasoning, _, answer = text.partition(_THINK_CLOSE_TAG) + return (reasoning.strip().removeprefix(_THINK_OPEN_TAG).strip(), answer.strip()) + + # A string THINKING level (e.g. "low") is forwarded verbatim as `reasoning_effort` # when the chat template declares that variable — this is the actual mechanism # Granite 4.2's chat template consumes (chat_template.jinja derives its boolean @@ -1989,6 +2034,10 @@ class used during generation, if any. if isinstance(hf_output, GenerateDecoderOnlyOutput) and mot._call.model_options: self._surface_logits(mot, hf_output) + # Built here (before the split below) because it needs `hf_output`'s KV cache/scores + # fields, which are cleared immediately after; cached under a key derived from + # `mot.value` further down, once the split has settled on the final string object. + cache_info: HFAloraCacheInfo | None = None if ( self._use_caches and isinstance(hf_output, GenerateDecoderOnlyOutput) @@ -2009,9 +2058,6 @@ class used during generation, if any. scores=hf_output.scores, ) - cache_key = id(mot.value) - self.cache_put(cache_key, cache_info) - # Clear KV cache and scores from HF output; retained via LRU cache above. # `ModelOutput` (`OrderedDict` subclass) does not sync `None` writes back # to the mapping, so plain attribute assignment leaves the dict entry — and @@ -2027,6 +2073,67 @@ class used during generation, if any. OrderedDict.__delitem__(hf_output, "logits") hf_output.logits = None + # Capture the raw text before any split below. Used for the stop-string check + # further down (a stop string could itself be think-tag-adjacent, so that check + # must stay on the pre-split text) and as the input to the split itself. Do not + # repoint this at `mot.value` after the split runs. + raw_value = mot.value + + # Gate on the template exposing a thinking var (some models think by default, so + # declaring the var isn't itself proof thinking is on) AND the resolved per-call + # value not being explicitly False (an answer that merely mentions "" on a + # template with thinking off must not be split). `None`/unset must still allow the + # split: Granite and Qwen3 both default `enable_thinking` to True in their own + # template source, so treating "unset" as "off" would under-split for the common case. + # Read from `mot._call.model_options` directly (not a value already filtered for the + # template) — it's the same dict `_filter_for_chat_template` resolves + # `ModelOption.THINKING` from when building the generation-time template kwargs, so + # this mirrors what the model was actually asked to do on this call. + # Skip for streaming: astream() assumes mot.value only grows, and shrinking it here + # would corrupt the final delta (see #1604 for proper incremental splitting later). + # + # Known limitation, deliberately not addressed here: this gate can under-split + # for a model that ignores an explicit ModelOption.THINKING=False, or that always + # emits blocks without declaring any of _CHAT_TEMPLATE_THINKING_VARS in its + # template. The resulting raw tags leak into mot.value and get replayed as ordinary + # `content` on the next turn — visible in the answer, and still read by requirement + # checks/judges, but not silently dropped anywhere. The alternative — splitting + # unconditionally — trades this for a worse failure: an over-split misfiles real + # answer text into mot.thinking, which gets attached as `reasoning_content` at + # replay time and is silently dropped from every subsequent prompt once a newer + # user turn exists, per Granite's own recency-based truncation (see the NOTE in + # mellea/backends/utils.py). Under-split stays diagnosable; over-split is invisible + # once replayed. Kept on the more conservative side deliberately; see #1604 for + # detecting always-thinking models without a declared var. + thinking_allowlist: frozenset[str] = getattr( + self, "_chat_template_allowlist", frozenset() + ) + resolved_thinking = ( + mot._call.model_options.get(ModelOption.THINKING) + if mot._call.model_options + else None + ) + if ( + not mot.generation.streaming + and resolved_thinking is not False + and thinking_allowlist.intersection(_CHAT_TEMPLATE_THINKING_VARS) + ): + thinking, answer = _split_think_tags(raw_value) + if thinking is not None: + mot.thinking = thinking + mot.value = answer + MelleaLogger.get_logger().debug( + "Split %d chars of thinking out of HF completion for %s.", + len(thinking), + self._model_id, + ) + + if cache_info is not None: + # Keyed after the split above has settled on the final `mot.value` object, so a + # later lookup by the same key against the same (now-split) thunk can find it. + cache_key = id(mot.value) + self.cache_put(cache_key, cache_info) + # Only scan for tools if we are not doing structured output and tool calls were provided to the model. if _format is None and tool_calls: mot.tool_calls = to_tool_calls(tools, mot.value) @@ -2077,8 +2184,8 @@ class used during generation, if any. stop_strings = ( mot._call.model_options.get(ModelOption.STOP_SEQUENCES) or [] ) - ends_with_stop_string = isinstance(mot.value, str) and any( - mot.value.endswith(s) for s in stop_strings + ends_with_stop_string = isinstance(raw_value, str) and any( + raw_value.endswith(s) for s in stop_strings ) if last_token in eos_set or ends_with_stop_string: mot.generation.finish_reasons = ["stop"] @@ -2134,6 +2241,7 @@ class used during generation, if any. "tools_available": tools, "tools_called": mot.tool_calls, "seed": seed, + "thinking": mot.thinking, } generate_log.action = mot._call.action generate_log.result = mot diff --git a/mellea/backends/utils.py b/mellea/backends/utils.py index 401e9ed79..d9fbaa0b9 100644 --- a/mellea/backends/utils.py +++ b/mellea/backends/utils.py @@ -99,10 +99,17 @@ def to_chat( # NOTE: `self.formatter.to_chat_messages` explicitly skips `Message` objects. However, we need # to print `Message`s to correctly serialize any documents with the message. Do the printing here. - # NOTE: reasoning is never replayed on the HF chat path — we serialize only `content` and never - # consult `should_replay_reasoning` (unlike the OpenAI/LiteLLM/Watsonx/Ollama chat paths). This is - # acceptable today because HF has a capture gap (per #1201) and never populates `Message.thinking` - # to begin with; when that gap is closed, replay must be wired in here. + # NOTE: `Message.thinking` is forwarded as `reasoning_content` (the key Granite/Qwen3 + # templates consume) for every assistant turn that has it. Granite's own chat template + # (not a turn-type check) then decides whether to keep or strip it: reasoning survives + # only for a turn at or after the most recent user message (`last_user_idx` / + # `truncate_history_thinking`, defaulted True and never overridden by mellea), and is + # stripped for anything from an earlier exchange, tool-call or not. In mellea's typical + # flow a tool-call turn has no intervening user message before its continuation, so it + # tends to survive, and a plain turn from a prior exchange tends not to — but that's a + # consequence of the recency rule, not a tool-call/plain-turn distinction Mellea enforces + # (unlike #1201's cross-backend `should_replay_reasoning`, which HF does not call). See + # test_rendered_prompt_*_turn in test_huggingface_thinking.py. ctx_as_conversation: list = [] for m in ctx_as_message_list: msg_dict: dict = {"role": m.role, "content": formatter.print(m)} @@ -113,6 +120,8 @@ def to_chat( msg_dict["tool_calls"] = m.tool_calls if m.tool_call_id: msg_dict["tool_call_id"] = m.tool_call_id + if m.role == "assistant" and m.thinking: + msg_dict["reasoning_content"] = m.thinking # Merge any author-declared provider fields (Mellea's known fields win; # a mismatched target raises). Must run after the known fields are set. msg_dict = merge_provider_fields(msg_dict, m.provider_fields, "huggingface") diff --git a/test/backends/test_huggingface_thinking.py b/test/backends/test_huggingface_thinking.py new file mode 100644 index 000000000..d602bc2b5 --- /dev/null +++ b/test/backends/test_huggingface_thinking.py @@ -0,0 +1,451 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for LocalHFBackend's ... tag splitting. + +No GPU or real model is needed — _split_think_tags is a pure string function. + +torch must be importable because importing huggingface.py triggers the top-level +`import torch`. Install mellea[hf] to satisfy this requirement. +""" + +import pytest + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") + +import mellea.backends.huggingface as hf_backend +from mellea.backends import ModelOption +from mellea.backends.huggingface import LocalHFBackend, _split_think_tags +from mellea.core.base import CBlock, ModelOutputThunk +from test.backends.test_huggingface_filter_options import ( + _GRANITE_THINKING_MODEL_ID, + _try_load_granite_tokenizer, +) + + +def test_split_think_tags_with_both_tags() -> None: + """Both tags present: reasoning and answer are split and stripped.""" + thinking, answer = _split_think_tags("reasoning herethe answer") + assert thinking == "reasoning here" + assert answer == "the answer" + + +def test_split_think_tags_missing_opening_tag() -> None: + """Only the closing tag present (e.g. granite-4.2's prompt-baked opening tag).""" + thinking, answer = _split_think_tags("reasoning herethe answer") + assert thinking == "reasoning here" + assert answer == "the answer" + + +def test_split_think_tags_no_closing_tag() -> None: + """No anywhere: text is returned unchanged, thinking is None.""" + thinking, answer = _split_think_tags("just a plain answer, no tags") + assert thinking is None + assert answer == "just a plain answer, no tags" + + +def test_split_think_tags_strips_surrounding_whitespace() -> None: + """Whitespace/newlines around tags (matches granite's \\n form) are stripped.""" + thinking, answer = _split_think_tags( + "\n reasoning here \n\n the answer \n" + ) + assert thinking == "reasoning here" + assert answer == "the answer" + + +def test_split_think_tags_empty_reasoning() -> None: + """Empty think block (e.g. granite's thinking-disabled form).""" + thinking, answer = _split_think_tags("the answer") + assert thinking == "" + assert answer == "the answer" + + +def test_split_think_tags_leading_whitespace_before_open_tag() -> None: + """A newline/whitespace before must not defeat removeprefix. + + Regression test: reasoning.removeprefix(_THINK_OPEN_TAG) only strips the tag + when it is the literal first character(s) of the string. Without stripping + first, "\n\n..." still starts with "\n", removeprefix is a no-op, and + the literal "" tag leaks into mot.thinking. + """ + thinking, answer = _split_think_tags("\n\nreasoning\n\nanswer") + assert thinking == "reasoning" + assert answer == "answer" + + +def test_split_think_tags_multiple_close_tags_uses_first() -> None: + """Multiple occurrences: only the first is treated as the boundary. + + Deliberate: Granite's own template uses the *last* occurrence when + truncating old reasoning out of replayed history, where dropping too much + is the safe direction. Here we're extracting a clean answer from a fresh + completion, where the safe direction is the opposite: if the model's + answer itself mentions the literal text "", splitting on the + first occurrence keeps the answer intact instead of corrupting it. + """ + thinking, answer = _split_think_tags("abc") + assert thinking == "a" + assert answer == "bc" + + +def _make_backend( + *, thinking_template_var: str | None = "think", use_caches: bool = False +) -> LocalHFBackend: + """Return a LocalHFBackend with __init__ bypassed, wired with the minimum + state post_processing needs when there is no real GenerateDecoderOnlyOutput + (i.e. every isinstance(hf_output, GenerateDecoderOnlyOutput) branch is skipped + unless use_caches=True and the caller also sets mot.raw.response). + + Args: + thinking_template_var: name of a thinking-related Jinja variable to bake + into the fake chat template (gates the think-split in post_processing), + or None for a template that does not reference any of them. + use_caches: whether to wire up a real `SimpleLRUCache` so the KV-cache + branch in post_processing actually runs. + """ + b: LocalHFBackend = LocalHFBackend.__new__(LocalHFBackend) + b._model_id = "test-org/test-model" + b.model_id = "test-org/test-model" + b._provider = "huggingface" + b._use_caches = use_caches + if use_caches: + from mellea.backends.cache import SimpleLRUCache + + object.__setattr__(b, "_cache", SimpleLRUCache(5)) + object.__setattr__(b, "_device", torch.device("cpu")) + + template = ( + f"{{{{ {thinking_template_var} }}}}" + if thinking_template_var + else "{{ messages }}" + ) + + class _FakeTokenizer: + chat_template = template + + object.__setattr__(b, "_tokenizer", _FakeTokenizer()) + return b + + +async def test_post_processing_splits_thinking_before_tool_scan(monkeypatch) -> None: + """post_processing must strip ... into mot.thinking and pass + only the post-split answer text to the tool-call scan, not the raw combined + output — verifies the ordering the inline comment in post_processing promises. + """ + recorded_text: list[str] = [] + + def fake_to_tool_calls(tools, text): + recorded_text.append(text) + return None + + monkeypatch.setattr(hf_backend, "to_tool_calls", fake_to_tool_calls) + + backend = _make_backend() + mot = ModelOutputThunk( + value="reasoning herecall get_weather(city='Boston')" + ) + mot._call.action = CBlock("What's the weather?") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=True, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking == "reasoning here" + assert mot.value == "call get_weather(city='Boston')" + assert recorded_text == ["call get_weather(city='Boston')"] + + +async def test_post_processing_skips_split_when_streaming() -> None: + """Streaming generations must not have mot.value shrunk by post_processing. + + Regression test for a real bug: ModelOutputThunk.astream() (mellea/core/base.py) + computes each delta from an offset captured before post_processing runs, assuming + mot._underlying_value only ever grows during streaming. If post_processing + replaces mot.value with the shorter post-split answer, that offset goes stale and + the final astream() delta is corrupted (truncated or empty) — verified separately + by simulating astream()'s delta math against this exact before/after state. + """ + backend = _make_backend() + mot = ModelOutputThunk(value="reasoning herethe answer") + mot.generation.streaming = True + mot._call.action = CBlock("test") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking is None + assert mot.value == "reasoning herethe answer" + + +async def test_post_processing_does_not_split_without_thinking_template_var() -> None: + """A chat template with no thinking-related variable must not trigger the split. + + Regression test: without this gate, any answer that happens to contain the + literal substring "" (e.g. a question about chat templates) would have + legitimate answer text incorrectly moved into mot.thinking. + """ + backend = _make_backend(thinking_template_var=None) + mot = ModelOutputThunk(value="Use the tag to close a reasoning block.") + mot._call.action = CBlock("How do reasoning tags work?") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking is None + assert mot.value == "Use the tag to close a reasoning block." + + +async def test_post_processing_does_not_split_when_thinking_explicitly_false() -> None: + """Regression test: a template declaring a thinking var is not proof thinking + was requested on this specific call. With ModelOption.THINKING explicitly + False, a literal "" in the answer must survive untouched, even though + the template declares "think" (the gate must check the resolved per-call value, + not just whether the template mentions the variable name at all). + """ + backend = _make_backend(thinking_template_var="think") + mot = ModelOutputThunk(value="Use the tag to close a reasoning block.") + mot._call.action = CBlock("How do reasoning tags work?") + mot._call.model_options = {ModelOption.THINKING: False} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking is None + assert mot.value == "Use the tag to close a reasoning block." + + +async def test_post_processing_splits_when_thinking_unset() -> None: + """Regression test: an unset/None ModelOption.THINKING must still allow the + split when the template declares a thinking var, since Granite and Qwen3 both + default thinking to True in their own template source — treating "unset" as + "off" would under-split the common case. + """ + backend = _make_backend(thinking_template_var="think") + mot = ModelOutputThunk(value="reasoning herethe answer") + mot._call.action = CBlock("test") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking == "reasoning here" + assert mot.value == "the answer" + + +async def test_post_processing_preserves_answer_mentioning_think_tag() -> None: + """False-positive risk test: thinking genuinely on, and the model's answer + itself explains what the tag does. The gate correctly fires (a + real reasoning block exists), and first-occurrence splitting keeps the + full answer intact rather than truncating it at the second, unrelated + occurrence — the risk jakelorocco raised in the original PR review. + """ + backend = _make_backend(thinking_template_var="think") + mot = ModelOutputThunk( + value=( + "the user wants an explanation" + "The tag marks the end of a reasoning block." + ) + ) + mot._call.action = CBlock("What does the tag do?") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=None, + ) + + assert mot.thinking == "the user wants an explanation" + assert mot.value == "The tag marks the end of a reasoning block." + + +async def test_post_processing_cache_key_findable_after_split() -> None: + """Regression test: the LRU cache key computed in post_processing must be + derived from mot.value *after* the split has run, so a later cache_get() call + against the same (now-split) thunk's value can find the entry. Before the + fix, the key was computed from the pre-split string's object identity, then + mot.value was reassigned to a new string a few lines later — orphaning the + cache entry under a key nothing would ever look up again. + """ + from transformers.generation.utils import GenerateDecoderOnlyOutput + + backend = _make_backend(thinking_template_var="think", use_caches=True) + sequences = torch.tensor([[1, 2, 3, 4]]) + hf_output = GenerateDecoderOnlyOutput( + sequences=sequences, scores=(torch.zeros(1, 10),) + ) + mot = ModelOutputThunk(value="reasoning herethe answer") + mot.raw.response = hf_output + mot._call.action = CBlock("test") + mot._call.model_options = {} + + await backend.post_processing( + mot, + conversation=[], + _format=None, + tool_calls=False, + tools={}, + seed=None, + input_ids=torch.tensor([[1, 2]]), + ) + + assert mot.thinking == "reasoning here" + assert mot.value == "the answer" + cached = backend.cache_get(id(mot.value)) + assert cached is not None + + +def _render_history(messages: list) -> str: + """Render a message list through the real granite-4.2-3b chat template. + + Loads the template only (no GPU, no model weights) via + `_try_load_granite_tokenizer`; skips if not locally cached. + """ + tok = _try_load_granite_tokenizer(_GRANITE_THINKING_MODEL_ID) + if tok is None: + pytest.skip( + f"{_GRANITE_THINKING_MODEL_ID} not in local HF cache — " + "run qualitative tests first" + ) + return tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + + +@pytest.mark.integration +def test_rendered_prompt_preserves_reasoning_on_tool_call_turn() -> None: + """On a tool-call turn, the `reasoning_content` forward in `to_chat()` + restores reasoning to the rendered prompt against the real template + (not a synthetic one). The plain multi-turn shape (next test) isn't + restored the same way. + """ + messages = [ + {"role": "user", "content": "What's the weather in Boston?"}, + { + "role": "assistant", + "content": "", + "reasoning_content": "I should call the weather tool.", + "tool_calls": [ + { + "id": "1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": {"city": "Boston"}, + }, + } + ], + }, + {"role": "tool", "content": "72F, sunny", "tool_call_id": "1"}, + ] + rendered = _render_history(messages) + assert "I should call the weather tool." in rendered + # The empty pair is what a dropped/truncated reasoning_content + # would leave behind; its absence here is the direct evidence reasoning survived. + assert "" not in rendered + + +@pytest.mark.integration +def test_rendered_prompt_drops_reasoning_on_plain_multi_turn() -> None: + """Deliberate limitation: on a plain multi-turn shape (assistant turn + with no tool call, followed by another user turn), Granite's own + `truncate_history_thinking` gate strips reasoning even though `to_chat()` + attaches `reasoning_content` — the reconstructed content carries both + tags, which is exactly what that gate matches on. Decided: keep HF + consistent with the #1201 cross-backend consensus (replay on tool-call + turns only) rather than extend replay to plain turns. + """ + messages = [ + {"role": "user", "content": "What is 2 + 2?"}, + { + "role": "assistant", + "content": "4", + "reasoning_content": "Two plus two equals four.", + }, + {"role": "user", "content": "And 3 + 3?"}, + ] + rendered = _render_history(messages) + assert "Two plus two equals four." not in rendered + + +@pytest.mark.parametrize("with_tool_call", [True, False]) +def test_parse_then_to_chat_round_trips_thinking_for_hf(with_tool_call: bool) -> None: + """Seam test: HF's `Message._parse()` (chat.py) carries `.thinking` onto the + parsed assistant Message in both branches (tool-call and plain), and that + Message then round-trips through `to_chat()` (utils.py) as `reasoning_content`. + + Closes the gap between the e2e tests (stop at `output.thinking`) and the + `to_chat` unit tests (hand-build `Message(..., thinking=...)` directly, + never exercising `_parse`) — this is the only test proving the link between + them for HF specifically. + """ + from typing import cast + + from mellea.backends.utils import to_chat + from mellea.core import ModelToolCall + from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter + from mellea.stdlib.components import Message + from mellea.stdlib.context import ChatContext + + mot = ModelOutputThunk(value="the answer") + mot.thinking = "reasoning trace" + mot.raw.provider = "huggingface" + mot.raw.response = None + if with_tool_call: + # Only non-None matters for _parse's branch selection; the placeholder is + # never read as a real ModelToolCall. + mot.tool_calls = [cast(ModelToolCall, None)] + + parsed = Message(role="assistant", content="placeholder")._parse(mot) + assert parsed.thinking == "reasoning trace" + + ctx = ChatContext() + ctx = ctx.add(Message("user", "hello")) + ctx = ctx.add(parsed) + action = Message("user", "next question") + formatter = ChatFormatter(model_id="test") + + result = to_chat(action, ctx, formatter, system_prompt=None) + assistant_msg = next(m for m in result if m["role"] == "assistant") + assert assistant_msg["reasoning_content"] == "reasoning trace" diff --git a/test/backends/test_huggingface_thinking_e2e.py b/test/backends/test_huggingface_thinking_e2e.py new file mode 100644 index 000000000..ccd70da01 --- /dev/null +++ b/test/backends/test_huggingface_thinking_e2e.py @@ -0,0 +1,92 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end tests for LocalHFBackend's ... tag splitting. + +Unlike test_huggingface_thinking.py (pure unit tests against synthetic strings), +these tests run real generation against granite-4.2-3b to verify the raw text +produced by `model.generate()` + `tokenizer.decode(skip_special_tokens=True)` +actually has the shape _split_think_tags assumes. This matters because +and are registered as *non-special* added tokens in granite-4.2's +tokenizer (special: false), so skip_special_tokens=True does not strip them — +but that is a property of this specific model/tokenizer, not something a +synthetic-string test can verify. +""" + +import re + +import pytest + +from test.predicates import require_gpu + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") + +pytestmark = [pytest.mark.huggingface, pytest.mark.e2e, require_gpu(min_vram_gb=8)] + +import mellea.backends.model_ids as model_ids +from mellea import MelleaSession +from mellea.backends import ModelOption +from mellea.backends.cache import SimpleLRUCache +from mellea.backends.huggingface import LocalHFBackend +from mellea.stdlib.context import ChatContext +from test.conftest import hf_skip + + +@pytest.fixture(scope="module") +def backend(): + """Shared granite-4.2-3b HuggingFace backend for all tests in this module.""" + with hf_skip(): + backend = LocalHFBackend( + model_id=model_ids.IBM_GRANITE_4_2_3B, cache=SimpleLRUCache(5) + ) + yield backend + + from test.conftest import cleanup_gpu_backend + + cleanup_gpu_backend(backend, "huggingface-thinking") + + +@pytest.fixture(scope="function") +def session(backend): + """Fresh HuggingFace session for each test.""" + session = MelleaSession(backend, ctx=ChatContext()) + yield session + session.reset() + + +@pytest.mark.qualitative +def test_thinking_enabled_populates_mot_thinking(session): + """ModelOption.THINKING=True: mot.thinking holds the real reasoning trace, + mot.value is the clean answer with no leftover tag, and the answer + itself still contains the expected content (not just an empty string that + would vacuously satisfy the tag-absence checks alone).""" + output = session.instruct( + "What is 2 + 2? Answer with just the number.", + model_options={ModelOption.THINKING: True, ModelOption.MAX_NEW_TOKENS: 400}, + ) + assert output.thinking, ( + f"Expected a non-empty reasoning trace, got: {output.thinking!r}" + ) + # granite-4.2's chat template bakes the opening into the generation + # prompt itself (see module docstring / _split_think_tags), so the model's own + # output never contains it — checking for its absence here would be vacuous. + # is the tag that actually appears in raw output and must be split out. + assert "" not in output.value + assert re.search(r"\b4\b", output.value), ( + f"Expected the digit 4 as its own token, got: {output.value!r}" + ) + + +@pytest.mark.qualitative +def test_thinking_disabled_leaves_mot_thinking_falsy(session): + """ModelOption.THINKING=False: no think block is generated, so mot.thinking + is falsy (None or empty string — both mean "no reasoning trace captured").""" + output = session.instruct( + "What is 2 + 2? Answer with just the number.", + model_options={ModelOption.THINKING: False, ModelOption.MAX_NEW_TOKENS: 100}, + ) + assert not output.thinking, f"Expected no reasoning trace, got: {output.thinking!r}" + assert "" not in output.value + assert re.search(r"\b4\b", output.value), ( + f"Expected the digit 4 as its own token, got: {output.value!r}" + ) diff --git a/test/backends/test_utils.py b/test/backends/test_utils.py index b226766bb..4889d1b68 100644 --- a/test/backends/test_utils.py +++ b/test/backends/test_utils.py @@ -151,6 +151,95 @@ def test_to_chat_basic_message(): assert result[1]["content"] == "next question" +def test_to_chat_attaches_reasoning_content_for_assistant_thinking(): + """An assistant Message carrying `.thinking` must have it forwarded as + `reasoning_content` on the wire dict. + """ + from mellea.backends.utils import to_chat + from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter + from mellea.stdlib.components import Message + from mellea.stdlib.context import ChatContext + + ctx = ChatContext() + ctx = ctx.add(Message("user", "hello")) + ctx = ctx.add(Message("assistant", "the answer", thinking="reasoning trace")) + action = Message("user", "next question") + formatter = ChatFormatter(model_id="test") + + result = to_chat(action, ctx, formatter, system_prompt=None) + assistant_msg = next(m for m in result if m["role"] == "assistant") + assert assistant_msg["reasoning_content"] == "reasoning trace" + + +def test_to_chat_known_reasoning_content_wins_over_provider_fields(): + """Regression/contract test: `reasoning_content` is set from `Message.thinking` + before `merge_provider_fields` runs (same known-fields-first pattern as + `tool_calls`/`tool_call_id`), so an author-declared `provider_fields` + collision on the same key is silently dropped (debug-logged, not raised) — + Mellea's own value always wins. + """ + from mellea.backends.utils import to_chat + from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter + from mellea.stdlib.components import Message + from mellea.stdlib.context import ChatContext + + ctx = ChatContext() + ctx = ctx.add(Message("user", "hello")) + ctx = ctx.add( + Message( + "assistant", + "the answer", + thinking="reasoning trace", + provider_fields={"huggingface": {"reasoning_content": "author override"}}, + ) + ) + action = Message("user", "next question") + formatter = ChatFormatter(model_id="test") + + result = to_chat(action, ctx, formatter, system_prompt=None) + assistant_msg = next(m for m in result if m["role"] == "assistant") + assert assistant_msg["reasoning_content"] == "reasoning trace" + + +def test_to_chat_omits_reasoning_content_when_no_thinking(): + """An assistant Message with no captured reasoning must not get a + `reasoning_content` key at all (not even an empty string).""" + from mellea.backends.utils import to_chat + from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter + from mellea.stdlib.components import Message + from mellea.stdlib.context import ChatContext + + ctx = ChatContext() + ctx = ctx.add(Message("user", "hello")) + ctx = ctx.add(Message("assistant", "the answer")) + action = Message("user", "next question") + formatter = ChatFormatter(model_id="test") + + result = to_chat(action, ctx, formatter, system_prompt=None) + assistant_msg = next(m for m in result if m["role"] == "assistant") + assert "reasoning_content" not in assistant_msg + + +def test_to_chat_ignores_thinking_on_non_assistant_message(): + """`.thinking` on a non-assistant Message (never set by real code paths, but + not type-guarded against) must not be forwarded as `reasoning_content` — the + Granite template only consumes that key on the assistant branch, so this is + otherwise inert, but the wire dict should not carry stray, unconsumed keys.""" + from mellea.backends.utils import to_chat + from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter + from mellea.stdlib.components import Message + from mellea.stdlib.context import ChatContext + + ctx = ChatContext() + ctx = ctx.add(Message("user", "hello", thinking="stray thinking")) + action = Message("user", "next question") + formatter = ChatFormatter(model_id="test") + + result = to_chat(action, ctx, formatter, system_prompt=None) + user_msg = next(m for m in result if m["content"] == "hello") + assert "reasoning_content" not in user_msg + + def test_to_chat_with_system_prompt(): from mellea.backends.utils import to_chat from mellea.formatters.template_formatter import TemplateFormatter as ChatFormatter