Update app.py

app.py

@@ -1,237 +1,195 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient # Keep for direct use if needed, though agent will use its own model
 import os
 import json
 import base64
 from PIL import Image
 import io
 
-# Import smolagents Tool
-from smolagents import Tool
+# Smolagents imports
+from smolagents import CodeAgent, Tool
+from smolagents.models import InferenceClientModel as SmolInferenceClientModel
+# We'll use PIL.Image directly for opening, AgentImage is for agent's internal typing if needed by a tool
+from smolagents.gradio_ui import pull_messages_from_step # For formatting agent steps
+from smolagents.memory import ActionStep, FinalAnswerStep, PlanningStep, MemoryStep # For type checking steps
+from smolagents.models import ChatMessageStreamDelta # For type checking stream deltas
+
 
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 print("Access token loaded.")
 
-# Initialize the image generation tool
-# This can be defined globally as it doesn't change per request
-try:
-    image_generation_tool = Tool.from_space(
-        "black-forest-labs/FLUX.1-schnell",
-        name="image_generator",
-        description="Generates an image from a text prompt. Use it when the user asks to 'generate an image of ...' or 'draw a picture of ...'. The input should be the descriptive prompt for the image."
-    )
-    print("Image generation tool loaded successfully.")
-except Exception as e:
-    print(f"Error loading image generation tool: {e}")
-    image_generation_tool = None
-
-# Function to encode image to base64
-def encode_image(image_path):
-    if not image_path:
-        print("No image path provided")
+# Function to encode image to base64 (remains useful if we ever need to pass base64 to a non-smolagent component)
+def encode_image(image_path_or_pil):
+    if not image_path_or_pil:
+        print("No image path or PIL Image provided")
         return None
     
     try:
-        print(f"Encoding image from path: {image_path}")
+        # print(f"Encoding image: {type(image_path_or_pil)}") # Debug
         
-        # If it's already a PIL Image
-        if isinstance(image_path, Image.Image):
-            image = image_path
-        else:
-            # Try to open the image file
-            image = Image.open(image_path)
+        if isinstance(image_path_or_pil, Image.Image):
+            image = image_path_or_pil
+        else: # Assuming it's a path
+            image = Image.open(image_path_or_pil)
         
-        # Convert to RGB if image has an alpha channel (RGBA)
         if image.mode == 'RGBA':
             image = image.convert('RGB')
         
-        # Encode to base64
         buffered = io.BytesIO()
-        image.save(buffered, format="JPEG") 
+        image.save(buffered, format="JPEG") # JPEG is generally smaller for transfer
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        print("Image encoded successfully")
+        # print("Image encoded successfully") # Debug
         return img_str
     except Exception as e:
         print(f"Error encoding image: {e}")
         return None
 
+# This function will now set up and run the smolagent
 def respond(
-    message_text, # Changed from 'message' to be explicit about text part
-    image_files,  # This will be a list of paths from gr.MultimodalTextbox
-    history: list[list[Any, str | None]], # History can now contain complex user messages
-    system_message,
+    message_text, # Text from MultimodalTextbox
+    image_file_paths,  # List of file paths from MultimodalTextbox
+    gradio_history: list[tuple[str, str]], # Gradio history (for context if needed, agent is stateless per call here)
+    system_message_for_agent, # System prompt for the main LLM agent
     max_tokens,
     temperature,
     top_p,
     frequency_penalty,
     seed,
-    provider,
-    custom_api_key,
-    custom_model,    
-    model_search_term,
-    selected_model
+    provider_for_agent_llm, 
+    api_key_for_agent_llm, 
+    model_id_for_agent_llm,
+    model_search_term, # Unused directly by agent logic
+    selected_model_for_agent_llm # Fallback model ID
 ):
