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..8963b1a68 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -1,12 +1,24 @@ import asyncio +import json +from collections.abc import Iterator +from http import HTTPStatus from types import SimpleNamespace +from typing import Any from unittest.mock import patch import pytest +from openai.types.chat import ChatCompletionChunk from openai.types.responses import ParsedResponseOutputMessage, ParsedResponseOutputText +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, +) +from opentelemetry.sdk.trace import TracerProvider from pydantic import BaseModel +from pytest_httpserver import HTTPServer +from pytest_httpserver.httpserver import RequestMatcher import langfuse.openai as lf_openai_module +from langfuse import Langfuse from langfuse._client.attributes import LangfuseOtelSpanAttributes from langfuse.openai import openai as lf_openai @@ -1399,3 +1411,152 @@ def test_with_raw_response_streaming_passes_through_untraced( span.name != "OpenAI-generation" for span in memory_exporter.get_finished_spans() ) + + +@pytest.fixture +def langfuse_http_client( + monkeypatch: pytest.MonkeyPatch, httpserver: HTTPServer +) -> Iterator[Langfuse]: + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "test-public-key") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "test-secret-key") + monkeypatch.setenv("LANGFUSE_BASE_URL", httpserver.url_for("").rstrip("/")) + httpserver.expect_request( + "/api/public/otel/v1/traces", method="POST" + ).respond_with_data(response_data=b"", content_type="application/x-protobuf") + tracer_provider = TracerProvider() + client = Langfuse(tracer_provider=tracer_provider, timeout=5) + yield client + client.shutdown() + tracer_provider.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_client", [False, True], ids=["sync", "async"]) +async def test_streaming_chat_completion_keeps_multiple_choices_separate( + langfuse_http_client: Langfuse, + httpserver: HTTPServer, + async_client: bool, +) -> None: + body = ( + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",' + '"created":1700000000,"model":"gpt-4o-mini",' + '"choices":[{"index":1,"delta":{"role":"assistant","content":"B"},' + '"finish_reason":null},{"index":0,"delta":{"role":"assistant","content":"A"},' + '"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: {"id":"chatcmpl-test","object":"chat.completion.chunk",' + '"created":1700000000,"model":"gpt-4o-mini","choices":[],' + '"usage":{"prompt_tokens":7,"completion_tokens":5,"total_tokens":12}}\n\n' + "data: [DONE]\n\n" + ) + args: dict[str, Any] = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "choose"}], + "n": 2, + "stream": True, + "stream_options": {"include_usage": True}, + } + httpserver.expect_request( + "/v1/chat/completions", method="POST", json=args + ).respond_with_data( + response_data=body, + content_type="text/event-stream", + ) + + chunks: list[ChatCompletionChunk] + if async_client: + async with lf_openai.AsyncOpenAI( + api_key="test", base_url=httpserver.url_for("/v1") + ) as async_provider: + async with await async_provider.chat.completions.create( + **args + ) as async_stream: + chunks = [chunk async for chunk in async_stream] + else: + with lf_openai.OpenAI( + api_key="test", base_url=httpserver.url_for("/v1") + ) as provider: + with provider.chat.completions.create(**args) as stream: + chunks = list(stream) + + assert [[choice.index for choice in chunk.choices] for chunk in chunks] == [ + [1, 0], + [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_http_client.flush() + ingestion = list( + httpserver.iter_matching_requests( + RequestMatcher(uri="/api/public/otel/v1/traces", method="POST") + ) + ) + assert len(ingestion) == 1 + ingestion_request, ingestion_response = ingestion[0] + assert ingestion_response.status_code == HTTPStatus.OK + assert ingestion_request.content_type == "application/x-protobuf" + payload = ExportTraceServiceRequest.FromString(ingestion_request.get_data()) + spans = [ + span + for resource in payload.resource_spans + for scope in resource.scope_spans + for span in scope.spans + ] + assert len(spans) == 1 + attributes = { + attribute.key: attribute.value.string_value for attribute in spans[0].attributes + } + output = json.loads(attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT]) + + assert isinstance(output, list), 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"}'), + ] + assert attributes["langfuse.observation.metadata.finish_reason"] == "tool_calls" + assert chunks[-1].usage is not None + assert chunks[-1].usage.model_dump(exclude_none=True) == { + "prompt_tokens": 7, + "completion_tokens": 5, + "total_tokens": 12, + } + assert ( + json.loads(attributes[LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS]) + == chunks[-1].usage.model_dump() + ) + print( + json.dumps( + { + "output": output, + "usage": json.loads( + attributes[LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS] + ), + }, + separators=(",", ":"), + ) + )