diff --git a/libs/partners/openrouter/langchain_openrouter/_stream_events.py b/libs/partners/openrouter/langchain_openrouter/_stream_events.py new file mode 100644 index 00000000000..3b314708e1c --- /dev/null +++ b/libs/partners/openrouter/langchain_openrouter/_stream_events.py @@ -0,0 +1,205 @@ +"""Native content-block streaming-event converter for openrouter. + +Builds text, reasoning, and tool-call blocks directly from openrouter's raw +OpenAI-shaped delta (openrouter streams tool-call args incrementally), feeding +the shared `BlockStreamTracker`. Unlike the compat bridge (which leaves +openrouter's `tool_calls`/`reasoning` in `additional_kwargs`), this surfaces +them as blocks. + +Differs from the groq converter in three ways: usage is a top-level +`chunk["usage"]` (carrying openrouter `cost`/`cost_details`), a no-choices +chunk may carry an `error` payload that must be propagated, and the reasoning +field is `delta.reasoning` (not `reasoning_content`). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from langchain_core.language_models.stream_events import ( + BlockStreamTracker, + accumulate_usage, + build_message_finish, +) + +from langchain_openrouter.chat_models import _create_usage_metadata + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterator + + from langchain_protocol.protocol import ( + MessageMetadata, + MessagesData, + MessageStartData, + ) + + +def _message_start(message_id: str | None, model: str | None) -> MessageStartData: + # Do not use a provider completion 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": "openrouter"} + if model: + metadata["model"] = model + return { + "event": "message-start", + "role": "ai", + "id": message_id or "", + "metadata": metadata, + } + + +def _feed_delta( + tracker: BlockStreamTracker, delta: dict[str, Any] +) -> Iterator[MessagesData]: + """Yield block events for one OpenAI-shaped openrouter delta.""" + if reasoning := delta.get("reasoning"): + yield from tracker.feed( + "reasoning", {"type": "reasoning", "reasoning": reasoning} + ) + if content := delta.get("content"): + yield from tracker.feed("text", {"type": "text", "text": content}) + # openrouter structured reasoning fragments (delta["reasoning_details"]) are + # not surfaced as blocks yet; tracked for a follow-up. + for tc in delta.get("tool_calls") or []: + idx = tc.get("index", 0) + fn = tc.get("function") or {} + yield from tracker.feed( + f"tool:{idx}", + { + "type": "tool_call_chunk", + "id": tc.get("id"), + "name": fn.get("name"), + "args": fn.get("arguments") or "", + "index": idx, + }, + ) + + +def _error_message(error: dict[str, Any]) -> str: + """Build the streaming-error message, matching `_stream`'s `ValueError`.""" + return ( + f"OpenRouter API returned an error during streaming: " + f"{error.get('message', str(error))} " + f"(code: {error.get('code', 'unknown')})" + ) + + +class _StreamState: + """Accumulates per-stream state shared by the sync and async converters.""" + + def __init__(self) -> None: + self.tracker = BlockStreamTracker() + self.started = False + self.usage: dict[str, Any] | None = None + self.response_metadata: dict[str, Any] = {"model_provider": "openrouter"} + self.model: str | None = None + + def prepare(self, chunk: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Coerce a raw chunk, fold in usage/cost, and return it plus choices. + + Raises: + ValueError: If a no-choices chunk carries an openrouter `error`. + """ + if not isinstance(chunk, dict): + chunk = chunk.model_dump(by_alias=True) + if self.model is None: + self.model = chunk.get("model") + if usage_payload := chunk.get("usage"): + self.usage = accumulate_usage( + self.usage, dict(_create_usage_metadata(usage_payload)) + ) + # Surface OpenRouter cost data, mirroring + # `_convert_chunk_to_message_chunk` so the v3 finish metadata keeps + # parity with the compat bridge's `response_metadata`. + if "cost" in usage_payload: + self.response_metadata["cost"] = usage_payload["cost"] + if "cost_details" in usage_payload: + self.response_metadata["cost_details"] = usage_payload["cost_details"] + choices = chunk.get("choices") or [] + if len(choices) == 0 and (error := chunk.get("error")): + raise ValueError(_error_message(error)) + return chunk, choices + + def record_finish(self, chunk: dict[str, Any], choice: dict[str, Any]) -> None: + """Capture finish-chunk response metadata, matching `_stream`. + + Mirrors the `generation_info` that `_stream`/`_astream` attach on the + final chunk so the v3 `message-finish` metadata equals the bridge's. + """ + self.response_metadata["finish_reason"] = choice["finish_reason"] + self.response_metadata["model_name"] = chunk.get("model") or self.model + if system_fingerprint := chunk.get("system_fingerprint"): + self.response_metadata["system_fingerprint"] = system_fingerprint + if native_finish_reason := choice.get("native_finish_reason"): + self.response_metadata["native_finish_reason"] = native_finish_reason + if response_id := chunk.get("id"): + self.response_metadata["id"] = response_id + if created := chunk.get("created"): + self.response_metadata["created"] = int(created) + if object_ := chunk.get("object"): + self.response_metadata["object"] = object_ + + +def convert_openrouter_stream( + raw: Iterator[Any], *, message_id: str | None = None +) -> Iterator[MessagesData]: + """Convert a raw openrouter chat stream to protocol events. + + Args: + raw: Raw openrouter chat-completion stream chunks (dicts or SDK + objects). + message_id: Overrides the provider message id on `message-start`. + + Yields: + Protocol `MessagesData` lifecycle events. + + Raises: + ValueError: If a chunk carries an openrouter `error` payload. + """ + state = _StreamState() + for chunk in raw: + chunk_dict, choices = state.prepare(chunk) + if len(choices) == 0: + continue + if not state.started: + state.started = True + yield _message_start(message_id, state.model) + choice = choices[0] + yield from _feed_delta(state.tracker, choice.get("delta") or {}) + if choice.get("finish_reason"): + state.record_finish(chunk_dict, choice) + + if not state.started: + return + yield from state.tracker.finish_all() + yield build_message_finish( + usage=state.usage, response_metadata=state.response_metadata + ) + + +async def aconvert_openrouter_stream( + raw: AsyncIterator[Any], *, message_id: str | None = None +) -> AsyncIterator[MessagesData]: + """Async twin of `convert_openrouter_stream`.""" + state = _StreamState() + async for chunk in raw: + chunk_dict, choices = state.prepare(chunk) + if len(choices) == 0: + continue + if not state.started: + state.started = True + yield _message_start(message_id, state.model) + choice = choices[0] + for ev in _feed_delta(state.tracker, choice.get("delta") or {}): + yield ev + if choice.get("finish_reason"): + state.record_finish(chunk_dict, choice) + + if not state.started: + return + for ev in state.tracker.finish_all(): + yield ev + yield build_message_finish( + usage=state.usage, response_metadata=state.response_metadata + ) diff --git a/libs/partners/openrouter/langchain_openrouter/chat_models.py b/libs/partners/openrouter/langchain_openrouter/chat_models.py index fcc98284852..e1a60438102 100644 --- a/libs/partners/openrouter/langchain_openrouter/chat_models.py +++ b/libs/partners/openrouter/langchain_openrouter/chat_models.py @@ -6,7 +6,7 @@ import json 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 langchain_core.callbacks import ( AsyncCallbackManagerForLLMRun, @@ -72,6 +72,9 @@ from typing_extensions import Self from langchain_openrouter._version import __version__ from langchain_openrouter.data._profiles import _PROFILES +if TYPE_CHECKING: + from langchain_protocol.protocol import MessagesData + _MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES) # LangChain-internal kwargs that must not be forwarded to the SDK. @@ -710,6 +713,72 @@ class ChatOpenRouter(BaseChatModel): ) 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 openrouter-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. + """ + from langchain_openrouter._stream_events import ( # noqa: PLC0415 + convert_openrouter_stream, + ) + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + if self.stream_usage: + params["stream_options"] = {"include_usage": True} + _strip_internal_kwargs(params) + sdk_messages = _wrap_messages_for_sdk(message_dicts) + raw = self.client.chat.send(messages=sdk_messages, **params) + for event in convert_openrouter_stream(raw, 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`.""" + from langchain_openrouter._stream_events import ( # noqa: PLC0415 + aconvert_openrouter_stream, + ) + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + if self.stream_usage: + params["stream_options"] = {"include_usage": True} + _strip_internal_kwargs(params) + sdk_messages = _wrap_messages_for_sdk(message_dicts) + raw = await self.client.chat.send_async(messages=sdk_messages, **params) + async for event in aconvert_openrouter_stream(raw, 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 + # # Internal methods # diff --git a/libs/partners/openrouter/tests/unit_tests/test_stream_events.py b/libs/partners/openrouter/tests/unit_tests/test_stream_events.py new file mode 100644 index 00000000000..cdbe1438c41 --- /dev/null +++ b/libs/partners/openrouter/tests/unit_tests/test_stream_events.py @@ -0,0 +1,392 @@ +"""Unit tests for the openrouter native stream-events converter.""" + +import os +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream + +from langchain_openrouter import ChatOpenRouter +from langchain_openrouter._stream_events import ( + aconvert_openrouter_stream, + convert_openrouter_stream, +) + +if "OPENROUTER_API_KEY" not in os.environ: + os.environ["OPENROUTER_API_KEY"] = "fake-key" + + +def _reasoning_text_tool() -> list[dict]: + m = "anthropic/claude-3.5-sonnet" + return [ + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "reasoning": "Let me "}, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"reasoning": "think."}, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"content": "Hi "}, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"content": "there"}, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "t1", + "function": { + "name": "get_weather", + "arguments": '{"city":', + }, + } + ] + }, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": ' "Paris"}'}} + ] + }, + "finish_reason": "tool_calls", + } + ], + }, + { + "model": m, + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + ] + + +def test_convert_openrouter_stream_lifecycle() -> None: + """Reasoning, text, and tool-call deltas surface as finished blocks.""" + events: list[Any] = list(convert_openrouter_stream(iter(_reasoning_text_tool()))) + assert_valid_event_stream(events) + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "openrouter" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == [ + "reasoning", + "text", + "tool_call", + ] + assert ( + cast("dict[str, Any]", finishes[0]["content"])["reasoning"] == "Let me think." + ) + assert cast("dict[str, Any]", finishes[1]["content"])["text"] == "Hi there" + tool = cast("dict[str, Any]", finishes[2]["content"]) + assert tool["name"] == "get_weather" + assert tool["args"] == {"city": "Paris"} + + assert events[-1]["event"] == "message-finish" + assert events[-1]["usage"] == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + } + + +def test_convert_openrouter_stream_error_chunk_raises() -> None: + """A no-choices chunk carrying an `error` payload raises `ValueError`.""" + chunks = [{"choices": [], "error": {"message": "rate limited", "code": 429}}] + with pytest.raises(ValueError, match="OpenRouter API returned an error"): + list(convert_openrouter_stream(iter(chunks))) + + +def test_convert_openrouter_stream_surfaces_cost() -> None: + """Cost/native-finish/model metadata reach `message-finish`, like the bridge. + + The final chunk carries both `finish_reason` metadata and usage/cost data; + both must land in the `message-finish` response metadata so switching to the + v3 path does not silently drop OpenRouter's documented cost surfacing. + """ + cost_details = { + "upstream_inference_cost": 7.745e-05, + "upstream_inference_prompt_cost": 8.95e-06, + "upstream_inference_completions_cost": 6.85e-05, + } + chunks: list[dict] = [ + { + "model": "openai/gpt-4o-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hi"}}], + }, + { + "model": "openai/gpt-4o-mini", + "id": "gen-cost-stream", + "system_fingerprint": "fp_abc", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + "native_finish_reason": "end_turn", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "cost": 7.5e-05, + "cost_details": cost_details, + }, + }, + ] + events: list[Any] = list(convert_openrouter_stream(iter(chunks))) + assert_valid_event_stream(events) + + finish = events[-1] + assert finish["event"] == "message-finish" + meta = finish["metadata"] + assert meta["cost"] == 7.5e-05 + assert meta["cost_details"] == cost_details + assert meta["finish_reason"] == "stop" + assert meta["native_finish_reason"] == "end_turn" + assert meta["model_name"] == "openai/gpt-4o-mini" + assert meta["system_fingerprint"] == "fp_abc" + assert meta["id"] == "gen-cost-stream" + assert finish["usage"] == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + } + + +def _parallel_tool_calls() -> list[dict]: + """Two tool calls interleaved at index 0 and 1 across chunks.""" + m = "anthropic/claude-3.5-sonnet" + return [ + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": "call_a", + "function": { + "name": "get_weather", + "arguments": '{"city":', + }, + }, + { + "index": 1, + "id": "call_b", + "function": { + "name": "get_time", + "arguments": '{"zone":', + }, + }, + ], + }, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": ' "Paris"}'}}, + {"index": 1, "function": {"arguments": ' "UTC"}'}}, + ] + }, + "finish_reason": "tool_calls", + } + ], + }, + ] + + +def test_convert_openrouter_stream_parallel_tool_calls() -> None: + """Args at distinct `tool:{idx}` keys finalize as separate tool calls.""" + events: list[Any] = list(convert_openrouter_stream(iter(_parallel_tool_calls()))) + assert_valid_event_stream(events) + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == ["tool_call", "tool_call"] + + first = cast("dict[str, Any]", finishes[0]["content"]) + second = cast("dict[str, Any]", finishes[1]["content"]) + assert first["name"] == "get_weather" + assert first["args"] == {"city": "Paris"} + assert second["name"] == "get_time" + assert second["args"] == {"zone": "UTC"} + + +async def test_aconvert_openrouter_stream_lifecycle() -> None: + """Async twin of `test_convert_openrouter_stream_lifecycle`.""" + + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + yield chunk + + events: list[Any] = [e async for e in aconvert_openrouter_stream(_araw())] + assert_valid_event_stream(events) + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "openrouter" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == [ + "reasoning", + "text", + "tool_call", + ] + assert ( + cast("dict[str, Any]", finishes[0]["content"])["reasoning"] == "Let me think." + ) + assert cast("dict[str, Any]", finishes[1]["content"])["text"] == "Hi there" + tool = cast("dict[str, Any]", finishes[2]["content"]) + assert tool["name"] == "get_weather" + assert tool["args"] == {"city": "Paris"} + + assert events[-1]["event"] == "message-finish" + assert events[-1]["usage"] == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + } + + +async def test_aconvert_openrouter_stream_error_chunk_raises() -> None: + """Async twin of `test_convert_openrouter_stream_error_chunk_raises`.""" + + async def _araw() -> Any: + yield {"choices": [], "error": {"message": "rate limited", "code": 429}} + + with pytest.raises(ValueError, match="OpenRouter API returned an error"): + [e async for e in aconvert_openrouter_stream(_araw())] + + +def test_openrouter_stream_events_v3_lifecycle() -> None: + """Drive `stream_events(version="v3")` through the native sync hook. + + Confirms the model-level path threads the LangChain run id onto + `message-start` (non-empty, unlike the converter's empty default) and + surfaces reasoning/text/tool blocks with `model_provider == "openrouter"`. + """ + llm = ChatOpenRouter(model="foo") + mock_client = MagicMock() + + def mock_send(*_a: Any, **_k: Any) -> Any: + for chunk in _reasoning_text_tool(): + mock = MagicMock() + mock.model_dump.return_value = chunk + yield mock + + mock_client.chat.send = mock_send + + with patch.object(llm, "client", mock_client): + events = list(llm.stream_events("Test query", version="v3")) + + assert_valid_event_stream(events) + + message_start = cast("dict[str, Any]", events[0]) + assert message_start["event"] == "message-start" + assert message_start["id"] # core seeds a non-empty LangChain run id + assert message_start["metadata"]["provider"] == "openrouter" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == [ + "reasoning", + "text", + "tool_call", + ] + tool = cast("dict[str, Any]", finishes[2]["content"]) + assert tool["name"] == "get_weather" + assert tool["args"] == {"city": "Paris"} + + message_finish = cast("dict[str, Any]", events[-1]) + assert message_finish["event"] == "message-finish" + assert message_finish["metadata"]["model_provider"] == "openrouter" + + +async def test_openrouter_astream_events_v3_lifecycle() -> None: + """Async twin of `test_openrouter_stream_events_v3_lifecycle`.""" + llm = ChatOpenRouter(model="foo") + + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + mock = MagicMock() + mock.model_dump.return_value = chunk + yield mock + + mock_client = MagicMock() + mock_client.chat.send_async = AsyncMock(return_value=_araw()) + + with patch.object(llm, "client", mock_client): + stream = await llm.astream_events("Test query", version="v3") + events = [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"] + assert message_start["metadata"]["provider"] == "openrouter" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == [ + "reasoning", + "text", + "tool_call", + ] + message_finish = cast("dict[str, Any]", events[-1]) + assert message_finish["metadata"]["model_provider"] == "openrouter"