mirror of
https://github.com/hwchase17/langchain.git
synced 2026-03-18 11:07:36 +00:00
126 lines
5.1 KiB
Python
126 lines
5.1 KiB
Python
"""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
|