Update app.py

app.py

@@ -1,13 +1,11 @@
 from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import os
-import datetime as dt
+import datetime
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 from smolagents.agent_types import AgentImage
-
-# ---------- Timezone tool ----------
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """Get the current local time for a given timezone.
@@ -25,22 +23,8 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {e}"
 
+final_answer = FinalAnswerTool()
 
-# ---------- Final answer wrapper (image-aware) ----------
-class FinalAnswerWithImages(FinalAnswerTool):
-    def __call__(self, value):
-        # If any tool returns a smolagents AgentImage, convert to PIL so Gradio can render it
-        if isinstance(value, AgentImage):
-            try:
-                value = value.to_pil()
-            except Exception:
-                pass
-        return super().__call__(value)
-
-final_answer = FinalAnswerWithImages()
-
-
-# ---------- LLM backend ----------
 model = HfApiModel(
     max_tokens=2096,
     temperature=0.5,
@@ -48,103 +32,33 @@ model = HfApiModel(
     custom_role_conversions=None,
 )
 
-# ---------- Hugging Face token for the Hub tool ----------
 HF_TOKEN = os.getenv("HF_TOKEN")
 assert HF_TOKEN, "HF_TOKEN is not set in the environment."
 
-# Load the remote image generation tool from the Hub
-_hf_image_tool = load_tool(
+image_generation_tool = load_tool(
     "agents-course/text-to-image",
     trust_remote_code=True,
     hf_token=HF_TOKEN,
 )
 
-# ---------- Helpers to normalize return values into a PIL image ----------
-def _to_pil(obj):
-    from PIL import Image
-    import io, base64, os
-
-    # smolagents AgentImage → PIL
-    if isinstance(obj, AgentImage) and hasattr(obj, "to_pil"):
-        return obj.to_pil()
-
-    # Already a PIL image
-    if hasattr(obj, "size") and callable(getattr(obj, "save", None)):
-        return obj  # looks like a PIL.Image.Image
-
-    # Raw bytes → PIL
-    if isinstance(obj, (bytes, bytearray)):
-        return Image.open(io.BytesIO(obj)).convert("RGB")
-
-    # File path → PIL
-    if isinstance(obj, str) and os.path.exists(obj):
-        return Image.open(obj).convert("RGB")
-
-    # base64 string → PIL
-    if isinstance(obj, str):
-        try:
-            return Image.open(io.BytesIO(base64.b64decode(obj))).convert("RGB")
-        except Exception:
-            pass
-
-    # Dict structures occasionally returned by tools (image / images / data)
-    if isinstance(obj, dict):
-        candidates = []
-        if "image" in obj:
-            candidates.append(obj["image"])
-        if "images" in obj and obj["images"]:
-            candidates.append(obj["images"][0])
-        if "data" in obj:
-            candidates.append(obj["data"])
-        for c in candidates:
-            try:
-                return _to_pil(c)
-            except Exception:
-                continue
-
-    raise ValueError(f"Unsupported image output type: {type(obj).__name__}")
-
-
-# ---------- Plain Python function the agent can call in python_interpreter ----------
-# IMPORTANT: no @tool decorator here.
-def generate_image(prompt: str):
-    """
-    Generate an image using the HF Hub tool and return a PIL image that Gradio can display.
-    The agent will call this function from python_interpreter like:
-        img = generate_image("a photorealistic cat")
-        final_answer(img)
-    """
-    # Most Hub tools accept just 'prompt'; passing extra kwargs can fail silently.
-    raw = _hf_image_tool(prompt=prompt)
-    pil = _to_pil(raw)
-    # Guard against empty images
-    if getattr(pil, "size", (0, 0)) == (0, 0):
-        raise RuntimeError("Image generation produced an empty image.")
-    return pil
-
-
-# (Optional) Also expose a tool version if the agent decides to use a tool instead of Python.
-@tool
-def make_image(prompt: str):
-    """Generate an image from text and return a displayable image (PIL)."""
-    return generate_image(prompt)
-
 
-# ---------- Prompts / Agent / UI ----------
+# if isinstance(final_answer, AgentImage):
+#     output = output.to_pil()  # Or output.save("cat.png")
+    
 with open("prompts.yaml", "r") as f:
     prompt_templates = yaml.safe_load(f)
 
 agent = CodeAgent(
     model=model,
     tools=[
-        final_answer,              # required to end turns
-        make_image,                # tool route (if the LLM picks a tool)
+        final_answer,
+        image_generation_tool,
         get_current_time_in_timezone,
         DuckDuckGoSearchTool(),
     ],
     max_steps=6,
-    verbosity_level=2,
+    verbosity_level=2,  # show tool traces
     prompt_templates=prompt_templates,
 )
 
-GradioUI(agent).launch()
+GradioUI(agent).launch()