mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-11 01:01:41 +00:00
feat(perplexity): native content-block streaming events
This commit is contained in:
Nick Hollon
194
libs/partners/perplexity/langchain_perplexity/_stream_events.py
Normal file
194
libs/partners/perplexity/langchain_perplexity/_stream_events.py
Normal file
@@ -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 `<think>…</think>` 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)
|
||||
@@ -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],
|
||||
|
||||
365
libs/partners/perplexity/tests/unit_tests/test_stream_events.py
Normal file
365
libs/partners/perplexity/tests/unit_tests/test_stream_events.py
Normal file
@@ -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",
|
||||
]
|
||||
Reference in New Issue
Block a user