trax/examples/research_agent_example.py

#!/usr/bin/env python3
"""Example usage of the Perplexity Research Agent.

This script demonstrates how to use the research agent programmatically
to conduct research using Perplexity's sonar-reasoning-pro model.
"""

import asyncio
import json
from datetime import datetime, timezone
from pathlib import Path

# Import the research agent components
from src.config import config
from src.services.protocols import ResearchQuery, ResearchResult
from src.services.research.service import OpenRouterResearchService
from src.services.research.config import ResearchConfig


async def conduct_research_example():
    """Example of conducting research using the agent."""
    
    # Check for API key
    if not config.OPENROUTER_API_KEY:
        print("❌ OPENROUTER_API_KEY not found in environment")
        print("Please set your OpenRouter API key in the environment")
        return
    
    try:
        # Initialize research service
        print("🔧 Initializing research service...")
        research_config = ResearchConfig.from_env(config.OPENROUTER_API_KEY)
        service = OpenRouterResearchService(research_config)
        
        # Example research queries
        queries = [
            "What are the latest developments in AI reasoning models like o1 and o3?",
            "How do vector databases compare for RAG applications in 2025?",
            "What are the best practices for fine-tuning large language models?"
        ]
        
        print(f"🧠 Conducting research on {len(queries)} topics...")
        print("=" * 60)
        
        results = []
        for i, query_text in enumerate(queries, 1):
            print(f"\n📚 Research #{i}: {query_text}")
            print("-" * 40)
            
            # Create research query
            research_query = ResearchQuery(
                query=query_text,
                context="Focus on recent developments and practical applications",
                max_tokens=4000,
                temperature=0.1,
                model="perplexity/sonar-reasoning-pro"
            )
            
            # Conduct research
            print("🔍 Researching...")
            result = await service.research(research_query)
            results.append(result)
            
            # Display results
            print(f"✅ Completed in {result.processing_time:.2f}s")
            print(f"🎯 Confidence: {result.confidence_score:.1%}")
            print(f"📊 Tokens used: {result.token_usage.get('total_tokens', 'N/A')}")
            print(f"📝 Answer preview: {result.answer[:200]}...")
            print(f"🔗 Sources found: {len(result.sources)}")
            
            if result.sources:
                print("   Sources:")
                for j, source in enumerate(result.sources[:3], 1):  # Show first 3
                    print(f"   {j}. {source}")
                if len(result.sources) > 3:
                    print(f"   ... and {len(result.sources) - 3} more")
        
        # Save results
        save_results(results)
        
        print("\n" + "=" * 60)
        print("🎉 Research completed successfully!")
        print(f"📁 Results saved to: examples/research_results/")
        
    except Exception as e:
        print(f"❌ Research failed: {e}")


def save_results(results: list[ResearchResult]):
    """Save research results to files."""
    
    # Create output directory
    output_dir = Path("examples/research_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    # Save individual results
    for i, result in enumerate(results, 1):
        # JSON format
        json_data = {
            'query': result.query,
            'answer': result.answer,
            'sources': result.sources,
            'confidence_score': result.confidence_score,
            'processing_time': result.processing_time,
            'model_used': result.model_used,
            'token_usage': result.token_usage,
            'timestamp': datetime.now(timezone.utc).isoformat()
        }
        
        json_file = output_dir / f"research_{i:02d}_{timestamp}.json"
        with open(json_file, 'w') as f:
            json.dump(json_data, f, indent=2)
        
        # Markdown format
        md_content = f"""# Research Report #{i}

## Query
{result.query}

## Answer
{result.answer}

## Sources
{chr(10).join(f"- {source}" for source in result.sources) if result.sources else "- Sources integrated in analysis"}

## Metadata
- Model: {result.model_used}
- Processing Time: {result.processing_time:.2f} seconds
- Confidence Score: {result.confidence_score:.1%}
- Tokens Used: {result.token_usage.get('total_tokens', 'N/A')}
- Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
        
        md_file = output_dir / f"research_{i:02d}_{timestamp}.md"
        with open(md_file, 'w') as f:
            f.write(md_content)
    
    # Save combined results
    combined_data = {
        'timestamp': datetime.now(timezone.utc).isoformat(),
        'total_queries': len(results),
        'results': [
            {
                'query': r.query,
                'answer': r.answer,
                'sources': r.sources,
                'confidence_score': r.confidence_score,
                'processing_time': r.processing_time,
                'model_used': r.model_used,
                'token_usage': r.token_usage
            }
            for r in results
        ]
    }
    
    combined_file = output_dir / f"batch_results_{timestamp}.json"
    with open(combined_file, 'w') as f:
        json.dump(combined_data, f, indent=2)


def main():
    """Main entry point."""
    print("🧠 Perplexity Research Agent Example")
    print("=" * 50)
    print("This example demonstrates using the research agent to conduct")
    print("research on multiple topics using Perplexity's sonar-reasoning-pro.")
    print()
    
    # Run the research example
    asyncio.run(conduct_research_example())


if __name__ == "__main__":
    main()