Spaces:
Build error
Build error
Enhance agent capabilities by integrating YAML-based prompt templates for web, data analysis, and media agents in agents.py. Update main.py to initialize agents with these templates, improving task handling and response accuracy. Introduce utility functions for extracting final answers and managing prompts, streamlining the overall agent workflow.
5c0be56
unverified
| import importlib | |
| import yaml | |
| from smolagents import CodeAgent | |
| from prompts import ( | |
| DATA_AGENT_SYSTEM_PROMPT, | |
| MEDIA_AGENT_SYSTEM_PROMPT, | |
| WEB_AGENT_SYSTEM_PROMPT, | |
| ) | |
| from tools import ( | |
| analyze_image, | |
| browse_webpage, | |
| extract_dates, | |
| find_in_page, | |
| parse_csv, | |
| perform_calculation, | |
| read_pdf, | |
| web_search, | |
| ) | |
| def create_web_agent(model): | |
| """ | |
| Create a specialized agent for web browsing tasks. | |
| Args: | |
| model: The model to use for the agent | |
| Returns: | |
| Configured CodeAgent for web browsing | |
| """ | |
| prompt_templates = yaml.safe_load( | |
| importlib.resources.files("smolagents.prompts") | |
| .joinpath("code_agent.yaml") | |
| .read_text() | |
| ) | |
| # prompt_templates["system_prompt"] = WEB_AGENT_SYSTEM_PROMPT | |
| web_agent = CodeAgent( | |
| tools=[web_search, browse_webpage, find_in_page, extract_dates], | |
| model=model, | |
| name="web_agent", | |
| description="Specialized agent for web browsing and searching. Use this agent to find information online, browse websites, and extract information from web pages.", | |
| add_base_tools=True, | |
| additional_authorized_imports=["requests", "bs4", "re", "json"], | |
| prompt_templates=prompt_templates, | |
| ) | |
| return web_agent | |
| def create_data_analysis_agent(model): | |
| """ | |
| Create a specialized agent for data analysis tasks. | |
| Args: | |
| model: The model to use for the agent | |
| Returns: | |
| Configured CodeAgent for data analysis | |
| """ | |
| prompt_templates = yaml.safe_load( | |
| importlib.resources.files("smolagents.prompts") | |
| .joinpath("code_agent.yaml") | |
| .read_text() | |
| ) | |
| # prompt_templates["system_prompt"] = DATA_AGENT_SYSTEM_PROMPT | |
| data_agent = CodeAgent( | |
| tools=[parse_csv, perform_calculation], | |
| model=model, | |
| name="data_agent", | |
| description="Specialized agent for data analysis. Use this agent to analyze data, perform calculations, and extract insights from structured data.", | |
| add_base_tools=True, | |
| additional_authorized_imports=["pandas", "numpy", "math", "csv", "io"], | |
| prompt_templates=prompt_templates, | |
| ) | |
| return data_agent | |
| def create_media_agent(model): | |
| """ | |
| Create a specialized agent for handling media (images, PDFs). | |
| Args: | |
| model: The model to use for the agent | |
| Returns: | |
| Configured CodeAgent for media handling | |
| """ | |
| prompt_templates = yaml.safe_load( | |
| importlib.resources.files("smolagents.prompts") | |
| .joinpath("code_agent.yaml") | |
| .read_text() | |
| ) | |
| # prompt_templates["system_prompt"] = MEDIA_AGENT_SYSTEM_PROMPT | |
| media_agent = CodeAgent( | |
| tools=[analyze_image, read_pdf], | |
| model=model, | |
| name="media_agent", | |
| description="Specialized agent for handling media files like images and PDFs. Use this agent to analyze images and extract text from PDF documents.", | |
| add_base_tools=True, | |
| additional_authorized_imports=["PIL", "io", "requests"], | |
| prompt_templates=prompt_templates, | |
| ) | |
| return media_agent | |