diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 354f7692fd5..685f85d04e9 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -30,6 +30,22 @@ class PerplexityChatConfig(OpenAIGPTConfig): dynamic_api_key = api_key or get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str("PERPLEXITY_API_KEY") return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, # mutable-ok: matches the base chat transform signature + model: str, + messages: list[AllMessageValues], # mutable-ok: matches the base chat transform signature + optional_params: dict, # mutable-ok: matches the base chat transform signature + litellm_params: dict, # mutable-ok: matches the base chat transform signature + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: matches the base chat transform signature + if not any(name.lower() == "x-pplx-integration" for name in headers): + headers["X-Pplx-Integration"] = "litellm" + return super().validate_environment( + headers, model, messages, optional_params, litellm_params, api_key, api_base + ) + def get_supported_openai_params(self, model: str) -> list: """ Perplexity supports a subset of OpenAI params diff --git a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py index 29d185b686d..5e835e44917 100644 --- a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py +++ b/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py @@ -5,12 +5,10 @@ Tests the response transformation to extract citation tokens and search queries from Perplexity API responses. """ -from unittest.mock import Mock - +import httpx import pytest # Add the project root to Python path - from litellm import ModelResponse from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig from litellm.types.utils import Usage @@ -19,6 +17,28 @@ from litellm.types.utils import Usage class TestPerplexityChatTransformation: """Test suite for Perplexity chat transformation functionality.""" + @pytest.mark.parametrize( + ("headers", "expected"), + [({}, "litellm"), ({"x-pplx-integration": "custom"}, "custom")], + ) + def test_integration_header_default_is_caller_overridable( + self, headers: dict[str, str], expected: str + ): + validated_headers = PerplexityChatConfig().validate_environment( + headers=headers, + model="sonar", + messages=[], + optional_params={}, + litellm_params={}, + ) + request = httpx.Request( + "POST", + "https://api.perplexity.ai/chat/completions", + headers=validated_headers, + ) + + assert request.headers.get_list("x-pplx-integration") == [expected] + def test_enhance_usage_with_citation_tokens(self): """Test extraction of citation tokens from API response.""" config = PerplexityChatConfig()