test_chat_models.py — langchain Source File

Source Code

"""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:
// ... (66 more lines)