Update app.py

app.py

@@ -1,19 +1,21 @@
 from smolagents import CodeAgent, InferenceClientModel, load_tool, tool
 import datetime
 import pytz
+import yaml
 from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
 
+# ────────────────────────────────────────────────
+# Custom Tools
+# ────────────────────────────────────────────────
+
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """Get the current local time in the requested timezone.
 
     Args:
         timezone: IANA timezone name (e.g. 'Asia/Kolkata', 'America/New_York')
-
-    Returns:
-        Human-readable string with current time or error message
     """
     try:
         tz = pytz.timezone(timezone)
@@ -31,6 +33,10 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
     return f"Tool called → arg1 = {arg1!r} | arg2 = {arg2} → What would you like to build?"
 
 
+# ────────────────────────────────────────────────
+# Main
+# ────────────────────────────────────────────────
+
 def main():
     final_answer_tool = FinalAnswerTool()
 
@@ -40,13 +46,34 @@ def main():
         max_tokens=2048,
     )
 
-    # Optional text-to-image tool
+    # Optional text-to-image tool from Hub
     image_tool = None
     try:
         image_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-        print("Text-to-image tool loaded")
+        print("✅ Text-to-image tool loaded from Hub")
+    except Exception as e:
+        print(f"⚠️ Text-to-image tool failed to load: {e}")
+
+    # ── Load your custom prompts.yaml + FIX the missing final_answer key ──
+    prompt_templates = None
+    try:
+        with open("prompts.yaml", encoding="utf-8") as f:
+            prompt_templates = yaml.safe_load(f)
+        print("✅ Custom prompt templates loaded")
+
+        # CRITICAL FIX for latest smolagents (v1.17+)
+        # The library now strictly requires the "final_answer" key
+        if prompt_templates and "final_answer" not in prompt_templates:
+            prompt_templates["final_answer"] = {
+                "pre_messages": "Here is the final answer:\n",
+                "post_messages": ""
+            }
+            print("✅ Added missing 'final_answer' template (required by latest smolagents)")
+
+    except FileNotFoundError:
+        print("⚠️ prompts.yaml not found → using default system prompt")
     except Exception as e:
-        print(f"Text-to-image tool failed to load: {e}")
+        print(f"⚠️ Failed to load prompts.yaml: {e}")
 
     # Tools list
     tools = [
@@ -57,22 +84,22 @@ def main():
     if image_tool is not None:
         tools.append(image_tool)
 
-    # Create agent – using library defaults for prompts
+    # Create the agent
     agent = CodeAgent(
         tools=tools,
         model=model,
-        max_steps=8,
+        max_steps=10,                    # increased for better performance
         verbosity_level=1,
         stream_outputs=True,
-        # planning_interval=4,                  # optional – uncomment if desired
-        # additional_authorized_imports=["math"],  # only if your code needs it
+        # planning_interval=4,           # uncomment if you want planning mode
+        prompt_templates=prompt_templates,   # now guaranteed to be valid
     )
 
-    # Launch UI
+    # Launch Gradio UI (perfect for HF Spaces / local)
     GradioUI(agent).launch(
         server_name="0.0.0.0",
         server_port=7860,
-        # share=True,                         # uncomment for public link (temporary)
+        # share=True,                    # uncomment for temporary public link
     )