agents.py

"""
Agent definitions for NBA data analysis.
"""
from crewai import Agent
from config import get_llm, NBA_DATA_PATH
from tools import get_agent_tools

# Get LLM
llm = get_llm()


def create_engineer_agent(csv_path: str = None) -> Agent:
    """
    Create the Engineer Agent for data processing and engineering tasks.
    
    Args:
        csv_path: Path to CSV file (defaults to NBA_DATA_PATH from config)
    
    Returns:
        Agent: Configured Engineer Agent
    """
    data_path = csv_path or NBA_DATA_PATH
    agent_tools = get_agent_tools(data_path)
    
    return Agent(
        role="Data Engineer",
        goal="Process, clean, and prepare data for analysis. Ensure data quality and create structured datasets.",
        backstory="""You are an expert data engineer with years of experience in sports analytics. 
        You specialize in processing large datasets, handling missing values, data validation, 
        and creating clean, analysis-ready datasets. You understand statistics deeply and 
        know how to structure data for optimal analysis.""",
        verbose=True,
        allow_delegation=False,
        llm=llm,
        tools=agent_tools,
    )


def create_analyst_agent(csv_path: str = None) -> Agent:
    """
    Create the Analyst Agent for data analysis and insights.
    
    Args:
        csv_path: Path to CSV file (defaults to NBA_DATA_PATH from config)
    
    Returns:
        Agent: Configured Analyst Agent
    """
    data_path = csv_path or NBA_DATA_PATH
    agent_tools = get_agent_tools(data_path)
    
    return Agent(
        role="Data Analyst",
        goal="Analyze data to extract meaningful insights, identify patterns, and provide actionable recommendations.",
        backstory="""You are a seasoned data analyst with a passion for analytics. 
        You excel at finding patterns in data, identifying trends, performing statistical analysis, 
        and translating complex data into clear, actionable insights. You understand performance 
        metrics and can provide strategic recommendations based on data.
        
        CRITICAL: When asked for aggregations, top N lists, totals, or statistical summaries:
        - ALWAYS use the 'analyze_nba_data' tool with pandas groupby operations
        - NEVER use semantic_search_nba_data for aggregation queries (it only returns individual records)
        - For "top 5 three-point shooters": use analyze_nba_data with groupby('Player')['3P'].sum()
        - Plan your analysis: understand what aggregation is needed, then write the appropriate pandas code""",
        verbose=True,
        allow_delegation=False,
        llm=llm,
        tools=agent_tools,
    )


def create_storyteller_agent() -> Agent:
    """
    Create the Storyteller Agent for creating engaging headlines and storylines from analysis results.
    
    Returns:
        Agent: Configured Storyteller Agent
    """
    return Agent(
        role="Sports Storyteller",
        goal="Transform data analysis results into engaging headlines and compelling storylines that bring statistics to life with narrative and context.",
        backstory="""You are a creative sports journalist and storyteller with a talent for turning 
        statistical analysis into captivating headlines and engaging storylines. You know how to make data come alive, 
        creating headlines that grab attention and writing compelling content that tells the story behind the numbers. 
        You understand what makes a great sports story and can transform complex analytics into memorable narratives 
        that connect with readers. You write with flair, using vivid language and storytelling techniques to make 
        statistics human and relatable. Your stories provide context, explain why the data matters, and bring the 
        performance metrics to life with engaging prose.""",
        verbose=True,
        allow_delegation=False,
        llm=llm,
        tools=[],  # Storyteller doesn't need data tools, just creates headlines and content from analysis
    )