agents.py

# agents.py
from pydantic import BaseModel
from typing import Dict, Any, List, Callable
from config import llm_config
from system_messages import USER_PROXY_MSG, DATA_SOURCE_MANAGER_MSG, DATA_QUALITY_AGENT_MSG, STATISTICAL_AGENT_MSG, QUALITATIVE_AGENT_MSG, VISUALIZATION_AGENT_MSG, REPORT_GENERATION_MSG
from utils import download_dataset_via_git, get_dataset_file_paths, assess_data_quality, perform_statistical_analysis, perform_qualitative_analysis, create_visualizations, generate_pdf_report

class DataAnalysisAgent(BaseModel):
    name: str
    system_message: str
    llm_config: Dict[str, Any]
    functions: List[Callable] = []
    
    class Config:
        arbitrary_types_allowed = True

    def model_dump(self) -> Dict[str, Any]:
        return {
            'name': self.name,
            'system_message': self.system_message,
            'llm_config': self.llm_config
        }

    def process_message(self, message: str, llm_config: Dict[str, Any]) -> Dict[str, Any]:
        """Process a message and return the response"""
        print(f"\n{self.name} processing message: {message}")
        
        try:
            # For functions assigned to this agent, try to execute them
            if self.functions:
                for func in self.functions:
                    try:
                        print(f"Executing function: {func.__name__}")
                        
                        # Special handling for report generation
                        if func.__name__ == 'generate_pdf_report' and isinstance(message, dict):
                            result = func(
                                message.get('analysis_results', {}),
                                message.get('visualization_paths', [])
                            )
                        else:
                            result = func(message)
                            
                        if result:
                            print(f"Function {func.__name__} succeeded with result: {str(result)}")
                            return {
                                "status": "success",
                                "message": result,  # Don't convert to string to preserve types
                                "next_agent": self._determine_next_agent(result)
                            }
                    except Exception as e:
                        print(f"Error executing function {func.__name__}: {str(e)}")
                        continue
            
            # If no functions succeeded or no functions assigned, process as conversation
            print(f"{self.name} completed processing")
            return {
                "status": "complete",
                "message": f"Processed by {self.name}: {message}",
                "next_agent": self._determine_next_agent(None)
            }
                
        except Exception as e:
            print(f"Error in {self.name}: {str(e)}")
            return {
                "status": "error",
                "message": f"Error in {self.name}: {str(e)}",
                "next_agent": None
            }
        
    def _determine_next_agent(self, result: Any) -> str:
        """Determine the next agent based on the result"""
        print(f"\nCurrent agent: {self.name}")
        print(f"Result type: {type(result)}")
        
        # Define the agent sequence
        agent_sequence = {
            "user_proxy": "data_source_manager",
            "data_source_manager": "data_quality_agent",
            "data_quality_agent": "statistical_agent",
            "statistical_agent": "qualitative_agent",
            "qualitative_agent": "visualization_agent",
            "visualization_agent": "report_generation_agent",
            "report_generation_agent": None
        }
        
        next_agent = agent_sequence.get(self.name)
        print(f"Moving to next agent: {next_agent}")
        return next_agent

def create_agents():
    # User Proxy Agent
    user_proxy = DataAnalysisAgent(
        name="user_proxy",
        system_message=USER_PROXY_MSG,
        llm_config=llm_config,
        functions=[download_dataset_via_git]
    )

    # Data Source Manager Agent
    data_source_manager = DataAnalysisAgent(
        name="data_source_manager",
        llm_config=llm_config,
        system_message=DATA_SOURCE_MANAGER_MSG,
        functions=[get_dataset_file_paths]
    )

    # Data Quality Assessment Agent
    data_quality_agent = DataAnalysisAgent(
        name="data_quality_agent",
        llm_config=llm_config,
        system_message=DATA_QUALITY_AGENT_MSG,
        functions=[assess_data_quality]
    )

    # Statistical Analysis Agent
    statistical_agent = DataAnalysisAgent(
        name="statistical_agent",
        llm_config=llm_config,
        system_message=STATISTICAL_AGENT_MSG,
        functions=[perform_statistical_analysis]
    )

    # Qualitative Analysis Agent
    qualitative_agent = DataAnalysisAgent(
        name="qualitative_agent",
        llm_config=llm_config,
        system_message=QUALITATIVE_AGENT_MSG,
        functions=[perform_qualitative_analysis]
    )

    # Visualization Agent
    visualization_agent = DataAnalysisAgent(
        name="visualization_agent",
        llm_config=llm_config,
        system_message=VISUALIZATION_AGENT_MSG,
        functions=[create_visualizations]
    )

    # Report Generation Agent
    report_generation_agent = DataAnalysisAgent(
        name="report_generation_agent",
        llm_config=llm_config,
        system_message=REPORT_GENERATION_MSG,
        functions=[generate_pdf_report]
    )

    # No need to register functions - they are passed in the constructor

    return {
        'user_proxy': user_proxy,
        'data_source_manager': data_source_manager,
        'data_quality_agent': data_quality_agent,
        'statistical_agent': statistical_agent,
        'qualitative_agent': qualitative_agent,
        'visualization_agent': visualization_agent,
        'report_generation_agent': report_generation_agent
    }