Feat: Image generation tool

app.py

@@ -66,6 +66,19 @@ def analyze_species_difference(species1: str, species2: str) -> str:
         return f"Based on the research, here's a comparison between {species1} and {species2}"
     else:
         return f"Sorry, I couldn't find information on the physiological differences between {species1} and {species2}. Please, try another query."
+
+@tool
+def generate_starship_image(starship_class: str) -> str:
+    """
+    Generate an image of a Star Trek starship based on the starship class.
+    """
+    prompt = f"Generate a digital image of Star Trek starship of the {starship_class} class."
+
+    try:
+        image_url = image_generation_tool(prompt)
+        return f"Here is the image of a {starship_class} class starship: {image_url}"
+    except Exception as e:
+        return f"Error generating image for {starship_class} class starship: {str(e)}"
         
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -104,7 +117,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_star_trek_quote, analyze_species_difference], ## add your tools here (don't remove final answer)
+    tools=[final_answer, get_star_trek_quote, analyze_species_difference, generate_starship_image], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,