Update Gradio_UI.py

Gradio_UI.py

@@ -1,5 +1,9 @@
 #!/usr/bin/env python
 # coding=utf-8
+"""
+Gradio UI for Travel Catalogue Creator
+Production-ready interface with streaming agent responses
+"""
 import os
 import re
 from typing import Optional, List
@@ -8,18 +12,27 @@ from smolagents.agents import ActionStep, MultiStepAgent
 from smolagents.memory import MemoryStep
 from smolagents.utils import _is_package_available
 
+
 def pull_messages_from_step(step_log: MemoryStep):
     """Extract ChatMessage objects from agent steps with proper nesting"""
+    if not _is_package_available("gradio"):
+        raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
+    
     import gradio as gr
+    
     if isinstance(step_log, ActionStep):
+        # Step header
         step_number = f"Step {step_log.step_number}" if step_log.step_number is not None else "Processing"
         yield gr.ChatMessage(role="assistant", content=f"**{step_number}**")
         
-        # Show LLM reasoning
+        # Show LLM reasoning/thinking
         if hasattr(step_log, "model_output") and step_log.model_output:
-            model_output = re.sub(r"```\s*<end_code>.*", "```", step_log.model_output.strip())
+            model_output = step_log.model_output.strip()
+            # Clean up code blocks
+            model_output = re.sub(r"```\s*<end_code>.*", "```", model_output)
             model_output = re.sub(r"<end_code>\s*```", "```", model_output)
-            if model_output.strip():
+            
+            if model_output:
                 yield gr.ChatMessage(role="assistant", content=model_output)
         
         parent_id = None
@@ -28,15 +41,18 @@ def pull_messages_from_step(step_log: MemoryStep):
         if hasattr(step_log, "tool_calls") and step_log.tool_calls:
             tool_call = step_log.tool_calls[0]
             parent_id = f"tool_{step_log.step_number}"
+            
+            # Format tool arguments
             args = tool_call.arguments
             if isinstance(args, dict):
                 content = "\n".join(f"• {k}: {v}" for k, v in args.items() if v and k != 'self')
             else:
                 content = str(args).strip()
+            
             metadata = {
                 "title": f"🛠️ Using: {tool_call.name}",
                 "id": parent_id,
-                "status": "pending",  # FIXED: "running" → "pending" (Gradio validation)
+                "status": "pending",  # Gradio requires "pending", not "running"
             }
             yield gr.ChatMessage(role="assistant", content=content, metadata=metadata)
         
@@ -44,29 +60,42 @@ def pull_messages_from_step(step_log: MemoryStep):
         if hasattr(step_log, "observations") and step_log.observations:
             obs = step_log.observations.strip()
             if obs and not obs.startswith("Execution logs:"):
-                metadata = {"title": "✅ Result", "status": "done"}  # FIXED: added status
+                metadata = {
+                    "title": "✅ Result",
+                    "status": "done"
+                }
                 if parent_id is not None:
                     metadata["parent_id"] = parent_id
                 yield gr.ChatMessage(role="assistant", content=obs, metadata=metadata)
         
         # Show errors
         if hasattr(step_log, "error") and step_log.error:
-            metadata = {"title": "⚠️ Warning", "status": "done"}  # FIXED: added status
+            metadata = {
+                "title": "⚠️ Warning",
+                "status": "done"
+            }
             if parent_id is not None:
                 metadata["parent_id"] = parent_id
             yield gr.ChatMessage(role="assistant", content=str(step_log.error), metadata=metadata)
         
-        # Step footer
+        # Step footer with timing and token info
         footer_parts = [step_number]
         if hasattr(step_log, "duration") and step_log.duration:
             footer_parts.append(f"⏱️ {float(step_log.duration):.1f}s")
         if hasattr(step_log, "input_token_count") and hasattr(step_log, "output_token_count"):
             footer_parts.append(f"💬 {step_log.input_token_count + step_log.output_token_count:,} tokens")
+        
         yield gr.ChatMessage(
             role="assistant",
             content=f'<span style="color: #888; font-size: 0.85em;">{" | ".join(footer_parts)}</span>',
         )
