context-engineering-intro/use-cases/agent-factory-with-subagents/examples/main_agent_reference/research_agent.py

"""
Research Agent that uses Brave Search and can invoke Email Agent.
"""

import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass

from pydantic_ai import Agent, RunContext

from .providers import get_llm_model
from .email_agent import email_agent, EmailAgentDependencies
from .tools import search_web_tool

logger = logging.getLogger(__name__)


SYSTEM_PROMPT = """
You are an expert research assistant with the ability to search the web and create email drafts. Your primary goal is to help users find relevant information and communicate findings effectively.

Your capabilities:
1. **Web Search**: Use Brave Search to find current, relevant information on any topic
2. **Email Creation**: Create professional email drafts through Gmail when requested

When conducting research:
- Use specific, targeted search queries
- Analyze search results for relevance and credibility
- Synthesize information from multiple sources
- Provide clear, well-organized summaries
- Include source URLs for reference

When creating emails:
- Use research findings to create informed, professional content
- Adapt tone and detail level to the intended recipient
- Include relevant sources and citations when appropriate
- Ensure emails are clear, concise, and actionable

Always strive to provide accurate, helpful, and actionable information.
"""


@dataclass
class ResearchAgentDependencies:
    """Dependencies for the research agent - only configuration, no tool instances."""
    brave_api_key: str
    gmail_credentials_path: str
    gmail_token_path: str
    session_id: Optional[str] = None


# Initialize the research agent
research_agent = Agent(
    get_llm_model(),
    deps_type=ResearchAgentDependencies,
    system_prompt=SYSTEM_PROMPT
)


@research_agent.tool
async def search_web(
    ctx: RunContext[ResearchAgentDependencies],
    query: str,
    max_results: int = 10
) -> List[Dict[str, Any]]:
    """
    Search the web using Brave Search API.
    
    Args:
        query: Search query
        max_results: Maximum number of results to return (1-20)
    
    Returns:
        List of search results with title, URL, description, and score
    """
    try:        
        # Ensure max_results is within valid range
        max_results = min(max(max_results, 1), 20)
        
        results = await search_web_tool(
            api_key=ctx.deps.brave_api_key,
            query=query,
            count=max_results
        )
        
        logger.info(f"Found {len(results)} results for query: {query}")
        return results
        
    except Exception as e:
        logger.error(f"Web search failed: {e}")
        return [{"error": f"Search failed: {str(e)}"}]


@research_agent.tool
async def create_email_draft(
    ctx: RunContext[ResearchAgentDependencies],
    recipient_email: str,
    subject: str,
    context: str,
    research_summary: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create an email draft based on research context using the Email Agent.
    
    Args:
        recipient_email: Email address of the recipient
        subject: Email subject line
        context: Context or purpose for the email
        research_summary: Optional research findings to include
    
    Returns:
        Dictionary with draft creation results
    """
    try:
        # Prepare the email content prompt
        if research_summary:
            email_prompt = f"""
Create a professional email to {recipient_email} with the subject "{subject}".

Context: {context}

Research Summary:
{research_summary}

Please create a well-structured email that:
1. Has an appropriate greeting
2. Provides clear context
3. Summarizes the key research findings professionally
4. Includes actionable next steps if appropriate
5. Ends with a professional closing

The email should be informative but concise, and maintain a professional yet friendly tone.
"""
        else:
            email_prompt = f"""
Create a professional email to {recipient_email} with the subject "{subject}".

Context: {context}

Please create a well-structured email that addresses the context provided.
"""
        
        # Create dependencies for email agent
        email_deps = EmailAgentDependencies(
            gmail_credentials_path=ctx.deps.gmail_credentials_path,
            gmail_token_path=ctx.deps.gmail_token_path,
            session_id=ctx.deps.session_id
        )
        
        # Run the email agent
        result = await email_agent.run(
            email_prompt,
            deps=email_deps,
            usage=ctx.usage  # Pass usage for token tracking
        )
        
        logger.info(f"Email agent invoked for recipient: {recipient_email}")
        
        return {
            "success": True,
            "agent_response": result.data,
            "recipient": recipient_email,
            "subject": subject,
            "context": context
        }
        
    except Exception as e:
        logger.error(f"Failed to create email draft via Email Agent: {e}")
        return {
            "success": False,
            "error": str(e),
            "recipient": recipient_email,
            "subject": subject
        }


@research_agent.tool
async def summarize_research(
    ctx: RunContext[ResearchAgentDependencies],
    search_results: List[Dict[str, Any]],
    topic: str,
    focus_areas: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create a comprehensive summary of research findings.
    
    Args:
        search_results: List of search result dictionaries
        topic: Main research topic
        focus_areas: Optional specific areas to focus on
    
    Returns:
        Dictionary with research summary
    """
    try:
        if not search_results:
            return {
                "summary": "No search results provided for summarization.",
                "key_points": [],
                "sources": []
            }
        
        # Extract key information
        sources = []
        descriptions = []
        
        for result in search_results:
            if "title" in result and "url" in result:
                sources.append(f"- {result['title']}: {result['url']}")
                if "description" in result:
                    descriptions.append(result["description"])
        
        # Create summary content
        content_summary = "\n".join(descriptions[:5])  # Limit to top 5 descriptions
        sources_list = "\n".join(sources[:10])  # Limit to top 10 sources
        
        focus_text = f"\nSpecific focus areas: {focus_areas}" if focus_areas else ""
        
        summary = f"""
Research Summary: {topic}{focus_text}

Key Findings:
{content_summary}

Sources:
{sources_list}
"""
        
        return {
            "summary": summary,
            "topic": topic,
            "sources_count": len(sources),
            "key_points": descriptions[:5]
        }
        
    except Exception as e:
        logger.error(f"Failed to summarize research: {e}")
        return {
            "summary": f"Failed to summarize research: {str(e)}",
            "key_points": [],
            "sources": []
        }


# Convenience function to create research agent with dependencies
def create_research_agent(
    brave_api_key: str,
    gmail_credentials_path: str,
    gmail_token_path: str,
    session_id: Optional[str] = None
) -> Agent:
    """
    Create a research agent with specified dependencies.
    
    Args:
        brave_api_key: Brave Search API key
        gmail_credentials_path: Path to Gmail credentials.json
        gmail_token_path: Path to Gmail token.json
        session_id: Optional session identifier
        
    Returns:
        Configured research agent
    """
    return research_agent