Fix extended thinking block round-trip for LiteLLM assistant messages (#85)

Fixes #71

Related downstream issue:
langchain-ai/deepagents#661

Author: Akshay-Dongare

langchain_litellm/chat_models/litellm.py

@@ -102,6 +102,27 @@ def _create_retry_decorator(
         error_types=errors, max_retries=llm.max_retries, run_manager=run_manager
     )
 
+def _inject_reasoning_content_into_content(
+    content: Any, reasoning_content: str
+) -> List[Dict[str, Any]]:
+    thinking_block = {"type": "thinking", "thinking": reasoning_content}
+    if isinstance(content, list):
+        has_thinking_block = any(
+            isinstance(block, dict)
+            and block.get("type") in ("thinking", "redacted_thinking")
+            for block in content
+        )
+        if has_thinking_block:
+            return content
+        return [thinking_block, *content]
+
+    if not content:
+        return [thinking_block]
+
+    if isinstance(content, str):
+        return [thinking_block, {"type": "text", "text": content}]
+
+    return [thinking_block, content]
 
 def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
     role = _dict["role"]
@@ -152,6 +173,9 @@ def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
 
         if _dict.get("reasoning_content"):
             additional_kwargs["reasoning_content"] = _dict["reasoning_content"]
+            content = _inject_reasoning_content_into_content(
+                content, _dict["reasoning_content"]
+            )
 
         # Check standard field first, then fallback to Vertex specific field
         provider_specific_fields = _dict.get("provider_specific_fields")
@@ -205,6 +229,9 @@ def _convert_delta_to_message_chunk(
         additional_kwargs["function_call"] = dict(function_call)
     if reasoning_content:
         additional_kwargs["reasoning_content"] = reasoning_content
+
+    if reasoning_content and (role == "assistant" or default_class == AIMessageChunk):
+        content = _inject_reasoning_content_into_content(content, reasoning_content)
 
     if provider_specific_fields is not None:
         additional_kwargs["provider_specific_fields"] = provider_specific_fields

tests/unit_tests/test_litellm.py

@@ -7,7 +7,11 @@
 from litellm.types.utils import ChatCompletionDeltaToolCall, Delta, Function
 
 from langchain_litellm.chat_models import ChatLiteLLM
-from langchain_litellm.chat_models.litellm import _convert_delta_to_message_chunk, _convert_dict_to_message
+from langchain_litellm.chat_models.litellm import (
+    _convert_delta_to_message_chunk,
+    _convert_dict_to_message,
+    _inject_reasoning_content_into_content,
+)
 
 
 class TestChatLiteLLMUnit(ChatModelUnitTests):
@@ -154,4 +158,41 @@ def test_provider_specific_fields_in_chat_result(self):
         result = llm._create_chat_result(mock_response)
 
         assert "provider_specific_fields" in result.llm_output
-        assert result.llm_output["provider_specific_fields"]["citations"][0]["source"] == "test"
+        assert result.llm_output["provider_specific_fields"]["citations"][0]["source"] == "test"
+
+
+def test_inject_reasoning_content_into_string_content() -> None:
+    result = _inject_reasoning_content_into_content("answer", "hidden chain")
+
+    assert result == [
+        {"type": "thinking", "thinking": "hidden chain"},
+        {"type": "text", "text": "answer"},
+    ]
+
+
+def test_inject_reasoning_content_into_empty_content() -> None:
+    result = _inject_reasoning_content_into_content("", "hidden chain")
+
+    assert result == [{"type": "thinking", "thinking": "hidden chain"}]
+
+
+def test_inject_reasoning_content_prepends_for_list_without_thinking() -> None:
+    content = [{"type": "text", "text": "answer"}]
+
+    result = _inject_reasoning_content_into_content(content, "hidden chain")
+
+    assert result == [
+        {"type": "thinking", "thinking": "hidden chain"},
+        {"type": "text", "text": "answer"},
+    ]
+
+
+def test_inject_reasoning_content_does_not_duplicate_existing_thinking() -> None:
+    content = [
+        {"type": "thinking", "thinking": "already there"},
+        {"type": "text", "text": "answer"},
+    ]
+
+    result = _inject_reasoning_content_into_content(content, "hidden chain")
+
+    assert result == content