From b7a60ac5f2acb97d99455e186f05d6ebfd0543bd Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 10 Jun 2026 12:13:42 -0400 Subject: [PATCH] feat(mistralai): native content-block streaming events --- .../langchain_mistralai/_stream_events.py | 153 +++++++++ .../langchain_mistralai/chat_models.py | 74 +++++ .../tests/unit_tests/test_stream_events.py | 297 ++++++++++++++++++ 3 files changed, 524 insertions(+) create mode 100644 libs/partners/mistralai/langchain_mistralai/_stream_events.py create mode 100644 libs/partners/mistralai/tests/unit_tests/test_stream_events.py diff --git a/libs/partners/mistralai/langchain_mistralai/_stream_events.py b/libs/partners/mistralai/langchain_mistralai/_stream_events.py new file mode 100644 index 0000000000000..00c906ef9b4d3 --- /dev/null +++ b/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) diff --git a/libs/partners/mistralai/langchain_mistralai/chat_models.py b/libs/partners/mistralai/langchain_mistralai/chat_models.py index be86042beb4cb..d379cb21a1543 100644 --- a/libs/partners/mistralai/langchain_mistralai/chat_models.py +++ b/libs/partners/mistralai/langchain_mistralai/chat_models.py @@ -85,6 +85,8 @@ 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 @@ async def _astream( ) 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], diff --git a/libs/partners/mistralai/tests/unit_tests/test_stream_events.py b/libs/partners/mistralai/tests/unit_tests/test_stream_events.py new file mode 100644 index 0000000000000..34011a43aef2b --- /dev/null +++ b/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, + }