feat: carry reasoning across turns with per-backend replay policy

mellea/backends/huggingface.py

@@ -613,6 +613,10 @@ async def _generate_from_intrinsic(
         conversation: list[dict] = []
         if system_prompt != "":
             conversation.append({"role": "system", "content": system_prompt})
+        # Intrinsic/adapter calls are single-shot evaluations over a rewritten
+        # conversation, not multi-turn generation, so reasoning is never replayed
+        # here (no `replay_reasoning=`). The HF chat path serializes via
+        # `to_chat`/`apply_chat_template`, not this helper.
         conversation.extend([message_to_openai_message(m) for m in ctx_as_message_list])
 
         docs = messages_to_docs(ctx_as_message_list)

mellea/backends/litellm.py

@@ -41,6 +41,7 @@
     get_current_event_loop,
     message_to_openai_message,
     send_to_queue,
+    should_replay_reasoning,
 )
 from ..stdlib.components import Message
 from ..stdlib.requirements import ALoraRequirement
@@ -355,8 +356,12 @@ async def _generate_from_chat_context_standard(
         system_prompt = model_opts.get(ModelOption.SYSTEM_PROMPT, "")
         if system_prompt != "":
             conversation.append({"role": "system", "content": system_prompt})
+        replay_flags = should_replay_reasoning(messages, self._provider)
         conversation.extend(
-            [message_to_openai_message(m, self.formatter) for m in messages]
+            [
+                message_to_openai_message(m, self.formatter, replay_reasoning=replay)
+                for m, replay in zip(messages, replay_flags)
+            ]
         )
 
         extra_params: dict[str, Any] = {}

mellea/backends/ollama.py

@@ -32,6 +32,7 @@
     ClientCache,
     get_current_event_loop,
     send_to_queue,
+    should_replay_reasoning,
 )
 from ..stdlib.components import Message
 from ..stdlib.requirements import ALoraRequirement
@@ -420,25 +421,29 @@ async def generate_from_chat_context(
 
         # 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.
-        for m in messages:
+        replay_flags = should_replay_reasoning(messages, self._provider)
+        for m, replay in zip(messages, replay_flags):
             if m.images is not None:
                 for img in m.images:
                     if isinstance(img, ImageUrlBlock):
                         raise ValueError(
                             "OllamaModelBackend does not support URL images (ImageUrlBlock). "
                             "Convert the image to a base64-encoded ImageBlock before passing it to Ollama."
                         )
-            conversation.append(
-                {
-                    "role": m.role,
-                    "content": self.formatter.print(m),
-                    "images": (
-                        _strip_data_uri_prefix([str(img.value) for img in m.images])
-                        if m.images
-                        else None
-                    ),
-                }
-            )
+            message_dict: dict[str, Any] = {
+                "role": m.role,
+                "content": self.formatter.print(m),
+                "images": (
+                    _strip_data_uri_prefix([str(img.value) for img in m.images])
+                    if m.images
+                    else None
+                ),
+            }
+            # Ollama's native SDK carries reasoning under the `thinking` key (see
+            # `chunk.message.thinking` on capture), not `reasoning_content`.
+            if replay and m.thinking:
+                message_dict["thinking"] = m.thinking
+            conversation.append(message_dict)
 
         # Append tool call information if applicable.
         tools: dict[str, AbstractMelleaTool] = dict()

mellea/backends/openai.py

@@ -43,6 +43,7 @@
     message_to_openai_message,
     messages_to_docs,
     send_to_queue,
+    should_replay_reasoning,
 )
 from ..stdlib.components import Intrinsic, Message
 from ..stdlib.requirements import LLMaJRequirement