-    print(f"Received message text: {message_text}")
-    print(f"Received {len(image_files) if image_files else 0} image files: {image_files}")
-    # print(f"History: {history}") # Can be very verbose
-    print(f"System message: {system_message}")
-    print(f"Max tokens: {max_tokens}, Temperature: {temperature}, Top-P: {top_p}")
-    print(f"Frequency Penalty: {frequency_penalty}, Seed: {seed}")
-    print(f"Selected provider: {provider}")         
-    print(f"Custom API Key provided: {bool(custom_api_key.strip())}")
-    print(f"Selected model (custom_model): {custom_model}")  
-    print(f"Model search term: {model_search_term}")
-    print(f"Selected model from radio: {selected_model}")
-
-    # Determine which token to use
-    token_to_use = custom_api_key if custom_api_key.strip() != "" else ACCESS_TOKEN
-    
-    if custom_api_key.strip() != "":
-        print("USING CUSTOM API KEY: BYOK token provided by user is being used for authentication")
-    else:
-        print("USING DEFAULT API KEY: Environment variable HF_TOKEN is being used for authentication")
-
-    user_text_message_lower = message_text.lower() if message_text else ""
-
-    image_keywords = ["generate image", "draw a picture of", "create an image of", "make an image of"]
-    is_image_generation_request = any(keyword in user_text_message_lower for keyword in image_keywords)
-
-    if is_image_generation_request and image_generation_tool:
-        print("Image generation request detected.")
-        image_prompt = message_text
-        for keyword in image_keywords:
-            if keyword in user_text_message_lower:
-                # Find the keyword in the original case-sensitive message text to split
-                keyword_start_index = user_text_message_lower.find(keyword)
-                image_prompt = message_text[keyword_start_index + len(keyword):].strip()
-                break
-        
-        print(f"Extracted image prompt: {image_prompt}")
-        if not image_prompt:
-            yield {"type": "text", "content": "Please provide a description for the image you want to generate."}
-            return
-
-        try:
-            generated_image_path = image_generation_tool(prompt=image_prompt)
-            print(f"Image generated by tool, path: {generated_image_path}")
-            yield {"type": "image", "path": str(generated_image_path)} # Ensure path is string
-            return
-        except Exception as e:
-            print(f"Error during image generation tool call: {e}")
-            yield {"type": "text", "content": f"Sorry, I couldn't generate the image. Error: {str(e)}"}
-            return
-    elif is_image_generation_request and not image_generation_tool:
-        yield {"type": "text", "content": "Image generation tool is not available or failed to load."}
-        return
+    print(f"Respond function called. Message: '{message_text}', Images: {image_file_paths}")
 
-    # If not an image generation request, proceed with text/multimodal LLM call
-    print("Proceeding with LLM call (text or multimodal).")
-    client = InferenceClient(token=token_to_use, provider=provider)
-    print(f"Hugging Face Inference Client initialized with {provider} provider.")
+    token_to_use = api_key_for_agent_llm if api_key_for_agent_llm.strip() != "" else ACCESS_TOKEN
+    model_to_use = model_id_for_agent_llm.strip() if model_id_for_agent_llm.strip() != "" else selected_model_for_agent_llm
 
-    if seed == -1:
-        seed = None
-
-    # Prepare messages for LLM
-    llm_user_content = []
-    if message_text and message_text.strip():
-        llm_user_content.append({"type": "text", "text": message_text})
-    
-    if image_files: # image_files is a list of paths from gr.MultimodalTextbox
-        for img_path in image_files:
-            if img_path:
-                try:
-                    encoded_image = encode_image(img_path) # img_path is already a path
-                    if encoded_image:
-                        llm_user_content.append({
-                            "type": "image_url",
-                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}
-                        })
-                except Exception as e:
-                    print(f"Error encoding image for LLM: {e}")
+    # --- Initialize the LLM for the CodeAgent ---
+    agent_llm_params = {
+        "model_id": model_to_use,
+        "token": token_to_use,
+        # smolagents's InferenceClientModel uses max_tokens for max_new_tokens
+        "max_tokens": max_tokens,
+        "temperature": temperature if temperature > 0.01 else None, # Some models require temp > 0
+        "top_p": top_p if top_p < 1.0 else None, # Often 1.0 means no top_p
+        "seed": seed if seed != -1 else None,
+    }
+    if provider_for_agent_llm and provider_for_agent_llm != "hf-inference":
+        agent_llm_params["provider"] = provider_for_agent_llm
     
