diff --git a/libs/partners/ollama/langchain_ollama/_stream_events.py b/libs/partners/ollama/langchain_ollama/_stream_events.py new file mode 100644 index 00000000000..3c03b1eefd8 --- /dev/null +++ b/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) diff --git a/libs/partners/ollama/langchain_ollama/chat_models.py b/libs/partners/ollama/langchain_ollama/chat_models.py index bdc7c67fd20..76cd4029e48 100644 --- a/libs/partners/ollama/langchain_ollama/chat_models.py +++ b/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], diff --git a/libs/partners/ollama/tests/unit_tests/test_stream_events.py b/libs/partners/ollama/tests/unit_tests/test_stream_events.py new file mode 100644 index 00000000000..3b2b5b67fd6 --- /dev/null +++ b/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"