-        yield gr.ChatMessage(role="assistant", content='<hr style="margin: 8px 0; border: 0; border-top: 1px solid #eee">')
+        
+        # Divider between steps
+        yield gr.ChatMessage(
+            role="assistant",
+            content='<hr style="margin: 8px 0; border: 0; border-top: 1px solid #eee">'
+        )
+
 
 def stream_to_gradio(
     agent: MultiStepAgent,
@@ -74,64 +103,166 @@ def stream_to_gradio(
     reset_agent_memory: bool = False,
     additional_args: Optional[dict] = None,
 ):
-    """Runs agent and streams messages as gradio ChatMessages"""
+    """
+    Runs agent and streams messages as gradio ChatMessages.
+    
+    Args:
+        agent: The MultiStepAgent instance to run
+        task: User's task/query string
+        reset_agent_memory: Whether to clear agent memory before running
+        additional_args: Optional additional arguments for the agent
+    
+    Yields:
+        Gradio ChatMessage objects for UI display
+    """
     if not _is_package_available("gradio"):
         raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
+    
     import gradio as gr
+    
     try:
+        # Run agent and stream steps
         for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
             if isinstance(step_log, ActionStep):
+                # Track token usage if available
                 if hasattr(agent.model, "last_input_token_count") and hasattr(agent.model, "last_output_token_count"):
                     step_log.input_token_count = agent.model.last_input_token_count
                     step_log.output_token_count = agent.model.last_output_token_count
+                
+                # Yield messages from this step
                 for message in pull_messages_from_step(step_log):
                     yield message
+        
+        # Handle final output
         final = handle_agent_output_types(step_log)
         if isinstance(final, AgentText):
             content = final.to_string()
-            yield gr.ChatMessage(role="assistant", content=content, metadata={"react": True, "status": "done"})
+            yield gr.ChatMessage(
+                role="assistant",
+                content=content,
+                metadata={"react": True, "status": "done"}
+            )
         else:
-            yield gr.ChatMessage(role="assistant", content=f"**Final Answer:** {str(final)}", metadata={"status": "done"})
+            yield gr.ChatMessage(
+                role="assistant",
+                content=f"**Final Answer:** {str(final)}",
+                metadata={"status": "done"}
+            )
+    
     except Exception as e:
+        # Handle errors gracefully
+        error_msg = str(e)
+        
+        # Provide helpful error messages
+        if "500 Internal Server Error" in error_msg or "Bad Request" in error_msg:
+            helpful_msg = (
+                "⚠️ **API Error:** The model service encountered an issue.\n\n"
+                "**Possible fixes:**\n"
+                "• Try rephrasing your request more clearly\n"
+                "• Include all required details: destination, dates, origin city, budget\n"
+                "• Example: *'5-day Barcelona trip from NYC, Oct 15-19, budget $1500 USD'*\n\n"
+                f"Technical details: {error_msg}"
+            )
+        elif "No results found" in error_msg:
+            helpful_msg = (
+                "⚠️ **Search Issue:** Couldn't find information about the destination.\n\n"
+                "Please try again with a different destination or check the spelling."
+            )
+        else:
+            helpful_msg = (
+                f"⚠️ **Error:** {error_msg}\n\n"
+                "Please try again with a clearer trip description including:\n"
+                "• Destination city\n"
+                "• Travel dates\n"
+                "• Origin city\n"
+                "• Budget amount + currency"
+            )
+        
         yield gr.ChatMessage(
             role="assistant",
-            content=f"⚠️ **Error during planning:** {str(e)}\nPlease try again with a clearer trip description (include destination, dates, origin city, and budget).",
+            content=helpful_msg,
             metadata={"status": "done"}
         )
 
+
 class GradioUI:
     """Production-ready Gradio interface for travel agent"""
+    
     def __init__(self, agent: MultiStepAgent, file_upload_folder: Optional[str] = None):
+        """
+        Initialize Gradio UI wrapper.
+        
+        Args:
+            agent: The MultiStepAgent instance to use
+            file_upload_folder: Optional folder path for file uploads
+        """
         if not _is_package_available("gradio"):
             raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
+        
         self.agent = agent
         self.file_upload_folder = file_upload_folder
+        
+        # Create upload folder if specified
         if self.file_upload_folder and not os.path.exists(file_upload_folder):
             os.makedirs(file_upload_folder, exist_ok=True)
     
     def interact_with_agent(self, prompt: str, history: List):
+        """
+        Handle user interaction with the agent.
+        
+        Args:
+            prompt: User's input text
+            history: Conversation history
+        
+        Yields:
+            Updated conversation history
+        """
         import gradio as gr
+        
+        # Add user message to history
         history.append(gr.ChatMessage(role="user", content=prompt))
         yield history
+        
+        # Stream agent responses
         for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
             history.append(msg)
             yield history
     
     def launch(self, **kwargs):
+        """
+        Launch the Gradio interface.
+        
+        Args:
+            **kwargs: Additional arguments passed to demo.launch()
+        """
         import gradio as gr
-        with gr.Blocks(title="✈️ Travel Catalogue Creator", fill_height=True, theme=gr.themes.Soft()) as demo:
+        
+        # Build UI
+        with gr.Blocks(
+            title="✈️ Travel Catalogue Creator",
+            fill_height=True,
+            theme=gr.themes.Soft()
+        ) as demo:
+            # Header
             gr.Markdown("# ✈️ Smart Travel Catalogue Creator")
-            gr.Markdown("Plan your trip with AI: weather forecasts, custom itineraries, packing lists & visual inspiration")
+            gr.Markdown(
+                "Plan your perfect trip with AI: weather forecasts, custom itineraries, "
+                "packing lists & visual inspiration"
+            )
+            
+            # Chat interface
             chatbot = gr.Chatbot(
                 label="Your Travel Assistant",
                 type="messages",
                 avatar_images=(
-                    None,
-                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",
+                    None,  # User avatar (default)
+                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",  # Agent avatar
                 ),
                 height=600,
                 show_copy_button=True,
             )
+            
+            # Input area
             with gr.Row():
                 text_input = gr.Textbox(
                     label="Describe your trip",
@@ -139,23 +270,33 @@ class GradioUI:
                     lines=2,
                 )
                 submit_btn = gr.Button("Plan My Trip", variant="primary")
+            
+            # Help text
             gr.Markdown("""
             ### 💡 Tips for best results:
-            • Include: destination, dates, origin city, budget amount + currency
-            • Example: *"5-day Lisbon trip from London, Sep 20-24, budget £800 GBP"*
+            • **Include:** destination, dates, origin city, budget amount + currency
+            • **Example:** *"5-day Lisbon trip from London, Sep 20-24, budget £800 GBP"*
+            • **Example:** *"Weekend Paris getaway from Amsterdam, March 10-12, €600 budget"*
             """)
+            
+            # Connect interactions
             submit_btn.click(
                 self.interact_with_agent,
                 inputs=[text_input, chatbot],
                 outputs=[chatbot],
                 show_progress="full",
             )
+            
             text_input.submit(
                 self.interact_with_agent,
                 inputs=[text_input, chatbot],
                 outputs=[chatbot],
                 show_progress="full",
             )
+        
+        # Launch
         demo.launch(**kwargs)
 
-__all__ = ["stream_to_gradio", "GradioUI"]
+
+# Export public API
+__all__ = ["stream_to_gradio", "GradioUI"]