-    if not llm_user_content: # Should not happen if user() function filters empty messages
-        print("No content for LLM, aborting.")
-        yield {"type": "text", "content": "Please provide some input."}
+    # HFIC specific params, add if not default and supported
+    if frequency_penalty != 0.0:
+         agent_llm_params["frequency_penalty"] = frequency_penalty
+         
+    agent_llm = SmolInferenceClientModel(**agent_llm_params)
+    print(f"Smolagents LLM for agent initialized: model='{model_to_use}', provider='{provider_for_agent_llm or 'default'}'")
+
+    # --- Define Tools for the Agent ---
+    agent_tools = []
+    try:
+        image_gen_tool = Tool.from_space(
+            space_id="black-forest-labs/FLUX.1-schnell",
+            name="image_generator",
+            description="Generates an image from a textual prompt. Input is a single string argument named 'prompt'. Output is an image file path.",
+            token=token_to_use 
+        )
+        agent_tools.append(image_gen_tool)
+        print("Image generation tool loaded: black-forest-labs/FLUX.1-schnell")
+    except Exception as e:
+        print(f"Error loading image generation tool: {e}")
+        yield f"Error: Could not load image generation tool. {e}"
         return
 
-    messages_for_llm = [{"role": "system", "content": system_message}]
-    print("Initial messages array constructed for LLM.")
-
-    for val in history: # history item is [user_content_list, assistant_response_str_or_dict]
-        user_content_list_hist = val[0]
-        assistant_response_hist = val[1]
-
-        if user_content_list_hist:
-            # user_content_list_hist is already in the correct format (list of dicts)
-            messages_for_llm.append({"role": "user", "content": user_content_list_hist})
-        
-        if assistant_response_hist:
-            # Assistant response could be text or an image dict from a previous tool call
-            if isinstance(assistant_response_hist, dict) and assistant_response_hist.get("type") == "image":
-                 messages_for_llm.append({"role": "assistant", "content": [{"type": "text", "text": f"Assistant previously displayed image: {assistant_response_hist.get('path')}"}]})
-            elif isinstance(assistant_response_hist, str):
-                 messages_for_llm.append({"role": "assistant", "content": assistant_response_hist})
-            # Else, if it's a dict but not an image type we understand for history, we might skip or log an error
+    # --- Initialize the CodeAgent ---
+    # If system_message_for_agent is empty, CodeAgent will use its default.
+    # The default is usually good as it explains how to use tools.
+    agent = CodeAgent(
+        tools=agent_tools,
+        model=agent_llm,
+        system_prompt=system_message_for_agent if system_message_for_agent and system_message_for_agent.strip() else None,
+        # add_base_tools=True, # Consider adding Python interpreter, etc.
+        stream_outputs=True # Important for Gradio streaming
+    )
+    print("Smolagents CodeAgent initialized.")
 
-    messages_for_llm.append({"role": "user", "content": llm_user_content})
-    # print(f"Full messages_for_llm: {messages_for_llm}") # Can be very verbose
+    # --- Prepare task and image inputs for the agent ---
+    agent_task_text = message_text
+    
+    pil_images_for_agent = []
+    if image_file_paths:
+        for file_path in image_file_paths:
+            try:
+                pil_images_for_agent.append(Image.open(file_path))
+            except Exception as e:
+                print(f"Error opening image file {file_path} for agent: {e}")
+    
+    print(f"Agent task: '{agent_task_text}'")
+    if pil_images_for_agent:
+        print(f"Passing {len(pil_images_for_agent)} image(s) to agent.")
 
