Update app.py

app.py

@@ -6,54 +6,68 @@ from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
 
+# Load HF image generation tool
+hf_image_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+
+@tool
+def generate_image(prompt: str) -> str:
+    """
+    Generate an image from a text prompt.
+
+    Args:
+        prompt (str): Description of the image.
+
+    Returns:
+        str: File path of the generated image.
+    """
+    img = hf_image_tool(prompt=prompt)
+    path = "/tmp/generated_image.png"
+    img.save(path)
+    return path
+
+
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """
-    Get the current local time in a specified timezone.
+    Get current time for a timezone.
 
     Args:
-        timezone (str): A valid timezone string such as 'Asia/Kolkata' or 'America/New_York'.
+        timezone (str): Example 'Asia/Kolkata'.
 
     Returns:
-        str: The current local time in the specified timezone.
+        str: Current time in that timezone.
     """
     try:
         tz = pytz.timezone(timezone)
         local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
-        return f"The current local time in {timezone} is: {local_time}"
+        return f"Time in {timezone}: {local_time}"
     except Exception as e:
-        return f"Error fetching time for timezone '{timezone}': {str(e)}"
+        return str(e)
 
 
-# Final answer tool (mandatory for agent)
+# Mandatory final answer tool
 final_answer = FinalAnswerTool()
 
 
-# LLM model for agent reasoning
+# LLM for reasoning
 model = HfApiModel(
     model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
-    max_tokens=2096,
-    temperature=0.5,
-)
-
-
-# Image generation tool from Hugging Face Agents course
-image_generation_tool = load_tool(
-    "agents-course/text-to-image",
-    trust_remote_code=True
+    max_tokens=2048,
+    temperature=0.5
 )
 
 
-# Load prompt templates
-with open("prompts.yaml", "r") as stream:
-    prompt_templates = yaml.safe_load(stream)
+# Prompt templates
+with open("prompts.yaml", "r") as f:
+    prompt_templates = yaml.safe_load(f)
 
 
-# Create agent
+# Agent
 agent = CodeAgent(
     model=model,
     tools=[
-        image_generation_tool,
+        generate_image,
         DuckDuckGoSearchTool(),
         get_current_time_in_timezone,
         final_answer
@@ -61,10 +75,10 @@ agent = CodeAgent(
     max_steps=6,
     verbosity_level=2,
     name="multitool_agent",
-    description="Agent that can search the web, generate images, and fetch timezone information.",
+    description="Agent that can generate images, search the web, and get timezone time.",
     prompt_templates=prompt_templates
 )
 
 
-# Launch Gradio UI
+# Launch UI
 GradioUI(agent).launch()