Update app.py

app.py

@@ -1,66 +1,89 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel, load_tool, tool
+from smolagents import CodeAgent, InferenceClientModel, tool, load_tool
 import datetime
-import requests
 import pytz
 import yaml
-from tools.final_answer import FinalAnswerTool
+
+# IMPORTANT: make sure this file exists or remove it
+try:
+    from tools.final_answer import FinalAnswerTool
+except ImportError:
+    class FinalAnswerTool:
+        def __call__(self, answer: str):
+            return answer
 
 from Gradio_UI import GradioUI
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity!
+
+# ------------------ TOOLS ------------------ #
 @tool
-def my_custom_tool(arg1:str, arg2:int)-> str: # it's important to specify the return type
-    # Keep this format for the tool description / args description but feel free to modify the tool
+def my_custom_tool(arg1: str, arg2: int) -> str:
     """A tool that does nothing yet 
     Args:
         arg1: the first argument
         arg2: the second argument
     """
-    return "What magic will you build ?"
+    return f"Received {arg1} and {arg2}"
+
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
-    """A tool that fetches the current local time in a specified timezone.
-    Args:
-        timezone: A string representing a valid timezone (e.g., 'America/New_York').
-    """
+    """Fetch current time in a timezone"""
     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}"
+        return f"{timezone}: {local_time}"
     except Exception as e:
-        return f"Error fetching time for timezone '{timezone}': {str(e)}"
+        return f"Error: {str(e)}"
 
 
-final_answer = FinalAnswerTool()
+# ------------------ MODEL ------------------ #
 model = InferenceClientModel(
-    max_tokens=2096,
+    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+    max_tokens=1024,  # safer for Spaces
     temperature=0.5,
-    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
-    custom_role_conversions=None,
 )
 
 
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# ------------------ LOAD PROMPTS ------------------ #
+try:
+    with open("prompts.yaml", "r") as f:
+        prompt_templates = yaml.safe_load(f)
+except FileNotFoundError:
+    prompt_templates = None
+
 
-# Load system prompt from prompt.yaml file
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
-    
+# ------------------ TOOLS ------------------ #
+final_answer = FinalAnswerTool()
+
+tools = [
+    final_answer,
+    my_custom_tool,
+    get_current_time_in_timezone,
+]
+
+# Optional (only if you actually use it)
+try:
+    image_generation_tool = load_tool(
+        "agents-course/text-to-image",
+        trust_remote_code=True
+    )
+    tools.append(image_generation_tool)
+except Exception:
+    pass
+
+
+# ------------------ AGENT ------------------ #
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], # add your tools here (don't remove final_answer)
-    max_steps=6,
+    tools=tools,
+    max_steps=4,  # lower = safer in Spaces
     verbosity_level=1,
-    grammar=None,
-    planning_interval=None,
-    name=None,
-    description=None,
-    prompt_templates=prompt_templates # Pass system prompt to CodeAgent
+    prompt_templates=prompt_templates,
 )
 
 
-GradioUI(agent).launch()
+# ------------------ UI ------------------ #
+app = GradioUI(agent)
+
+# IMPORTANT for Spaces
+app.launch(server_name="0.0.0.0", server_port=7860)