From aadf5971a35dde09bec09cc5c4d01040ad383df4 Mon Sep 17 00:00:00 2001 From: Sujit Nirmal Date: Tue, 20 Oct 2026 15:28:19 +0000 Subject: [PATCH] fix(perplexity): keep citations and search_results on assembled streaming responses Every Perplexity stream chunk carries the top-level citations and search_results arrays, and a non-streaming response exposes both, but the response rebuilt from a stream dropped them. Success callbacks and spend logs for streamed Perplexity calls therefore never saw the sources. Copy the last streamed values onto the assembled response through the provider streaming metadata hook --- .../llms/perplexity/chat/transformation.py | 29 ++++++++- .../test_perplexity_chat_transformation.py | 60 ++++++++++++++++++- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 017a74cd51e..5af1da2004d 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -2,12 +2,12 @@ Translate from OpenAI's `/v1/chat/completions` to Perplexity's `/v1/chat/completions` """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -24,6 +24,19 @@ _PERPLEXITY_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"X _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +class _StreamedSearchSources(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + citations: tuple[str, ...] | None = None + search_results: tuple[Mapping[str, object], ...] | None = None + + +def _streamed_search_sources(chunk: object) -> _StreamedSearchSources: + try: + return _StreamedSearchSources.model_validate(chunk) + except ValidationError: + return _StreamedSearchSources() + + class PerplexityChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -120,6 +133,18 @@ class PerplexityChatConfig(OpenAIGPTConfig): return model_response + def apply_assembled_streaming_response_metadata(self, response: ModelResponse, chunks: Sequence[object]) -> None: + sources: Final = tuple(map(_streamed_search_sources, reversed(chunks))) + citations: Final = next((source.citations for source in sources if source.citations), ()) + search_results: Final = next((source.search_results for source in sources if source.search_results), ()) + fields: Final = ( + ("citations", list(citations)), + ("search_results", [dict(result) for result in search_results]), + ) + for field_name, value in fields: + if value: + setattr(response, field_name, value) + def _enhance_usage_with_perplexity_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ Extract citation tokens and search queries from Perplexity API response diff --git a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py index 2c95265890d..3860ed44b89 100644 --- a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py +++ b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py @@ -13,10 +13,10 @@ import pytest # Add the project root to Python path -from litellm import ModelResponse +from litellm import ModelResponse, stream_chunk_builder from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig -from litellm.types.utils import Usage +from litellm.types.utils import ModelResponseStream, Usage class TestPerplexityChatTransformation: @@ -917,3 +917,59 @@ def test_transform_response_keeps_body_values_of_every_json_type_as_the_provider assert response.search_results == search_results assert [type(result) for result in response.search_results] == [dict, str, int] assert repr(response.usage.citation_tokens) == "18" + + +def _perplexity_stream_logging() -> Logging: + logging_obj = Logging( + model="sonar", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="perplexity-stream-test", + function_id="perplexity-stream-test", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}, custom_llm_provider="perplexity") + return logging_obj + + +def _perplexity_stream_chunk(delta: dict, finish_reason: str | None, sources: dict) -> ModelResponseStream: + return ModelResponseStream( + id="cmpl-1", + object="chat.completion.chunk", + created=1700000000, + model="sonar", + choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason}], + **sources, + ) + + +@pytest.mark.parametrize("dumped", [False, True]) +def test_stream_chunk_builder_keeps_perplexity_citations_and_search_results(dumped: bool): + sources = { + "citations": ["https://a.example", "https://b.example"], + "search_results": [{"url": "https://a.example", "title": "A", "date": None}], + } + chunks = [ + _perplexity_stream_chunk({"role": "assistant", "content": "Paris [1]"}, None, sources), + _perplexity_stream_chunk({"content": " is in France [2]."}, "stop", sources), + ] + + response = stream_chunk_builder( + chunks=[chunk.model_dump() if dumped else chunk for chunk in chunks], logging_obj=_perplexity_stream_logging() + ) + + assert response.choices[0].message.content == "Paris [1] is in France [2]." + assert response.citations == sources["citations"] + assert response.search_results == sources["search_results"] + + +def test_stream_chunk_builder_adds_no_perplexity_sources_when_the_stream_had_none(): + response = stream_chunk_builder( + chunks=[_perplexity_stream_chunk({"role": "assistant", "content": "Hi"}, "stop", {})], + logging_obj=_perplexity_stream_logging(), + ) + + assert response.choices[0].message.content == "Hi" + assert not hasattr(response, "citations") + assert not hasattr(response, "search_results")