Update app.py

app.py

@@ -1,69 +1,174 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+import os
 import datetime
-import requests
-import pytz
 import yaml
-from tools.final_answer import FinalAnswerTool
+import pytz
+
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
+from tools.final_answer import FinalAnswerTool  # keep your existing file
+from Gradio_UI import GradioUI  # our UI that unwraps FinalAnswerStep and shows images
 
-from Gradio_UI import GradioUI
+# ---------------------------
+# Tools
+# ---------------------------
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
-def my_custom_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
-    """A tool that does nothing yet 
+def my_custom_tool(arg1: str, arg2: int) -> str:
+    """A toy tool that echoes args.
     Args:
         arg1: the first argument
         arg2: the second argument
     """
-    return "What magic will you build ?"
+    return f"my_custom_tool got: arg1={arg1}, arg2={arg2}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
-    """A tool that fetches the current local time in a specified timezone.
+    """Fetch the current local time in a specified timezone.
     Args:
         timezone: A string representing a valid timezone (e.g., 'America/New_York').
     """
     try:
-        # Create timezone object
         tz = pytz.timezone(timezone)
-        # Get current time in that 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}"
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-
 final_answer = FinalAnswerTool()
 
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
-
+# ---------------------------
+# Model (works on Spaces)
+# ---------------------------
+# If the model is gated or you use a private Inference Endpoint, add a Space secret "HF_TOKEN".
 model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
+    max_tokens=2096,
+    temperature=0.5,
+    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+    custom_role_conversions=None,
 )
 
+# ---------------------------
+# Text-to-image (Hub tool) + wrapper that returns a local file path
+# ---------------------------
+_raw_text_to_image = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+@tool
+def generate_image(prompt: str) -> str:
+    """Generate an image from text and save it locally. Always returns `IMAGE:<abs_path>`.
+    Args:
+        prompt: The image description to generate.
+    """
+    result = _raw_text_to_image(prompt=prompt)
+
+    import base64, io
+    from PIL import Image, ImageDraw
+
+    out_dir = os.path.abspath("generated_images")
+    os.makedirs(out_dir, exist_ok=True)
+    out_path = os.path.join(out_dir, f"img_{int(datetime.datetime.now().timestamp())}.png")
+
+    def _save_pil(img):
+        img.save(out_path)
+        return out_path
 
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+    try:
+        # PIL Image?
+        if "PIL.Image" in str(type(result)):
+            return f"IMAGE:{_save_pil(result)}"
+
+        # dict variants seen in the wild
+        if isinstance(result, dict):
+            for key in ["image", "pil_image"]:
+                if key in result and "PIL.Image" in str(type(result[key])):
+                    return f"IMAGE:{_save_pil(result[key])}"
+            for key in ["image_base64", "b64", "base64"]:
+                if key in result and isinstance(result[key], str):
+                    data = base64.b64decode(result[key], validate=False)
+                    Image.open(io.BytesIO(data)).save(out_path)
+                    return f"IMAGE:{out_path}"
+            for key in ["path", "image_path"]:
+                if key in result and isinstance(result[key], str):
+                    p = result[key]
+                    if os.path.isfile(p):
+                        return f"IMAGE:{os.path.abspath(p)}"
+                    try:
+                        Image.open(p).save(out_path)
+                        return f"IMAGE:{out_path}"
+                    except Exception:
+                        pass
+
+        # raw bytes
+        if isinstance(result, (bytes, bytearray)):
+            Image.open(io.BytesIO(result)).save(out_path)
+            return f"IMAGE:{out_path}"
+
+        # string: file path or base64 or text
+        if isinstance(result, str):
+            if os.path.isfile(result):
+                return f"IMAGE:{os.path.abspath(result)}"
+            try:
+                data = base64.b64decode(result, validate=False)
+                Image.open(io.BytesIO(data)).save(out_path)
+                return f"IMAGE:{out_path}"
+            except Exception:
+                img = Image.new("RGB", (1024, 512), color=(245, 245, 245))
+                d = ImageDraw.Draw(img)
+                d.multiline_text((20, 20), result[:4000], fill=(0, 0, 0))
+                img.save(out_path)
+                return f"IMAGE:{out_path}"
+
+        # ultimate fallback: render the prompt as an image
+        img = Image.new("RGB", (1024, 512), color=(245, 245, 245))
+        d = ImageDraw.Draw(img)
+        d.multiline_text((20, 20), prompt, fill=(0, 0, 0))
+        img.save(out_path)
+        return f"IMAGE:{out_path}"
+
+    except Exception as e:
+        return f"Error generating image: {e}"
+
+search = DuckDuckGoSearchTool()
 
-with open("prompts.yaml", 'r') as stream:
+# ---------------------------
+# Prompts
+# ---------------------------
+with open("prompts.yaml", "r", encoding="utf-8") as stream:
     prompt_templates = yaml.safe_load(stream)
-    
+
+IMAGE_RULE = (
+    "If you generate an image using the `generate_image` tool, "
+    "include the exact returned line `IMAGE:<abs_path>` on a new line in your final answer. "
+    "Always finish by calling final_answer(<plain text> with any IMAGE lines)."
+)
+if isinstance(prompt_templates, dict):
+    if "rules" in prompt_templates and isinstance(prompt_templates["rules"], str):
+        prompt_templates["rules"] += "\n" + IMAGE_RULE
+    else:
+        prompt_templates["rules"] = IMAGE_RULE
+
+# ---------------------------
+# Agent
+# ---------------------------
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
+    tools=[
+        final_answer,
+        get_current_time_in_timezone,
+        my_custom_tool,
+        search,
+        generate_image,
+    ],
     max_steps=6,
     verbosity_level=1,
     grammar=None,
     planning_interval=None,
     name=None,
     description=None,
-    prompt_templates=prompt_templates
+    prompt_templates=prompt_templates,
 )
 
-
-GradioUI(agent).launch()
+# ---------------------------
+# Launch Gradio (Spaces-friendly)
+# ---------------------------
+if __name__ == "__main__":
+    # Gradio will bind correctly on Spaces; no need to set host/port manually.
+    GradioUI(agent).launch()