fix(openai): keep n>1 stream choices separate in recorded output

langfuse/openai.py

@@ -760,7 +760,11 @@ def _extract_streamed_response_api_response(chunks: Any) -> Any:
 
 
 def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
-    completion: Any = defaultdict(lambda: None) if resource.type == "chat" else ""
+    # Per-choice accumulators keyed by the provider's choice index: n>1
+    # streams interleave choices within a single chunk, so a shared
+    # accumulator would merge content and tool calls across choices.
+    chat_completions: dict[int, defaultdict] = {}
+    completion_texts: dict[int, str] = {}
     model, usage, finish_reason, service_tier = None, None, None, None
 
     for chunk in chunks:
@@ -778,7 +782,15 @@ def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
         for choice in choices:
             if _is_openai_v1():
                 choice = choice.__dict__
+
+            choice_index = choice.get("index", None)
+            if not isinstance(choice_index, int):
+                choice_index = 0
+
             if resource.type == "chat":
+                completion = chat_completions.setdefault(
+                    choice_index, defaultdict(lambda: None)
+                )
                 delta = choice.get("delta", None)
                 choice_finish_reason = choice.get("finish_reason", None)
                 if choice_finish_reason is not None:
@@ -868,9 +880,11 @@ def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
                             ) + tool_arguments
 
             if resource.type == "completion":
-                completion += choice.get("text", "")
+                completion_texts[choice_index] = completion_texts.get(
+                    choice_index, ""
+                ) + (choice.get("text", "") or "")
 
-    def get_response_for_chat() -> Any:
+    def get_response_for_chat(completion: Any) -> Any:
         content = completion["content"]
 
         if completion["tool_calls"]:
@@ -897,9 +911,28 @@ def get_response_for_chat() -> Any:
 
         return content or None
 
+    response: Any
+    if resource.type == "chat":
+        if len(chat_completions) > 1:
+            response = [
+                get_response_for_chat(chat_completions[index])
+                for index in sorted(chat_completions)
+            ]
+        elif chat_completions:
+            response = get_response_for_chat(next(iter(chat_completions.values())))
+        else:
+            response = get_response_for_chat(defaultdict(lambda: None))
+    else:
+        if len(completion_texts) > 1:
+            response = [completion_texts[index] for index in sorted(completion_texts)]
+        elif completion_texts:
+            response = next(iter(completion_texts.values()))
+        else:
+            response = ""
+
     return (
         model,
-        get_response_for_chat() if resource.type == "chat" else completion,
+        response,
         usage,
         {"finish_reason": finish_reason} if finish_reason is not None else None,
         service_tier,

tests/unit/test_openai.py

@@ -251,6 +251,55 @@ def _make_single_chunk_stream():
     )
 
 
+def _make_chat_stream_chunks_two_choices():
+    # Interleaved n=2 stream: choice indexes arrive as [0, 1], [1, 0], [0, 1].
+    # Choice 0 streams content "A" then "0"; choice 1 streams "B" then "1".
+    # Each choice carries one tool call at tool index 0.
+    def _choice(index, content, tool_args, finish_reason=None):
+        return SimpleNamespace(
+            index=index,
+            delta=SimpleNamespace(
+                role="assistant" if content is not None else None,
+                content=content,
+                function_call=None,
+                tool_calls=[
+                    SimpleNamespace(
+                        index=0,
+                        id=f"call_{index}",
+                        type="function",
+                        function=SimpleNamespace(
+                            name=f"tool_{index}", arguments=tool_args
+                        ),
+                    )
+                ]
+                if tool_args is not None
+                else None,
+            ),
+            finish_reason=finish_reason,
+        )
+
+    return [
+        SimpleNamespace(
+            model="gpt-4o-mini",
+            choices=[_choice(0, "A", '{"x":'), _choice(1, "B", '{"y":')],
+            usage=None,
+        ),
+        SimpleNamespace(
+            model="gpt-4o-mini",
+            choices=[_choice(1, "1", "1}"), _choice(0, "0", "0}")],
+            usage=None,
+        ),
+        SimpleNamespace(
+            model="gpt-4o-mini",
+            choices=[
+                _choice(0, None, None, finish_reason="stop"),
+                _choice(1, None, None, finish_reason="stop"),
+            ],
+            usage=SimpleNamespace(prompt_tokens=3, completion_tokens=2, total_tokens=5),
+        ),
+    ]
+
+
 def test_chat_completion_exports_generation_span(
     langfuse_memory_client, get_span, json_attr
 ):
@@ -401,6 +450,92 @@ def test_streaming_chat_completion_preserves_tool_calls_after_content():
     assert metadata == {"finish_reason": "tool_calls"}
 
 
+def test_streaming_chat_completion_multiple_choices_are_not_merged():
+    model, completion, usage, metadata, _service_tier = (
+        lf_openai_module._extract_streamed_openai_response(
+            SimpleNamespace(type="chat"),
+            _make_chat_stream_chunks_two_choices(),
+        )
+    )
+
+    assert model == "gpt-4o-mini"
+    assert completion == [
+        {
+            "role": "assistant",
+            "content": "A0",
+            "tool_calls": [
+                {
+                    "id": "call_0",
+                    "type": "function",
+                    "function": {"name": "tool_0", "arguments": '{"x":0}'},
+                }
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": "B1",
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {"name": "tool_1", "arguments": '{"y":1}'},
+                }
+            ],
+        },
+    ]
+    assert usage.total_tokens == 5
+    assert metadata == {"finish_reason": "stop"}
+
+
+def test_streaming_chat_completion_single_choice_output_unchanged():
+    model, completion, usage, metadata, _service_tier = (
+        lf_openai_module._extract_streamed_openai_response(
+            SimpleNamespace(type="chat"),
+            _make_chat_stream_chunks(),
+        )
+    )
+
+    assert model == "gpt-4o-mini"
+    assert completion == "2"
+    assert usage.total_tokens == 4
+    assert metadata == {"finish_reason": "stop"}
+
+
+def test_streaming_chat_completion_single_nonzero_choice_index_not_lost():
+    # An n>1 stream that only ever yields output for choice 1 (for example
+    # after partial consumption) must still record that choice's response
+    # instead of reading index 0 and recording None.
+    chunks = _make_chat_stream_chunks_two_choices()
+    only_choice_one = [
+        SimpleNamespace(
+            model="gpt-4o-mini",
+            choices=[choice for choice in chunk.choices if choice.index == 1],
+            usage=chunk.usage,
+        )
+        for chunk in chunks
+    ]
+
+    model, completion, _usage, _metadata, _service_tier = (
+        lf_openai_module._extract_streamed_openai_response(
+            SimpleNamespace(type="chat"),
+            only_choice_one,
+        )
+    )
+
+    assert model == "gpt-4o-mini"
+    assert completion == {
+        "role": "assistant",
+        "content": "B1",
+        "tool_calls": [
+            {
+                "id": "call_1",
+                "type": "function",
+                "function": {"name": "tool_1", "arguments": '{"y":1}'},
+            }
+        ],
+    }
+
+
 def test_response_api_output_serializes_openai_parsed_response_objects():
     class ParsedOutput(BaseModel):
         name: str