fix(ai): keep the choices of an n>1 OpenAI chat stream apart

.sampo/changesets/openai-stream-n-choices.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Keep the choices of a streamed OpenAI chat completion with `n > 1` apart in `$ai_output_choices`, as a non-streamed one already does. Their text and tool calls were interleaved into one choice.

posthog/ai/openai/_streaming.py

@@ -3,13 +3,14 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Union
 
-from ..types import StreamingEventData, TokenUsage
+from ..types import FormattedMessage, StreamingEventData, TokenUsage
 from ..utils import merge_usage_stats, _responses_stop_reason
 from .openai_converter import (
     accumulate_openai_tool_calls,
     extract_openai_content_from_chunk,
     extract_openai_tool_calls_from_chunk,
     extract_openai_usage_from_chunk,
+    format_openai_streaming_output,
 )
 
 
@@ -46,13 +47,26 @@ def process_chunk(self, chunk: Any) -> None:
                 self.stop_reason = stop_reason
 
 
+@dataclass
+class _ChoiceStream:
+    """One choice's view of a chunk, so the per-chunk helpers read that choice."""
+
+    choices: List[Any]
+
+
 @dataclass
 class _ChatCompletionsStreamState:
-    """Accumulates state specific to a Chat Completions stream."""
+    """Accumulates state specific to a Chat Completions stream.
+
+    With ``n > 1`` each chunk carries the delta of one choice, named by its
+    ``index``, so content and tool calls are kept per choice.
+    """
 
     usage_stats: TokenUsage = field(default_factory=lambda: TokenUsage())
-    output: List[Any] = field(default_factory=list)
-    _tool_calls: Dict[Union[int, str], Dict[str, Any]] = field(default_factory=dict)
+    _outputs: Dict[int, List[Any]] = field(default_factory=dict)
+    _tool_calls: Dict[int, Dict[Union[int, str], Dict[str, Any]]] = field(
+        default_factory=dict
+    )
     model: Optional[str] = None
     stop_reason: Optional[str] = None
     service_tier: Optional[str] = None
@@ -67,23 +81,53 @@ def process_chunk(self, chunk: Any) -> None:
         if chunk_usage:
             merge_usage_stats(self.usage_stats, chunk_usage)
 
-        content = extract_openai_content_from_chunk(chunk, "chat")
-        if content is not None:
-            self.output.append(content)
-
-        chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk)
-        if chunk_tool_calls:
-            accumulate_openai_tool_calls(self._tool_calls, chunk_tool_calls)
-
-        choices = getattr(chunk, "choices", None)
-        if choices:
-            finish_reason = getattr(choices[0], "finish_reason", None)
-            if finish_reason is not None:
+        for choice in getattr(chunk, "choices", None) or []:
+            index = getattr(choice, "index", None)
+            if not isinstance(index, int):
+                index = 0
+            view = _ChoiceStream([choice])
+
+            content = extract_openai_content_from_chunk(view, "chat")
+            if content is not None:
+                self._outputs.setdefault(index, []).append(content)
+
+            chunk_tool_calls = extract_openai_tool_calls_from_chunk(view)
+            if chunk_tool_calls:
+                accumulate_openai_tool_calls(
+                    self._tool_calls.setdefault(index, {}), chunk_tool_calls
+                )
+
+            # Like a non-streamed completion, the stop reason is the first choice's.
+            finish_reason = getattr(choice, "finish_reason", None)
+            if finish_reason is not None and index == 0:
                 self.stop_reason = finish_reason
 
     @property
     def tool_calls(self) -> Optional[List[Dict[str, Any]]]:
-        return list(self._tool_calls.values()) if self._tool_calls else None
+        calls = [
+            call
+            for index in sorted(self._tool_calls)
+            for call in self._tool_calls[index].values()
+        ]
+        return calls or None
+
+    def _indexes(self) -> List[int]:
+        return sorted(set(self._outputs) | set(self._tool_calls))
+
+    def formatted_output(self) -> List[FormattedMessage]:
+        """One assistant message with each choice's content in index order, the
+        shape a non-streamed completion is captured with."""
+        content: List[Any] = []
+        for index in self._indexes():
+            tool_calls = self._tool_calls.get(index)
+            content.extend(
+                format_openai_streaming_output(
+                    self._outputs.get(index, []),
+                    "chat",
+                    list(tool_calls.values()) if tool_calls else None,
+                )[0]["content"]
+            )
+        return [{"role": "assistant", "content": content}]
 
 
 def _build_streaming_event_data(

posthog/ai/openai/openai.py

@@ -439,9 +439,7 @@ def _capture_streaming_event(
             base_url=self._client.base_url,
             kwargs=kwargs,
             formatted_input=_format_openai_streaming_input(kwargs, "chat"),
-            formatted_output=_format_openai_streaming_output(
-                state.output, "chat", state.tool_calls
-            ),
+            formatted_output=state.formatted_output(),
             usage_stats=state.usage_stats,
             latency=latency,
             distinct_id=posthog_distinct_id,

posthog/ai/openai/openai_async.py

@@ -443,9 +443,7 @@ async def _capture_streaming_event(
             base_url=self._client.base_url,
             kwargs=kwargs,
             formatted_input=format_openai_streaming_input(kwargs, "chat"),
-            formatted_output=format_openai_streaming_output(
-                state.output, "chat", state.tool_calls
-            ),
+            formatted_output=state.formatted_output(),
             usage_stats=state.usage_stats,
             latency=latency,
             distinct_id=posthog_distinct_id,

posthog/test/ai/openai/test_openai.py

@@ -2028,6 +2028,59 @@ def test_streaming_chat_extracts_model_from_chunk_when_not_in_kwargs(mock_client
         assert props["$ai_model"] == "gpt-4o-stored-prompt"
 
 
+def test_streaming_chat_keeps_n_choices_apart(mock_client):
+    # With n=2 each chunk carries one choice's delta, named by its index.
+    def chunk(index, content=None, finish_reason=None):
+        return ChatCompletionChunk(
+            id="chunk",
+            model="gpt-4o-mini",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=index,
+                    delta=ChoiceDelta(content=content),
+                    finish_reason=finish_reason,
+                )
+            ],
+        )
+
+    chunks = [
+        chunk(0, "Hello"),
+        chunk(1, "Bonjour"),
+        chunk(0, " world"),
+        chunk(1, " le monde"),
+        chunk(1, finish_reason="length"),
+        chunk(0, finish_reason="stop"),
+    ]
+
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        mock_create.return_value = chunks
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        list(
+            client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": "Hello"}],
+                n=2,
+                stream=True,
+                posthog_distinct_id="test-id",
+            )
+        )
+
+    props = mock_client.capture.call_args[1]["properties"]
+    # The same shape a non-streamed n=2 completion is captured with.
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "Hello world"},
+                {"type": "text", "text": "Bonjour le monde"},
+            ],
+        }
+    ]
+    assert props["$ai_stop_reason"] == "stop"
+
+
 def test_streaming_chat_prefers_kwargs_model_over_chunk_model(mock_client):
     """Test that model from kwargs takes precedence over model from chunk."""
     chunks = [