-    model_to_use = custom_model.strip() if custom_model.strip() != "" else selected_model
-    print(f"Model selected for LLM inference: {model_to_use}")
+    # --- Run the agent and stream response ---
+    # Agent is reset each turn. For conversational memory, agent instance
+    # would need to be stored in session_state and agent.run(..., reset=False) used.
+    
+    current_agent_response_text = ""
+    try:
+        # The agent.run method returns a generator when stream=True
+        for step_item in agent.run(
+            task=agent_task_text, 
+            images=pil_images_for_agent, 
+            stream=True, 
+            reset=True # Explicitly reset for stateless operation per call
+        ):
+            if isinstance(step_item, ChatMessageStreamDelta):
+                if step_item.content:
+                    current_agent_response_text += step_item.content
+                    yield current_agent_response_text # Yield accumulated text
+            
+            elif isinstance(step_item, (ActionStep, PlanningStep, FinalAnswerStep)):
+                # A structured step. Format it for Gradio.
+                # pull_messages_from_step yields gr.ChatMessage objects.
+                for gradio_chat_msg in pull_messages_from_step(step_item, skip_model_outputs=agent.stream_outputs):
+                    # The 'bot' function will handle these gr.ChatMessage objects.
+                    yield gradio_chat_msg # Yield the gr.ChatMessage object directly
+                current_agent_response_text = "" # Reset text buffer after a structured step
+            
+            # else:
+                # print(f"Unhandled stream item type: {type(step_item)}") # Debug
 
-    response_text = ""
-    print(f"Sending request to {provider} provider for LLM.")
+        # If there's any remaining text not part of a gr.ChatMessage, yield it.
+        # This usually shouldn't happen if stream_to_gradio logic is followed,
+        # as text deltas should be part of the last gr.ChatMessage or yielded before it.
+        # However, if the agent's final textual answer comes as pure deltas after all steps.
+        if current_agent_response_text and not isinstance(step_item, FinalAnswerStep):
+             # Check if the last yielded item already contains this text
+            if not (isinstance(step_item, gr.ChatMessage) and step_item.content == current_agent_response_text):
+                 yield current_agent_response_text
 
-    parameters = {
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "frequency_penalty": frequency_penalty,
-    }
-    
-    if seed is not None:
-        parameters["seed"] = seed
 
-    try:
-        stream = client.chat_completion(
-            model=model_to_use,
-            messages=messages_for_llm,
-            stream=True,
-            **parameters
-        )
-        
-        print("Received LLM tokens: ", end="", flush=True)
-        
-        for chunk in stream:
-            if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
-                if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
-                    token_text = chunk.choices[0].delta.content
-                    if token_text:
-                        print(token_text, end="", flush=True)
-                        response_text += token_text
-                        yield {"type": "text", "content": response_text}
-        
-        print()
     except Exception as e:
-        print(f"Error during LLM inference: {e}")
-        response_text += f"\nError: {str(e)}"
-        yield {"type": "text", "content": response_text}
+        error_message = f"Error during agent execution: {str(e)}"
+        print(error_message)
+        yield error_message # Yield the error message to be displayed in UI
 
-    print("Completed LLM response generation.")
+    print("Agent run completed.")
 
+
+# Function to validate provider selection based on BYOK
 def validate_provider(api_key, provider):
     if not api_key.strip() and provider != "hf-inference":
         return gr.update(value="hf-inference")
     return gr.update(value=provider)
 
+# GRADIO UI
 with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
     chatbot = gr.Chatbot(
         height=600, 
         show_copy_button=True, 
-        placeholder="Select a model and begin chatting. Now supports multiple inference providers and multimodal inputs. Try 'generate image of a cat playing chess'.",
+        placeholder="Select a model and begin chatting. Now uses smolagents with tools!",
         layout="panel",
-        bubble_full_width=False
+        bubble_full_width=False # For better display of images/files
     )
     print("Chatbot interface created.")
     
@@ -247,164 +205,199 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
     
     with gr.Accordion("Settings", open=False):
         system_message_box = gr.Textbox(
-            value="You are a helpful AI assistant that can understand images and text. If asked to generate an image, respond by saying you will call the image_generator tool.", 
-            placeholder="You are a helpful assistant.",
-            label="System Prompt"
+            value="You are a helpful AI assistant. You can generate images if asked. Be precise with your prompts for image generation.", 
+            placeholder="You are a helpful AI assistant.",
+            label="System Prompt for Agent"
         )
         
         with gr.Row():
             with gr.Column():