@@ -633,6 +634,10 @@ async def _generate_from_intrinsic(
         conversation: list[dict] = []
         if system_prompt != "":
             conversation.append({"role": "system", "content": system_prompt})
+        # Intrinsic/adapter calls are single-shot evaluations over a rewritten
+        # conversation, not multi-turn generation, so reasoning is never replayed
+        # here (no `replay_reasoning=`) — unlike the chat path in
+        # `_generate_from_context`, which applies `should_replay_reasoning`.
         conversation.extend([message_to_openai_message(m) for m in messages])
 
         docs = messages_to_docs(messages)
@@ -859,8 +864,12 @@ async def _generate_from_chat_context_standard(
         system_prompt = model_opts.get(ModelOption.SYSTEM_PROMPT, "")
         if system_prompt != "":
             conversation.append({"role": "system", "content": system_prompt})
+        replay_flags = should_replay_reasoning(messages, self._provider)
         conversation.extend(
-            [message_to_openai_message(m, self.formatter) for m in messages]
+            [
+                message_to_openai_message(m, self.formatter, replay_reasoning=replay)
+                for m, replay in zip(messages, replay_flags)
+            ]
         )
 
         extra_params: dict[str, Any] = {}

mellea/backends/watsonx.py

@@ -43,6 +43,7 @@
     extract_model_tool_requests,
     get_current_event_loop,
     send_to_queue,
+    should_replay_reasoning,
 )
 from ..stdlib.components import Message
 from ..stdlib.requirements import ALoraRequirement
@@ -402,9 +403,15 @@ async def generate_from_chat_context(
 
         # 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.
-        conversation.extend(
-            [{"role": m.role, "content": self.formatter.print(m)} for m in messages]
-        )
+        replay_flags = should_replay_reasoning(messages, self._provider)
+        for m, replay in zip(messages, replay_flags):
+            message_dict: dict[str, Any] = {
+                "role": m.role,
+                "content": self.formatter.print(m),
+            }
+            if replay and m.thinking:
+                message_dict["reasoning_content"] = m.thinking
+            conversation.append(message_dict)
 
         if _format is not None:
             model_opts["response_format"] = {

mellea/helpers/__init__.py

@@ -22,6 +22,7 @@
     extract_model_tool_requests,
     message_to_openai_message,
     messages_to_docs,
+    should_replay_reasoning,
 )
 from .server_type import (
     _server_type,
@@ -42,5 +43,6 @@
     "message_to_openai_message",
     "messages_to_docs",
     "send_to_queue",
+    "should_replay_reasoning",
     "wait_for_all_mots",
 ]

mellea/helpers/openai_compatible_helpers.py

@@ -177,19 +177,63 @@ def chat_completion_delta_merge(
     return merged
 
 
-def message_to_openai_message(msg: Message, formatter: Formatter | None = None) -> dict:
+def should_replay_reasoning(
+    messages: list[Message], provider: str | None
+) -> list[bool]:
+    """Decide, per message, whether its reasoning trace should be replayed to the provider.
+
+    Implements the cross-provider consensus rule from issue #1201: an assistant
+    message's reasoning is round-tripped only when that turn issued a tool call —
+    detected by a `tool`-role message immediately following it — and stripped on
+    plain follow-up turns. Non-assistant messages and assistant messages without
+    reasoning always return `False`.
+
+    Args:
+        messages: The conversation in order, as it will be serialised.
+        provider: The backend provider name (e.g. `"openai"`, `"ollama"`).
+            Currently unused — every provider follows the consensus rule above.
+            It is a reserved hook for a provider-specific deviation (e.g. a model
+            that must replay reasoning on plain turns, or must not after a tool
+            call); add a keyed branch here once such a case is verified live.
+
+    Returns:
+        A list of booleans, one per message in `messages`, indicating whether that
+        message's reasoning should be included in the serialised payload.
+    """
+    flags: list[bool] = []
+    for i, msg in enumerate(messages):
+        if msg.role != "assistant" or not msg.thinking:
+            flags.append(False)
+            continue
+        # The prior turn "had a tool call" iff the next message is a tool result.
+        prior_turn_had_tool_call = (
+            i + 1 < len(messages) and messages[i + 1].role == "tool"
+        )
+        flags.append(prior_turn_had_tool_call)
+    return flags
+
+
+def message_to_openai_message(
+    msg: Message, formatter: Formatter | None = None, *, replay_reasoning: bool = False
+) -> dict:
     """Serialise a Mellea `Message` to the format required by OpenAI-compatible API providers.
 
     Args:
         msg: The `Message` object to serialise.
         formatter: Optional formatter used to render the message content (including
             documents) through the template system. When `None`, uses the raw
             `msg.content` string without document rendering.
+        replay_reasoning: When `True` and `msg.thinking` is a non-empty string,
+            the reasoning trace is emitted under the `"reasoning_content"` key so
+            the provider receives the model's prior reasoning. Defaults to `False`
+            (reasoning is stripped), preserving the historical behaviour; callers
+            decide per-turn via their replay policy (see `should_replay_reasoning`).
 
     Returns:
         A dict with `"role"` and `"content"` fields. When the message carries
         images, `"content"` is a list of text and image-URL dicts; otherwise it
-        is a plain string.
+        is a plain string. When `replay_reasoning` is `True` and reasoning is
+        present, the dict also carries a `"reasoning_content"` field.
     """
     # 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.
@@ -205,28 +249,32 @@ def message_to_openai_message(msg: Message, formatter: Formatter | None = None)
                 url = raw if raw.startswith("data:") else f"data:image/png;base64,{raw}"
             img_list.append({"type": "image_url", "image_url": {"url": url}})
 
-        return {
+        result: dict[str, Any] = {
             "role": msg.role,
             "content": [{"type": "text", "text": content}, *img_list],
         }
     else:
-        return {"role": msg.role, "content": content}
-        # Target format:
-        # {
-        #     "role": "user",
-        #     "content": [
-        #       {
-        #         "type": "text",
-        #         "text": "What's in this picture?"
-        #       },
-        #       {
-        #         "type": "image_url",
-        #         "image_url": {
-        #           "url": "data:image/jpeg;base64,<base64_string>"
-        #         }
-        #       }
-        #     ]
-        #   }
+        result = {"role": msg.role, "content": content}
+
+    if replay_reasoning and msg.thinking:
+        result["reasoning_content"] = msg.thinking
+    return result
+    # Target format:
+    # {
+    #     "role": "user",
+    #     "content": [
+    #       {
+    #         "type": "text",
+    #         "text": "What's in this picture?"
+    #       },
+    #       {
+    #         "type": "image_url",
+    #         "image_url": {
+    #           "url": "data:image/jpeg;base64,<base64_string>"
+    #         }
+    #       }
+    #     ]
+    #   }
 
 
 def messages_to_docs(msgs: list[Message]) -> list[dict[str, str]]:

mellea/stdlib/components/chat.py

@@ -42,6 +42,13 @@ class Message(Component["Message"]):
             backends that require base64 will raise a ``ValueError``).
         documents (list[Document] | None): Optional documents associated with
             the message.
+        thinking (str | None): Optional reasoning trace produced by a thinking
+            model on the turn that generated this message. Populated by `_parse`
+            from `ModelOutputThunk.thinking`; carried through `as_chat_history`
+            so backends can round-trip it on subsequent turns per their replay
+            policy. `None` or empty for messages that carry no reasoning (e.g.
+            user turns, or assistant turns from non-thinking models); the replay
+            policy and serializers treat both falsy cases identically.
 
     Attributes:
         Role (type): Type alias for the allowed role literals: `"system"`,
@@ -57,14 +64,16 @@ def __init__(
         *,
         images: None | list[ImageBlock | ImageUrlBlock] = None,
         documents: None | Iterable[str | Document] = None,
+        thinking: str | None = None,
     ):
-        """Initialize a Message with a role, text content, and optional images and documents."""
+        """Initialize a Message with a role, content, optional images/documents, and an optional reasoning trace."""
         if role not in get_args(Message.Role):
             raise ValueError(
                 f"Invalid role {role!r}. Must be one of: {list(get_args(Message.Role))}"
             )
         self.role = role
         self.content = content  # TODO this should be private.
+        self.thinking = thinking
         self._content_cblock = CBlock(self.content)
         self._images = images
         self._docs = _coerce_to_documents(documents)
@@ -130,24 +139,49 @@ def _parse(self, computed: ModelOutputThunk) -> "Message":
             # A tool was successfully requested.
             # Assistant responses for tool calling differ by backend. For the default formatter,
             # we put all of the function data into the content field in the same format we received it.
+            # Carry any captured reasoning onto the tool-issuing assistant Message: this is the
+            # exact turn the replay policy round-trips (see `should_replay_reasoning`), so the
+            # reasoning must survive parsing here or the round-trip has nothing to replay.
             if provider == "ollama" and response is not None:
+                thinking = (
+                    computed.thinking if response.message.role == "assistant" else None
+                )
                 return Message(
-                    role=response.message.role, content=str(response.message.tool_calls)
+                    role=response.message.role,
+                    content=str(response.message.tool_calls),
+                    thinking=thinking,
                 )
             if provider in ("openai", "watsonx", "litellm") and isinstance(
                 response, dict
             ):
                 choice = response["choices"][0] if "choices" in response else response
                 msg = choice["message"]
-                return Message(role=msg["role"], content=str(msg.get("tool_calls", [])))
+                thinking = computed.thinking if msg["role"] == "assistant" else None
+                return Message(
+                    role=msg["role"],
+                    content=str(msg.get("tool_calls", [])),
+                    thinking=thinking,
+                )
             # Hugging Face (or others). There are no guarantees on how the model represented the function calls.
             # Output it in the same format we received the tool call request.
             assert computed.value is not None
-            return Message(role="assistant", content=computed.value)
+            return Message(
+                role="assistant", content=computed.value, thinking=computed.thinking
+            )
 
+        # No tool call on this turn: carry any captured reasoning onto the parsed
+        # assistant Message so it can round-trip on subsequent turns. `thinking` is
+        # None for non-thinking models and for role="tool" recoveries.
         if provider == "ollama" and response is not None:
             # Ollama can return role="tool"; preserve role recovery from the response.
-            return Message(role=response.message.role, content=response.message.content)
+            thinking = (
+                computed.thinking if response.message.role == "assistant" else None
+            )
+            return Message(
+                role=response.message.role,
+                content=response.message.content,
+                thinking=thinking,
+            )
         if provider in ("openai", "watsonx", "litellm") and isinstance(response, dict):
             choices = response.get("choices") or [{}]
             msg = choices[0].get("message", {})
@@ -157,12 +191,15 @@ def _parse(self, computed: ModelOutputThunk) -> "Message":
             content = msg.get("content") or response.get("message", {}).get(
                 "content", ""
             )
-            return Message(role=role, content=content)
+            thinking = computed.thinking if role == "assistant" else None
+            return Message(role=role, content=content, thinking=thinking)
 
         # Hugging Face: raw.response is token tensors with no role/content to parse.
         # Unknown provider: nothing to switch on. Both fall back to the decoded text.
         assert computed.value is not None
-        return Message(role="assistant", content=computed.value)
+        return Message(
+            role="assistant", content=computed.value, thinking=computed.thinking
+        )
 
 
 class ToolMessage(Message):