Final_Assignment / agents.py
BMukhtar's picture
score 35
c79086f
# 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()