diff --git a/langfuse/openai.py b/langfuse/openai.py index ea9dcd68c..e54cca1e1 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -32,7 +32,7 @@ from inspect import isawaitable, isclass from typing import Any, Optional, cast -from openai._types import NotGiven +from openai._types import NotGiven, Omit from packaging.version import Version from pydantic import BaseModel from pydantic_core import to_jsonable_python @@ -137,6 +137,24 @@ class OpenAiDefinition: min_version="1.50.0", max_version="1.92.0", ), + OpenAiDefinition( + module="openai.resources.beta.chat.completions", + object="Completions", + method="stream", + type="chat", + sync=True, + min_version="1.40.0", + max_version="1.92.0", + ), + OpenAiDefinition( + module="openai.resources.beta.chat.completions", + object="AsyncCompletions", + method="stream", + type="chat", + sync=False, + min_version="1.40.0", + max_version="1.92.0", + ), OpenAiDefinition( module="openai.resources.chat.completions", object="Completions", @@ -153,6 +171,22 @@ class OpenAiDefinition: sync=False, min_version="1.92.0", ), + OpenAiDefinition( + module="openai.resources.chat.completions", + object="Completions", + method="stream", + type="chat", + sync=True, + min_version="1.92.0", + ), + OpenAiDefinition( + module="openai.resources.chat.completions", + object="AsyncCompletions", + method="stream", + type="chat", + sync=False, + min_version="1.92.0", + ), OpenAiDefinition( module="openai.resources.responses", object="Responses", @@ -185,6 +219,22 @@ class OpenAiDefinition: sync=False, min_version="1.66.0", ), + OpenAiDefinition( + module="openai.resources.responses", + object="Responses", + method="stream", + type="chat", + sync=True, + min_version="1.66.0", + ), + OpenAiDefinition( + module="openai.resources.responses", + object="AsyncResponses", + method="stream", + type="chat", + sync=False, + min_version="1.66.0", + ), OpenAiDefinition( module="openai.resources.embeddings", object="Embeddings", @@ -204,10 +254,18 @@ class OpenAiDefinition: _RESPONSES_PROMPT_FIELDS = ("tools", "tool_choice", "parallel_tool_calls") _STRUCTURED_OUTPUT_METADATA_FIELDS = ("response_format", "text_format") +_LANGFUSE_STREAM_ARG_NAMES = ( + "name", + "langfuse_prompt", + "langfuse_public_key", + "trace_id", + "parent_observation_id", +) +_LANGFUSE_OPENAI_STREAM_ARGS = object() def _is_not_given(value: Any) -> bool: - return isinstance(value, NotGiven) + return isinstance(value, (NotGiven, Omit)) def _get_attr_or_item(value: Any, key: str, default: Any = None) -> Any: @@ -321,6 +379,20 @@ def __init__( self.args["trace_id"] = trace_id self.args["parent_observation_id"] = parent_observation_id + extra_body = kwargs.get("extra_body") + if isinstance(extra_body, dict) and _LANGFUSE_OPENAI_STREAM_ARGS in extra_body: + request_extra_body = extra_body.copy() + stream_args = dict(request_extra_body.pop(_LANGFUSE_OPENAI_STREAM_ARGS)) + kwargs["extra_body"] = request_extra_body + + if "metadata" in stream_args: + self.metadata = stream_args.pop("metadata") + self.args["metadata"] = _get_structured_output_metadata( + self.metadata, kwargs + ) + + self.args.update(stream_args) + self.kwargs = kwargs def get_langfuse_args(self) -> Any: @@ -356,13 +428,13 @@ def _extract_responses_prompt(kwargs: Any) -> Any: for key in _RESPONSES_PROMPT_FIELDS: value = kwargs.get(key, None) - if value is not None and not isinstance(value, NotGiven): + if value is not None and not _is_not_given(value): prompt_fields[key] = _serialize_openai_value(value) - if isinstance(input_value, NotGiven): + if _is_not_given(input_value): input_value = None - if isinstance(instructions, NotGiven): + if _is_not_given(instructions): instructions = None if instructions is None: @@ -395,13 +467,15 @@ def _extract_chat_prompt(kwargs: Any) -> Any: """Extracts the user input from prompts. Returns an array of messages or dict with messages and functions""" prompt = {} - if kwargs.get("functions") is not None: + if kwargs.get("functions") is not None and not _is_not_given(kwargs["functions"]): prompt.update({"functions": kwargs["functions"]}) - if kwargs.get("function_call") is not None: + if kwargs.get("function_call") is not None and not _is_not_given( + kwargs["function_call"] + ): prompt.update({"function_call": kwargs["function_call"]}) - if kwargs.get("tools") is not None: + if kwargs.get("tools") is not None and not _is_not_given(kwargs["tools"]): prompt.update({"tools": kwargs["tools"]}) if prompt: @@ -531,11 +605,9 @@ def _get_langfuse_data_from_kwargs(resource: OpenAiDefinition, kwargs: Any) -> A raise ValueError("parent_observation_id requires trace_id to be set") metadata = kwargs.get("metadata", {}) - if ( - metadata is not None - and not isinstance(metadata, NotGiven) - and not isinstance(metadata, dict) - ): + if _is_not_given(metadata): + metadata = {} + elif metadata is not None and not isinstance(metadata, dict): if isinstance(metadata, BaseModel): metadata = _serialize_openai_value(metadata) else: @@ -556,63 +628,61 @@ def _get_langfuse_data_from_kwargs(resource: OpenAiDefinition, kwargs: Any) -> A parsed_temperature = ( kwargs.get("temperature", 1) - if not isinstance(kwargs.get("temperature", 1), NotGiven) + if not _is_not_given(kwargs.get("temperature", 1)) else 1 ) parsed_max_tokens = ( kwargs.get("max_tokens", float("inf")) - if not isinstance(kwargs.get("max_tokens", float("inf")), NotGiven) + if not _is_not_given(kwargs.get("max_tokens", float("inf"))) else float("inf") ) parsed_max_completion_tokens = ( kwargs.get("max_completion_tokens", None) - if not isinstance(kwargs.get("max_completion_tokens", float("inf")), NotGiven) + if not _is_not_given(kwargs.get("max_completion_tokens", float("inf"))) else None ) parsed_top_p = ( - kwargs.get("top_p", 1) - if not isinstance(kwargs.get("top_p", 1), NotGiven) - else 1 + kwargs.get("top_p", 1) if not _is_not_given(kwargs.get("top_p", 1)) else 1 ) parsed_frequency_penalty = ( kwargs.get("frequency_penalty", 0) - if not isinstance(kwargs.get("frequency_penalty", 0), NotGiven) + if not _is_not_given(kwargs.get("frequency_penalty", 0)) else 0 ) parsed_presence_penalty = ( kwargs.get("presence_penalty", 0) - if not isinstance(kwargs.get("presence_penalty", 0), NotGiven) + if not _is_not_given(kwargs.get("presence_penalty", 0)) else 0 ) parsed_seed = ( kwargs.get("seed", None) - if not isinstance(kwargs.get("seed", None), NotGiven) + if not _is_not_given(kwargs.get("seed", None)) else None ) - parsed_n = kwargs.get("n", 1) if not isinstance(kwargs.get("n", 1), NotGiven) else 1 + parsed_n = kwargs.get("n", 1) if not _is_not_given(kwargs.get("n", 1)) else 1 parsed_service_tier = ( kwargs.get("service_tier", None) - if not isinstance(kwargs.get("service_tier", None), NotGiven) + if not _is_not_given(kwargs.get("service_tier", None)) else None ) if resource.type == "embedding": parsed_dimensions = ( kwargs.get("dimensions", None) - if not isinstance(kwargs.get("dimensions", None), NotGiven) + if not _is_not_given(kwargs.get("dimensions", None)) else None ) parsed_encoding_format = ( kwargs.get("encoding_format", "float") - if not isinstance(kwargs.get("encoding_format", "float"), NotGiven) + if not _is_not_given(kwargs.get("encoding_format", "float")) else "float" ) @@ -1077,7 +1147,7 @@ def _instrument_openai_stream( raw_iterator = response._iterator completion_start_time: Optional[datetime] = None is_finalized = False - close = response.close + close = response.response.close def finalize_once() -> None: nonlocal is_finalized @@ -1115,7 +1185,7 @@ def traced_close() -> Any: finalize_once() response._iterator = traced_iterator() - response.close = traced_close + response.response.close = traced_close return response @@ -1139,7 +1209,7 @@ def _instrument_openai_async_stream( raw_iterator = response._iterator completion_start_time: Optional[datetime] = None is_finalized = False - close = response.close + close = response.response.aclose async def finalize_once() -> None: nonlocal is_finalized @@ -1180,7 +1250,7 @@ async def traced_aclose() -> Any: return await traced_close() response._iterator = traced_iterator() - response.close = traced_close + response.response.aclose = traced_close response.aclose = traced_aclose return response @@ -1195,7 +1265,7 @@ def _get_raw_response_mode(kwargs: Any) -> Optional[str]: """ extra_headers = kwargs.get("extra_headers", None) - if extra_headers is None or isinstance(extra_headers, NotGiven): + if extra_headers is None or _is_not_given(extra_headers): return None try: @@ -1244,6 +1314,36 @@ def _unwrap_raw_response(openai_response: Any) -> Any: return openai_response +@_langfuse_wrapper +def _wrap_stream( + open_ai_resource: OpenAiDefinition, wrapped: Any, args: Any, kwargs: Any +) -> Any: + """Forward Langfuse arguments to the create call owned by the stream manager.""" + arg_names: tuple[str, ...] = _LANGFUSE_STREAM_ARG_NAMES + if open_ai_resource.module == "openai.resources.beta.chat.completions": + arg_names += ("metadata",) + + langfuse_args = {name: kwargs[name] for name in arg_names if name in kwargs} + + if not langfuse_args: + return wrapped(**kwargs) + + if open_ai_resource.module == "openai.resources.responses" and any( + name in kwargs and not _is_not_given(kwargs[name]) + for name in ("response_id", "starting_after") + ): + return wrapped(**kwargs) + + for name in langfuse_args: + kwargs.pop(name) + + extra_body = dict(kwargs.get("extra_body") or {}) + extra_body[_LANGFUSE_OPENAI_STREAM_ARGS] = langfuse_args + kwargs["extra_body"] = extra_body + + return wrapped(**kwargs) + + @_langfuse_wrapper def _wrap( open_ai_resource: OpenAiDefinition, wrapped: Any, args: Any, kwargs: Any @@ -1436,10 +1536,18 @@ def register_tracing() -> None: ): continue + wrapper = ( + _wrap_stream(resource) + if resource.method == "stream" + else _wrap(resource) + if resource.sync + else _wrap_async(resource) + ) + wrap_function_wrapper( resource.module, f"{resource.object}.{resource.method}", - _wrap(resource) if resource.sync else _wrap_async(resource), + wrapper, ) diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index 681be4fbf..9f275b783 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -1,4 +1,5 @@ import asyncio +import json from types import SimpleNamespace from unittest.mock import patch @@ -8,9 +9,15 @@ import langfuse.openai as lf_openai_module from langfuse._client.attributes import LangfuseOtelSpanAttributes +from langfuse.api import Prompt_Text +from langfuse.model import TextPromptClient from langfuse.openai import openai as lf_openai +class StreamAnswer(BaseModel): + answer: int + + class DummySyncResponse: def __init__(self) -> None: self.closed = False @@ -1261,24 +1268,92 @@ def _chat_completion_payload(): } -def _chat_completion_chunk_sse_body(): +def _chat_completion_chunk_sse_body(content: str = "2"): + chunks = [ + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "sample-model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": content}, + "finish_reason": None, + } + ], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "sample-model", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "sample-model", + "choices": [], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + }, + }, + ] + + return "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + ( + "data: [DONE]\n\n" + ) + + +def _response_stream_sse_body(): return ( - 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",' - '"created":1700000000,"model":"gpt-4o-mini-2024-07-18",' - '"choices":[{"index":0,"delta":{"role":"assistant","content":"2"},' - '"finish_reason":null}]}\n\n' + 'data: {"type":"response.created","sequence_number":0,"response":' + '{"id":"resp-test","object":"response","created_at":1700000000,' + '"model":"sample-model","output":[],"parallel_tool_calls":true,' + '"tool_choice":"auto","tools":[],"status":"in_progress"}}\n\n' + 'data: {"type":"response.completed","sequence_number":1,"response":' + '{"id":"resp-test","object":"response","created_at":1700000000,' + '"model":"sample-model","output":[{"id":"msg-test","type":"message",' + '"status":"completed","role":"assistant","content":[{"type":"output_text",' + '"text":"Hello","annotations":[]}]}],"parallel_tool_calls":true,' + '"tool_choice":"auto","tools":[],"status":"completed","usage":' + '{"input_tokens":3,"input_tokens_details":{"cached_tokens":0},' + '"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},' + '"total_tokens":4}}}\n\n' "data: [DONE]\n\n" ) -def _mock_transport_openai_client(async_client: bool = False): +def _mock_transport_openai_client( + async_client: bool = False, request_bodies=None, chat_stream_content: str = "2" +): import httpx def handler(request: httpx.Request) -> httpx.Response: + if request_bodies is not None: + request_bodies.append(json.loads(request.content)) + + if request.url.path.endswith("/responses"): + return httpx.Response( + 200, + content=_response_stream_sse_body().encode(), + headers={"content-type": "text/event-stream"}, + ) + if b'"stream": true' in request.content or b'"stream":true' in request.content: return httpx.Response( 200, - content=_chat_completion_chunk_sse_body().encode(), + content=_chat_completion_chunk_sse_body(chat_stream_content).encode(), headers={"content-type": "text/event-stream"}, ) @@ -1296,6 +1371,388 @@ def handler(request: httpx.Request) -> httpx.Response: ) +def _stream_manager(openai_client, resource, name): + if resource == "chat": + return openai_client.chat.completions.stream( + name=name, + model="sample-model", + messages=[{"role": "user", "content": "Hello"}], + ) + + return openai_client.responses.stream( + name=name, + model="sample-model", + input="Hello", + ) + + +def test_openai_stream_helpers_accept_langfuse_arguments( + langfuse_memory_client, find_spans, get_span, json_attr +): + request_bodies = [] + openai_client = _mock_transport_openai_client( + request_bodies=request_bodies, + chat_stream_content='{"answer":2}', + ) + prompt = TextPromptClient( + Prompt_Text( + name="stream-helper-prompt", + version=3, + prompt="Hello", + type="text", + labels=[], + config={}, + tags=[], + ) + ) + trace_id = "1" * 32 + parent_observation_id = "2" * 16 + extra_body = {"custom_field": "kept"} + + with patch.object( + lf_openai_module, "get_client", return_value=langfuse_memory_client + ) as get_client_mock: + with openai_client.chat.completions.stream( + name="unit-openai-chat-stream-helper", + langfuse_prompt=prompt, + langfuse_public_key="stream-public-key", + trace_id=trace_id, + parent_observation_id=parent_observation_id, + metadata={"stream": "helper"}, + extra_body=extra_body, + model="sample-model", + messages=[{"role": "user", "content": "Hello"}], + response_format=StreamAnswer, + stream_options={"include_usage": True}, + ) as stream: + completion = stream.get_final_completion() + + get_client_mock.assert_called_once_with(public_key="stream-public-key") + + assert completion.model == "sample-model" + assert completion.choices[0].message.content == '{"answer":2}' + assert completion.choices[0].message.parsed == StreamAnswer(answer=2) + assert extra_body == {"custom_field": "kept"} + assert request_bodies[0]["custom_field"] == "kept" + assert all( + name not in request_bodies[0] + for name in ( + "name", + "langfuse_prompt", + "langfuse_public_key", + "trace_id", + "parent_observation_id", + ) + ) + + with openai_client.responses.stream( + name="unit-openai-responses-stream-helper", + model="sample-model", + input="Hello", + ) as stream: + response = stream.get_final_response() + + assert response.model == "sample-model" + assert response.status == "completed" + assert response.output_text == "Hello" + + langfuse_memory_client.flush() + assert len(find_spans("unit-openai-chat-stream-helper")) == 1 + assert len(find_spans("unit-openai-responses-stream-helper")) == 1 + assert find_spans("OpenAI-generation") == [] + + chat_span = get_span("unit-openai-chat-stream-helper") + assert ( + chat_span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME] + == "stream-helper-prompt" + ) + assert ( + chat_span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION] == 3 + ) + assert chat_span.attributes["langfuse.observation.metadata.stream"] == "helper" + assert json_attr(chat_span, LangfuseOtelSpanAttributes.OBSERVATION_INPUT) == [ + {"role": "user", "content": "Hello"} + ] + assert json_attr( + chat_span, LangfuseOtelSpanAttributes.OBSERVATION_MODEL_PARAMETERS + ) == { + "temperature": 1, + "max_tokens": "Infinity", + "top_p": 1, + "frequency_penalty": 0, + "presence_penalty": 0, + } + assert ( + chat_span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] + == '{"answer":2}' + ) + assert ( + chat_span.attributes[ + LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME + ] + is not None + ) + assert json_attr( + chat_span, LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS + ) == { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + "prompt_tokens_details": None, + "completion_tokens_details": None, + } + assert format(chat_span.context.trace_id, "032x") == trace_id + assert chat_span.parent is not None + assert format(chat_span.parent.span_id, "016x") == parent_observation_id + + responses_span = get_span("unit-openai-responses-stream-helper") + assert ( + responses_span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_INPUT] + == "Hello" + ) + assert json_attr( + responses_span, LangfuseOtelSpanAttributes.OBSERVATION_MODEL_PARAMETERS + ) == { + "temperature": 1, + "max_tokens": "Infinity", + "top_p": 1, + "frequency_penalty": 0, + "presence_penalty": 0, + } + assert json_attr(responses_span, LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT) == { + "id": "msg-test", + "content": [ + { + "annotations": [], + "text": "Hello", + "type": "output_text", + "logprobs": None, + } + ], + "role": "assistant", + "status": "completed", + "type": "message", + "phase": None, + } + assert json_attr( + responses_span, LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS + ) == { + "input_tokens": 3, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 4, + } + + +@pytest.mark.asyncio +async def test_async_openai_stream_helpers_accept_langfuse_arguments( + langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client(async_client=True) + + async with openai_client.chat.completions.stream( + name="unit-openai-async-chat-stream-helper", + model="sample-model", + messages=[{"role": "user", "content": "Hello"}], + ) as stream: + completion = await stream.get_final_completion() + + assert completion.model == "sample-model" + assert completion.choices[0].message.content == "2" + + async with openai_client.responses.stream( + name="unit-openai-async-responses-stream-helper", + model="sample-model", + input="Hello", + ) as stream: + response = await stream.get_final_response() + + assert response.model == "sample-model" + assert response.status == "completed" + + langfuse_memory_client.flush() + assert len(find_spans("unit-openai-async-chat-stream-helper")) == 1 + assert len(find_spans("unit-openai-async-responses-stream-helper")) == 1 + assert find_spans("OpenAI-generation") == [] + + +@pytest.mark.parametrize("omitted_argument", ["response_id", "starting_after"]) +@pytest.mark.asyncio +async def test_async_responses_stream_treats_openai_omit_as_absent( + omitted_argument, langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client(async_client=True) + name = f"unit-openai-responses-stream-omit-{omitted_argument}" + + async with openai_client.responses.stream( + name=name, + model="sample-model", + input="Hello", + **{omitted_argument: lf_openai.omit}, + ) as stream: + response = await stream.get_final_response() + + assert response.status == "completed" + + langfuse_memory_client.flush() + assert len(find_spans(name)) == 1 + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize( + "consume_one_event", [False, True], ids=["unconsumed", "partial"] +) +def test_stream_helpers_finalize_on_context_exit( + resource, consume_one_event, langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client() + name = f"unit-openai-{resource}-stream-context-exit-{consume_one_event}" + manager = _stream_manager(openai_client, resource, name) + + with manager as stream: + if consume_one_event: + next(stream) + + langfuse_memory_client.flush() + assert len(find_spans(name)) == 1 + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize( + "consume_one_event", [False, True], ids=["unconsumed", "partial"] +) +@pytest.mark.asyncio +async def test_async_stream_helpers_finalize_on_context_exit( + resource, consume_one_event, langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client(async_client=True) + name = f"unit-openai-async-{resource}-stream-context-exit-{consume_one_event}" + manager = _stream_manager(openai_client, resource, name) + + async with manager as stream: + if consume_one_event: + await stream.__anext__() + + langfuse_memory_client.flush() + assert len(find_spans(name)) == 1 + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +def test_stream_manager_reuse_preserves_langfuse_arguments( + resource, langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client() + name = f"unit-openai-reused-{resource}-stream-manager" + manager = _stream_manager(openai_client, resource, name) + + with manager as stream: + list(stream) + + with manager as stream: + list(stream) + + langfuse_memory_client.flush() + assert len(find_spans(name)) == 2 + + +@pytest.mark.parametrize( + ("openai_version", "expected_registrations"), + [ + ("1.39.9", set()), + ( + "1.40.0", + { + ( + "openai.resources.beta.chat.completions", + "Completions.stream", + ), + ( + "openai.resources.beta.chat.completions", + "AsyncCompletions.stream", + ), + }, + ), + ( + "1.92.0", + { + ("openai.resources.chat.completions", "Completions.stream"), + ("openai.resources.chat.completions", "AsyncCompletions.stream"), + }, + ), + ], +) +def test_chat_stream_registration_version_boundaries( + monkeypatch, openai_version, expected_registrations +): + registrations = set() + + def record_registration(module, name, _wrapper): + if module in { + "openai.resources.beta.chat.completions", + "openai.resources.chat.completions", + } and name.endswith(".stream"): + registrations.add((module, name)) + + monkeypatch.setattr(lf_openai_module.openai, "__version__", openai_version) + monkeypatch.setattr(lf_openai_module, "wrap_function_wrapper", record_registration) + + lf_openai_module.register_tracing() + + assert registrations == expected_registrations + + +def test_legacy_beta_stream_forwards_langfuse_metadata(): + captured_args = {} + + def legacy_stream(*, model, messages, extra_body): + extractor = lf_openai_module.OpenAiArgsExtractor( + model=model, + messages=messages, + extra_body=extra_body, + ) + captured_args.update(extractor.get_langfuse_args()) + return "stream-manager" + + resource = lf_openai_module.OpenAiDefinition( + module="openai.resources.beta.chat.completions", + object="Completions", + method="stream", + type="chat", + sync=True, + ) + wrapper = lf_openai_module._wrap_stream(resource) + + result = wrapper( + legacy_stream, + None, + (), + { + "metadata": {"suite": "legacy-beta"}, + "model": "sample-model", + "messages": [{"role": "user", "content": "Hello"}], + }, + ) + + assert result == "stream-manager" + assert captured_args["metadata"] == {"suite": "legacy-beta"} + + +def test_responses_stream_retrieval_rejects_langfuse_arguments( + langfuse_memory_client, find_spans +): + openai_client = _mock_transport_openai_client() + + with pytest.raises(TypeError): + openai_client.responses.stream( + response_id="resp-test", + name="unit-openai-responses-stream-retrieval", + ) + + langfuse_memory_client.flush() + assert find_spans("unit-openai-responses-stream-retrieval") == [] + + def test_with_raw_response_chat_completion_captures_output_and_usage( langfuse_memory_client, get_span, json_attr ): diff --git a/tests/unit/test_openai_stream_compatibility.py b/tests/unit/test_openai_stream_compatibility.py new file mode 100644 index 000000000..dae8ff519 --- /dev/null +++ b/tests/unit/test_openai_stream_compatibility.py @@ -0,0 +1,96 @@ +import httpx +import pytest +from packaging.version import Version + +from langfuse.openai import openai as lf_openai + +_OPENAI_VERSION = Version(lf_openai.__version__) +_STREAM_RESOURCES = [ + "chat", + pytest.param( + "responses", + marks=pytest.mark.skipif( + _OPENAI_VERSION < Version("1.66.0"), + reason="Responses stream helpers require OpenAI 1.66 or newer", + ), + ), +] + + +def _stream_response(_request): + return httpx.Response( + 200, + content=b"data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +def _chat_completions(client): + if _OPENAI_VERSION < Version("1.92.0"): + return client.beta.chat.completions + + return client.chat.completions + + +def _stream_request(client, resource): + if resource == "chat": + return _chat_completions(client), { + "messages": [{"role": "user", "content": "Hello"}] + } + + return client.responses, {"input": "Hello"} + + +@pytest.mark.parametrize("resource", _STREAM_RESOURCES) +def test_stream_helper_compatibility(resource, langfuse_memory_client, get_span): + client = lf_openai.OpenAI( + api_key="test", + http_client=httpx.Client(transport=httpx.MockTransport(_stream_response)), + ) + endpoint, request = _stream_request(client, resource) + name = f"unit-openai-{resource}-stream-compatibility" + + with endpoint.stream( + name=name, + metadata={"compatibility": lf_openai.__version__}, + model="sample-model", + **request, + ): + pass + + client.close() + langfuse_memory_client.flush() + span = get_span(name) + assert ( + span.attributes["langfuse.observation.metadata.compatibility"] + == lf_openai.__version__ + ) + + +@pytest.mark.parametrize("resource", _STREAM_RESOURCES) +@pytest.mark.asyncio +async def test_async_stream_helper_compatibility( + resource, langfuse_memory_client, get_span +): + client = lf_openai.AsyncOpenAI( + api_key="test", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_stream_response)), + ) + endpoint, request = _stream_request(client, resource) + name = f"unit-openai-async-{resource}-stream-compatibility" + + async with endpoint.stream( + name=name, + metadata={"compatibility": lf_openai.__version__}, + model="sample-model", + **request, + ): + pass + + await client.close() + langfuse_memory_client.flush() + span = get_span(name) + assert ( + span.attributes["langfuse.observation.metadata.compatibility"] + == lf_openai.__version__ + )