-                max_tokens_slider = gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max tokens")
-                temperature_slider = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-                top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+                max_tokens_slider = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max New Tokens")
+                temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.01, label="Temperature")
+                top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-P")
             with gr.Column():
                 frequency_penalty_slider = gr.Slider(minimum=-2.0, maximum=2.0, value=0.0, step=0.1, label="Frequency Penalty")
                 seed_slider = gr.Slider(minimum=-1, maximum=65535, value=-1, step=1, label="Seed (-1 for random)")
         
-        providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
-        provider_radio = gr.Radio(choices=providers_list, value="hf-inference", label="Inference Provider")
-        byok_textbox = gr.Textbox(value="", label="BYOK (Bring Your Own Key)", info="Enter a custom Hugging Face API key here. When empty, only 'hf-inference' provider can be used.", placeholder="Enter your Hugging Face API token", type="password")
-        custom_model_box = gr.Textbox(value="", label="Custom Model", info="(Optional) Provide a custom Hugging Face model path. Overrides any selected featured model.", placeholder="meta-llama/Llama-3.3-70B-Instruct")
-        model_search_box = gr.Textbox(label="Filter Models", placeholder="Search for a featured model...", lines=1)
+        providers_list = [
+            "hf-inference", "cerebras", "together", "sambanova", "novita", 
+            "cohere", "fireworks-ai", "hyperbolic", "nebius",
+        ]
+        provider_radio = gr.Radio(choices=providers_list, value="hf-inference", label="Inference Provider for Agent's LLM")
+        byok_textbox = gr.Textbox(value="", label="BYOK (Your HF Token or Provider API Key)", info="Enter API key for the selected provider. Uses HF_TOKEN if empty.", placeholder="Enter your API token", type="password")
+        custom_model_box = gr.Textbox(value="", label="Custom Model ID for Agent's LLM", info="(Optional) Provide a custom model ID. Overrides featured model.", placeholder="meta-llama/Llama-3.3-70B-Instruct")
+        model_search_box = gr.Textbox(label="Filter Featured Models", placeholder="Search for a featured model...", lines=1)
         
         models_list = [
-            "meta-llama/Llama-3.2-11B-Vision-Instruct", "meta-llama/Llama-3.3-70B-Instruct", "meta-llama/Llama-3.1-70B-Instruct",
-            "meta-llama/Llama-3.0-70B-Instruct", "meta-llama/Llama-3.2-3B-Instruct", "meta-llama/Llama-3.2-1B-Instruct",
+            "meta-llama/Llama-3.3-70B-Instruct", "meta-llama/Llama-3.1-70B-Instruct", "meta-llama/Llama-3.0-70B-Instruct",
+            "meta-llama/Llama-3.2-11B-Vision-Instruct", "meta-llama/Llama-3.2-3B-Instruct", "meta-llama/Llama-3.2-1B-Instruct",
             "meta-llama/Llama-3.1-8B-Instruct", "NousResearch/Hermes-3-Llama-3.1-8B", "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
             "mistralai/Mistral-Nemo-Instruct-2407", "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.3",
-            "mistralai/Mistral-7B-Instruct-v0.2", "Qwen/Qwen3-235B-A22B", "Qwen/Qwen3-32B", "Qwen/Qwen2.5-72B-Instruct",
-            "Qwen/Qwen2.5-3B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct", "Qwen/QwQ-32B", "Qwen/Qwen2.5-Coder-32B-Instruct",
-            "microsoft/Phi-3.5-mini-instruct", "microsoft/Phi-3-mini-128k-instruct", "microsoft/Phi-3-mini-4k-instruct",
+            "Qwen/Qwen3-235B-A22B", "Qwen/Qwen3-32B", "Qwen/Qwen2.5-72B-Instruct", "Qwen/Qwen2.5-3B-Instruct",
+            "Qwen/Qwen2.5-Coder-32B-Instruct", "microsoft/Phi-3.5-mini-instruct", "microsoft/Phi-3-mini-128k-instruct",
         ]
-        featured_model_radio = gr.Radio(label="Select a model below", choices=models_list, value="meta-llama/Llama-3.2-11B-Vision-Instruct", interactive=True)
+        featured_model_radio = gr.Radio(label="Select a Featured Model for Agent's LLM", choices=models_list, value="meta-llama/Llama-3.3-70B-Instruct", interactive=True)
         
         gr.Markdown("[View all Text-to-Text models](https://huggingface.co/models?inference_provider=all&pipeline_tag=text-generation&sort=trending) | [View all multimodal models](https://huggingface.co/models?inference_provider=all&pipeline_tag=image-text-to-text&sort=trending)")
 
