feat(ollama): native content-block streaming events

libs/partners/ollama/langchain_ollama/_stream_events.py

@@ -0,0 +1,230 @@
+"""Native content-block streaming-event converter for Ollama.
+
+Maps the raw Ollama chat stream (dicts with `message.content`,
+`message.thinking`, `message.tool_calls`, and final `done`/usage fields) to the
+protocol `MessagesData` lifecycle, feeding the shared `BlockStreamTracker`.
+Unlike the compat bridge (which buries reasoning in `additional_kwargs`), this
+surfaces thinking as real `reasoning` content blocks.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, cast
+
+from langchain_core.language_models.stream_events import (
+    BlockStreamTracker,
+    accumulate_usage,
+    build_message_finish,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterator
+
+    from langchain_core.messages import ToolCall
+    from langchain_protocol.protocol import (
+        MessageMetadata,
+        MessagesData,
+        MessageStartData,
+    )
+
+# Bound `_get_tool_calls_from_response`.
+GetToolCalls = Callable[[Any], "list[ToolCall]"]
+
+# Stable per-block keys for the tracker (Ollama gives no indices).
+_TEXT_KEY = "text"
+_REASONING_KEY = "reasoning"
+
+
+def _message_start(message_id: str | None, model: str | None) -> MessageStartData:
+    metadata: MessageMetadata = {"provider": "ollama"}
+    if model:
+        metadata["model"] = model
+    # `cast` rather than a bare TypedDict literal: the strict `ty` checker
+    # rejects the literal against the external `MessageStartData` TypedDict.
+    return cast(
+        "MessageStartData",
+        {
+            "event": "message-start",
+            "role": "ai",
+            "id": message_id or "",
+            "metadata": metadata,
+        },
+    )
+
+
+def convert_ollama_stream(
+    raw: Iterator[Any],
+    get_tool_calls: GetToolCalls,
+    *,
+    reasoning: bool | None = None,
+    message_id: str | None = None,
+) -> Iterator[MessagesData]:
+    """Convert a raw Ollama chat stream to protocol events.
+
+    Args:
+        raw: Raw Ollama stream items (dicts; non-dicts skipped).
+        get_tool_calls: `ChatOllama`'s `_get_tool_calls_from_response`, injected
+            so the converter stays pure.
+        reasoning: When truthy, surface `message.thinking` as reasoning blocks.
+        message_id: Left empty by default so the v3 stream's seeded run id stands.
+
+    Yields:
+        Protocol `MessagesData` lifecycle events.
+    """
+    # Local import to avoid a circular import: `chat_models` imports this module.
+    from langchain_ollama.chat_models import (  # noqa: PLC0415
+        _get_usage_metadata_from_generation_info,
+    )
+
+    tracker = BlockStreamTracker()
+    started = False
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": "ollama"}
+    tool_idx = 0
+
+    for resp in raw:
+        if not isinstance(resp, dict):
+            continue
+
+        message = resp.get("message") or {}
+
+        # Skip "load" responses with empty content, matching the compat bridge
+        # (`_iterate_over_stream`): the model was loaded but generated nothing,
+        # so emitting an empty message-start/finish would diverge from the bridge.
+        content = message.get("content") or ""
+        if (
+            resp.get("done") is True
+            and resp.get("done_reason") == "load"
+            and not content.strip()
+        ):
+            continue
+
+        if not started:
+            started = True
+            yield _message_start(message_id, resp.get("model"))
+
+        thinking = message.get("thinking")
+        if reasoning and thinking:
+            yield from tracker.feed(
+                _REASONING_KEY, {"type": "reasoning", "reasoning": thinking}
+            )
+
+        if content:
+            yield from tracker.feed(_TEXT_KEY, {"type": "text", "text": content})
+
+        if message.get("tool_calls"):
+            for tc in get_tool_calls(resp):
+                yield from tracker.feed(
+                    f"tool:{tool_idx}",
+                    {
+                        "type": "tool_call_chunk",
+                        "id": tc.get("id"),
+                        "name": tc.get("name"),
+                        "args": json.dumps(tc.get("args") or {}),
+                    },
+                )
+                tool_idx += 1
+
+        if resp.get("done") is True:
+            usage = accumulate_usage(
+                usage, _get_usage_metadata_from_generation_info(resp)
+            )
+            done_meta = {
+                k: v for k, v in resp.items() if k != "message" and v is not None
+            }
+            if "model" in done_meta:
+                done_meta["model_name"] = done_meta["model"]
+            response_metadata.update(done_meta)
+
+    if not started:
+        return
+    yield from tracker.finish_all()
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)
+
+
+async def aconvert_ollama_stream(
+    raw: AsyncIterator[Any],
+    get_tool_calls: GetToolCalls,
+    *,
+    reasoning: bool | None = None,
+    message_id: str | None = None,
+) -> AsyncIterator[MessagesData]:
+    """Async twin of `convert_ollama_stream`.
+
+    `get_tool_calls` and the usage helper stay sync.
+    """
+    # Local import to avoid a circular import: `chat_models` imports this module.
+    from langchain_ollama.chat_models import (  # noqa: PLC0415
+        _get_usage_metadata_from_generation_info,
+    )
+
+    tracker = BlockStreamTracker()
+    started = False
+    usage: dict[str, Any] | None = None
+    response_metadata: dict[str, Any] = {"model_provider": "ollama"}
+    tool_idx = 0
+
+    async for resp in raw:
+        if not isinstance(resp, dict):
+            continue
+
+        message = resp.get("message") or {}
+
+        # Skip "load" responses with empty content, matching the compat bridge
+        # (`_aiterate_over_stream`): the model was loaded but generated nothing,
+        # so emitting an empty message-start/finish would diverge from the bridge.
+        content = message.get("content") or ""
+        if (
+            resp.get("done") is True
+            and resp.get("done_reason") == "load"
+            and not content.strip()
+        ):
+            continue
+
+        if not started:
+            started = True
+            yield _message_start(message_id, resp.get("model"))
+
+        thinking = message.get("thinking")
+        if reasoning and thinking:
+            for ev in tracker.feed(
+                _REASONING_KEY, {"type": "reasoning", "reasoning": thinking}
+            ):
+                yield ev
+
+        if content:
+            for ev in tracker.feed(_TEXT_KEY, {"type": "text", "text": content}):
+                yield ev
+
+        if message.get("tool_calls"):
+            for tc in get_tool_calls(resp):
+                for ev in tracker.feed(
+                    f"tool:{tool_idx}",
+                    {
+                        "type": "tool_call_chunk",
+                        "id": tc.get("id"),
+                        "name": tc.get("name"),
+                        "args": json.dumps(tc.get("args") or {}),
+                    },
+                ):
+                    yield ev
+                tool_idx += 1
+
+        if resp.get("done") is True:
+            usage = accumulate_usage(
+                usage, _get_usage_metadata_from_generation_info(resp)
+            )
+            done_meta = {
+                k: v for k, v in resp.items() if k != "message" and v is not None
+            }
+            if "model" in done_meta:
+                done_meta["model_name"] = done_meta["model"]
+            response_metadata.update(done_meta)
+
+    if not started:
+        return
+    for ev in tracker.finish_all():
+        yield ev
+    yield build_message_finish(usage=usage, response_metadata=response_metadata)

libs/partners/ollama/langchain_ollama/chat_models.py

@@ -47,7 +47,7 @@ import logging
 import warnings
 from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
 from operator import itemgetter
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 from uuid import uuid4
 
 from langchain_core.callbacks import CallbackManagerForLLMRun
@@ -97,6 +97,9 @@ from langchain_ollama._utils import (
 )
 from langchain_ollama._version import __version__
 
+if TYPE_CHECKING:
+    from langchain_protocol.protocol import MessagesData
+
 log = logging.getLogger(__name__)
 
 
@@ -1307,6 +1310,45 @@ class ChatOllama(BaseChatModel):
                 )
             yield 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 Ollama-native content-block protocol events.
+
+        Detected by `langchain-core`'s `_iter_v2_events`; powers
+        `stream_events(version="v3")`. Surfaces `thinking` as reasoning blocks,
+        which the compat-bridge path does not. `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
+        # `_get_usage_metadata_from_generation_info` from this module.
+        from langchain_ollama._stream_events import (  # noqa: PLC0415
+            convert_ollama_stream,
+        )
+
+        reasoning = kwargs.get("reasoning", self.reasoning)
+        for event in convert_ollama_stream(
+            self._create_chat_stream(messages, stop, **kwargs),
+            _get_tool_calls_from_response,
+            reasoning=reasoning,
+            message_id=message_id,
+        ):
+            if run_manager is not None and event["event"] == "content-block-delta":
+                # Widen to a plain dict for the optional `delta` access: TypedDict
+                # union narrowing on the discriminator differs across `ty`
+                # versions, so neither a typed key access nor a typed cast is
+                # portable here.
+                delta = cast("dict[str, Any]", event).get("delta") or {}
+                if delta.get("type") == "text-delta":
+                    run_manager.on_llm_new_token(str(delta.get("text", "")))
+            yield event
+
     async def _aiterate_over_stream(
         self,
         messages: list[BaseMessage],
@@ -1389,6 +1431,37 @@ class ChatOllama(BaseChatModel):
                 )
             yield chunk
 
+    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
+        # `_get_usage_metadata_from_generation_info` from this module.
+        from langchain_ollama._stream_events import (  # noqa: PLC0415
+            aconvert_ollama_stream,
+        )
+
+        reasoning = kwargs.get("reasoning", self.reasoning)
+        async for event in aconvert_ollama_stream(
+            self._acreate_chat_stream(messages, stop, **kwargs),
+            _get_tool_calls_from_response,
+            reasoning=reasoning,
+            message_id=message_id,
+        ):
+            if run_manager is not None and event["event"] == "content-block-delta":
+                # See sync twin: widen to a plain dict for the optional `delta`
+                # access (portable across `ty` TypedDict-narrowing differences).
+                delta = cast("dict[str, Any]", event).get("delta") or {}
+                if delta.get("type") == "text-delta":
+                    await run_manager.on_llm_new_token(str(delta.get("text", "")))
+            yield event
+
     async def _agenerate(
         self,
         messages: list[BaseMessage],

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

@@ -0,0 +1,244 @@
+"""Unit tests for the Ollama 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_ollama import ChatOllama
+from langchain_ollama._stream_events import (
+    aconvert_ollama_stream,
+    convert_ollama_stream,
+)
+from langchain_ollama.chat_models import _get_tool_calls_from_response
+
+
+def _thinking_then_text_then_tool() -> list[dict]:
+    return [
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "", "thinking": "Let me "},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "", "thinking": "think."},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "The answer"},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": " is 42."},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"city": "Paris"},
+                        }
+                    }
+                ],
+            },
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": ""},
+            "done": True,
+            "done_reason": "stop",
+            "prompt_eval_count": 10,
+            "eval_count": 7,
+        },
+    ]
+
+
+def test_convert_ollama_stream_lifecycle() -> None:
+    events: list[Any] = list(
+        convert_ollama_stream(
+            iter(_thinking_then_text_then_tool()),
+            _get_tool_calls_from_response,
+            reasoning=True,
+        )
+    )
+    assert_valid_event_stream(events)
+
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"] == ""  # empty → core's seeded run id stands
+    assert events[0]["metadata"]["provider"] == "ollama"
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    types = [f["content"]["type"] for f in finishes]
+    assert types == ["reasoning", "text", "tool_call"]
+    reasoning = cast("dict[str, Any]", finishes[0]["content"])
+    assert reasoning["reasoning"] == "Let me think."
+    text = cast("dict[str, Any]", finishes[1]["content"])
+    assert text["text"] == "The answer is 42."
+    tool = cast("dict[str, Any]", finishes[2]["content"])
+    assert tool["name"] == "get_weather"
+    assert tool["args"] == {"city": "Paris"}
+
+    message_finish = events[-1]
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["usage"] == {
+        "input_tokens": 10,
+        "output_tokens": 7,
+        "total_tokens": 17,
+    }
+
+
+async def test_aconvert_ollama_stream_lifecycle() -> None:
+    async def _araw() -> Any:
+        for chunk in _thinking_then_text_then_tool():
+            yield chunk
+
+    events: list[Any] = [
+        ev
+        async for ev in aconvert_ollama_stream(
+            _araw(),
+            _get_tool_calls_from_response,
+            reasoning=True,
+        )
+    ]
+    assert_valid_event_stream(events)
+
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"] == ""
+    assert events[0]["metadata"]["provider"] == "ollama"
+
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    types = [f["content"]["type"] for f in finishes]
+    assert types == ["reasoning", "text", "tool_call"]
+    reasoning = cast("dict[str, Any]", finishes[0]["content"])
+    assert reasoning["reasoning"] == "Let me think."
+    text = cast("dict[str, Any]", finishes[1]["content"])
+    assert text["text"] == "The answer is 42."
+    tool = cast("dict[str, Any]", finishes[2]["content"])
+    assert tool["name"] == "get_weather"
+    assert tool["args"] == {"city": "Paris"}
+
+    message_finish = events[-1]
+    assert message_finish["event"] == "message-finish"
+    assert message_finish["usage"] == {
+        "input_tokens": 10,
+        "output_tokens": 7,
+        "total_tokens": 17,
+    }
+
+
+def test_convert_ollama_stream_reasoning_disabled() -> None:
+    """With reasoning off, thinking is not surfaced as a block."""
+    chunks = [
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "", "thinking": "hidden"},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "hi"},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": ""},
+            "done": True,
+            "prompt_eval_count": 1,
+            "eval_count": 1,
+        },
+    ]
+    events: list[Any] = list(
+        convert_ollama_stream(
+            iter(chunks), _get_tool_calls_from_response, reasoning=False
+        )
+    )
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    assert [f["content"]["type"] for f in finishes] == ["text"]
+
+
+def test_convert_ollama_stream_skips_leading_load_response() -> None:
+    """A leading `done_reason="load"` empty chunk is skipped, like the bridge."""
+    chunks = [
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": ""},
+            "done": True,
+            "done_reason": "load",
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": "hi"},
+            "done": False,
+        },
+        {
+            "model": "qw3",
+            "message": {"role": "assistant", "content": ""},
+            "done": True,
+            "done_reason": "stop",
+            "prompt_eval_count": 1,
+            "eval_count": 1,
+        },
+    ]
+    events: list[Any] = list(
+        convert_ollama_stream(
+            iter(chunks), _get_tool_calls_from_response, reasoning=True
+        )
+    )
+    assert_valid_event_stream(events)
+    # message-start carries the model from the first *non-load* chunk, not the
+    # skipped load chunk, and the stream is not empty.
+    assert events[0]["event"] == "message-start"
+    finishes = [e for e in events if e["event"] == "content-block-finish"]
+    assert [f["content"]["type"] for f in finishes] == ["text"]
+
+
+def test_ollama_stream_events_v3_model_level() -> None:
+    """End-to-end: `stream_events(version="v3")` drives the native hook."""
+    llm = ChatOllama(model="qw3", reasoning=True)
+    with patch.object(
+        ChatOllama,
+        "_create_chat_stream",
+        return_value=iter(_thinking_then_text_then_tool()),
+    ):
+        events: list[Any] = list(llm.stream_events("hi", version="v3"))
+
+    assert_valid_event_stream(events)
+    # message-start id is the LC run id (threaded by core), not empty.
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"]
+    finish_types = [
+        e["content"]["type"] for e in events if e["event"] == "content-block-finish"
+    ]
+    assert finish_types == ["reasoning", "text", "tool_call"]
+    assert events[-1]["metadata"]["model_provider"] == "ollama"
+
+
+async def test_ollama_astream_events_v3_model_level() -> None:
+    """Async end-to-end: `astream_events(version="v3")` drives the native hook."""
+
+    async def _araw() -> Any:
+        for chunk in _thinking_then_text_then_tool():
+            yield chunk
+
+    llm = ChatOllama(model="qw3", reasoning=True)
+    with patch.object(ChatOllama, "_acreate_chat_stream", return_value=_araw()):
+        stream = await llm.astream_events("hi", version="v3")
+        events: list[Any] = [ev async for ev in stream]
+
+    assert_valid_event_stream(events)
+    assert events[0]["event"] == "message-start"
+    assert events[0]["id"]
+    finish_types = [
+        e["content"]["type"] for e in events if e["event"] == "content-block-finish"
+    ]
+    assert finish_types == ["reasoning", "text", "tool_call"]
+    assert events[-1]["metadata"]["model_provider"] == "ollama"