Update app.py

app.py

@@ -8,87 +8,99 @@ from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 import os
 
-# ----- Custom Tools -----
+# ------------------------
+# Custom tools
+# ------------------------
 @tool
 def my_custom_tool(arg1: str, arg2: int) -> str:
     """
-    Simple demo tool.
-    
+    A demo custom tool.
+
     Args:
         arg1: The first argument
         arg2: The second argument
+
+    Returns:
+        A string summarizing the inputs.
     """
     return f"Your inputs were arg1={arg1} and arg2={arg2}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """
-    Returns current local time in a given timezone.
-    
-    Args:
-        timezone: A valid timezone string.
+    Returns the current local time in a specified timezone.
     """
     try:
         tz = pytz.timezone(timezone)
-        now = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
-        return f"The current local time in {timezone} is {now}."
+        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
+        return f"The current local time in {timezone} is {local_time}."
     except Exception as e:
-        return f"Error with timezone '{timezone}': {str(e)}"
+        return f"Error fetching time for '{timezone}': {str(e)}"
 
-# ----- Final Answer -----
+# ------------------------
+# Final answer tool
+# ------------------------
 final_answer = FinalAnswerTool()
 
-# ----- Model Setup -----
+# ------------------------
+# Agent model setup
+# ------------------------
 model = HfApiModel(
     model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
     max_tokens=2096,
     temperature=0.5
 )
 
-# ----- Image Generation via HF Inference API -----
+# ------------------------
+# Image generation using HF Inference API
+# ------------------------
 @tool
 def generate_image(prompt: str) -> AgentImage:
     """
-    Generates a photorealistic image using the HF Inference API.
-    
-    Args:
-        prompt: text prompt to describe the image.
-    
-    Returns:
-        AgentImage that Gradio can render.
+    Generates an image from a text prompt using SDXL (Stable Diffusion XL).
     """
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
-        raise ValueError("HF_TOKEN not set in Space settings!")
+        raise ValueError("HF_TOKEN not set in Space secrets")
 
-    # example HF image endpoint
-    api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1"
+    # SDXL model
+    api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
     headers = {"Authorization": f"Bearer {hf_token}"}
+
+    # Request body for image generation
     payload = {"inputs": prompt}
 
-    resp = requests.post(api_url, headers=headers, json=payload)
-    resp.raise_for_status()
+    response = requests.post(api_url, headers=headers, json=payload)
+    response.raise_for_status()
+
+    # The response will be raw image bytes (PNG)
+    image_bytes = response.content
 
-    image_bytes = resp.content
     return AgentImage.from_bytes(image_bytes)
 
-# ----- Load Prompt Templates -----
+# ------------------------
+# Load prompt templates
+# ------------------------
 with open("prompts.yaml", "r") as f:
     prompt_templates = yaml.safe_load(f)
 
-# ----- Create Agent -----
+# ------------------------
+# Create the agent
+# ------------------------
 agent = CodeAgent(
     model=model,
     tools=[
         final_answer,
         generate_image,
         get_current_time_in_timezone,
-        my_custom_tool,
+        my_custom_tool
     ],
     max_steps=6,
     verbosity_level=1,
-    prompt_templates=prompt_templates,
+    prompt_templates=prompt_templates
 )
 
-# ----- Launch UI -----
+# ------------------------
+# Start the UI
+# ------------------------
 GradioUI(agent).launch()