Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
@@ -40,13 +41,14 @@ def calculator(a: int, b: int) -> int:
     """
     return int(a) * int(b)
 
-# (Optional) simple weather tool using Open-Meteo (no key needed)
+# Simple weather tool using Open-Meteo (no key needed)
 CITY_COORDS = {
     "new york": (40.7128, -74.0060),
     "paris": (48.8566, 2.3522),
     "london": (51.5074, -0.1278),
     "tokyo": (35.6762, 139.6503),
 }
+
 @tool
 def get_weather(city: str) -> str:
     """Get current temperature and humidity for a known city (New York, Paris, London, Tokyo).
@@ -74,56 +76,181 @@ def get_weather(city: str) -> str:
     except Exception as e:
         return f"Weather fetch failed: {e}"
 
+# -------------------------
+# Simple image generation tool (fallback)
+# -------------------------
+@tool
+def simple_image_description(prompt: str) -> str:
+    """Generate a detailed description of an image based on the prompt.
+    This is a fallback when actual image generation is not available.
+    Args:
+        prompt: Description of the image to generate
+    """
+    return f"Image description: {prompt}. Note: This is a text description as image generation is currently unavailable. To enable actual image generation, please add your Hugging Face token to the space settings."
+
 # -------------------------
 # Required final answer tool
 # -------------------------
-final_answer = FinalAnswerTool()
+try:
+    final_answer = FinalAnswerTool()
+    print("✓ Final answer tool loaded")
+except Exception as e:
+    print(f"✗ Error loading final answer tool: {e}")
+    # Create a fallback
+    @tool
+    def final_answer(answer: str) -> str:
+        """Provide the final answer to the user."""
+        return answer
 
 # -------------------------
-# Model (you can swap to a different one if overloaded)
+# Model setup with error handling
 # -------------------------
-model = HfApiModel(
-    max_tokens=2096,
-    temperature=0.5,
-    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # change if needed
-    custom_role_conversions=None,
-)
+try:
+    model = HfApiModel(
+        max_tokens=2096,
+        temperature=0.5,
+        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
+        custom_role_conversions=None,
+    )
+    print("✓ Model loaded successfully")
+except Exception as e:
+    print(f"✗ Error loading model: {e}")
+    # Try a fallback model
+    try:
+        model = HfApiModel(
+            max_tokens=1024,
+            temperature=0.5,
+            model_id='microsoft/DialoGPT-medium',
+            custom_role_conversions=None,
+        )
+        print("✓ Fallback model loaded")
+    except Exception as e2:
+        print(f"✗ Fallback model also failed: {e2}")
+        raise e2
 
 # -------------------------
 # Pre-made tools from Hub / smolagents
 # -------------------------
-# DuckDuckGo web search (no key)
-duck_search = DuckDuckGoSearchTool()
-# Text-to-image tool from the course (remote tool, trust code)
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# DuckDuckGo web search (no key) - with error handling
+try:
+    duck_search = DuckDuckGoSearchTool()
+    print("✓ DuckDuckGo search tool loaded")
+except Exception as e:
+    print(f"✗ Error loading DuckDuckGo tool: {e}")
+    duck_search = None
+
+# Text-to-image tool - with proper authentication handling
+image_generation_tool = None
+hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+
+if hf_token:
+    try:
+        # Set the token for huggingface_hub
+        from huggingface_hub import login
+        login(token=hf_token)
+        
+        image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+        print("✓ Image generation tool loaded successfully")
+    except Exception as e:
+        print(f"✗ Failed to load image generation tool: {e}")
+        print("  Using fallback image description tool instead")
+        image_generation_tool = simple_image_description
+else:
+    print("! No HF_TOKEN found - using fallback image description tool")
+    print("  To enable real image generation, add your Hugging Face token in Space settings")
+    image_generation_tool = simple_image_description
 
 # -------------------------
-# Prompts
+# Load prompts with error handling
 # -------------------------
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
+prompt_templates = {}
+try:
+    if os.path.exists("prompts.yaml"):
+        with open("prompts.yaml", 'r') as stream:
+            prompt_templates = yaml.safe_load(stream)
+        print("✓ Prompts loaded from prompts.yaml")
+    else:
+        print("! prompts.yaml not found, using default prompts")
+        prompt_templates = {
+            "system_prompt": """You are a helpful AI assistant with access to various tools.
+            
+Available tools:
+- calculator: For mathematical operations
+- get_weather: For weather information in major cities
+- get_current_time_in_timezone: For time information in different timezones
+- duck_search: For web searches (if available)
+- simple_image_description or image_generation_tool: For image-related requests
+- final_answer: To provide your final response
+
+When users request image generation:
+1. If image_generation_tool is available, use it to create actual images
+2. If only simple_image_description is available, provide detailed descriptions
+
+Always try to use the most appropriate tool for the task and provide helpful responses."""
+        }
+except Exception as e:
+    print(f"✗ Error loading prompts: {e}")
+    prompt_templates = {}
+
+# -------------------------
+# Build tools list
+# -------------------------
+tools_list = [
+    final_answer,
+    calculator,
+    get_current_time_in_timezone,
+    get_weather,
+    my_custom_tool,
+    image_generation_tool,  # This will be either the real tool or the fallback
+]
+
+if duck_search:
+    tools_list.append(duck_search)
+
+print(f"✓ Loaded {len(tools_list)} tools")
 
 # -------------------------
 # Agent with tools registered
 # -------------------------
-agent = CodeAgent(
-    model=model,
-    tools=[
-        final_answer,               # keep this
-        duck_search,                # web search
-        image_generation_tool,      # text -> image
-        calculator,                 # local calc tool
-        get_current_time_in_timezone,
-        get_weather,
-        my_custom_tool,             # your playground tool
-    ],
-    max_steps=6,
-    verbosity_level=1,
-    grammar=None,
-    planning_interval=None,
-    name=None,
-    description=None,
-    prompt_templates=prompt_templates
-)
-
-GradioUI(agent).launch()
+try:
+    agent = CodeAgent(
+        model=model,
+        tools=tools_list,
+        max_steps=10,
+        verbosity_level=2,
+        grammar=None,
+        planning_interval=None,
+        name=None,
+        description=None,
+        prompt_templates=prompt_templates
+    )
+    print("✓ Agent created successfully")
+except Exception as e:
+    print(f"✗ Error creating agent: {e}")
+    raise e
+
+# -------------------------
+# Launch UI
+# -------------------------
+try:
+    ui = GradioUI(agent)
+    print("✓ UI created successfully")
+    ui.launch()
+except Exception as e:
+    print(f"✗ Error launching UI: {e}")
+    # Fallback: create a simple Gradio interface
+    import gradio as gr
+    
+    def chat_function(message):
+        try:
+            return agent.run(message)
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    demo = gr.Interface(
+        fn=chat_function,
+        inputs=gr.Textbox(label="Your message"),
+        outputs=gr.Textbox(label="Agent response"),
+        title="AI Agent"
+    )
+    demo.launch()