feat(openai): native content-block streaming for chat completions

Author: nick-hollon-lc

libs/partners/openai/langchain_openai/__init__.py

@@ -3,6 +3,10 @@
 from langchain_openai._version import __version__
 from langchain_openai.chat_models import AzureChatOpenAI, ChatOpenAI
 from langchain_openai.chat_models._client_utils import StreamChunkTimeoutError
+from langchain_openai.chat_models._stream_events import (
+    aconvert_openai_completions_stream,
+    convert_openai_completions_stream,
+)
 from langchain_openai.embeddings import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from langchain_openai.llms import AzureOpenAI, OpenAI
 from langchain_openai.tools import custom_tool
@@ -16,5 +20,7 @@
     "OpenAIEmbeddings",
     "StreamChunkTimeoutError",
     "__version__",
+    "aconvert_openai_completions_stream",
+    "convert_openai_completions_stream",
     "custom_tool",
 ]

libs/partners/openai/langchain_openai/chat_models/_stream_events.py

@@ -0,0 +1,162 @@
+"""Native content-block streaming-event converter for OpenAI Chat Completions.
+
+Drives raw OpenAI Chat Completions chunks into the shared `BlockStreamTracker`,
+reusing `BaseChatOpenAI._convert_chunk_to_generation_chunk` for content
+extraction (it already yields indexed content blocks + tool-call chunks). This
+converter is the reuse seam for OpenAI-compatible providers (deepseek, groq,
+fireworks, xai, openrouter), which adapt their chunk shape to OpenAI's and call
+it with a different `provider`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from langchain_core.language_models.stream_events import (
+    BlockStreamTracker,
+    accumulate_usage,
+    build_message_finish,
+    iter_protocol_blocks,
+)
+from langchain_core.messages import AIMessageChunk
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterator
+
+    from langchain_core.outputs import ChatGenerationChunk
+    from langchain_protocol.protocol import (
+        MessageMetadata,
+        MessagesData,
+        MessageStartData,
+    )
+
+# Bound `BaseChatOpenAI._convert_chunk_to_generation_chunk`.
+MakeChunk = Callable[..., "ChatGenerationChunk | None"]
+
+
+def _message_start(
+    message_id: str | None, model: str | None, provider: str
+) -> MessageStartData:
+    metadata: MessageMetadata = {"provider": provider}
+    if model:
+        metadata["model"] = model
+    return {
+        "event": "message-start",
+        "role": "ai",
+        "id": message_id or "",
+        "metadata": metadata,
+    }
+
+
+def convert_openai_completions_stream(
+    raw: Iterator[Any],
+    make_chunk: MakeChunk,
+    *,
+    base_generation_info: dict[str, Any] | None = None,
+    message_id: str | None = None,
+    provider: str = "openai",
+) -> Iterator[MessagesData]:
+    """Convert a raw OpenAI Chat Completions chunk stream to protocol events.
+
+    Args:
+        raw: Raw OpenAI chunks (dicts or SDK objects with `model_dump`).
+        make_chunk: `BaseChatOpenAI._convert_chunk_to_generation_chunk`, injected
+            so the converter stays pure and unit-testable.
+        base_generation_info: Passed to `make_chunk` for the first chunk only
+            (mirrors `_stream`), `{}` thereafter.
+        message_id: Message id for `message-start`. Left empty by default so
+            the v3 stream's seeded LangChain run id stands (matching the compat
+            bridge); the provider completion id is deliberately not used here.
+        provider: `model_provider` id for downstream reuse (groq, deepseek, ...).
+
+    Yields:
+        Protocol `MessagesData` lifecycle events.
+    """
+    tracker = BlockStreamTracker()
+    started = False
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": provider}
+    model: str | None = None
+    first = True
+
+    for chunk in raw:
+        if not isinstance(chunk, dict):
+            chunk = chunk.model_dump()
+        if model is None and chunk.get("model"):
+            model = chunk["model"]
+        gen = make_chunk(chunk, AIMessageChunk, base_generation_info if first else {})
+        first = False
+        if gen is None:
+            continue
+        msg = gen.message
+        if not started:
+            started = True
+            yield _message_start(message_id, model, provider)
+        for key, block in iter_protocol_blocks(msg):
+            yield from tracker.feed(key, block)
+        usage_metadata = getattr(msg, "usage_metadata", None)
+        if usage_metadata:
+            usage = accumulate_usage(usage, usage_metadata)
+        merged = {**(gen.generation_info or {}), **(msg.response_metadata or {})}
+        if merged:
+            response_metadata.update(merged)
+            # `_convert_chunk_to_generation_chunk` hardcodes
+            # `model_provider="openai"`; re-apply the caller's `provider` so
+            # OpenAI-compatible reuse (groq, deepseek, ...) isn't mislabeled.
+            response_metadata["model_provider"] = provider
+
+    if not started:
+        return
+    yield from tracker.finish_all()
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)
+
+
+async def aconvert_openai_completions_stream(
+    raw: AsyncIterator[Any],
+    make_chunk: MakeChunk,
+    *,
+    base_generation_info: dict[str, Any] | None = None,
+    message_id: str | None = None,
+    provider: str = "openai",
+) -> AsyncIterator[MessagesData]:
+    """Async twin of `convert_openai_completions_stream`. `make_chunk` is sync."""
+    tracker = BlockStreamTracker()
+    started = False
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": provider}
+    model: str | None = None
+    first = True
+
+    async for chunk in raw:
+        if not isinstance(chunk, dict):
+            chunk = chunk.model_dump()
+        if model is None and chunk.get("model"):
+            model = chunk["model"]
+        gen = make_chunk(chunk, AIMessageChunk, base_generation_info if first else {})
+        first = False
+        if gen is None:
+            continue
+        msg = gen.message
+        if not started:
+            started = True
+            yield _message_start(message_id, model, provider)
+        for key, block in iter_protocol_blocks(msg):
+            for ev in tracker.feed(key, block):
+                yield ev
+        usage_metadata = getattr(msg, "usage_metadata", None)
+        if usage_metadata:
+            usage = accumulate_usage(usage, usage_metadata)
+        merged = {**(gen.generation_info or {}), **(msg.response_metadata or {})}
+        if merged:
+            response_metadata.update(merged)
+            # `_convert_chunk_to_generation_chunk` hardcodes
+            # `model_provider="openai"`; re-apply the caller's `provider` so
+            # OpenAI-compatible reuse (groq, deepseek, ...) isn't mislabeled.
+            response_metadata["model_provider"] = provider
+
+    if not started:
+        return
+    for ev in tracker.finish_all():
+        yield ev
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)

libs/partners/openai/langchain_openai/chat_models/base.py

@@ -56,6 +56,10 @@
     LanguageModelInput,
     ModelProfileRegistry,
 )
+from langchain_core.language_models._compat_bridge import (
+    achunks_to_events,
+    chunks_to_events,
+)
 from langchain_core.language_models.chat_models import (
     BaseChatModel,
     LangSmithParams,
@@ -149,11 +153,16 @@
     _convert_from_v1_to_responses,
     _convert_to_v03_ai_message,
 )
+from langchain_openai.chat_models._stream_events import (
+    aconvert_openai_completions_stream,
+    convert_openai_completions_stream,
+)
 from langchain_openai.data._profiles import _PROFILES
 
 if TYPE_CHECKING:
     import httpx
     from langchain_core.language_models import ModelProfile
+    from langchain_protocol.protocol import MessagesData
     from openai.types.responses import Response
 
 logger = logging.getLogger(__name__)
@@ -1911,6 +1920,147 @@ async def _astream(
                 )
             yield generation_chunk
 
+    def _stream_chat_model_events(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        *,
+        message_id: str | None = None,
+        **kwargs: Any,
+    ) -> Iterator[MessagesData]:
+        """Emit OpenAI-native content-block events for the Chat Completions path.
+
+        Defers to the compat bridge for cases this converter does not yet
+        specialize: the Responses API, structured output (`response_format`),
+        and raw-header mode. Detected by core's `_iter_v2_events`.
+        """
+        # Responses API / structured output / raw headers: bridge over `_stream`,
+        # which (on `ChatOpenAI`) routes to the Responses path when applicable.
+        # `response_format` may arrive via call kwargs or be baked into
+        # `model_kwargs`; both fold into the payload, so check both.
+        if (
+            self._use_responses_api({**kwargs, **self.model_kwargs})
+            or kwargs.get("response_format") is not None
+            or self.model_kwargs.get("response_format") is not None
+            or self.include_response_headers
+        ):
+            # Forward kwargs untouched (as core's `_iter_v2_events` would):
+            # `_stream` handles `stream_usage` itself, and the Responses path
+            # rejects a stray `stream_usage` kwarg, so we must not inject one.
+            yield from chunks_to_events(
+                self._stream(
+                    messages,
+                    stop=stop,
+                    run_manager=run_manager,
+                    **kwargs,
+                ),
+                message_id=message_id,
+            )
+            return
+
+        self._ensure_sync_client_available()
+        kwargs["stream"] = True
+        stream_usage = self._should_stream_usage(
+            kwargs.pop("stream_usage", None), **kwargs
+        )
+        if stream_usage:
+            kwargs["stream_options"] = {"include_usage": stream_usage}
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        try:
+            with self.client.create(**payload) as response:
+                for event in convert_openai_completions_stream(
+                    response,
+                    self._convert_chunk_to_generation_chunk,
+                    message_id=message_id,
+                ):
+                    if (
+                        run_manager is not None
+                        and event["event"] == "content-block-delta"
+                        and event["delta"].get("type") == "text-delta"
+                    ):
+                        # Text-only by design on the v3 events path: the events
+                        # themselves carry block/usage detail, so the legacy
+                        # `chunk=`/`logprobs=` callback args are not threaded.
+                        run_manager.on_llm_new_token(
+                            str(event["delta"].get("text", ""))
+                        )
+                    yield event
+        except openai.BadRequestError as e:
+            _handle_openai_bad_request(e)
+        except openai.APIError as e:
+            _handle_openai_api_error(e)
+
+    async def _astream_chat_model_events(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        *,
+        message_id: str | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[MessagesData]:
+        """Async twin of `_stream_chat_model_events`."""
+        if (
+            self._use_responses_api({**kwargs, **self.model_kwargs})
+            or kwargs.get("response_format") is not None
+            or self.model_kwargs.get("response_format") is not None
+            or self.include_response_headers
+        ):
+            # Forward kwargs untouched (as core's `_aiter_v2_events` would):
+            # `_astream` handles `stream_usage` itself, and the Responses path
+            # rejects a stray `stream_usage` kwarg, so we must not inject one.
+            async for event in achunks_to_events(
+                self._astream(
+                    messages,
+                    stop=stop,
+                    run_manager=run_manager,
+                    **kwargs,
+                ),
+                message_id=message_id,
+            ):
+                yield event
+            return
+
+        kwargs["stream"] = True
+        stream_usage = self._should_stream_usage(
+            kwargs.pop("stream_usage", None), **kwargs
+        )
+        if stream_usage:
+            kwargs["stream_options"] = {"include_usage": stream_usage}
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        try:
+            response = await self.async_client.create(**payload)
+            async with response as stream:
+                # Mirror `_astream`: apply per-chunk stall protection before the
+                # converter consumes the stream.
+                timed_stream = _astream_with_chunk_timeout(
+                    stream,
+                    self.stream_chunk_timeout,
+                    model_name=self.model_name,
+                )
+                async for event in aconvert_openai_completions_stream(
+                    timed_stream,
+                    self._convert_chunk_to_generation_chunk,
+                    message_id=message_id,
+                ):
+                    if (
+                        run_manager is not None
+                        and event["event"] == "content-block-delta"
+                        and event["delta"].get("type") == "text-delta"
+                    ):
+                        # Text-only by design on the v3 events path: the events
+                        # themselves carry block/usage detail, so the legacy
+                        # `chunk=`/`logprobs=` callback args are not threaded.
+                        await run_manager.on_llm_new_token(
+                            str(event["delta"].get("text", ""))
+                        )
+                    yield event
+        except openai.BadRequestError as e:
+            _handle_openai_bad_request(e)
+        except openai.APIError as e:
+            _handle_openai_api_error(e)
+
     async def _agenerate(
         self,
         messages: list[BaseMessage],

libs/partners/openai/tests/unit_tests/chat_models/test_base.py

@@ -635,13 +635,38 @@ def mock_create(*args: Any, **kwargs: Any) -> MockSyncContextManager:
         events = list(llm.stream_events("你的名字叫什么？只回答名字", version="v3"))
 
     assert_valid_event_stream(events)
+    # `message-start` carries the stream's LangChain run id (threaded from core),
+    # not the provider completion id and not an empty string.
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"]
+    assert not events[0]["id"].startswith("chatcmpl")
     # At minimum, a text block with the accumulated answer.
     finishes = [e for e in events if e["event"] == "content-block-finish"]
     assert len(finishes) >= 1
     text_finishes = [f for f in finishes if f["content"]["type"] == "text"]
     assert len(text_finishes) == 1
 
 
+async def test_openai_astream_events_v3_lifecycle(mock_openai_completion: list) -> None:
+    """Async twin of `test_openai_stream_events_v3_lifecycle`."""
+    from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
+
+    llm = ChatOpenAI(model="gpt-4o", api_key=SecretStr("test"))
+    mock_client = MagicMock()
+
+    async def mock_acreate(*args: Any, **kwargs: Any) -> MockAsyncContextManager:
+        return MockAsyncContextManager(mock_openai_completion)
+
+    mock_client.create = mock_acreate
+    with patch.object(llm, "async_client", mock_client):
+        stream = await llm.astream_events("test", version="v3")
+        events = [e async for e in stream]
+
+    assert_valid_event_stream(events)
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    assert len([f for f in finishes if f["content"]["type"] == "text"]) == 1
+
+
 @pytest.fixture
 def mock_completion() -> dict:
     return {