fix: preserve reasoning and thinking metadata in Message.fit (#327)

**Key Changes:**

- Updated `Message.fit` to retain `thinking_blocks` and
`reasoning_content`
- Added tests to ensure preservation of these metadata fields
- Improved robustness for handling objects with optional metadata fields

**Added:**

- Unit tests for `Message.fit` verifying that `thinking_blocks` and
`reasoning_content` are correctly preserved or ignored as appropriate in
  various scenarios - `tests/test_chat.py`

**Changed:**

- Enhanced `Message.fit` method to detect and copy `thinking_blocks` and
  `reasoning_content` from input objects into the resulting `Message`'s
  metadata only if they are present and non-empty - `rigging/message.py`
- Updated logic to ensure that empty or missing `thinking_blocks` and
`reasoning_content` are not included in metadata, improving message
object
  cleanliness and consistency

Fixes ENG-3882

Author: l50

rigging/message.py

@@ -1360,11 +1360,19 @@ def fit(cls, message: t.Union["Message", MessageDict, Content, str]) -> "Message
         """Helper function to convert various common types to a Message object."""
         if isinstance(message, (str, *ContentTypes)):
             return cls(role="user", content=[message])
-        return (
-            cls.model_validate(message)
-            if isinstance(message, dict)
-            else message.model_copy(deep=True)
-        )
+        if isinstance(message, dict):
+            return cls.model_validate(message)
+
+        metadata: dict[str, t.Any] = {}
+        if hasattr(message, "thinking_blocks") and message.thinking_blocks:
+            metadata["thinking_blocks"] = message.thinking_blocks
+        if hasattr(message, "reasoning_content") and message.reasoning_content:
+            metadata["reasoning_content"] = message.reasoning_content
+
+        result = message.model_copy(deep=True)
+        if metadata:
+            result.metadata.update(metadata)
+        return result
 
     @classmethod
     def apply_to_list(cls, messages: t.Sequence["Message"], **kwargs: str) -> list["Message"]:

tests/test_chat.py

@@ -44,6 +44,100 @@ def test_message_from_str() -> None:
     assert msg.content == "Please say hello."
 
 
+def test_message_fit_preserves_thinking_blocks() -> None:
+    from pydantic import BaseModel
+
+    class LiteLLMMessage(BaseModel):
+        model_config = {"extra": "allow"}
+        role: str
+        content: str
+        thinking_blocks: list[dict[str, t.Any]] | None = None
+        metadata: dict[str, t.Any] = {}
+
+        def model_copy(self, deep: bool = False) -> "Message":  # noqa: ARG002
+            return Message(role=self.role, content=self.content, metadata=self.metadata.copy())  # type: ignore[return-value]
+
+    original = LiteLLMMessage(
+        role="assistant",
+        content="Hello!",
+        thinking_blocks=[{"type": "thinking", "content": "Let me think..."}],
+    )
+    fitted = Message.fit(original)  # type: ignore[arg-type]
+    assert fitted.metadata.get("thinking_blocks") == [{"type": "thinking", "content": "Let me think..."}]
+
+
+def test_message_fit_preserves_reasoning_content() -> None:
+    from pydantic import BaseModel
+
+    class LiteLLMMessage(BaseModel):
+        model_config = {"extra": "allow"}
+        role: str
+        content: str
+        reasoning_content: str | None = None
+        metadata: dict[str, t.Any] = {}
+
+        def model_copy(self, deep: bool = False) -> "Message":  # noqa: ARG002
+            return Message(role=self.role, content=self.content, metadata=self.metadata.copy())  # type: ignore[return-value]
+
+    original = LiteLLMMessage(
+        role="assistant",
+        content="Hello!",
+        reasoning_content="I reasoned about this carefully.",
+    )
+    fitted = Message.fit(original)  # type: ignore[arg-type]
+    assert fitted.metadata.get("reasoning_content") == "I reasoned about this carefully."
+
+
+def test_message_fit_preserves_both_thinking_fields() -> None:
+    from pydantic import BaseModel
+
+    class LiteLLMMessage(BaseModel):
+        model_config = {"extra": "allow"}
+        role: str
+        content: str
+        thinking_blocks: list[dict[str, t.Any]] | None = None
+        reasoning_content: str | None = None
+        metadata: dict[str, t.Any] = {}
+
+        def model_copy(self, deep: bool = False) -> "Message":  # noqa: ARG002
+            return Message(role=self.role, content=self.content, metadata=self.metadata.copy())  # type: ignore[return-value]
+
+    original = LiteLLMMessage(
+        role="assistant",
+        content="Hello!",
+        thinking_blocks=[{"type": "thinking", "content": "Thinking..."}],
+        reasoning_content="Reasoning...",
+    )
+    fitted = Message.fit(original)  # type: ignore[arg-type]
+    assert fitted.metadata.get("thinking_blocks") == [{"type": "thinking", "content": "Thinking..."}]
+    assert fitted.metadata.get("reasoning_content") == "Reasoning..."
+
+
+def test_message_fit_ignores_empty_thinking_blocks() -> None:
+    from pydantic import BaseModel
+
+    class LiteLLMMessage(BaseModel):
+        model_config = {"extra": "allow"}
+        role: str
+        content: str
+        thinking_blocks: list[dict[str, t.Any]] | None = None
+        reasoning_content: str | None = None
+        metadata: dict[str, t.Any] = {}
+
+        def model_copy(self, deep: bool = False) -> "Message":  # noqa: ARG002
+            return Message(role=self.role, content=self.content, metadata=self.metadata.copy())  # type: ignore[return-value]
+
+    original = LiteLLMMessage(
+        role="assistant",
+        content="Hello!",
+        thinking_blocks=[],
+        reasoning_content="",
+    )
+    fitted = Message.fit(original)  # type: ignore[arg-type]
+    assert "thinking_blocks" not in fitted.metadata
+    assert "reasoning_content" not in fitted.metadata
+
+
 def test_message_str_representation() -> None:
     msg = Message("assistant", "I am an AI assistant.")
     assert str(msg) == "[assistant]: I am an AI assistant."