diff --git a/langfuse/openai.py b/langfuse/openai.py index 3e480f541..d22904e67 100644 --- a/langfuse/openai.py +++ b/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,13 +875,13 @@ 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: @@ -884,10 +889,10 @@ def get_response_for_chat() -> Any: 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, diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index 681be4fbf..91cc9f017 100644 --- a/tests/unit/test_openai.py +++ b/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"}'), + ]