From 9f01e8cdde6ea50f4c14985a7e5bff84c1cb32ad Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Thu, 11 Jun 2026 09:28:35 -0400 Subject: [PATCH] feat(perplexity): native content-block streaming events --- .../langchain_perplexity/_stream_events.py | 194 ++++++++++ .../langchain_perplexity/chat_models.py | 69 +++- .../tests/unit_tests/test_stream_events.py | 365 ++++++++++++++++++ 3 files changed, 627 insertions(+), 1 deletion(-) create mode 100644 libs/partners/perplexity/langchain_perplexity/_stream_events.py create mode 100644 libs/partners/perplexity/tests/unit_tests/test_stream_events.py diff --git a/libs/partners/perplexity/langchain_perplexity/_stream_events.py b/libs/partners/perplexity/langchain_perplexity/_stream_events.py new file mode 100644 index 00000000000..ae715c5983b --- /dev/null +++ b/libs/partners/perplexity/langchain_perplexity/_stream_events.py @@ -0,0 +1,194 @@ +"""Native content-block streaming-event converter for Perplexity. + +Builds text and tool-call blocks directly from Perplexity's OpenAI-shaped +delta, feeding the shared `BlockStreamTracker`. Unlike the compat bridge (which +buries tool_calls in `additional_kwargs` and loses citations on the v3 path), +this surfaces tool calls as proper blocks and puts search extras (citations, +search_results, images, etc.) in `message-finish` `response_metadata`. + +Usage is **cumulative**: each chunk's `usage` field is a running total, so the +message total is the value from the *last* chunk that carries usage — not an +accumulation across chunks. + +Note: Perplexity reasoning is inline `…` text inside `content`; +there is no separate `reasoning_content` field, so it surfaces as plain text. + +Citations and other search extras are emitted by the legacy `_stream` path via +`additional_kwargs`. On the v3 path `additional_kwargs` is dropped by the compat +bridge, so this converter puts them in `response_metadata` instead — +native v3 output is a strict superset of what the bridge provides. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from langchain_core.language_models.stream_events import ( + BlockStreamTracker, + build_message_finish, +) + +from langchain_perplexity.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": "perplexity"} + 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 Perplexity delta.""" + if content := delta.get("content"): + yield from tracker.feed("text", {"type": "text", "text": content}) + for tc in delta.get("tool_calls") or []: + idx = tc.get("index", 0) + fn = tc.get("function") or {} + args = fn.get("arguments") + 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 _collect_extras(chunk: dict[str, Any]) -> dict[str, Any]: + """Build response_metadata extras from the first chunk. + + Mirrors the first-chunk block in `_stream`: always includes `citations` + (default `[]`), the present-keyed `images`/`related_questions`/ + `search_results`, and the truthy-keyed `videos`/`reasoning_steps`. + """ + extras: dict[str, Any] = {"citations": chunk.get("citations", [])} + for key in ("images", "related_questions", "search_results"): + if key in chunk: + extras[key] = chunk[key] + for key in ("videos", "reasoning_steps"): + if chunk.get(key): + extras[key] = chunk[key] + return extras + + +def convert_perplexity_stream( + raw: Iterator[Any], *, message_id: str | None = None +) -> Iterator[MessagesData]: + """Convert a raw Perplexity chat stream to protocol events. + + Args: + raw: Raw Perplexity 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 + latest_usage: dict[str, Any] | None = None + response_metadata: dict[str, Any] = {"model_provider": "perplexity"} + model: str | None = None + first_chunk = True + + for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() + if model is None: + model = chunk.get("model") + if usage := chunk.get("usage"): + # Usage is cumulative; track the latest total — do NOT accumulate. + latest_usage = dict(_create_usage_metadata(usage)) + if "num_search_queries" not in response_metadata: + if num_sq := usage.get("num_search_queries"): + response_metadata["num_search_queries"] = num_sq + if "search_context_size" not in response_metadata: + if scs := usage.get("search_context_size"): + response_metadata["search_context_size"] = scs + choices = chunk.get("choices") or [] + if len(choices) == 0: + continue + if first_chunk: + response_metadata.update(_collect_extras(chunk)) + first_chunk = False + 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=latest_usage, response_metadata=response_metadata) + + +async def aconvert_perplexity_stream( + raw: AsyncIterator[Any], *, message_id: str | None = None +) -> AsyncIterator[MessagesData]: + """Async twin of `convert_perplexity_stream`.""" + tracker = BlockStreamTracker() + started = False + latest_usage: dict[str, Any] | None = None + response_metadata: dict[str, Any] = {"model_provider": "perplexity"} + model: str | None = None + first_chunk = True + + async for chunk in raw: + if not isinstance(chunk, dict): + chunk = chunk.model_dump() + if model is None: + model = chunk.get("model") + if usage := chunk.get("usage"): + # Usage is cumulative; track the latest total — do NOT accumulate. + latest_usage = dict(_create_usage_metadata(usage)) + if "num_search_queries" not in response_metadata: + if num_sq := usage.get("num_search_queries"): + response_metadata["num_search_queries"] = num_sq + if "search_context_size" not in response_metadata: + if scs := usage.get("search_context_size"): + response_metadata["search_context_size"] = scs + choices = chunk.get("choices") or [] + if len(choices) == 0: + continue + if first_chunk: + response_metadata.update(_collect_extras(chunk)) + first_chunk = False + 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=latest_usage, response_metadata=response_metadata) diff --git a/libs/partners/perplexity/langchain_perplexity/chat_models.py b/libs/partners/perplexity/langchain_perplexity/chat_models.py index b45853d5a76..15e4484690a 100644 --- a/libs/partners/perplexity/langchain_perplexity/chat_models.py +++ b/libs/partners/perplexity/langchain_perplexity/chat_models.py @@ -6,7 +6,10 @@ import json import logging from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from operator import itemgetter -from typing import Any, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast + +if TYPE_CHECKING: + from langchain_protocol.protocol import MessagesData from langchain_core.callbacks import ( AsyncCallbackManagerForLLMRun, @@ -1369,6 +1372,70 @@ class ChatPerplexity(BaseChatModel): await run_manager.on_llm_new_token(chunk.text, chunk=chunk) 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 Perplexity-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_perplexity._stream_events import convert_perplexity_stream + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs} + params.pop("stream", None) + if stop: + params["stop_sequences"] = stop + raw = self.client.chat.completions.create( + messages=message_dicts, stream=True, **params + ) + for event in convert_perplexity_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_perplexity._stream_events import aconvert_perplexity_stream + + message_dicts, params = self._create_message_dicts(messages, stop) + params = {**params, **kwargs} + params.pop("stream", None) + if stop: + params["stop_sequences"] = stop + raw = await self.async_client.chat.completions.create( + messages=message_dicts, stream=True, **params + ) + async for event in aconvert_perplexity_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 + def _generate( self, messages: list[BaseMessage], diff --git a/libs/partners/perplexity/tests/unit_tests/test_stream_events.py b/libs/partners/perplexity/tests/unit_tests/test_stream_events.py new file mode 100644 index 00000000000..a1f40e26a52 --- /dev/null +++ b/libs/partners/perplexity/tests/unit_tests/test_stream_events.py @@ -0,0 +1,365 @@ +"""Unit tests for the Perplexity 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_perplexity import ChatPerplexity +from langchain_perplexity._stream_events import ( + aconvert_perplexity_stream, + convert_perplexity_stream, +) + +if "PPLX_API_KEY" not in os.environ: + os.environ["PPLX_API_KEY"] = "fake-key" + + +def _text_with_citations() -> list[dict]: + """Fixture: plain text response with citations and search_results. + + Usage is cumulative — each chunk's `usage` is a running total. + Includes a final `choices: []` chunk carrying only the cumulative total. + """ + m = "sonar" + return [ + { + "model": m, + "citations": ["https://example.com/1", "https://example.com/2"], + "search_results": [{"title": "Example", "url": "https://example.com/1"}], + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hello "}, + "finish_reason": None, + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 2, + "total_tokens": 12, + }, + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"content": "world"}, + "finish_reason": None, + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": {"content": "!"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 6, + "total_tokens": 16, + }, + }, + # Usage-only chunk (choices: []) — final cumulative total. Strictly + # larger than the prior chunk so the test proves message-finish uses + # THIS total (last-wins), not the last choices-chunk's total and not a + # sum of per-chunk deltas. + { + "model": m, + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 7, + "total_tokens": 17, + "num_search_queries": 2, + "search_context_size": "high", + }, + }, + ] + + +def _tool_call_chunks() -> list[dict]: + """Fixture: single tool call streamed across two chunks.""" + m = "sonar" + return [ + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": "call_abc", + "function": { + "name": "get_weather", + "arguments": '{"city":', + }, + } + ], + }, + "finish_reason": None, + } + ], + "usage": {"prompt_tokens": 15, "completion_tokens": 3, "total_tokens": 18}, + }, + { + "model": m, + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": ' "Paris"}'}} + ] + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 15, "completion_tokens": 8, "total_tokens": 23}, + }, + ] + + +def test_convert_perplexity_stream_lifecycle() -> None: + events: list[Any] = list(convert_perplexity_stream(iter(_text_with_citations()))) + assert_valid_event_stream(events) + + # message-start: empty id (LangChain run id slot), correct provider + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "perplexity" + + # Block order: exactly one text block + finishes = [e for e in events if e["event"] == "content-block-finish"] + types = [f["content"]["type"] for f in finishes] + assert types == ["text"] + assert cast("dict[str, Any]", finishes[0]["content"])["text"] == "Hello world!" + + # message-finish: usage == LAST cumulative total (the no-choices final + # chunk's 7/17, not the prior choices-chunk's 6/16 nor a sum of deltas) + message_finish = events[-1] + assert message_finish["event"] == "message-finish" + assert message_finish["usage"] == { + "input_tokens": 10, + "output_tokens": 7, + "total_tokens": 17, + } + + # Extras present in response_metadata + assert message_finish["metadata"]["model_provider"] == "perplexity" + assert message_finish["metadata"]["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + assert message_finish["metadata"]["search_results"] == [ + {"title": "Example", "url": "https://example.com/1"} + ] + + +async def test_aconvert_perplexity_stream_lifecycle() -> None: + async def _araw() -> Any: + for chunk in _text_with_citations(): + yield chunk + + events: list[Any] = [e async for e in aconvert_perplexity_stream(_araw())] + assert_valid_event_stream(events) + + assert events[0]["event"] == "message-start" + assert events[0]["id"] == "" + assert events[0]["metadata"]["provider"] == "perplexity" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + types = [f["content"]["type"] for f in finishes] + assert types == ["text"] + assert cast("dict[str, Any]", finishes[0]["content"])["text"] == "Hello world!" + + message_finish = events[-1] + assert message_finish["event"] == "message-finish" + # Usage must equal the LAST cumulative total (7/17), not a sum of deltas + assert message_finish["usage"] == { + "input_tokens": 10, + "output_tokens": 7, + "total_tokens": 17, + } + assert message_finish["metadata"]["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + assert message_finish["metadata"]["search_results"] == [ + {"title": "Example", "url": "https://example.com/1"} + ] + + +def test_convert_perplexity_stream_tool_call() -> None: + """Tool calls are surfaced as `tool_call` blocks keyed `tool:{idx}`.""" + events: list[Any] = list(convert_perplexity_stream(iter(_tool_call_chunks()))) + 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 = cast("dict[str, Any]", finishes[0]["content"]) + assert tool["name"] == "get_weather" + assert tool["args"] == {"city": "Paris"} + + +def test_perplexity_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 text blocks with `model_provider == "perplexity"`. + """ + llm = ChatPerplexity(model="sonar") + mock_client = MagicMock() + + def mock_create(*_a: Any, **_k: Any) -> Any: + return iter(_text_with_citations()) + + mock_client.chat.completions.create = mock_create + + with patch.object(llm, "client", mock_client): + stream = llm.stream_events("Test query", version="v3") + events = list(stream) + + 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"] == "perplexity" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == ["text"] + + message_finish = cast("dict[str, Any]", events[-1]) + assert message_finish["event"] == "message-finish" + assert message_finish["metadata"]["model_provider"] == "perplexity" + # Citations ride on message-finish metadata (not additional_kwargs — the v3 + # path has no additional_kwargs channel; this is an intentional relocation). + assert message_finish["metadata"]["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + + # Round-trip: extras and model_name reach the assembled message's + # response_metadata (the user-facing guarantee of the de-risk design). + output = stream.output + assert output.response_metadata["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + assert output.response_metadata["model_name"] == "sonar" + + +async def test_perplexity_astream_events_v3_lifecycle() -> None: + """Async twin of `test_perplexity_stream_events_v3_lifecycle`.""" + llm = ChatPerplexity(model="sonar") + + async def _araw() -> Any: + for chunk in _text_with_citations(): + yield chunk + + async def mock_create(*_a: Any, **_k: Any) -> Any: + return _araw() + + mock_async_client = MagicMock() + mock_async_client.chat.completions.create = mock_create + + with patch.object(llm, "async_client", mock_async_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"] # non-empty LC run id + assert message_start["metadata"]["provider"] == "perplexity" + + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [f["content"]["type"] for f in finishes] == ["text"] + + message_finish = cast("dict[str, Any]", events[-1]) + assert message_finish["event"] == "message-finish" + assert message_finish["metadata"]["model_provider"] == "perplexity" + # Citations ride on message-finish metadata (not additional_kwargs — the v3 + # path has no additional_kwargs channel; this is an intentional relocation). + assert message_finish["metadata"]["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + + output = await stream.output + assert output.response_metadata["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ] + assert output.response_metadata["model_name"] == "sonar" + + +def _no_usage_text() -> list[dict]: + """Fixture: text chunks with no `usage` field at all.""" + m = "sonar" + return [ + { + "model": m, + "choices": [ + {"index": 0, "delta": {"content": "Hi"}, "finish_reason": "stop"} + ], + } + ] + + +def test_convert_perplexity_stream_no_usage() -> None: + """No `usage` on any chunk → message-finish omits the usage field.""" + events: list[Any] = list(convert_perplexity_stream(iter(_no_usage_text()))) + assert_valid_event_stream(events) + message_finish = events[-1] + assert message_finish["event"] == "message-finish" + assert message_finish.get("usage") is None + + +def test_convert_perplexity_stream_usage_only_yields_nothing() -> None: + """A stream with only `choices: []` usage chunks yields no events.""" + usage_only = [ + {"model": "sonar", "choices": [], "usage": {"total_tokens": 5}}, + {"model": "sonar", "choices": [], "usage": {"total_tokens": 9}}, + ] + assert list(convert_perplexity_stream(iter(usage_only))) == [] + + +def test_convert_perplexity_stream_accepts_sdk_objects() -> None: + """Non-dict chunks are normalized via `model_dump()`.""" + + class _Chunk: + def __init__(self, data: dict) -> None: + self._data = data + + def model_dump(self) -> dict: + return self._data + + raw = [_Chunk(c) for c in _text_with_citations()] + events: list[Any] = list(convert_perplexity_stream(iter(raw))) + assert_valid_event_stream(events) + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert cast("dict[str, Any]", finishes[0]["content"])["text"] == "Hello world!" + assert events[-1]["metadata"]["citations"] == [ + "https://example.com/1", + "https://example.com/2", + ]