fix(openai): keep streamed choices separate

langfuse/openai.py

@@ -760,7 +760,10 @@ 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 ""
+    chat_completions: defaultdict[int, defaultdict[str, Any]] = defaultdict(
+        lambda: defaultdict(lambda: None)
+    )
+    completion: Any = ""
     model, usage, finish_reason, service_tier = None, None, None, None
 
     for chunk in chunks:
@@ -775,10 +778,12 @@ def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
 
         choices = chunk.get("choices") or []
 
-        for choice in choices:
+        for choice_position, choice in enumerate(choices):
             if _is_openai_v1():
                 choice = choice.__dict__
             if resource.type == "chat":
+                choice_index = cast(int, choice.get("index", choice_position))
+                completion = chat_completions[choice_index]
                 delta = choice.get("delta", None)
                 choice_finish_reason = choice.get("finish_reason", None)
                 if choice_finish_reason is not None:
@@ -870,24 +875,24 @@ def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
             if resource.type == "completion":
                 completion += choice.get("text", "")
 
-    def get_response_for_chat() -> Any:
-        content = completion["content"]
+    def get_response_for_chat(chat_completion: Any) -> Any:
+        content = chat_completion["content"]
 
-        if completion["tool_calls"]:
+        if chat_completion["tool_calls"]:
             response = {
                 "role": "assistant",
-                "tool_calls": completion["tool_calls"],
+                "tool_calls": chat_completion["tool_calls"],
             }
 
             if content is not None:
                 response["content"] = content
 
             return response
 
-        if completion["function_call"]:
+        if chat_completion["function_call"]:
             response = {
                 "role": "assistant",
-                "function_call": completion["function_call"],
+                "function_call": chat_completion["function_call"],
             }
 
             if content is not None:
@@ -897,9 +902,15 @@ def get_response_for_chat() -> Any:
 
         return content or None
 
+    chat_outputs = [
+        get_response_for_chat(chat_completion)
+        for _, chat_completion in sorted(chat_completions.items())
+    ]
+    chat_output = chat_outputs[0] if len(chat_outputs) == 1 else chat_outputs or None
+
     return (
         model,
-        get_response_for_chat() if resource.type == "chat" else completion,
+        chat_output if resource.type == "chat" else completion,
         usage,
         {"finish_reason": finish_reason} if finish_reason is not None else None,
         service_tier,

tests/unit/test_openai.py

@@ -1,7 +1,9 @@
 import asyncio
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import patch
 
+import httpx
 import pytest
 from openai.types.responses import ParsedResponseOutputMessage, ParsedResponseOutputText
 from pydantic import BaseModel
@@ -1399,3 +1401,77 @@ def test_with_raw_response_streaming_passes_through_untraced(
         span.name != "OpenAI-generation"
         for span in memory_exporter.get_finished_spans()
     )
+
+
+def test_streaming_chat_completion_keeps_multiple_choices_separate(
+    langfuse_memory_client: Any, get_span: Any, json_attr: Any
+) -> None:
+    body = (
+        'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",'
+        '"created":1700000000,"model":"gpt-4o-mini",'
+        '"choices":[{"index":0,"delta":{"role":"assistant","content":"A"},'
+        '"finish_reason":null},{"index":1,"delta":{"role":"assistant","content":"B"},'
+        '"finish_reason":null}]}\n\n'
+        'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",'
+        '"created":1700000000,"model":"gpt-4o-mini",'
+        '"choices":[{"index":1,"delta":{"content":"1","tool_calls":[{"index":0,'
+        '"id":"call-1","type":"function","function":{"name":"lookup",'
+        '"arguments":"{\\"value\\":\\"B"}}]},"finish_reason":null},{"index":0,'
+        '"delta":{"content":"0","tool_calls":[{"index":0,"id":"call-0",'
+        '"type":"function","function":{"name":"lookup",'
+        '"arguments":"{\\"value\\":\\"A"}}]},"finish_reason":null}]}\n\n'
+        'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",'
+        '"created":1700000000,"model":"gpt-4o-mini",'
+        '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,'
+        '"function":{"arguments":"0\\"}"}}]},"finish_reason":"tool_calls"},'
+        '{"index":1,"delta":{"tool_calls":[{"index":0,'
+        '"function":{"arguments":"1\\"}"}}]},"finish_reason":"tool_calls"}]}\n\n'
+        "data: [DONE]\n\n"
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            status_code=200,
+            content=body.encode(),
+            headers={"content-type": "text/event-stream"},
+        )
+
+    openai_client = lf_openai.OpenAI(
+        api_key="test",
+        http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    chunks = list(
+        openai_client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "choose"}],
+            n=2,
+            stream=True,
+        )
+    )
+
+    assert [[choice.index for choice in chunk.choices] for chunk in chunks] == [
+        [0, 1],
+        [1, 0],
+        [0, 1],
+    ]
+    assert [
+        (choice.index, tool_call.index)
+        for chunk in chunks
+        for choice in chunk.choices
+        for tool_call in choice.delta.tool_calls or []
+    ] == [(1, 0), (0, 0), (0, 0), (1, 0)]
+
+    langfuse_memory_client.flush()
+    span = get_span("OpenAI-generation")
+    output = json_attr(span, LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT)
+
+    assert [choice["content"] for choice in output] == ["A0", "B1"]
+    assert [
+        (tool_call["id"], tool_call["function"]["arguments"])
+        for choice in output
+        for tool_call in choice["tool_calls"]
+    ] == [
+        ("call-0", '{"value":"A0"}'),
+        ("call-1", '{"value":"B1"}'),
+    ]