From 01d9cf78f0ada8a7a8764e0972c62201c67f79b9 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 10 Jun 2026 14:49:32 -0400 Subject: [PATCH] feat(groq): native content-block streaming events --- .../groq/langchain_groq/_stream_events.py | 151 +++++++++ .../groq/langchain_groq/chat_models.py | 59 +++- .../tests/unit_tests/test_stream_events.py | 307 ++++++++++++++++++ 3 files changed, 516 insertions(+), 1 deletion(-) create mode 100644 libs/partners/groq/langchain_groq/_stream_events.py create mode 100644 libs/partners/groq/tests/unit_tests/test_stream_events.py diff --git a/libs/partners/groq/langchain_groq/_stream_events.py b/libs/partners/groq/langchain_groq/_stream_events.py new file mode 100644 index 0000000000000..32edf8532f77c --- /dev/null +++ b/libs/partners/groq/langchain_groq/_stream_events.py @@ -0,0 +1,151 @@ +"""Native content-block streaming-event converter for groq. + +Builds text, reasoning, and tool-call blocks directly from groq's raw +OpenAI-shaped delta (groq streams tool-call args incrementally), feeding the +shared `BlockStreamTracker`. Unlike the compat bridge (which leaves groq's +`tool_calls`/`reasoning` in `additional_kwargs`), this surfaces them as blocks. +""" + +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_groq.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": "groq"} + 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 groq 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}) + # groq server-tool results (delta["executed_tools"]) 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 {} + args = fn.get("arguments") + if args == "null": # groq sends JSON null for no-arg tools; normalize + args = "{}" + yield from tracker.feed( + f"tool:{idx}", + { + "type": "tool_call_chunk", + "id": tc.get("id"), + "name": fn.get("name"), + "args": args or "", + "index": idx, + }, + ) + + +def convert_groq_stream( + raw: Iterator[Any], *, message_id: str | None = None +) -> Iterator[MessagesData]: + """Convert a raw groq chat stream to protocol events. + + Args: + raw: Raw groq chat-completion stream chunks (dicts or SDK objects). + message_id: Overrides the provider message id on `message-start`. + + Yields: + Protocol `MessagesData` lifecycle events. + """ + tracker = BlockStreamTracker() + started = False + usage: dict[str, Any] | None = None + response_metadata: dict[str, Any] = {"model_provider": "groq"} + model: str | None = None + + for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() # noqa: PLW2901 + if model is None: + model = chunk.get("model") + if groq_usage := (chunk.get("x_groq") or {}).get("usage"): + usage = accumulate_usage(usage, dict(_create_usage_metadata(groq_usage))) + choices = chunk.get("choices") or [] + if len(choices) == 0: + continue + if not started: + started = True + yield _message_start(message_id, model) + choice = choices[0] + yield from _feed_delta(tracker, choice.get("delta") or {}) + if finish_reason := choice.get("finish_reason"): + response_metadata["finish_reason"] = finish_reason + + if not started: + return + yield from tracker.finish_all() + yield build_message_finish(usage=usage, response_metadata=response_metadata) + + +async def aconvert_groq_stream( + raw: AsyncIterator[Any], *, message_id: str | None = None +) -> AsyncIterator[MessagesData]: + """Async twin of `convert_groq_stream`.""" + tracker = BlockStreamTracker() + started = False + usage: dict[str, Any] | None = None + response_metadata: dict[str, Any] = {"model_provider": "groq"} + model: str | None = None + + async for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() # noqa: PLW2901 + if model is None: + model = chunk.get("model") + if groq_usage := (chunk.get("x_groq") or {}).get("usage"): + usage = accumulate_usage(usage, dict(_create_usage_metadata(groq_usage))) + choices = chunk.get("choices") or [] + if len(choices) == 0: + continue + if not started: + started = True + yield _message_start(message_id, model) + choice = choices[0] + for ev in _feed_delta(tracker, choice.get("delta") or {}): + yield ev + if finish_reason := choice.get("finish_reason"): + response_metadata["finish_reason"] = finish_reason + + 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/groq/langchain_groq/chat_models.py b/libs/partners/groq/langchain_groq/chat_models.py index 9cfaa31042a56..a6c723c553f18 100644 --- a/libs/partners/groq/langchain_groq/chat_models.py +++ b/libs/partners/groq/langchain_groq/chat_models.py @@ -6,7 +6,7 @@ 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, @@ -74,6 +74,9 @@ from langchain_groq._version import __version__ from langchain_groq.data._profiles import _PROFILES +if TYPE_CHECKING: + from langchain_protocol.protocol import MessagesData + _MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES) _STRICT_STRUCTURED_OUTPUT_MODELS = frozenset( { @@ -762,6 +765,60 @@ async def _astream( ) 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 groq-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_groq._stream_events import convert_groq_stream # noqa: PLC0415 + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + raw = self.client.create(messages=message_dicts, **params) + for event in convert_groq_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_groq._stream_events import aconvert_groq_stream # noqa: PLC0415 + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + raw = await self.async_client.create(messages=message_dicts, **params) + async for event in aconvert_groq_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/groq/tests/unit_tests/test_stream_events.py b/libs/partners/groq/tests/unit_tests/test_stream_events.py new file mode 100644 index 0000000000000..2a76e3318e6c1 --- /dev/null +++ b/libs/partners/groq/tests/unit_tests/test_stream_events.py @@ -0,0 +1,307 @@ +"""Unit tests for the groq native stream-events converter.""" + +import os +from typing import Any, cast +from unittest.mock import MagicMock, patch + +from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream + +from langchain_groq import ChatGroq +from langchain_groq._stream_events import aconvert_groq_stream, convert_groq_stream + +if "GROQ_API_KEY" not in os.environ: + os.environ["GROQ_API_KEY"] = "fake-key" + + +def _reasoning_text_tool() -> list[dict]: + m = "llama-3.3-70b" + 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": [], + "x_groq": { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + } + }, + }, + ] + + +def test_convert_groq_stream_lifecycle() -> None: + events: list[Any] = list(convert_groq_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"] == "groq" + + 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"] + 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"} + + message_finish = events[-1] + assert message_finish["event"] == "message-finish" + assert message_finish["usage"] == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + } + + +def test_groq_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 == "groq"`. + """ + llm = ChatGroq(model="foo") + mock_client = MagicMock() + + def mock_create(*_a: Any, **_k: Any) -> Any: + return iter(_reasoning_text_tool()) + + mock_client.create = mock_create + + 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"] == "groq" + + 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"] == "groq" + + +async def test_groq_astream_events_v3_lifecycle() -> None: + """Async twin of `test_groq_stream_events_v3_lifecycle`.""" + llm = ChatGroq(model="foo") + + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + yield chunk + + async def mock_create(*a: Any, **k: Any) -> Any: + return _araw() + + mock_client = MagicMock() + mock_client.create = mock_create + + with patch.object(llm, "async_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"] == "groq" + + 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"] == "groq" + + +async def test_aconvert_groq_stream_lifecycle() -> None: + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + yield chunk + + events: list[Any] = [e async for e in aconvert_groq_stream(_araw())] + assert_valid_event_stream(events) + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "groq" + + 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"] + 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"} + + message_finish = events[-1] + assert message_finish["event"] == "message-finish" + assert message_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 = "llama-3.3-70b" + 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_groq_stream_parallel_tool_calls() -> None: + """Args at distinct `tool:{idx}` keys finalize as separate tool calls.""" + events: list[Any] = list(convert_groq_stream(iter(_parallel_tool_calls()))) + assert_valid_event_stream(events) + + finishes = [e for e in events if e["event"] == "content-block-finish"] + types = [f["content"]["type"] for f in finishes] + assert types == ["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"}