models - litellm - stop reasoning

Author: pgrayy

src/strands/models/litellm.py

@@ -139,22 +139,22 @@ async def stream(
 
         logger.debug("got response from model")
         yield self.format_chunk({"chunk_type": "message_start"})
-        yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"})
 
         tool_calls: dict[int, list[Any]] = {}
+        started_reasoning = False
+        started_text = False
 
         async for event in response:
             # Defensive: skip events with empty or missing choices
             if not getattr(event, "choices", None):
                 continue
             choice = event.choices[0]
 
-            if choice.delta.content:
-                yield self.format_chunk(
-                    {"chunk_type": "content_delta", "data_type": "text", "data": choice.delta.content}
-                )
-
             if hasattr(choice.delta, "reasoning_content") and choice.delta.reasoning_content:
+                if not started_reasoning:
+                    yield self.format_chunk({"chunk_type": "content_start", "data_type": "reasoning_content"})
+                    started_reasoning = True
+
                 yield self.format_chunk(
                     {
                         "chunk_type": "content_delta",
@@ -163,14 +163,27 @@ async def stream(
                     }
                 )
 
+            if choice.delta.content:
+                if started_reasoning:
+                    yield self.format_chunk({"chunk_type": "content_stop", "data_type": "reasoning_content"})
+                    started_reasoning = False
+
+                if not started_text:
+                    yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"})
+                    started_text = True
+
+                yield self.format_chunk(
+                    {"chunk_type": "content_delta", "data_type": "text", "data": choice.delta.content}
+                )
+
             for tool_call in choice.delta.tool_calls or []:
                 tool_calls.setdefault(tool_call.index, []).append(tool_call)
 
             if choice.finish_reason:
+                if started_text:
+                    yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"})
                 break
 
-        yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"})
-
         for tool_deltas in tool_calls.values():
             yield self.format_chunk({"chunk_type": "content_start", "data_type": "tool", "data": tool_deltas[0]})
 

tests/strands/models/test_litellm.py

@@ -182,6 +182,8 @@ async def test_stream(litellm_acompletion, api_key, model_id, model, agenerator,
         {"messageStart": {"role": "assistant"}},
         {"contentBlockStart": {"start": {}}},
         {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "\nI'm thinking"}}}},
+        {"contentBlockStop": {}},
+        {"contentBlockStart": {"start": {}}},
         {"contentBlockDelta": {"delta": {"text": "I'll calculate"}}},
         {"contentBlockDelta": {"delta": {"text": "that for you"}}},
         {"contentBlockStop": {}},
@@ -251,8 +253,6 @@ async def test_stream_empty(litellm_acompletion, api_key, model_id, model, agene
     tru_events = await alist(response)
     exp_events = [
         {"messageStart": {"role": "assistant"}},
-        {"contentBlockStart": {"start": {}}},
-        {"contentBlockStop": {}},
         {"messageStop": {"stopReason": "end_turn"}},
     ]