# import smolagents.models as sm_models

# _orig_roles = sm_models.MessageRole.roles

# @classmethod
# def _roles_with_control(cls):
#     return _orig_roles() + ["control"]

# sm_models.MessageRole.roles = _roles_with_control



from smolagents import (CodeAgent, 
                        GradioUI)
from smolagents.default_tools import (DuckDuckGoSearchTool, 
                                      VisitWebpageTool, 
                                      WikipediaSearchTool, 
                                      SpeechToTextTool,
                                      PythonInterpreterTool)
import yaml
from final_answer import FinalAnswerTool, check_reasoning, ensure_formatting
from tools import (use_vision_model, youtube_frames_to_images,
                         read_file, 
                         extract_text_from_image, analyze_csv_file, 
                         analyze_excel_file, youtube_transcribe,
                         transcribe_audio, review_youtube_video)
import os
from model_provider import create_react_model
import time

# Load prompts from YAML file
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
def get_manager_agent():
    return CodeAgent(
        model=create_react_model(),
        tools=[FinalAnswerTool(), 
            DuckDuckGoSearchTool(), 
            VisitWebpageTool(max_output_length=500000), 
            WikipediaSearchTool(extract_format='HTML'),
            SpeechToTextTool(),
            youtube_transcribe,
            #    use_vision_model,
            #    youtube_frames_to_images,
            #    review_youtube_video,
            read_file, 
            extract_text_from_image, 
            analyze_csv_file, analyze_excel_file,
            transcribe_audio,
            ],
        managed_agents=[],
        additional_authorized_imports=['os', 'pandas', 'numpy', 'PIL', 'tempfile', 'PIL.Image'],
        max_steps=10,
        verbosity_level=1,
        planning_interval=10,
        name="Manager",
        description="The manager of the team, responsible for overseeing and guiding the team's work.",
        final_answer_checks=[
            check_reasoning, 
            ensure_formatting,
        ],
        prompt_templates=prompt_templates
    )

manager_agent = get_manager_agent()



if __name__ == "__main__":
    GradioUI(manager_agent).launch()