-    chat_history = gr.State([])
-    
-    def filter_models(search_term):
-        print(f"Filtering models with search term: {search_term}")
-        filtered = [m for m in models_list if search_term.lower() in m.lower()]
-        print(f"Filtered models: {filtered}")
-        return gr.update(choices=filtered)
-
-    def set_custom_model_from_radio(selected):
-        print(f"Featured model selected: {selected}")
-        return selected
-
-    def user(user_multimodal_input, history):
-        print(f"User input (raw from gr.MultimodalTextbox): {user_multimodal_input}")
-        
-        text_content = user_multimodal_input.get("text", "").strip()
-        files = user_multimodal_input.get("files", []) # These are temp file paths from Gradio
-
-        if not text_content and not files:
-            print("Empty input, skipping history append.")
-            # Optionally, could raise gr.Error("Please enter a message or upload an image.")
-            # For now, let's allow the bot to respond if history is not empty,
-            # or do nothing if history is also empty.
-            return history
+    # Chat history state (using gr.State to manage it properly)
+    # The chatbot's value itself will be the history display.
+    # We might need a separate gr.State if agent needs to be conversational across turns.
+    # For now, agent is stateless per turn.
 
-        # Prepare content for history: a list of dicts for multimodal display
-        history_user_entry_content = []
-        if text_content:
-            history_user_entry_content.append({"type": "text", "text": text_content})
+    # Function for the chat interface
+    def user(user_multimodal_input_dict, history):
+        print(f"User input: {user_multimodal_input_dict}")
+        text_content = user_multimodal_input_dict.get("text", "")
+        files = user_multimodal_input_dict.get("files", [])
         
-        for file_path_obj in files: # file_path_obj is a FileData object from Gradio
-            if file_path_obj and hasattr(file_path_obj, 'name') and file_path_obj.name:
-                # Gradio's Chatbot can display images directly from file paths
-                # We store it in a format that `respond` can also understand
-                # The path is temporary, Gradio handles making it accessible for display
-                history_user_entry_content.append({"type": "image_url", "image_url": {"url": file_path_obj.name}})
-                print(f"Adding image to history entry: {file_path_obj.name}")
+        user_display_parts = []
+        if text_content and text_content.strip():
+            user_display_parts.append(text_content)
+        for file_path_obj in files: # file_path_obj is a tempfile._TemporaryFileWrapper
+            user_display_parts.append((file_path_obj.name, os.path.basename(file_path_obj.name)))
             
-        if history_user_entry_content:
-            history.append([history_user_entry_content, None]) # User part, Bot part (initially None)
-        
+        if not user_display_parts:
+            return history
+            
+        # Append the user's multimodal message to history for display
+        # The actual data (dict) is passed to `bot` function separately.
+        history.append([user_display_parts if len(user_display_parts) > 1 else user_display_parts[0], None])
         return history
-    
+
     def bot(history, system_msg, max_tokens, temperature, top_p, freq_penalty, seed, provider, api_key, custom_model, search_term, selected_model):
-        if not history or not history[-1][0]: # If no user message or empty user message content
-            print("No user message to process in bot function or user message content is empty.")
-            yield history # Return current history without processing
+        if not history or not history[-1][0]: # If no user input
+            yield history
             return
 
-        user_content_list = history[-1][0] # This is now a list of content dicts
+        # The user's input (text and list of file paths) is in history[-1][0]
+        # If `user` function stores the dict:
+        raw_user_input_dict = history[-1][0] if isinstance(history[-1][0], dict) else {"text": str(history[-1][0]), "files": []}
         
