feat(mistralai): native content-block streaming events

libs/partners/mistralai/langchain_mistralai/_stream_events.py

@@ -0,0 +1,153 @@
+"""Native content-block streaming-event converter for MistralAI.
+
+Mirrors `ChatMistralAI._stream`: threads `(index, index_type, default_class)`
+through `_convert_chunk_to_message_chunk` (injected to avoid a circular import)
+and feeds each resulting `AIMessageChunk`'s content blocks into the shared
+`BlockStreamTracker`.
+"""
+
+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.messages import BaseMessageChunk
+    from langchain_protocol.protocol import (
+        MessageMetadata,
+        MessagesData,
+        MessageStartData,
+    )
+
+# The module-level `_convert_chunk_to_message_chunk`, injected so the converter
+# stays pure and avoids a circular import. It takes a raw chunk plus the running
+# `default_class`, `index`, and `index_type`, returning the built chunk and the
+# updated index/index_type.
+ConvertChunk = Callable[..., "tuple[BaseMessageChunk, int, str]"]
+
+
+def _message_start(message_id: str | None, model: str | None) -> MessageStartData:
+    # Do not use the provider chunk id here: on the v3 path core seeds the
+    # stream with the LangChain run id, and an empty id lets that stand
+    # (matching the compat bridge). Only an explicit `message_id` wins.
+    metadata: MessageMetadata = {"provider": "mistralai"}
+    if model:
+        metadata["model"] = model
+    return {
+        "event": "message-start",
+        "role": "ai",
+        "id": message_id or "",
+        "metadata": metadata,
+    }
+
+
+def convert_mistral_stream(
+    raw: Iterator[Any],
+    convert_chunk: ConvertChunk,
+    *,
+    output_version: str | None = None,
+    message_id: str | None = None,
+) -> Iterator[MessagesData]:
+    """Convert a raw Mistral chat stream to protocol events.
+
+    Args:
+        raw: Raw Mistral chat-completion chunks (OpenAI-shaped dicts).
+        convert_chunk: `_convert_chunk_to_message_chunk`, injected so the
+            converter stays pure and avoids a circular import.
+        output_version: Forwarded to `convert_chunk`; reasoning blocks only
+            surface under `"v1"`.
+        message_id: Overrides the id on `message-start`.
+
+    Yields:
+        Protocol `MessagesData` lifecycle events.
+    """
+    tracker = BlockStreamTracker()
+    started = False
+    index = -1
+    index_type = ""
+    default_class: type[BaseMessageChunk] = AIMessageChunk
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": "mistralai"}
+    model: str | None = None
+
+    for chunk in raw:
+        if len(chunk.get("choices", [])) == 0:
+            continue
+        if model is None:
+            model = chunk.get("model")
+        new_chunk, index, index_type = convert_chunk(
+            chunk, default_class, index, index_type, output_version
+        )
+        # Make future chunks the same type as the first chunk.
+        default_class = new_chunk.__class__
+        if not started:
+            started = True
+            yield _message_start(message_id, model)
+        if isinstance(new_chunk, AIMessageChunk):
+            for key, block in iter_protocol_blocks(new_chunk):
+                yield from tracker.feed(key, block)
+            if new_chunk.usage_metadata:
+                usage = accumulate_usage(usage, new_chunk.usage_metadata)
+            if new_chunk.response_metadata:
+                response_metadata.update(new_chunk.response_metadata)
+
+    if not started:
+        return
+    yield from tracker.finish_all()
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)
+
+
+async def aconvert_mistral_stream(
+    raw: AsyncIterator[Any],
+    convert_chunk: ConvertChunk,
+    *,
+    output_version: str | None = None,
+    message_id: str | None = None,
+) -> AsyncIterator[MessagesData]:
+    """Async twin of `convert_mistral_stream`. `convert_chunk` is sync."""
+    tracker = BlockStreamTracker()
+    started = False
+    index = -1
+    index_type = ""
+    default_class: type[BaseMessageChunk] = AIMessageChunk
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": "mistralai"}
+    model: str | None = None
+
+    async for chunk in raw:
+        if len(chunk.get("choices", [])) == 0:
+            continue
+        if model is None:
+            model = chunk.get("model")
+        new_chunk, index, index_type = convert_chunk(
+            chunk, default_class, index, index_type, output_version
+        )
+        # Make future chunks the same type as the first chunk.
+        default_class = new_chunk.__class__
+        if not started:
+            started = True
+            yield _message_start(message_id, model)
+        if isinstance(new_chunk, AIMessageChunk):
+            for key, block in iter_protocol_blocks(new_chunk):
+                for ev in tracker.feed(key, block):
+                    yield ev
+            if new_chunk.usage_metadata:
+                usage = accumulate_usage(usage, new_chunk.usage_metadata)
+            if new_chunk.response_metadata:
+                response_metadata.update(new_chunk.response_metadata)
+
+    if not started:
+        return
+    for ev in tracker.finish_all():
+        yield ev
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)

