Fix gradio UI and add image saving tool and .gitignore

.gitattributes

@@ -33,3 +33,14 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+# Images (store in Git LFS)
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.tif filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+*.ico filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,35 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual environments
+venv/
+.venv/
+env/
+.env/
+
+# Packaging/build
+build/
+dist/
+*.egg-info/
+
+# Tooling caches
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+coverage.xml
+
+# OS/editor
+.DS_Store
+Thumbs.db
+.idea/
+.vscode/
+
+# Gradio runtime artifacts
+.gradio/
+
+# App outputs
+generated_images/
+uploads/

Gradio_UI.py

@@ -25,7 +25,7 @@ from typing import Optional
 import requests
 from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
 from smolagents.agents import ActionStep, MultiStepAgent
-from smolagents.memory import MemoryStep
+from smolagents.memory import FinalAnswerStep, MemoryStep
 from smolagents.utils import _is_package_available
 
 
@@ -215,10 +215,28 @@ def stream_to_gradio(agent, task: str, reset_agent_memory: bool = False, additio
         for message in pull_messages_from_step(step_log):
             yield message
 
-    final_answer = handle_agent_output_types(step_log)
+    raw_final_answer = step_log.final_answer if isinstance(step_log, FinalAnswerStep) else step_log
+
+    # If a tool returns a local image path (e.g. via `save_image`), render it inline in the chat.
+    if isinstance(raw_final_answer, str):
+        candidate_path = raw_final_answer.strip()
+        if candidate_path and os.path.exists(candidate_path):
+            mime_type, _ = mimetypes.guess_type(candidate_path)
+            if mime_type and mime_type.startswith("image/"):
+                yield gr.ChatMessage(role="assistant", content={"path": candidate_path, "mime_type": mime_type})
+                return
+
+    final_answer = handle_agent_output_types(raw_final_answer)
 
     if isinstance(final_answer, AgentText):
-        yield gr.ChatMessage(role="assistant", content=f"**Final answer:**\n{final_answer.to_string()}\n")
+        # If the text is actually a local image path, render the image.
+        text = final_answer.to_string().strip()
+        if text and os.path.exists(text):
+            mime_type, _ = mimetypes.guess_type(text)
+            if mime_type and mime_type.startswith("image/"):
+                yield gr.ChatMessage(role="assistant", content={"path": text, "mime_type": mime_type})
+                return
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:**\n{text}\n")
 
     elif isinstance(final_answer, AgentImage):
         img_path = _save_agent_image(final_answer)
@@ -314,7 +332,6 @@ class GradioUI:
             file_uploads_log = gr.State([])
             chatbot = gr.Chatbot(
                 label="Agent",
-                type="messages",
                 avatar_images=(
                     None,
                     "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",

app.py

@@ -1,9 +1,11 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
+from tools.save_image import SaveImageTool
+from tools.web_search import WebSearchTool
 
 from Gradio_UI import GradioUI
 
@@ -15,7 +17,7 @@ def name_meaning(name: str) -> str:
         name: A name to look up.
     """
     return (
-        f"CALL DuckDuckGoSearchTool with query: '{name} name meaning origin'. "
+        f"CALL web_search with query: '{name} name meaning origin'. "
         f"Prefer sources like BehindTheName, Nameberry, Oxford Reference, Britannica. "
         f"Return: origin, meaning, variants, and 1-2 links."
     )
@@ -76,43 +78,51 @@ def get_weather(city: str) -> str:
         return f"Weather lookup failed: {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 = 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,
-)
-
-
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
-    
-agent = CodeAgent(
-    model=model,
-    tools=[
-        final_answer,
-        DuckDuckGoSearchTool(),     # Web search
-        image_generation_tool,      # HF Hub tool
-        name_meaning,             # Your custom tool
-        get_current_time_in_timezone,
-        get_weather                 # Extra creative tool!
-    ],
-    max_steps=6,
-    verbosity_level=1,
-    grammar=None,
-    planning_interval=None,
-    name=None,
-    description=None,
-    prompt_templates=prompt_templates
-)
-
-
-GradioUI(agent).launch()
+def build_agent() -> CodeAgent:
+    final_answer = FinalAnswerTool()
+    save_image = SaveImageTool()
+    web_search = WebSearchTool()
+
+    # 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 = 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,
+    )
+
+    # Import tool from Hub
+    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+    with open("prompts.yaml", "r") as stream:
+        prompt_templates = yaml.safe_load(stream)
+
+    return CodeAgent(
+        model=model,
+        tools=[
+            final_answer,
+            save_image,
+            web_search,  # Web search
+            image_generation_tool,  # HF Hub tool
+            name_meaning,  # Your custom tool
+            get_current_time_in_timezone,
+            get_weather,  # Extra creative tool!
+        ],
+        max_steps=6,
+        verbosity_level=1,
+        grammar=None,
+        planning_interval=None,
+        name=None,
+        description=None,
+        prompt_templates=prompt_templates,
+    )
+
+
+def main():
+    agent = build_agent()
+    GradioUI(agent).launch()
+
+
+if __name__ == "__main__":
+    main()

prompts.yaml

@@ -24,12 +24,8 @@
   Thought: I will now generate an image showcasing the oldest person, save it locally, and return the path.
   Code:
   ```py
-  import os
-  os.makedirs("generated_images", exist_ok=True)
-
   image = image_generator(prompt="A portrait of John Doe, a 55-year-old man living in Canada.")
-  path = "generated_images/john_doe.png"
-  image.save(path)
+  path = save_image(image=image, filename="john_doe.png")
 
   final_answer(path)
   ```<end_code>

tools/save_image.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import re
+import shutil
+import uuid
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Optional
+
+from smolagents.tools import Tool
+
+
+class SaveImageTool(Tool):
+    name = "save_image"
+    description = (
+        "Save an image to the local `generated_images/` folder and return the saved file path. "
+        "Use this instead of importing `os` for filesystem operations."
+    )
+    inputs = {
+        "image": {"type": "any", "description": "An image object (PIL, AgentImage), bytes, or an existing file path."},
+        "filename": {
+            "type": "string",
+            "description": "Optional output filename (e.g. `cat.png`). Defaults to a random `.png` name.",
+            "nullable": True,
+        },
+    }
+    output_type = "string"
+
+    def forward(self, image: Any, filename: Optional[str] = None) -> str:
+        base_dir = Path("generated_images")
+        base_dir.mkdir(parents=True, exist_ok=True)
+        base_dir_resolved = base_dir.resolve()
+
+        safe_name = self._sanitize_filename(filename) if filename else f"image_{uuid.uuid4().hex[:8]}.png"
+        out_path = (base_dir / safe_name).resolve()
+        if not out_path.is_relative_to(base_dir_resolved):
+            raise ValueError("Refusing to write outside `generated_images/`.")
+
+        # If `image` is already a path on disk, just copy it.
+        if isinstance(image, (str, Path)):
+            src = Path(image)
+            if src.exists() and src.is_file():
+                shutil.copyfile(src, out_path)
+                return str(out_path)
+
+        pil_img = self._to_pil(image)
+        pil_img.save(out_path)
+        return str(out_path)
+
+    @staticmethod
+    def _sanitize_filename(filename: str) -> str:
+        name = Path(filename).name  # drop any path parts
+        name = re.sub(r"[^A-Za-z0-9._-]", "_", name).strip("._")
+        if not name:
+            name = f"image_{uuid.uuid4().hex[:8]}.png"
+        if "." not in name:
+            name += ".png"
+        return name
+
+    @staticmethod
+    def _to_pil(image: Any):
+        from PIL import Image
+
+        if hasattr(image, "save"):
+            return image
+
+        if hasattr(image, "to_pil"):
+            pil = image.to_pil()
+            if pil is not None and hasattr(pil, "save"):
+                return pil
+
+        if isinstance(image, (bytes, bytearray)):
+            return Image.open(BytesIO(image))
+
+        if hasattr(image, "to_string"):
+            as_str = image.to_string()
+            if isinstance(as_str, str):
+                p = Path(as_str)
+                if p.exists() and p.is_file():
+                    return Image.open(p)
+
+        raise TypeError(f"Unsupported image type for saving: {type(image)}")
+

tools/web_search.py

@@ -1,27 +1,44 @@
+from __future__ import annotations
+
 from typing import Any, Optional
+
 from smolagents.tools import Tool
-import duckduckgo_search
 
-class DuckDuckGoSearchTool(Tool):
+
+class WebSearchTool(Tool):
     name = "web_search"
-    description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
-    inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
-    output_type = "string"
+    description = "Search the web with DuckDuckGo and return a short list of results."
+    inputs = {
+        "query": {"type": "string", "description": "Search query."},
+        "max_results": {
+            "type": "integer",
+            "description": "Maximum number of results to return (1-10). Defaults to 5.",
+            "nullable": True,
+        },
+    }
+    output_type = "any"
+
+    def forward(self, query: str, max_results: Optional[int] = 5) -> Any:
+        if not isinstance(query, str) or not query.strip():
+            raise ValueError("`query` must be a non-empty string.")
+
+        limit = 5 if max_results is None else int(max_results)
+        limit = max(1, min(10, limit))
 
-    def __init__(self, max_results=10, **kwargs):
-        super().__init__()
-        self.max_results = max_results
         try:
-            from duckduckgo_search import DDGS
-        except ImportError as e:
-            raise ImportError(
-                "You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`."
-            ) from e
-        self.ddgs = DDGS(**kwargs)
-
-    def forward(self, query: str) -> str:
-        results = self.ddgs.text(query, max_results=self.max_results)
-        if len(results) == 0:
-            raise Exception("No results found! Try a less restrictive/shorter query.")
-        postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
-        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)
+            from ddgs import DDGS  # type: ignore
+        except Exception as e:  # pragma: no cover
+            raise ModuleNotFoundError("Missing dependency: `ddgs` (pip install ddgs).") from e
+
+        results: list[dict[str, str]] = []
+        with DDGS() as ddgs:
+            for r in ddgs.text(query, max_results=limit):
+                title = (r.get("title") or "").strip()
+                url = (r.get("href") or r.get("url") or "").strip()
+                snippet = (r.get("body") or r.get("snippet") or "").strip()
+                if not (title or url or snippet):
+                    continue
+                results.append({"title": title, "url": url, "snippet": snippet})
+
+        return results
+