-        # Extract text and image file paths from the user_content_list for the `respond` function
-        text_for_respond = ""
-        image_files_for_respond = []
+        # If `user` function stores formatted display parts:
+        # We need to reconstruct or rely on msg input to bot.
+        # For now, assuming msg.submit passes the raw dict.
+        # Let's adjust the Gradio flow to pass `msg` directly to `bot` as well.
+
+        # The `msg` variable in `msg.submit` holds the raw MultimodalTextbox output.
+        # We need to pass this raw dict to `respond`.
+        # The `history` is for display.
+        
+        # This part is tricky as `bot` gets `history` which is already formatted for display.
+        # A common pattern is to pass `msg` (raw input) also to `bot`.
+        # Let's assume `history[-1][0]` contains enough info or we adjust `user` fn.
+        # For simplicity, let's assume `user` stores the raw dict if needed,
+        # or `bot` can parse `history[-1][0]` if it's a string/list of tuples.
+
+        # Let's assume `history[-1][0]` is the raw `user_multimodal_input_dict`
+        # This means the `user` function must append it like: `history.append([user_multimodal_input_dict, None])`
+        # And the chatbot will display `str(user_multimodal_input_dict)`.
+        # This is what the current `user` function does.
+
+        user_input_data = history[-1][0] # This should be the dict from MultimodalTextbox
+        text_input_for_agent = user_input_data.get("text", "")
+        # Files from MultimodalTextbox are temp file paths
+        image_file_paths_for_agent = [f.name for f in user_input_data.get("files", []) if hasattr(f, 'name')]
 
-        for item in user_content_list:
-            if item["type"] == "text":
-                text_for_respond = item["text"]
-            elif item["type"] == "image_url": 
-                image_files_for_respond.append(item["image_url"]["url"])
 
-        history[-1][1] = "" # Clear placeholder for bot response / Initialize bot response
+        history[-1][1] = "" # Initialize assistant's part for streaming
         
