diff --git a/.sampo/changesets/openai-stream-n-choices.md b/.sampo/changesets/openai-stream-n-choices.md new file mode 100644 index 00000000..d2184d13 --- /dev/null +++ b/.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. diff --git a/posthog/ai/openai/_streaming.py b/posthog/ai/openai/_streaming.py index cc6baf13..3a7d059b 100644 --- a/posthog/ai/openai/_streaming.py +++ b/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( diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 1bdcdeb6..998e3fd1 100644 --- a/posthog/ai/openai/openai.py +++ b/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, diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 1282e1dd..a37aebb6 100644 --- a/posthog/ai/openai/openai_async.py +++ b/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, diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index dcbfde24..edbe3c9d 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/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 = [