Update app.py

app.py

@@ -1,18 +1,19 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
+import re
+from typing import Optional
+from duckduckgo_search import DDGS
 from tools.final_answer import FinalAnswerTool
-
 from Gradio_UI import GradioUI
 
-#my first try
 @tool
-def dog_average_longevity(breed: str)-> str: 
+def dog_average_longevity(breed: str) -> str: 
     """A tool that uses DuckDuckGo to search the web and returns the average lifespan for a breed of dogs. 
         Args: 
-           breed(str) : a string representing the breed of the dog (e.g., "Labrador").    
+           breed(str): A string representing the breed of the dog (e.g., "Labrador").    
     """
     query = f"average lifespan for {breed} dog in years"
     
@@ -31,28 +32,25 @@ def extract_lifespan(text: str) -> Optional[str]:
     """
     Extracts the first occurrence of a lifespan range or single number in years.
     
-    arg1:  text: The input text from search results.
-    
-    return: Extracted lifespan range or single value (e.g., '10-12 years' or '13 years') or None.
+    :param text: The input text from search results.
+    :return: Extracted lifespan range or single value (e.g., '10-12 years' or '13 years') or None.
     """
-    pattern = r"(\d{1,2}(-\d{1,2})?\s?(years|yrs)?)"
+    pattern = r"(\d{1,2}\s?(-|to)\s?\d{1,2}|\d{1,2})\s?(years|yrs)?"
     match = re.search(pattern, text, re.IGNORECASE)
     
     if match:
         return match.group(1)
     return None   
 
-
 final_answer = FinalAnswerTool()
 model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',# it is possible that this model may be overloaded
-custom_role_conversions=None,
+    max_tokens=2096,
+    temperature=0.5,
+    model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',
+    custom_role_conversions=None,
 )
 
-
-# Import tool from Hub
+# Import tool from Hugging Face Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
 with open("prompts.yaml", 'r') as stream:
@@ -60,7 +58,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, dog_average_longevity], ## add your tools here (don't remove final answer)
+    tools=[final_answer, dog_average_longevity, image_generation_tool],  # Now includes image generation
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -70,5 +68,4 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
-
-GradioUI(agent).launch()
+GradioUI(agent).launch()