src/agents/research_agent.py

"""
Research Agent Module
======================
Agent responsible for web search and research tasks.
"""

from dataclasses import dataclass, field
from typing import Any
from .base import BaseAgent

# Import the searcher tool
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from tools.searcher import search as web_search


@dataclass
class ResearchAgent(BaseAgent):
    """
    Agent that performs research/search tasks.
    
    Uses the searcher tool for web search with fallback to simulated results.
    """
    role: str = "research"
    tools: list[str] = field(default_factory=lambda: ["web_search", "document_fetch"])
    use_real_search: bool = True  # Flag to enable/disable real search
    
    async def run(self, input: dict[str, Any]) -> dict[str, Any]:
        """
        Execute a research query.
        
        Args:
            input: Dictionary with 'query' key containing the search query
            
        Returns:
            Dictionary with search results
        """
        query = input.get("query", "")
        self.log(f"Researching: {query}")
        
        # Use real search or simulated based on flag
        if self.use_real_search:
            try:
                search_results = await web_search(query, max_results=5)
            except Exception as e:
                self.log(f"Search failed, using simulation: {e}", level="warning")
                search_results = self._simulate_search(query)
        else:
            search_results = self._simulate_search(query)
        
        return {
            "agent": "research",
            "query": query,
            "results": search_results
        }
    
    def _simulate_search(self, query: str) -> list[dict[str, str]]:
        """
        Generate simulated search results for testing.
        
        Args:
            query: The search query
            
        Returns:
            List of simulated search result dictionaries
        """
        # Generate mock results based on query keywords
        base_results = [
            {
                "title": f"Research findings on {query}",
                "url": f"https://example.com/research/{query.replace(' ', '-')}",
                "snippet": f"This article explores the latest developments in {query}. "
                           f"Key insights and analysis are provided."
            },
            {
                "title": f"A comprehensive guide to {query}",
                "url": f"https://example.org/guide/{query.replace(' ', '-')}",
                "snippet": f"Learn everything you need to know about {query}. "
                           f"Expert opinions and case studies included."
            },
            {
                "title": f"Latest news: {query}",
                "url": f"https://news.example.com/{query.replace(' ', '-')}",
                "snippet": f"Breaking developments in {query}. "
                           f"Stay updated with the most recent information."
            }
        ]
        
        return base_results