ShayanP's picture
Update app.py
0606701 verified
Raw History Blame Contribute Delete
2.59 kB
from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml
import logging
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@tool
def my_custom_tool(arg1: int, arg2: int) -> int:
"""Add two numbers.
Args:
arg1: The first number.
arg2: The second number.
"""
return arg1 + arg2
@tool
def get_current_time_in_timezone(timezone: str) -> str:
"""Fetch the current local time in a specified timezone.
Args:
timezone: A valid timezone, such as Asia/Dubai or America/New_York.
"""
try:
tz = pytz.timezone(timezone)
local_time = datetime.datetime.now(tz).strftime(
"%Y-%m-%d %H:%M:%S %Z"
)
return f"The current local time in {timezone} is: {local_time}"
except pytz.exceptions.UnknownTimeZoneError:
return (
f"Invalid timezone: {timezone}. "
"Try a valid timezone such as Asia/Dubai."
)
except Exception as exc:
logger.exception("Failed to fetch timezone")
return f"Error fetching time: {exc}"
def main():
# Configure the model
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
custom_role_conversions=None,
)
# Load the image tool
image_generation_tool = load_tool(
"agents-course/text-to-image",
trust_remote_code=True,
)
# Load prompt templates
with open("prompts.yaml", "r", encoding="utf-8") as stream:
prompt_templates = yaml.safe_load(stream)
if not isinstance(prompt_templates, dict):
raise ValueError("prompts.yaml must contain a YAML mapping.")
final_answer = FinalAnswerTool()
# Initialize the agent with all tools
agent = CodeAgent(
model=model,
tools=[
get_current_time_in_timezone,
my_custom_tool,
image_generation_tool,
final_answer,
],
max_steps=6,
verbosity_level=1,
grammar=None,
planning_interval=None,
name="AI Assistant",
description=(
"An AI assistant that can calculate, retrieve local time, "
"and generate images."
),
prompt_templates=prompt_templates,
)
# Do not pass share=True here; your Gradio_UI.py
# already sets share=True internally.
GradioUI(agent).launch()
if __name__ == "__main__":
main()