diff --git a/libs/partners/fireworks/langchain_fireworks/_stream_events.py b/libs/partners/fireworks/langchain_fireworks/_stream_events.py new file mode 100644 index 00000000000..7c08508d807 --- /dev/null +++ b/libs/partners/fireworks/langchain_fireworks/_stream_events.py @@ -0,0 +1,160 @@ +"""Native content-block streaming-event converter for Fireworks. + +Builds text, reasoning, and tool-call blocks directly from Fireworks' raw +OpenAI-shaped delta (Fireworks streams tool-call args incrementally), feeding +the shared `BlockStreamTracker`. Unlike the compat bridge (which leaves +Fireworks' `tool_calls`/`reasoning_content` in `additional_kwargs`), this +surfaces them as first-class blocks. + +Reasoning delta key is ``reasoning_content`` (not ``reasoning``). +""" + +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_fireworks.chat_models import _usage_to_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": "fireworks"} + 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 Fireworks delta.""" + if reasoning := delta.get("reasoning_content"): + yield from tracker.feed( + "reasoning", {"type": "reasoning", "reasoning": reasoning} + ) + if content := delta.get("content"): + yield from tracker.feed("text", {"type": "text", "text": content}) + # Deferred-as-scope-cut: function_call (legacy) and logprobs are not + # surfaced as blocks; the bridge already drops them. + 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": # normalize JSON null for no-arg tools + 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_fireworks_stream( + raw: Iterator[Any], *, message_id: str | None = None +) -> Iterator[MessagesData]: + """Convert a raw Fireworks chat stream to protocol events. + + Args: + raw: Raw Fireworks 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": "fireworks"} + model: str | None = None + + for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() + if model is None: + model = chunk.get("model") + if service_tier := chunk.get("service_tier"): + response_metadata["service_tier"] = service_tier + # Usage arrives on the final no-choices chunk when + # `stream_options.include_usage=True`; also tolerate it on other chunks. + if fw_usage := chunk.get("usage"): + usage = accumulate_usage(usage, dict(_usage_to_metadata(fw_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_fireworks_stream( + raw: AsyncIterator[Any], *, message_id: str | None = None +) -> AsyncIterator[MessagesData]: + """Async twin of `convert_fireworks_stream`.""" + tracker = BlockStreamTracker() + started = False + usage: dict[str, Any] | None = None + response_metadata: dict[str, Any] = {"model_provider": "fireworks"} + model: str | None = None + + async for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() + if model is None: + model = chunk.get("model") + if service_tier := chunk.get("service_tier"): + response_metadata["service_tier"] = service_tier + if fw_usage := chunk.get("usage"): + usage = accumulate_usage(usage, dict(_usage_to_metadata(fw_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/fireworks/langchain_fireworks/chat_models.py b/libs/partners/fireworks/langchain_fireworks/chat_models.py index 79440a6af78..b5e129ad89a 100644 --- a/libs/partners/fireworks/langchain_fireworks/chat_models.py +++ b/libs/partners/fireworks/langchain_fireworks/chat_models.py @@ -8,6 +8,7 @@ import logging from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from operator import itemgetter from typing import ( + TYPE_CHECKING, Any, Literal, NoReturn, @@ -103,6 +104,9 @@ from langchain_fireworks._compat import _convert_from_v1_to_chat_completions from langchain_fireworks._version import __version__ from langchain_fireworks.data._profiles import _PROFILES +if TYPE_CHECKING: + from langchain_protocol.protocol import MessagesData + logger = logging.getLogger(__name__) @@ -1115,6 +1119,74 @@ class ChatFireworks(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 Fireworks-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_fireworks._stream_events import ( + convert_fireworks_stream, + ) + + kwargs.pop("stream_usage", None) # never leak into the Fireworks payload + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + if self.stream_usage and "stream_options" not in params: + params["stream_options"] = {"include_usage": True} + raw = _completion_with_retry( + self, run_manager=run_manager, messages=message_dicts, **params + ) + for event in convert_fireworks_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_fireworks._stream_events import ( + aconvert_fireworks_stream, + ) + + kwargs.pop("stream_usage", None) # never leak into the Fireworks payload + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs, "stream": True} + if self.stream_usage and "stream_options" not in params: + params["stream_options"] = {"include_usage": True} + raw = await _acompletion_with_retry( + self, run_manager=run_manager, messages=message_dicts, **params + ) + async for event in aconvert_fireworks_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 + async def _agenerate( self, messages: list[BaseMessage], diff --git a/libs/partners/fireworks/tests/unit_tests/test_stream_events.py b/libs/partners/fireworks/tests/unit_tests/test_stream_events.py new file mode 100644 index 00000000000..5ed9aa09e85 --- /dev/null +++ b/libs/partners/fireworks/tests/unit_tests/test_stream_events.py @@ -0,0 +1,317 @@ +"""Unit tests for the Fireworks 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_fireworks import ChatFireworks +from langchain_fireworks._stream_events import ( + aconvert_fireworks_stream, + convert_fireworks_stream, +) + +if "FIREWORKS_API_KEY" not in os.environ: + os.environ["FIREWORKS_API_KEY"] = "fake-key" + + +def _reasoning_text_tool() -> list[dict]: + m = "accounts/fireworks/models/deepseek-r1" + return [ + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "reasoning_content": "Let me "}, + "finish_reason": None, + } + ], + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"reasoning_content": "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", + } + ], + }, + # Final no-choices chunk carrying usage (Fireworks, not under x_groq) + { + "model": m, + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + ] + + +def test_convert_fireworks_stream_lifecycle() -> None: + events: list[Any] = list(convert_fireworks_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"] == "fireworks" + assert events[0]["metadata"]["model"] == "accounts/fireworks/models/deepseek-r1" + + 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_fireworks_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 == "fireworks"`. + """ + llm = ChatFireworks(model="accounts/fireworks/models/deepseek-r1") + 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"] == "fireworks" + + 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"] == "fireworks" + + +async def test_fireworks_astream_events_v3_lifecycle() -> None: + """Async twin of `test_fireworks_stream_events_v3_lifecycle`.""" + llm = ChatFireworks(model="accounts/fireworks/models/deepseek-r1") + + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + yield chunk + + # `_acompletion_with_retry` does `await llm.async_client.create(**kwargs)` + # (1.x SDK: `create()` is a coroutine resolving to an `AsyncStream`), so + # the mock must be a coroutine returning an async iterator. + 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"] == "fireworks" + + 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"] == "fireworks" + + +async def test_aconvert_fireworks_stream_lifecycle() -> None: + async def _araw() -> Any: + for chunk in _reasoning_text_tool(): + yield chunk + + events: list[Any] = [e async for e in aconvert_fireworks_stream(_araw())] + assert_valid_event_stream(events) + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "fireworks" + + 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 = "accounts/fireworks/models/deepseek-r1" + 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_fireworks_stream_parallel_tool_calls() -> None: + """Args at distinct `tool:{idx}` keys finalize as separate tool calls.""" + events: list[Any] = list(convert_fireworks_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"} + + # No usage chunk in this fixture: message-finish should carry no usage. + assert events[-1]["event"] == "message-finish" + assert events[-1].get("usage") is None