diff --git a/langfuse/openai.py b/langfuse/openai.py index 3e480f541..de200dc3e 100644 --- a/langfuse/openai.py +++ b/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, diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index 681be4fbf..455b19567 100644 --- a/tests/unit/test_openai.py +++ b/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