Update app.py

app.py

@@ -18,26 +18,27 @@ hub_stats_tool = HubStatsTool()
 # Load the guest dataset and initialize the guest info tool
 guest_info_tool = load_guest_dataset()
 
-# Create Alfred with all the tools
-alfred = CodeAgent(
-    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
+# Custom agent class to format responses
+class FormattedCodeAgent(CodeAgent):
+    def run(self, prompt, reset=True, **kwargs):
+        # Call the parent class's run method
+        response = super().run(prompt, reset=reset, **kwargs)
+        
+        # Convert response to a string if it’s a dictionary
+        if isinstance(response, dict):
+            response_str = "\n".join([f"{key}: {value}" for key, value in response.items()])
+        else:
+            response_str = str(response)
+        
+        return response_str
+
+# Create Alfred with formatted responses
+alfred = FormattedCodeAgent(
+    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool],
     model=model,
     add_base_tools=True,  # Add any additional base tools
     planning_interval=3   # Enable planning every 3 steps
 )
 
-# Wrapper function to handle agent responses and format for GradioUI
-def process_agent_response(user_input):
-    # Run the agent with reset=False to maintain conversation memory
-    response = alfred.run(user_input, reset=False)
-    
-    # Convert response to a string if it’s a dictionary
-    if isinstance(response, dict):
-        response_str = "\n".join([f"{key}: {value}" for key, value in response.items()])
-    else:
-        response_str = str(response)
-    
-    return response_str
-
 # Launch the GradioUI interface
-GradioUI(alfred, process_function=process_agent_response).launch()
+GradioUI(alfred).launch()