Update app.py

app.py

@@ -1,68 +1,86 @@
 #!/usr/bin/env python
 # coding=utf-8
 """
-Robust startup file that guarantees the prompt template dict matches what
-smolagents.CodeAgent expects, without importing internal constants.
+Robust startup file for your SmolAgents Space.
+– Merges custom prompts.yaml with safe defaults
+– Defines tools, model, CodeAgent, and Gradio UI
 """
 
-import datetime
 import copy
+import datetime
 import pytz
 import yaml
 
-# ------------- smolagents -------------
+# --------- SmolAgents imports ---------
 from smolagents import CodeAgent, HfApiModel, tool
-# Local Gradio wrapper
+
+# --------- Local Gradio wrapper --------
 from Gradio_UI import GradioUI
 
-# ------------ local tools -------------
+# --------- Local tools -----------------
 from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 from tools.web_search import DuckDuckGoSearchTool
 
-# ---------- example custom tool -------
+# ---------------------------------------------------------------------------
+# Example custom tool (now with full doc-string so SmolAgents can JSON-schema it)
 @tool
 def my_custom_tool(arg1: str, arg2: int) -> str:
-    """A placeholder tool to illustrate custom tool wiring."""
+    """Example Tool: a simple placeholder.
+
+    Args:
+        arg1 (str): Any string input (unused – demo only).
+        arg2 (int): Any integer input (unused – demo only).
+
+    Returns:
+        str: A constant demo message.
+    """
     return "What magic will you build ?"
 
 
+# A small utility tool: get local time in a given timezone
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
-    """Return current local time in a given timezone."""
+    """Return the current local time in the specified timezone.
+
+    Args:
+        timezone (str): IANA timezone name, e.g. 'America/New_York'.
+
+    Returns:
+        str: Formatted datetime string or an error message.
+    """
     try:
         tz = pytz.timezone(timezone)
         return datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
     except Exception as exc:
-        return f"Invalid timezone ({exc})"
+        return f"Error: {exc}"
 
 
-# ---------------- model ----------------
+# ---------------------------------------------------------------------------
+# Model configuration
 model = HfApiModel(
-    model_id="Qwen/Qwen3-32B",
+    model_id="Qwen/Qwen3-32B",   # Adjust if overloaded
     max_tokens=2096,
     temperature=0.5,
     custom_role_conversions=None,
 )
 
-# ------------- load YAML --------------
+# ---------------------------------------------------------------------------
+# 1) Load prompts.yaml (may be imperfect)
 try:
     with open("prompts.yaml", "r") as f:
-        user_templates = yaml.safe_load(f) or {}
+        custom_prompts = yaml.safe_load(f) or {}
 except Exception as exc:
-    print("⚠️  Could not load prompts.yaml — using bare-bones defaults.", exc)
-    user_templates = {}
+    print("⚠️  Could not load prompts.yaml – using bare defaults.", exc)
+    custom_prompts = {}
 
-# --------- repair / fill gaps ---------
-def ensure_block(mapping, key, default):
-    """If key missing or wrong type, replace with default."""
-    if key not in mapping or not isinstance(mapping[key], type(default)):
-        mapping[key] = copy.deepcopy(default)
-
-# Minimal viable defaults
+# 2) Build minimal required template structure
 MIN_SYSTEM = "You are a helpful agent."
-MIN_FINAL = "Thought: I now know the answer.\nCode:\n```py\nfinal_answer({{answer}})\n```<end_code>"
-MIN_PLAN_DICT = {
+MIN_FINAL = (
+    "Thought: I now know the answer.\n"
+    "Code:\n```py\nfinal_answer({{answer}})\n```<end_code>"
+)
+MIN_PLAN = {
     "initial_facts": "|-\n  List any facts.",
     "initial_plan": "|-\n  Write a high-level plan.",
     "update_facts_pre_messages": "|-\n  Pre facts.",
@@ -70,24 +88,30 @@ MIN_PLAN_DICT = {
     "update_plan_pre_messages": "|-\n  Pre plan.",
     "update_plan_post_messages": "|-\n  Post plan.",
 }
-MIN_AGENT_DICT = {
+MIN_MGR = {
     "task": "|-\n  You are {{name}}.",
     "report": "|-\n  {{final_answer}}",
 }
 
-# Top-level fix-ups
-ensure_block(user_templates, "system_prompt", MIN_SYSTEM)
-ensure_block(user_templates, "final_answer", MIN_FINAL)
-ensure_block(user_templates, "planning", MIN_PLAN_DICT)
-ensure_block(user_templates, "managed_agent", MIN_AGENT_DICT)
+def ensure(mapping, key, default):
+    if key not in mapping or not isinstance(mapping[key], type(default)):
+        mapping[key] = copy.deepcopy(default)
+
+ensure(custom_prompts, "system_prompt", MIN_SYSTEM)
+ensure(custom_prompts, "final_answer", MIN_FINAL)
+ensure(custom_prompts, "planning",    MIN_PLAN)
+ensure(custom_prompts, "managed_agent", MIN_MGR)
+
+# If planning / managed_agent came in as strings, convert to dict default
+if isinstance(custom_prompts["planning"], str):
+    custom_prompts["planning"] = copy.deepcopy(MIN_PLAN)
+if isinstance(custom_prompts["managed_agent"], str):
+    custom_prompts["managed_agent"] = copy.deepcopy(MIN_MGR)
 
-# If planning/managed_agent came in as strings, convert to dict with default
-if isinstance(user_templates["planning"], str):
-    user_templates["planning"] = copy.deepcopy(MIN_PLAN_DICT)
-if isinstance(user_templates["managed_agent"], str):
-    user_templates["managed_agent"] = copy.deepcopy(MIN_AGENT_DICT)
+prompt_templates = custom_prompts  # now guaranteed correct structure
 
-# ------------- create agent -----------
+# ---------------------------------------------------------------------------
+# Build the CodeAgent
 agent = CodeAgent(
     model=model,
     tools=[
@@ -99,8 +123,10 @@ agent = CodeAgent(
     ],
     max_steps=6,
     verbosity_level=1,
-    prompt_templates=user_templates,
+    prompt_templates=prompt_templates,
 )
 
-# ------------- launch UI --------------
-GradioUI(agent).launch()
+# ---------------------------------------------------------------------------
+# Launch the Gradio UI
+if __name__ == "__main__":
+    GradioUI(agent).launch()