-        # Call the respond function which is now a generator
-        for response_chunk in respond(
-            text_for_respond, 
-            image_files_for_respond, 
-            history[:-1], # Pass previous history
-            system_msg, max_tokens, temperature, top_p, freq_penalty, seed,
-            provider, api_key, custom_model, search_term, selected_model
+        # Buffer for current text stream from agent
+        # Handles both pure text deltas and text content from gr.ChatMessage
+        current_text_for_turn = ""
+
+        for item in respond(
+            message_text=text_input_for_agent,
+            image_file_paths=image_file_paths_for_agent,
+            gradio_history=history[:-1], # Pass previous turns for context if agent uses it
+            system_message_for_agent=system_msg,
+            max_tokens=max_tokens, temperature=temperature, top_p=top_p,
+            frequency_penalty=freq_penalty, seed=seed,
+            provider_for_agent_llm=provider, api_key_for_agent_llm=api_key,
+            model_id_for_agent_llm=custom_model,
+            model_search_term=search_term, # unused
+            selected_model_for_agent_llm=selected_model
         ):
-            current_bot_response = history[-1][1]
-            if isinstance(response_chunk, dict):
-                if response_chunk["type"] == "text":
-                    # If current bot response is already an image dict, we can't append text.
-                    # This indicates a new text response after an image, or just text.
-                    if isinstance(current_bot_response, dict) and current_bot_response.get("type") == "image":
-                        # This case should ideally not happen if an image is the final response from a tool.
-                        # If it does, we might need to start a new bot message in history.
-                        # For now, we'll overwrite if the new chunk is text.
-                        history[-1][1] = response_chunk["content"]
-                    elif isinstance(current_bot_response, str):
-                         history[-1][1] = response_chunk["content"] # Accumulate text
-                    else: # current_bot_response is likely "" or None
-                         history[-1][1] = response_chunk["content"]
-
-                elif response_chunk["type"] == "image":
-                    # Image response from tool. Gradio Chatbot displays this as an image.
-                    # The path should be accessible by Gradio.
-                    # If there was prior text content for this turn, it's now overwritten by the image.
-                    # This means a tool call that produces an image is considered the primary response for that turn.
-                    history[-1][1] = {"path": response_chunk["path"], "mime_type": "image/jpeg"} # Assuming JPEG, could be PNG
+            if isinstance(item, str): # LLM text delta from agent's thought or textual answer
+                current_text_for_turn = item 
+                history[-1][1] = current_text_for_turn
+            elif isinstance(item, gr.ChatMessage):
+                # This is a structured step (thought, tool output, image, etc.)
+                # We need to append this to the history as a new message or part of current message.
+                # For simplicity, let's append its string content to the current turn's assistant message.
+                # If it's an image/file, we'll represent it as a markdown link.
+                if isinstance(item.content, str):
+                    current_text_for_turn = item.content # Replace if it's a full message
+                elif isinstance(item.content, dict) and "path" in item.content:
+                    # This is typically an image or audio file
+                    file_path = item.content["path"]
+                    # We need to make this file accessible to Gradio if it's temporary from agent
+                    # For now, just put a placeholder.
+                    # If it's an output from a tool, the path might be relative to where smolagents saves it.
+                    # Gradio needs an absolute path or a URL.
+                    # A common pattern is to copy temp files to a static dir served by Gradio or use gr.File.
+                    # For now, let's assume Gradio can handle local paths if they are in a folder it knows.
+                    # We'll display it as a tuple for Gradio Chatbot.
+                    # This means history[-1][1] needs to become a list.
+                    
+                    # If current_text_for_turn is not empty, make history[-1][1] a list
+                    if current_text_for_turn and not isinstance(history[-1][1], list):
+                        history[-1][1] = [current_text_for_turn]
+                    elif not current_text_for_turn and not isinstance(history[-1][1], list):
+                         history[-1][1] = []
+
+
+                    alt_text = item.metadata.get("title", os.path.basename(file_path)) if item.metadata else os.path.basename(file_path)
+                    
+                    # Add as new component to the list for current assistant message
+                    if isinstance(history[-1][1], list):
+                        history[-1][1].append((file_path, alt_text))
+                    else: # Should have been made a list above
+                        history[-1][1] = [(file_path, alt_text)]
+                    
+                    current_text_for_turn = "" # Reset text buffer after a file
+                
+                # If it's not a delta, but a full message, replace the current text
+                if not isinstance(history[-1][1], list): # if it hasn't become a list due to file
+                    history[-1][1] = current_text_for_turn
+
             yield history
-            
+
+    # Event handlers
+    # `msg.submit`'s first argument is the function to call.
+    # Its `inputs` are the Gradio components whose values are passed to the function.
+    # Its `outputs` are the Gradio components that are updated by the function's return value.
+    # The `user` function now appends the raw dict from MultimodalTextbox to history.
+    # The `bot` function takes this history.
+    
+    # When msg is submitted:
+    # 1. Call `user` to update history with user's input. Output is `chatbot`.
+    # 2. Then call `bot` with the updated history. Output is `chatbot`.
+    # 3. Then clear `msg`
     msg.submit(
         user,
         [msg, chatbot],
-        [chatbot],
+        [chatbot], # `user` returns the new history, updating the chatbot display
         queue=False
     ).then(
         bot,
         [chatbot, system_message_box, max_tokens_slider, temperature_slider, top_p_slider, 
          frequency_penalty_slider, seed_slider, provider_radio, byok_textbox, custom_model_box, 
          model_search_box, featured_model_radio],
-        [chatbot]
+        [chatbot] # `bot` yields history updates, streaming to chatbot
     ).then(
-        lambda: {"text": "", "files": []}, # Clear MultimodalTextbox
+        lambda: {"text": "", "files": []},  # Clear MultimodalTextbox
         None,
         [msg]
     )
     
     model_search_box.change(fn=filter_models, inputs=model_search_box, outputs=featured_model_radio)
-    print("Model search box change event linked.")
-
     featured_model_radio.change(fn=set_custom_model_from_radio, inputs=featured_model_radio, outputs=custom_model_box)
-    print("Featured model radio button change event linked.")
-    
     byok_textbox.change(fn=validate_provider, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
-    print("BYOK textbox change event linked.")
-
     provider_radio.change(fn=validate_provider, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
-    print("Provider radio button change event linked.")
 
 print("Gradio interface initialized.")
 
 if __name__ == "__main__":
     print("Launching the demo application.")
-    demo.launch(show_api=True)
+    demo.launch(show_api=False) # show_api=False for cleaner launch, True for API docs