libs/partners/mistralai/langchain_mistralai/chat_models.py

@@ -85,6 +85,8 @@ if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterator
     from contextlib import AbstractAsyncContextManager
 
+    from langchain_protocol.protocol import MessagesData
+
 logger = logging.getLogger(__name__)
 
 # Mistral enforces a specific pattern for tool call IDs
@@ -830,6 +832,78 @@ class ChatMistralAI(BaseChatModel):
                 )
             yield gen_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 Mistral-native content-block protocol events.
+
+        Detected by `langchain-core`'s `_iter_v2_events`; powers
+        `stream_events(version="v3")`. Falls through to the compat bridge
+        only if this method is absent. `message_id` is threaded from the
+        stream so `message-start` matches the bridge's LangChain run id.
+        """
+        # Local import avoids a circular import: `_stream_events` imports
+        # `_convert_chunk_to_message_chunk` from this module.
+        from langchain_mistralai._stream_events import convert_mistral_stream
+
+        message_dicts, params = self._create_message_dicts(messages, stop)
+        params = {**params, **kwargs, "stream": True}
+        raw = self.completion_with_retry(
+            messages=message_dicts, run_manager=run_manager, **params
+        )
+        for event in convert_mistral_stream(
+            raw,
+            _convert_chunk_to_message_chunk,
+            output_version=self.output_version,
+            message_id=message_id,
+        ):
+            if (
+                run_manager is not None
+                and event["event"] == "content-block-delta"
+                and event["delta"].get("type") == "text-delta"
+            ):
+                run_manager.on_llm_new_token(str(event["delta"].get("text", "")))
+            yield event
+
+    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`."""
+        # Local import avoids a circular import: `_stream_events` imports
+        # `_convert_chunk_to_message_chunk` from this module.
+        from langchain_mistralai._stream_events import aconvert_mistral_stream
+
+        message_dicts, params = self._create_message_dicts(messages, stop)
+        params = {**params, **kwargs, "stream": True}
+        raw = await acompletion_with_retry(
+            self, messages=message_dicts, run_manager=run_manager, **params
+        )
+        async for event in aconvert_mistral_stream(
+            raw,
+            _convert_chunk_to_message_chunk,
+            output_version=self.output_version,
+            message_id=message_id,
+        ):
+            if (
+                run_manager is not None
+                and event["event"] == "content-block-delta"
+                and event["delta"].get("type") == "text-delta"
+            ):
+                await run_manager.on_llm_new_token(str(event["delta"].get("text", "")))
+            yield event
+
     async def _agenerate(
         self,
         messages: list[BaseMessage],

libs/partners/mistralai/tests/unit_tests/test_stream_events.py

@@ -0,0 +1,297 @@
+"""Unit tests for the MistralAI native stream-events converter."""
+
+from typing import Any, cast
+from unittest.mock import patch
+
+from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
+
+from langchain_mistralai import ChatMistralAI
+from langchain_mistralai._stream_events import convert_mistral_stream
+from langchain_mistralai.chat_models import _convert_chunk_to_message_chunk
+
+
+def _text_then_tool() -> list[dict]:
+    cid, model = "cmpl-1", "mistral-large"
+    return [
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": "Hello"},
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": " world"},
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "id": "t1",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": '{"city": "Paris"}',
+                                },
+                            }
+                        ],
+                    },
+                    "finish_reason": "tool_calls",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 9,
+                "completion_tokens": 3,
+                "total_tokens": 12,
+            },
+        },
+    ]
+
+
+def _reasoning_then_text_v1() -> list[dict]:
+    """Mistral ``output_version="v1"`` chunks: a thinking block then text.
+
+    Under v1 `delta.content` is a list of typed blocks. A `thinking` block
+    carries its text in a `thinking` sub-block list; `_convert_chunk_to_message_chunk`
+    maps it to a `reasoning` content block. When the block `type` changes
+    (`thinking` -> `text`) the converter's threaded `index`/`index_type`
+    advance, splitting the stream into two distinct blocks.
+    """
+    cid, model = "cmpl-1", "magistral-medium"
+    return [
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "thinking",
+                                "thinking": [{"type": "text", "text": "Let me "}],
+                            }
+                        ],
+                    },
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "content": [
+                            {
+                                "type": "thinking",
+                                "thinking": [{"type": "text", "text": "think."}],
+                            }
+                        ]
+                    },
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "id": cid,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": [{"type": "text", "text": "Hi Bob."}]},
+                    "finish_reason": "stop",
+                }
+            ],
+        },
+    ]
+
+
+def test_convert_mistral_stream_v1_reasoning() -> None:
+    """v1 reasoning path: index/index_type threading splits thinking from text.
+
+    Guards the bespoke converter's core motivation — that the
+    `index`/`index_type` returned by `_convert_chunk_to_message_chunk` are
+    threaded back in so a type change (`thinking` -> `text`) opens a new
+    block rather than merging. The reasoning-as-blocks behavior against a
+    live model is covered by `test_reasoning_v1` in the integration tests.
+    """
+    events: list[Any] = list(
+        convert_mistral_stream(
+            iter(_reasoning_then_text_v1()),
+            _convert_chunk_to_message_chunk,
+            output_version="v1",
+        )
+    )
+    assert_valid_event_stream(events)
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    # Two distinct blocks: the thinking deltas accumulate into one reasoning
+    # block, then the type change to `text` advances index/index_type.
+    assert [f["content"]["type"] for f in finishes] == ["reasoning", "text"]
+    assert [f["index"] for f in finishes] == [0, 1]
+
+    reasoning = cast("dict[str, Any]", finishes[0]["content"])
+    text = cast("dict[str, Any]", finishes[1]["content"])
+    assert reasoning["reasoning"] == "Let me think."
+    assert text["text"] == "Hi Bob."
+
+
+def test_convert_mistral_stream_lifecycle() -> None:
+    events: list[Any] = list(
+        convert_mistral_stream(
+            iter(_text_then_tool()),
+            _convert_chunk_to_message_chunk,
+            output_version="v0",
+        )
+    )
+    assert_valid_event_stream(events)
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"] == ""
+    assert events[0]["metadata"]["provider"] == "mistralai"
+
+    text = "".join(
+        e["delta"].get("text", "")
+        for e in events
+        if e["event"] == "content-block-delta"
+        and e["delta"].get("type") == "text-delta"
+    )
+    assert text == "Hello world"
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    tool_finishes = [f for f in finishes if f["content"]["type"] == "tool_call"]
+    assert len(tool_finishes) == 1
+    tc = cast("dict[str, Any]", tool_finishes[0]["content"])
+    assert tc["name"] == "get_weather"
+    assert tc["args"] == {"city": "Paris"}
+
+    message_finish = events[-1]
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["usage"] == {
+        "input_tokens": 9,
+        "output_tokens": 3,
+        "total_tokens": 12,
+    }
+
+
+def test_mistral_stream_events_v3_lifecycle() -> None:
+    """Validate `stream_events(version="v3")` over a text + tool_call stream.
+
+    Threads a realistic chunk sequence through `_stream_chat_model_events`
+    via a mocked raw client and asserts a spec-conformant event stream.
+    """
+    llm = ChatMistralAI(api_key="test")  # type: ignore[arg-type]
+
+    with patch.object(
+        ChatMistralAI,
+        "completion_with_retry",
+        return_value=iter(_text_then_tool()),
+    ):
+        events: list[Any] = list(llm.stream_events("Test query", version="v3"))
+
+    assert_valid_event_stream(events)
+
+    # `message-start` must carry the stream's LangChain run id (threaded from
+    # core), not the empty converter default.
+    message_start = cast("dict[str, Any]", events[0])
+    assert message_start["event"] == "message-start"
+    assert message_start["id"]
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    tool_finishes = [f for f in finishes if f["content"]["type"] == "tool_call"]
+    assert len(tool_finishes) == 1
+    tc = cast("dict[str, Any]", tool_finishes[0]["content"])
+    assert tc["name"] == "get_weather"
+    assert tc["args"] == {"city": "Paris"}
+
+    message_finish = cast("dict[str, Any]", events[-1])
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["metadata"]["model_provider"] == "mistralai"
+
+
+async def test_mistral_astream_events_v3_lifecycle() -> None:
+    """Async twin of `test_mistral_stream_events_v3_lifecycle`."""
+    llm = ChatMistralAI(api_key="test")  # type: ignore[arg-type]
+
+    async def _acompletion(*args: Any, **kwargs: Any) -> Any:
+        async def _gen() -> Any:
+            for chunk in _text_then_tool():
+                yield chunk
+
+        return _gen()
+
+    with patch(
+        "langchain_mistralai.chat_models.acompletion_with_retry",
+        new=_acompletion,
+    ):
+        stream = await llm.astream_events("Test query", version="v3")
+        events: list[Any] = [e async for e in stream]
+
+    assert_valid_event_stream(events)
+    message_start = cast("dict[str, Any]", events[0])
+    assert message_start["event"] == "message-start"
+    assert message_start["id"]
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    tool_finishes = [f for f in finishes if f["content"]["type"] == "tool_call"]
+    assert len(tool_finishes) == 1
+    tc = cast("dict[str, Any]", tool_finishes[0]["content"])
+    assert tc["name"] == "get_weather"
+    assert tc["args"] == {"city": "Paris"}
+
+    message_finish = cast("dict[str, Any]", events[-1])
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["metadata"]["model_provider"] == "mistralai"
+
+
+async def test_aconvert_mistral_stream_lifecycle() -> None:
+    from langchain_mistralai._stream_events import aconvert_mistral_stream
+
+    async def _araw() -> Any:
+        for chunk in _text_then_tool():
+            yield chunk
+
+    events: list[Any] = [
+        e
+        async for e in aconvert_mistral_stream(
+            _araw(), _convert_chunk_to_message_chunk, output_version="v0"
+        )
+    ]
+    assert_valid_event_stream(events)
+    assert events[0]["event"] == "message-start"
+    assert events[0]["metadata"]["provider"] == "mistralai"
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    tool_finishes = [f for f in finishes if f["content"]["type"] == "tool_call"]
+    assert len(tool_finishes) == 1
+    tc = cast("dict[str, Any]", tool_finishes[0]["content"])
+    assert tc["name"] == "get_weather"
+    assert tc["args"] == {"city": "Paris"}
+
+    message_finish = events[-1]
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["usage"] == {
+        "input_tokens": 9,
+        "output_tokens": 3,
+        "total_tokens": 12,
+    }