Update app.py

app.py

@@ -1,17 +1,13 @@
 """Run hf agent."""
-
 # pylint: disable=line-too-long, broad-exception-caught
 import datetime
-
-
-# import requests
 import pytz
 import yaml
+import re
 from Gradio_UI import GradioUI
 from loguru import logger
 from smolagents import (
     CodeAgent,
-    # DuckDuckGoSearchTool,
     HfApiModel,
     load_tool,
     tool,
@@ -19,32 +15,25 @@ from smolagents import (
 from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 
-
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
+# Custom tool example
 @tool
-def my_cutom_tool(arg1: str, arg2: int) -> str:  # it's import to specify the return type
-    # Keep this format for the description / args / args description but feel free to modify the tool
+def my_custom_tool(arg1: str, arg2: int) -> str:  # it's important to specify the return type
     """
     Run a tool that does nothing yet.
-
     Args:
         arg1: the first argument
         arg2: the second argument
-
     """
     return f"What magic will you build? {arg1=}, {arg2=}"
 
-
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """
     Run a tool that fetches the current local time in a specified timezone.
-
     Args:
         timezone: A string representing a valid timezone (e.g., 'America/New_York').
-
     """
-    logger.debug(f" get_current_time_in_timezone {timezone=}")
+    logger.debug(f"get_current_time_in_timezone {timezone=}")
     try:
         # Create timezone object
         tz = pytz.timezone(timezone)
@@ -55,27 +44,32 @@ def get_current_time_in_timezone(timezone: str) -> str:
         logger.error(e)
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-
+# Initialize tools
 final_answer = FinalAnswerTool()
 visit_webpage = VisitWebpageTool()
-
 model = HfApiModel()
 
-
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
+# Load prompt templates
 with open("prompts.yaml", "r", encoding="utf8") as stream:
     prompt_templates = yaml.safe_load(stream)
 
+# Add the missing 'final_answer' template if it doesn't exist
+if 'final_answer' not in prompt_templates:
+    prompt_templates['final_answer'] = "{{final_answer}}"
+
+# Create the agent
 agent = CodeAgent(
     model=model,
     tools=[
-        my_cutom_tool,
+        my_custom_tool,
         get_current_time_in_timezone,
         final_answer,
         visit_webpage,
-    ],  ## add your tools here (don't remove final answer)
+        image_generation_tool,  # Added the image generation tool
+    ],  # add your tools here
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -85,11 +79,10 @@ agent = CodeAgent(
     prompt_templates=prompt_templates,
 )
 
-# debug = True
-# share = True
+# Launch the Gradio UI
 ssr = True
 GradioUI(agent).launch(
-    # debug=debug,  # already set in GradioUI
-    # share=share,
-    # ssr=ssr,
-)
+    # debug=True,
+    # share=True,
+    ssr=ssr,
+)