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")