Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from smolagents import GradioUI, CodeAgent, HfApiModel
-import re
+import ast
 
 # Import our custom tools from their modules
 from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
@@ -21,100 +21,73 @@ hub_stats_tool = HubStatsTool()
 # Load the guest dataset and initialize the guest info tool
 guest_tool = load_guest_dataset()
 
-class FormattedCodeAgent(CodeAgent):
+# Make guest tool description VERY clear
+guest_tool.description = """USE THIS TOOL FOR GALA GUEST QUESTIONS!
+Official guest database with names, relations, descriptions, emails.
+Query examples: "all guests", "Lady Ada Lovelace", "scientists".
+Returns formatted guest information."""
+
+# Simple agent that formats output
+class SimpleAgent(CodeAgent):
     def run(self, query: str):
-        # Get the raw response
-        raw_response = super().run(query)
-        
-        print(f"📦 Raw response from Alfred:\n{raw_response}")
+        response = super().run(query)
         
-        # Parse the response to extract meaningful content
-        if "Final answer:" in raw_response:
-            # Extract just the final answer
-            parts = raw_response.split("Final answer:")
-            if len(parts) > 1:
-                final_answer = parts[-1].strip()
-                
-                # Clean up Python list formatting
-                if final_answer.startswith('[') and final_answer.endswith(']'):
-                    # It's a Python list - format as bullet points
-                    try:
-                        import ast
-                        guest_list = ast.literal_eval(final_answer)
-                        if isinstance(guest_list, list):
-                            formatted = "🎩 **Gala Guests Found:**\n\n"
-                            for guest in guest_list:
-                                formatted += f"• {guest}\n"
-                            return formatted
-                    except:
-                        pass
-                
-                # Return cleaned answer
-                return final_answer.replace("['", "").replace("']", "").replace("', '", "\n• ")
+        # Handle list responses
+        if isinstance(response, list):
+            if response:
+                formatted = "🎩 **Gala Guests:**\n\n"
+                for guest in response:
+                    formatted += f"• {guest}\n"
+                return formatted
+            return "No guests found."
         
-        # Check for printed output
-        if "Attending the gala are:" in raw_response:
-            # Extract the guest list from print statement
-            lines = raw_response.split('\n')
-            for line in lines:
-                if "Attending the gala are:" in line:
-                    # Try to extract list
-                    import re
-                    match = re.search(r"\[(.*?)\]", line)
-                    if match:
-                        guest_str = match.group(1)
-                        guests = [g.strip().strip("'\"") for g in guest_str.split(',')]
+        # Handle string responses that contain lists
+        if isinstance(response, str):
+            # Check if it's a string representation of a list
+            response = response.strip()
+            if (response.startswith('[') and response.endswith(']')) or \
+               (response.startswith("['") and response.endswith("']")):
+                try:
+                    guest_list = ast.literal_eval(response)
+                    if isinstance(guest_list, list):
                         formatted = "🎩 **Gala Guests:**\n\n"
-                        for guest in guests:
+                        for guest in guest_list:
                             formatted += f"• {guest}\n"
                         return formatted
-        
-        # Default: clean up code execution traces
-        lines = raw_response.split('\n')
-        cleaned = []
-        skip = False
-        
-        for line in lines:
-            # Skip execution headers
-            if line.startswith('━━━━') or line.startswith('─ Executing'):
-                skip = True
-                continue
-            if skip and (line.startswith('─') or line.strip() == ''):
-                continue
-            skip = False
+                except:
+                    pass
             
-            # Skip code blocks and execution logs
-            if line.startswith('```') or line.startswith('Execution logs:'):
-                continue
-                
-            cleaned.append(line)
+            # Clean up code execution traces
+            lines = response.split('\n')
+            clean_lines = []
+            for line in lines:
+                if not any(line.startswith(x) for x in ['━━━━', '─ Executing', '```', 'Execution logs:']):
+                    clean_lines.append(line)
+            return '\n'.join(clean_lines).strip()
         
-        return '\n'.join(cleaned).strip()
+        return str(response)
 
-# Create Alfred with custom formatting
-alfred = FormattedCodeAgent(
+# Create Alfred
+alfred = SimpleAgent(
     tools=[guest_tool, weather_info_tool, hub_stats_tool, search_tool], 
     model=model,
     add_base_tools=True,
-    planning_interval=3,
-    max_steps=5
+    planning_interval=2,
+    max_steps=3
 )
 
 # Test
-print("\n🧪 TESTING ALFRED...")
-test_query = "Who is coming to the gala?"
-print(f"Query: '{test_query}'")
-response = alfred.run(test_query)
-print(f"Response:\n{response}")
+print("\n🎩 Testing Alfred...")
+test_response = alfred.run("Who is coming to the gala?")
+print(f"Test response:\n{test_response}")
 
-# Create a custom Gradio interface
+# Create Gradio interface
 with gr.Blocks(title="🎩 Alfred Assistant") as demo:
     gr.Markdown("# 🎩 Alfred Assistant")
-    gr.Markdown("Your AI assistant with multiple capabilities including guest information retrieval.")
+    gr.Markdown("Ask about gala guests, weather, model stats, or search the web.")
     
-    # Main chat interface
     chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Ask Alfred anything...")
+    msg = gr.Textbox(label="Your question...")
     clear = gr.ClearButton([msg, chatbot])
     
     def respond(message, chat_history):
@@ -124,26 +97,19 @@ with gr.Blocks(title="🎩 Alfred Assistant") as demo:
     
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
     
-    # Test Section
-    with gr.Accordion("🧪 Test Queries", open=False):
-        gr.Markdown("Try these example queries:")
-        
-        test_queries = [
-            "Tell me about 'Lady Ada Lovelace'",
-            "What's the weather in Paris?",
-            "Show me stats for the gpt2 model",
-            "Search for latest AI news",
-            "Who is coming to the gala?"
-        ]
-        
-        with gr.Row():
-            for query in test_queries:
-                test_btn = gr.Button(query, size="sm")
-                test_btn.click(
-                    fn=lambda q=query: q,
-                    inputs=[],
-                    outputs=[msg]
-                )
+    # Examples
+    with gr.Accordion("📋 Example Questions", open=True):
+        gr.Examples(
+            examples=[
+                "Who is coming to the gala?",
+                "Tell me about Lady Ada Lovelace",
+                "What's the weather in London?",
+                "Show me stats for GPT-2",
+                "Search for Python programming news"
+            ],
+            inputs=msg,
+            label="Click any example:"
+        )
 
 if __name__ == "__main__":
     demo.launch()