src/gaia/tools/perplexity_tool.py

"""
Perplexity API tool for the GAIA agent.

This module provides a tool for querying the Perplexity API with advanced models:
- sonar-reasoning: Balanced model for general reasoning and search tasks
- sonar-reasoning-pro: Enhanced version with stronger analytical capabilities
- sonar-deep-research: Specialized model for in-depth research and complex queries

The tool uses a fixed temperature of 0.1 for maximum accuracy and allows selecting
the appropriate model for different types of tasks.

The tool handles API responses and errors, and extracts content and citations from responses.

Models Reference:
- sonar-reasoning: Balanced accuracy and speed for general queries
- sonar-reasoning-pro: Higher accuracy, deeper analysis, suitable for complex reasoning tasks
- sonar-deep-research: Maximum depth and research capabilities, ideal for scholarly research
"""

import os
import json
import logging
import urllib.request
import urllib.error
from typing import Dict, Any, List, Optional, Union

logger = logging.getLogger("gaia_agent.tools.perplexity")

if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

class PerplexityTool:
    """Tool for querying the Perplexity API."""
    
    # Valid models supported by the Perplexity API
    VALID_MODELS = ["sonar", "sonar-reasoning", "sonar-reasoning-pro", "sonar-deep-research"]
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize the Perplexity API tool.
        
        Args:
            config: Optional configuration dictionary that may include:
                - api_key: Perplexity API key (otherwise read from environment)
                - model: Model to use (default: "sonar-reasoning")
                - temperature: Temperature setting (default: 0.1)
                - timeout: Timeout in seconds (default: 30)
        """
        self.config = config or {}
        
        self.api_key = self.config.get("api_key") or os.getenv("PERPLEXITY_API_KEY")
        if not self.api_key:
            logger.warning("Perplexity API key not found in config or environment variables")
        
        self.api_url = "https://api.perplexity.ai"
        
        # Get model from config, environment variable, or use default
        model = self.config.get("model") or os.getenv("PERPLEXITY_MODEL", "sonar-reasoning")
        if model not in self.VALID_MODELS:
            logger.warning(f"Invalid model '{model}'. Using default 'sonar-reasoning' instead.")
            model = "sonar-reasoning"
        self.model = model
        
        # Set temperature to 0.1 for maximum accuracy unless otherwise specified
        self.temperature = self.config.get("temperature", 0.1)
        self.timeout = self.config.get("timeout", 30)
        
        self.endpoint = f"{self.api_url}/chat/completions"
        
        logger.info(f"Initialized Perplexity tool with model: {self.model}")
        
    
    def query(self,
               question: str,
               max_tokens: int = 500,
               system_prompt: Optional[str] = None,
               model: Optional[str] = None) -> Dict[str, Any]:
        """
        Send a query to the Perplexity API.
        
        Args:
            question: The question to ask
            max_tokens: Maximum number of tokens in the response
            system_prompt: Optional system prompt
            model: Override the model for this specific query
            
        Returns:
            Dictionary containing the query results
            
        Raises:
            Exception: If an error occurs during the query
        """
        
        if not self.api_key:
            raise Exception("Perplexity API key not available")
        
        try:
            # Determine which model to use (either the override or the default)
            use_model = model if model else self.model
            
            # Validate the model
            if use_model not in self.VALID_MODELS:
                logger.warning(f"Invalid model specified: '{use_model}'. Using '{self.model}' instead.")
                use_model = self.model
            
            messages = []
            
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
                
            messages.append({"role": "user", "content": question})
            
            # Model parameter must be included in the request body
            payload = {
                "model": use_model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": self.temperature
            }
            
            logger.info(f"Querying Perplexity API with model: {use_model}")
            
            payload_bytes = json.dumps(payload).encode('utf-8')
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            req = urllib.request.Request(
                self.endpoint,
                data=payload_bytes,
                headers=headers,
                method="POST"
            )
            
            with urllib.request.urlopen(req, timeout=self.timeout) as response:
                response_data = response.read().decode('utf-8')
                api_response = json.loads(response_data)
                
                content = self._extract_content(api_response)
                citations = self._extract_citations(api_response)
                usage = self._get_usage(api_response)
                
                result = {
                    "query": question,
                    "content": content,
                    "citations": citations,
                    "usage": usage,
                    "model": use_model,  # Use the actual model used for the query
                    "raw_response": api_response
                }
                
                
                return result
                
        except urllib.error.HTTPError as e:
            logger.error(f"HTTP Error {e.code} when querying Perplexity API: {e.reason}")
            try:
                error_data = e.read().decode('utf-8')
                logger.error(f"Error details: {error_data}")
            except:
                logger.error("Could not read error details")
            raise Exception(f"Perplexity API HTTP error: {e.code} {e.reason}")
            
        except urllib.error.URLError as e:
            logger.error(f"URL Error when querying Perplexity API: {e.reason}")
            raise Exception(f"Perplexity API connection error: {e.reason}")
            
        except Exception as e:
            logger.error(f"Error querying Perplexity API: {str(e)}")
            raise Exception(f"Perplexity API error: {str(e)}")
    
    def search(self,
               query: str,
               max_tokens: int = 500,
               model: Optional[str] = None,
               depth: str = "standard") -> Dict[str, Any]:
        """
        Search the web using the Perplexity API.
        
        Args:
            query: The search query
            max_tokens: Maximum number of tokens in the response
            model: Optional model override (uses instance model by default)
            depth: Research depth - "standard", "deep", or "comprehensive"
            
        Returns:
            Dictionary containing the search results
            
        Raises:
            Exception: If an error occurs during the search
        """
        # Select the appropriate model based on depth if not explicitly specified
        if not model:
            if depth == "deep":
                model = "sonar-reasoning-pro"
            elif depth == "comprehensive":
                model = "sonar-deep-research"
        
        # Adjust the system prompt based on the requested depth
        if depth == "standard":
            system_prompt = """You are a helpful web search assistant.
            Search the web for accurate and up-to-date information to answer the user's query.
            Provide a comprehensive answer with relevant facts and details.
            Include citations to your sources."""
        elif depth == "deep":
            system_prompt = """You are an advanced web search and analysis assistant.
            Perform a deep search for accurate, detailed, and up-to-date information to answer the user's query.
            Analyze the information critically, identifying key patterns, insights, and connections.
            Provide a thorough answer with comprehensive analysis.
            Include citations to your sources and note the reliability of each source."""
        elif depth == "comprehensive":
            system_prompt = """You are an expert research assistant conducting comprehensive analysis.
            Perform exhaustive research on the user's query, finding detailed information from diverse sources.
            Critically analyze all information, evaluating source credibility and identifying consensus vs. disagreements.
            Synthesize findings into a comprehensive research summary addressing all aspects of the query.
            Include thorough citations to your sources, noting the authority and reliability of each source.
            Identify any gaps in available information or areas where further research would be beneficial."""
        else:
            # Default to standard if depth is not recognized
            system_prompt = """You are a helpful web search assistant.
            Search the web for accurate and up-to-date information to answer the user's query.
            Provide a comprehensive answer with relevant facts and details.
            Include citations to your sources."""
        
        logger.info(f"Performing {depth} search with model: {model or self.model}")
        return self.query(query, max_tokens, system_prompt, model)
    
    def _extract_content(self, response: Dict[str, Any]) -> Optional[str]:
        """
        Extract content from a Perplexity API response.
        
        Args:
            response: The API response dictionary
            
        Returns:
            The extracted content, or None if not found
        """
        if "choices" in response and len(response["choices"]) > 0:
            if "message" in response["choices"][0]:
                return response["choices"][0]["message"].get("content")
        return None
    
    def _extract_citations(self, response: Dict[str, Any]) -> List[str]:
        """
        Extract citations from a Perplexity API response.
        
        Args:
            response: The API response dictionary
            
        Returns:
            List of citation URLs
        """
        return response.get("citations", [])
    
    def _get_usage(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """
        Extract usage information from a Perplexity API response.
        
        Args:
            response: The API response dictionary
            
        Returns:
            Dictionary containing usage information
        """
        return response.get("usage", {})


def create_perplexity_tool(model: str = "sonar-reasoning") -> PerplexityTool:
    """
    Create a Perplexity tool instance with the specified model.
    
    Args:
        model: The Perplexity model to use. Options:
            - "sonar-reasoning": Balanced model for general reasoning and search (default)
            - "sonar-reasoning-pro": Enhanced model with stronger analytical capabilities
            - "sonar-deep-research": Specialized model for in-depth research and complex queries
    
    Returns:
        A configured PerplexityTool instance
    """
    config = {"model": model}
    return PerplexityTool(config)


def create_perplexity_tool_pro() -> PerplexityTool:
    """
    Create a Perplexity tool instance using the sonar-reasoning-pro model.
    
    This model provides enhanced analysis capabilities for complex reasoning tasks.
    
    Returns:
        A PerplexityTool instance configured with the sonar-reasoning-pro model
    """
    return create_perplexity_tool("sonar-reasoning-pro")


def create_perplexity_tool_research() -> PerplexityTool:
    """
    Create a Perplexity tool instance using the sonar-deep-research model.
    
    This model provides maximum depth and research capabilities, ideal for
    scholarly research and complex investigative queries.
    
    Returns:
        A PerplexityTool instance configured with the sonar-deep-research model
    """
    return create_perplexity_tool("sonar-deep-research")


# Additional specialized search methods

def deep_search(query: str, max_tokens: int = 750) -> Dict[str, Any]:
    """
    Perform a deep search using the sonar-reasoning-pro model.
    
    This is a convenience function for performing a search with enhanced depth
    and analysis. It uses the sonar-reasoning-pro model which offers improved
    analytical capabilities compared to the standard model.
    
    Args:
        query: The search query
        max_tokens: Maximum number of tokens in the response (default: 750)
        
    Returns:
        Dictionary containing the search results with enhanced analysis
    """
    tool = create_perplexity_tool_pro()
    return tool.search(query, max_tokens, depth="deep")


def comprehensive_research(query: str, max_tokens: int = 1000) -> Dict[str, Any]:
    """
    Perform comprehensive research using the sonar-deep-research model.
    
    This is a convenience function for performing exhaustive research on complex topics.
    It uses the sonar-deep-research model which provides the maximum depth and
    research capabilities available.
    
    Args:
        query: The research query
        max_tokens: Maximum number of tokens in the response (default: 1000)
        
    Returns:
        Dictionary containing thorough research results with comprehensive analysis
    """
    tool = create_perplexity_tool_research()
    return tool.search(query, max_tokens, depth="comprehensive")