fix: keep finish_reason in create_stream when the finish chunk also carries usage

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -945,7 +945,7 @@ async def create_stream(
 
             # for liteLLM chunk usage, do the following hack keeping the pervious chunk.stop_reason (if set).
             # set the stop_reason for the usage chunk to the prior stop_reason
-            stop_reason = choice.finish_reason if chunk.usage is None and stop_reason is None else stop_reason
+            stop_reason = choice.finish_reason if stop_reason is None else stop_reason
             maybe_model = chunk.model
 
             reasoning_content: str | None = None

python/packages/autogen-ext/tests/models/test_openai_model_client.py

@@ -362,6 +362,39 @@ async def test_openai_chat_completion_client_create_stream_no_usage_explicit(mon
     assert chunks[2] == " Yet Another Hello"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("finish_reason", ["stop", "length"])
+async def test_openai_chat_completion_client_create_stream_finish_reason_with_usage(
+    monkeypatch: pytest.MonkeyPatch, finish_reason: Literal["stop", "length"]
+) -> None:
+    # Some OpenAI-compatible servers (e.g. the Amazon Bedrock runtime /openai/v1 endpoint with GPT-6 models) put
+    # usage on the chunk that carries finish_reason, then send the usual choices=[] usage chunk.
+    usage = CompletionUsage(prompt_tokens=3, completion_tokens=3, total_tokens=6)
+
+    async def _mock_stream_usage(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]:
+        for choices, chunk_usage in [
+            ([ChunkChoice(finish_reason=None, index=0, delta=ChoiceDelta(content="Hello", role="assistant"))], None),
+            ([ChunkChoice(finish_reason=finish_reason, index=0, delta=ChoiceDelta())], usage),
+            ([], usage),
+        ]:
+            yield ChatCompletionChunk(
+                id="id", choices=choices, created=0, model="gpt-4o", object="chat.completion.chunk", usage=chunk_usage
+            )
+
+    async def _mock_create_usage(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]:
+        return _mock_stream_usage(*args, **kwargs)
+
+    monkeypatch.setattr(AsyncCompletions, "create", _mock_create_usage)
+    client = OpenAIChatCompletionClient(model="gpt-4o", api_key="api_key")
+    chunks: List[str | CreateResult] = []
+    async for chunk in client.create_stream(messages=[UserMessage(content="Hello", source="user")], include_usage=True):
+        chunks.append(chunk)
+    assert isinstance(chunks[-1], CreateResult)
+    assert chunks[-1].content == "Hello"
+    assert chunks[-1].usage == RequestUsage(prompt_tokens=3, completion_tokens=3)
+    assert chunks[-1].finish_reason == finish_reason
+
+
 @pytest.mark.asyncio
 async def test_openai_chat_completion_client_none_usage(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test that completion_tokens and prompt_tokens handle None usage correctly.