Update app.py

app.py

@@ -21,7 +21,21 @@ from huggingface_hub import login
 login(hf_token)
 
 
-
+@tool
+def generate_image(prompt: str) -> str:
+    """A tool that generates an image based on the given prompt.
+    
+    Args:
+        prompt: A text description of the image to generate.
+    
+    Returns:
+        The URL or path to the generated image.
+    """
+    try:
+        image_url = image_generation_tool(prompt)
+        return image_url
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
 
 @tool
 def web_search(query:str)-> str: 
@@ -73,7 +87,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, web_search], ## add your tools here (don't remove final answer)
+    tools=[final_answer, web_search, generate_image], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,