trax/tests/test_research_agent.py

163 lines
6.4 KiB
Python

"""Unit tests for the Perplexity research agent using OpenRouter."""
import pytest
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
from datetime import datetime, timezone
from typing import Dict, Any
from src.services.protocols import ResearchQuery, ResearchResult
from src.services.research.service import OpenRouterResearchService
from src.services.research.config import ResearchConfig
class TestPerplexityResearchAgent:
"""Test suite for Perplexity sonar-reasoning-pro research agent."""
@pytest.fixture
def mock_openrouter_service(self):
"""Create a mock OpenRouter service."""
service = AsyncMock(spec=OpenRouterResearchService)
service.research = AsyncMock()
service.get_available_models = MagicMock(return_value=[
"perplexity/sonar-reasoning-pro",
"perplexity/sonar-medium-online",
"openai/gpt-4"
])
return service
@pytest.fixture
def sample_sonar_query(self):
"""Create a sample query for sonar-reasoning-pro."""
return ResearchQuery(
query="What are the latest developments in AI reasoning models?",
context="Focus on models like o1, o3, and reasoning capabilities",
max_tokens=4000,
temperature=0.1,
model="perplexity/sonar-reasoning-pro"
)
@pytest.fixture
def sample_sonar_result(self):
"""Create a sample result from sonar-reasoning-pro."""
return ResearchResult(
query="What are the latest developments in AI reasoning models?",
answer="Recent developments include OpenAI's o1 model which demonstrates...",
sources=[
"https://arxiv.org/abs/2024.12345",
"https://openai.com/blog/o1-reasoning"
],
confidence_score=0.92,
processing_time=3.2,
model_used="perplexity/sonar-reasoning-pro",
token_usage={
"prompt_tokens": 200,
"completion_tokens": 1200,
"total_tokens": 1400
}
)
@pytest.mark.asyncio
async def test_sonar_research_execution(self, mock_openrouter_service, sample_sonar_query, sample_sonar_result):
"""Test sonar-reasoning-pro research execution."""
mock_openrouter_service.research.return_value = sample_sonar_result
result = await mock_openrouter_service.research(sample_sonar_query)
mock_openrouter_service.research.assert_called_once_with(sample_sonar_query)
assert result.model_used == "perplexity/sonar-reasoning-pro"
assert result.confidence_score > 0.9 # Sonar should have high confidence
assert len(result.sources) > 0
def test_sonar_query_validation(self):
"""Test sonar-specific query validation."""
query = ResearchQuery(
query="Test reasoning query",
model="perplexity/sonar-reasoning-pro",
max_tokens=4000,
temperature=0.1
)
assert query.model == "perplexity/sonar-reasoning-pro"
assert query.max_tokens <= 4000 # Sonar limit
assert 0.0 <= query.temperature <= 1.0
def test_sonar_result_quality(self, sample_sonar_result):
"""Test sonar result quality metrics."""
assert sample_sonar_result.confidence_score >= 0.8 # High confidence expected
assert sample_sonar_result.processing_time < 10.0 # Reasonable speed
assert len(sample_sonar_result.sources) >= 1 # Should have sources
assert sample_sonar_result.token_usage["total_tokens"] > 0
@pytest.mark.asyncio
async def test_sonar_error_handling(self, mock_openrouter_service, sample_sonar_query):
"""Test sonar error handling."""
mock_openrouter_service.research.side_effect = Exception("OpenRouter API Error")
with pytest.raises(Exception, match="OpenRouter API Error"):
await mock_openrouter_service.research(sample_sonar_query)
def test_sonar_model_availability(self, mock_openrouter_service):
"""Test sonar model availability."""
models = mock_openrouter_service.get_available_models()
assert "perplexity/sonar-reasoning-pro" in models
assert "perplexity/sonar-medium-online" in models
@pytest.mark.asyncio
async def test_sonar_performance_benchmark(self, mock_openrouter_service, sample_sonar_query):
"""Test sonar performance benchmarks."""
result = ResearchResult(
query=sample_sonar_query.query,
answer="Comprehensive reasoning analysis...",
sources=["https://example.com"],
confidence_score=0.95,
processing_time=2.8,
model_used="perplexity/sonar-reasoning-pro",
token_usage={"total_tokens": 1500}
)
mock_openrouter_service.research.return_value = result
research_result = await mock_openrouter_service.research(sample_sonar_query)
# Performance expectations for sonar-reasoning-pro
assert research_result.processing_time < 5.0 # Should be fast
assert research_result.confidence_score >= 0.8 # High confidence
assert research_result.token_usage["total_tokens"] < 2000 # Reasonable token usage
class TestOpenRouterIntegration:
"""Integration tests for OpenRouter with Perplexity."""
@pytest.fixture
def openrouter_config(self):
"""Create OpenRouter configuration."""
return ResearchConfig(
api_key="test-openrouter-key",
base_url="https://openrouter.ai/api/v1",
default_model="perplexity/sonar-reasoning-pro",
max_tokens=4000,
temperature=0.1
)
def test_openrouter_config_validation(self, openrouter_config):
"""Test OpenRouter configuration."""
assert openrouter_config.api_key == "test-openrouter-key"
assert openrouter_config.default_model == "perplexity/sonar-reasoning-pro"
openrouter_config.validate()
@pytest.mark.asyncio
async def test_openrouter_service_initialization(self, openrouter_config):
"""Test OpenRouter service initialization."""
service = OpenRouterResearchService(openrouter_config)
assert hasattr(service, 'research')
assert hasattr(service, 'get_available_models')
assert callable(service.research)
assert callable(service.get_available_models)
if __name__ == "__main__":
pytest.main([__file__, "-v"])