Update app.py

app.py

@@ -1,9 +1,11 @@
-"""Run hf agent."""
+"""Run enhanced hf agent with improved capabilities."""
 # pylint: disable=line-too-long, broad-exception-caught
 import datetime
 import pytz
 import yaml
 import re
+import os
+from pathlib import Path
 from Gradio_UI import GradioUI
 from loguru import logger
 from smolagents import (
@@ -15,16 +17,45 @@ from smolagents import (
 from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 
-# Custom tool example
+
+# Enhanced configuration with environment variables and defaults
+class AgentConfig:
+    def __init__(self):
+        self.max_steps = int(os.environ.get("AGENT_MAX_STEPS", "8"))
+        self.verbosity_level = int(os.environ.get("AGENT_VERBOSITY", "1"))
+        self.planning_interval = int(os.environ.get("AGENT_PLANNING_INTERVAL", "3"))
+        self.model_name = os.environ.get("AGENT_MODEL", "default")
+        self.debug = os.environ.get("AGENT_DEBUG", "false").lower() == "true"
+        self.share = os.environ.get("AGENT_SHARE", "false").lower() == "true"
+        self.ssr = os.environ.get("AGENT_SSR", "true").lower() == "true"
+        
+        # Create log directory
+        log_dir = Path("logs")
+        log_dir.mkdir(exist_ok=True)
+        now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.log_file = log_dir / f"agent_run_{now}.log"
+        
+        # Setup logging
+        logger.add(self.log_file, rotation="100 MB")
+        logger.info(f"Agent configuration initialized: {self.__dict__}")
+
+
+# Custom tool example with better error handling
 @tool
-def my_custom_tool(arg1: str, arg2: int) -> str:  # it's important to specify the return type
+def my_custom_tool(arg1: str, arg2: int) -> str:
     """
-    Run a tool that does nothing yet.
+    Run a tool that demonstrates custom functionality.
     Args:
-        arg1: the first argument
-        arg2: the second argument
+        arg1: the first argument (string)
+        arg2: the second argument (integer)
     """
-    return f"What magic will you build? {arg1=}, {arg2=}"
+    logger.debug(f"my_custom_tool called with {arg1=}, {arg2=}")
+    try:
+        return f"Custom tool processed: {arg1} with value {arg2}"
+    except Exception as e:
+        logger.error(f"Error in my_custom_tool: {e}")
+        return f"Error processing request: {str(e)}"
+
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -40,49 +71,126 @@ def get_current_time_in_timezone(timezone: str) -> str:
         # 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}"
+    except pytz.exceptions.UnknownTimeZoneError:
+        logger.warning(f"Unknown timezone: {timezone}")
+        return f"Error: '{timezone}' is not a recognized timezone. Please use a valid timezone like 'America/New_York' or 'Asia/Tokyo'."
     except Exception as e:
-        logger.error(e)
+        logger.error(f"Error in get_current_time_in_timezone: {e}")
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-# Initialize tools
-final_answer = FinalAnswerTool()
-visit_webpage = VisitWebpageTool()
-model = HfApiModel()
 
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# Helper function to ensure all required templates exist
+def ensure_complete_templates(templates):
+    """Ensure all required templates exist and provide defaults if needed."""
+    required_templates = [
+        "system_prompt", "final_answer",
+        "planning.initial_facts", "planning.initial_plan",
+        "planning.update_facts_pre_messages", "planning.update_facts_post_messages",
+        "planning.update_plan_pre_messages", "planning.update_plan_post_messages"
+    ]
+    
+    # Check if nested keys exist and create them if not
+    for template in required_templates:
+        if "." in template:
+            parent, child = template.split(".", 1)
+            if parent not in templates:
+                templates[parent] = {}
+            if child not in templates[parent]:
+                logger.warning(f"Missing template {template}, using default")
+                # Set a minimal default
+                templates[parent][child] = f"{{{{ {child.replace('_', ' ')} }}}}"
+        elif template not in templates:
+            logger.warning(f"Missing template {template}, using default")
+            # Set a minimal default
+            templates[template] = f"{{{{ {template.replace('_', ' ')} }}}}"
+    
+    return templates
 
-# Load prompt templates
-with open("prompts.yaml", "r", encoding="utf8") as stream:
-    prompt_templates = yaml.safe_load(stream)
 
-# Add the missing 'final_answer' template if it doesn't exist
-if 'final_answer' not in prompt_templates:
-    prompt_templates['final_answer'] = "{{final_answer}}"
+def main():
+    """Main function to set up and run the agent."""
+    config = AgentConfig()
+    
+    # Initialize tools with better error handling
+    final_answer = FinalAnswerTool()
+    visit_webpage = VisitWebpageTool()
+    
+    # Use try-except for more resilient model loading
+    try:
+        model = HfApiModel(model_name=config.model_name)
+        logger.info(f"Successfully loaded model: {config.model_name}")
+    except Exception as e:
+        logger.error(f"Error loading model {config.model_name}: {e}")
+        logger.info("Falling back to default model")
+        model = HfApiModel()
+    
+    # Import tool from Hub with error handling
+    try:
+        image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+        logger.info("Successfully loaded image generation tool")
+    except Exception as e:
+        logger.error(f"Error loading image generation tool: {e}")
+        image_generation_tool = None
+
+    # Load prompt templates with fallback
+    prompt_templates = {}
+    try:
+        with open("prompts.yaml", "r", encoding="utf8") as stream:
+            prompt_templates = yaml.safe_load(stream) or {}
+        logger.info("Successfully loaded prompts.yaml")
+    except Exception as e:
+        logger.error(f"Error loading prompts.yaml: {e}")
+        logger.info("Using default templates")
+    
+    # Ensure all required templates exist
+    prompt_templates = ensure_complete_templates(prompt_templates)
 
-# Create the agent
-agent = CodeAgent(
-    model=model,
-    tools=[
+    # Prepare tools list with available tools
+    tools = [
         my_custom_tool,
         get_current_time_in_timezone,
         final_answer,
         visit_webpage,
-        image_generation_tool,  # Added the image generation tool
-    ],  # add your tools here
-    max_steps=6,
-    verbosity_level=1,
-    grammar=None,
-    planning_interval=None,
-    name=None,
-    description=None,
-    prompt_templates=prompt_templates,
-)
+    ]
+    
+    if image_generation_tool:
+        tools.append(image_generation_tool)
+    
+    # Create the agent with enhanced configuration
+    agent = CodeAgent(
+        model=model,
+        tools=tools,
+        max_steps=config.max_steps,
+        verbosity_level=config.verbosity_level,
+        grammar=None,
+        planning_interval=config.planning_interval,  # Enable periodic replanning
+        name="EnhancedAgent",  # Give the agent a name
+        description="An enhanced CodeAgent with improved capabilities",
+        prompt_templates=prompt_templates,
+    )
+    
+    # Add hooks for better monitoring (available in smolagents)
+    @agent.add_hook("before_step")
+    def log_step(agent, step_number):
+        logger.info(f"Starting step {step_number}/{agent.max_steps}")
+    
+    @agent.add_hook("after_step")
+    def analyze_step(agent, step_number, thought, code, observation):
+        logger.info(f"Completed step {step_number}")
+        if "error" in observation.lower():
+            logger.warning(f"Potential error in step {step_number}: {observation[:100]}...")
+    
+    # Launch the Gradio UI with configured settings
+    try:
+        GradioUI(agent).launch(
+            debug=config.debug,
+            share=config.share,
+            ssr=config.ssr,
+        )
+    except Exception as e:
+        logger.critical(f"Failed to launch UI: {e}")
+        print(f"Critical error: Failed to launch UI: {e}")
+
 
-# Launch the Gradio UI
-ssr = True
-GradioUI(agent).launch(
-    # debug=True,
-    # share=True,
-    ssr=ssr,
-)
+if __name__ == "__main__":
+    main()