Update app.py

app.py

@@ -1,12 +1,9 @@
-"""Run enhanced hf agent with improved capabilities."""
+"""Modified Gradio UI integration to fix type errors."""
 # 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 (
     CodeAgent,
@@ -18,43 +15,16 @@ from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 
 
-# 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
+# Custom tools
 @tool
 def my_custom_tool(arg1: str, arg2: int) -> str:
     """
-    Run a tool that demonstrates custom functionality.
+    Run a tool that does nothing yet.
     Args:
-        arg1: the first argument (string)
-        arg2: the second argument (integer)
+        arg1: the first argument
+        arg2: the second argument
     """
-    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)}"
+    return f"What magic will you build? {arg1=}, {arg2=}"
 
 
 @tool
@@ -71,81 +41,83 @@ 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(f"Error in get_current_time_in_timezone: {e}")
+        logger.error(e)
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
 
-# 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
+# Define our default prompt templates with the correct structure
+DEFAULT_PROMPTS = {
+    "system_prompt": "You are an expert assistant who can solve any task using code blobs.",
+    "final_answer": {
+        "pre_messages": "Here is the final answer to your question:",
+        "post_messages": "Hope this helps!"
+    },
+    "planning": {
+        "initial_facts": "First, let's identify what we know and what we need to find out.",
+        "initial_plan": "Let's develop a step-by-step plan to solve this task.",
+        "update_facts_pre_messages": "Let's update what we know based on our progress.",
+        "update_facts_post_messages": "Here's our updated knowledge about the task.",
+        "update_plan_pre_messages": "Let's review our plan based on what we've learned.",
+        "update_plan_post_messages": "Here's our revised plan to complete the task."
+    }
+}
 
 
-def main():
-    """Main function to set up and run the agent."""
-    config = AgentConfig()
-    
-    # Initialize tools with better error handling
+def create_agent():
+    """Create and configure the agent with proper error handling."""
+    # Initialize tools
     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()
+        logger.info("Successfully loaded model")
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        raise RuntimeError(f"Failed to initialize model: {e}")
     
-    # Import tool from Hub with error handling
+    # Import tool from Hub with proper 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 = {}
+    
+    # Load prompt templates or use defaults
+    prompt_templates = DEFAULT_PROMPTS.copy()
     try:
         with open("prompts.yaml", "r", encoding="utf8") as stream:
-            prompt_templates = yaml.safe_load(stream) or {}
-        logger.info("Successfully loaded prompts.yaml")
+            loaded_prompts = yaml.safe_load(stream)
+            if loaded_prompts:
+                # Update system prompt if available
+                if "system_prompt" in loaded_prompts:
+                    prompt_templates["system_prompt"] = loaded_prompts["system_prompt"]
+                
+                # Handle final_answer format
+                if "final_answer" in loaded_prompts:
+                    if isinstance(loaded_prompts["final_answer"], dict):
+                        if "pre_messages" in loaded_prompts["final_answer"]:
+                            prompt_templates["final_answer"]["pre_messages"] = loaded_prompts["final_answer"]["pre_messages"]
+                        if "post_messages" in loaded_prompts["final_answer"]:
+                            prompt_templates["final_answer"]["post_messages"] = loaded_prompts["final_answer"]["post_messages"]
+                    else:
+                        # String format, use as pre_messages
+                        prompt_templates["final_answer"]["pre_messages"] = loaded_prompts["final_answer"]
+                
+                # Handle planning section
+                if "planning" in loaded_prompts and isinstance(loaded_prompts["planning"], dict):
+                    for key in prompt_templates["planning"]:
+                        if key in loaded_prompts["planning"]:
+                            prompt_templates["planning"][key] = loaded_prompts["planning"][key]
+                
+                logger.info("Successfully merged prompt templates")
     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)
-
-    # Prepare tools list with available tools
+    # Prepare tools list
     tools = [
         my_custom_tool,
         get_current_time_in_timezone,
@@ -156,31 +128,80 @@ def main():
     if image_generation_tool:
         tools.append(image_generation_tool)
     
-    # Create the agent with enhanced configuration
-    agent = CodeAgent(
+    # Create and return the agent
+    return CodeAgent(
         model=model,
         tools=tools,
-        max_steps=config.max_steps,
-        verbosity_level=config.verbosity_level,
+        max_steps=6,
+        verbosity_level=1,
         grammar=None,
-        planning_interval=config.planning_interval,  # Enable periodic replanning
-        name="EnhancedAgent",  # Give the agent a name
-        description="An enhanced CodeAgent with improved capabilities",
+        planning_interval=None,
+        name="Agent",
+        description="A CodeAgent with improved capabilities",
         prompt_templates=prompt_templates,
     )
+
+
+def run_app():
+    """
+    Run the application with custom Gradio interface to avoid schema validation issues.
+    This circumvents the error in json_schema_to_python_type.
+    """
+    # Import gradio only when needed to customize the interface
+    try:
+        import gradio as gr
+        from Gradio_UI import GradioUI
+    except ImportError as e:
+        logger.error(f"Error importing Gradio: {e}")
+        raise
     
-     
-    # Launch the Gradio UI with configured settings
     try:
-        GradioUI(agent).launch(
-            #debug=config.debug,
-            #share=config.share,
-            #ssr=config.ssr,
+        # Create the agent
+        agent = create_agent()
+        
+        # Basic interface without relying on complex schema validation
+        with gr.Blocks() as demo:
+            gr.Markdown("# Code Agent Interface")
+            
+            with gr.Row():
+                with gr.Column(scale=4):
+                    input_text = gr.Textbox(
+                        label="Your task",
+                        placeholder="Describe your task here...",
+                        lines=5
+                    )
+                    
+                    submit_btn = gr.Button("Submit")
+                    
+                with gr.Column(scale=6):
+                    output_text = gr.Markdown(label="Agent Response")
+            
+            # Handle submissions manually to avoid schema issues
+            def process_request(task):
+                try:
+                    result = agent(task)
+                    return result
+                except Exception as e:
+                    logger.error(f"Error in agent processing: {e}")
+                    return f"Error processing your request: {str(e)}"
+            
+            submit_btn.click(
+                fn=process_request,
+                inputs=input_text,
+                outputs=output_text
+            )
+        
+        # Launch with minimal options
+        demo.launch(
+            server_name="0.0.0.0",
+            share=False,
+            debug=False
         )
+        
     except Exception as e:
-        logger.critical(f"Failed to launch UI: {e}")
-        print(f"Critical error: Failed to launch UI: {e}")
+        logger.error(f"Failed to start application: {e}")
+        print(f"Critical error: {e}")
 
 
 if __name__ == "__main__":
-    main()
+    run_app()