feat(perplexity): overhaul integration with official SDK and Search API (#34412)

libs/partners/perplexity/tests/integration_tests/test_chat_models.py

@@ -0,0 +1,125 @@
+"""Integration tests for ChatPerplexity."""
+
+import os
+
+import pytest
+from langchain_core.messages import HumanMessage
+
+from langchain_perplexity import ChatPerplexity, MediaResponse, WebSearchOptions
+
+
+@pytest.mark.skipif(not os.environ.get("PPLX_API_KEY"), reason="PPLX_API_KEY not set")
+class TestChatPerplexityIntegration:
+    def test_standard_generation(self) -> None:
+        """Test standard generation."""
+        chat = ChatPerplexity(model="sonar", temperature=0)
+        message = HumanMessage(content="Hello! How are you?")
+        response = chat.invoke([message])
+        assert response.content
+        assert isinstance(response.content, str)
+
+    async def test_async_generation(self) -> None:
+        """Test async generation."""
+        chat = ChatPerplexity(model="sonar", temperature=0)
+        message = HumanMessage(content="Hello! How are you?")
+        response = await chat.ainvoke([message])
+        assert response.content
+        assert isinstance(response.content, str)
+
+    def test_pro_search(self) -> None:
+        """Test Pro Search (reasoning_steps extraction)."""
+        # Pro search is available on sonar-pro
+        chat = ChatPerplexity(
+            model="sonar-pro",
+            temperature=0,
+            web_search_options=WebSearchOptions(search_type="pro"),
+            streaming=True,
+        )
+        message = HumanMessage(content="Who won the 2024 US election and why?")
+
+        # We need to collect chunks to check reasoning steps
+        chunks = list(chat.stream([message]))
+        full_content = "".join(c.content for c in chunks if isinstance(c.content, str))
+        assert full_content
+
+        # Check if any chunk has reasoning_steps
+        has_reasoning = any("reasoning_steps" in c.additional_kwargs for c in chunks)
+        if has_reasoning:
+            assert True
+        else:
+            # Fallback assertion if no reasoning steps returned
+            assert len(chunks) > 0
+
+    async def test_streaming(self) -> None:
+        """Test streaming."""
+        chat = ChatPerplexity(model="sonar", temperature=0)
+        message = HumanMessage(content="Count to 5")
+        async for chunk in chat.astream([message]):
+            assert isinstance(chunk.content, str)
+
+    def test_citations_and_search_results(self) -> None:
+        """Test that citations and search results are returned."""
+        chat = ChatPerplexity(model="sonar", temperature=0)
+        message = HumanMessage(content="Who is the CEO of OpenAI?")
+        response = chat.invoke([message])
+
+        # Citations are usually in additional_kwargs
+        assert "citations" in response.additional_kwargs
+        # Search results might be there too
+        # Note: presence depends on whether search was performed
+        if response.additional_kwargs.get("citations"):
+            assert len(response.additional_kwargs["citations"]) > 0
+
+    def test_search_control(self) -> None:
+        """Test search control parameters."""
+        # Test disabled search (should complete without citations)
+        chat = ChatPerplexity(model="sonar", disable_search=True)
+        message = HumanMessage(content="What is 2+2?")
+        response = chat.invoke([message])
+        assert response.content
+
+        # Test search classifier
+        chat_classifier = ChatPerplexity(model="sonar", enable_search_classifier=True)
+        response_classifier = chat_classifier.invoke([message])
+        assert response_classifier.content
+
+    def test_search_recency_filter(self) -> None:
+        """Test search_recency_filter parameter."""
+        chat = ChatPerplexity(model="sonar", search_recency_filter="month")
+        message = HumanMessage(content="Latest AI news")
+        response = chat.invoke([message])
+        assert response.content
+
+    def test_search_domain_filter(self) -> None:
+        """Test search_domain_filter parameter."""
+        chat = ChatPerplexity(model="sonar", search_domain_filter=["wikipedia.org"])
+        message = HumanMessage(content="Python programming language")
+        response = chat.invoke([message])
+
+        # Verify citations come from wikipedia if any
+        if citations := response.additional_kwargs.get("citations"):
+            assert any("wikipedia.org" in c for c in citations)
+
+    def test_media_and_metadata(self) -> None:
+        """Test related questions and images."""
+        chat = ChatPerplexity(
+            model="sonar-pro",
+            return_related_questions=True,
+            return_images=True,
+            # Media response overrides for video
+            media_response=MediaResponse(overrides={"return_videos": True}),
+        )
+        message = HumanMessage(content="Apollo 11 moon landing")
+        response = chat.invoke([message])
+
+        # Check related questions
+        if related := response.additional_kwargs.get("related_questions"):
+            assert len(related) > 0
+
+        # Check images
+        if images := response.additional_kwargs.get("images"):
+            assert len(images) > 0
+
+        # Check videos (might not always be present but structure should handle it)
+        if videos := response.additional_kwargs.get("videos"):
+            assert len(videos) > 0