fix(openai): support Langfuse arguments in stream helpers

langfuse/openai.py

@@ -145,6 +145,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",
@@ -161,6 +179,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",
@@ -193,6 +227,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",
@@ -212,6 +262,14 @@ 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:
@@ -329,6 +387,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:
@@ -1080,7 +1152,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
@@ -1118,7 +1190,7 @@ def traced_close() -> Any:
             finalize_once()
 
     response._iterator = traced_iterator()
-    response.close = traced_close
+    response.response.close = traced_close
 
     return response
 
@@ -1142,7 +1214,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
@@ -1183,7 +1255,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
@@ -1247,6 +1319,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
@@ -1439,10 +1541,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,
         )