Update app.py

app.py

@@ -6,41 +6,24 @@ import base64
 from PIL import Image
 import io
 
-# Import smolagents components
-from smolagents import CodeAgent, Tool
-from smolagents.models import InferenceClientModel as SmolInferenceClientModel # Alias to avoid conflict
+# Import smolagents Tool
+from smolagents import Tool
 
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 print("Access token loaded.")
 
-# --- Smolagents Setup for Image Generation ---
-print("Initializing smolagents components for image generation...")
+# 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", # The Space ID of the image generation tool
+        "black-forest-labs/FLUX.1-schnell",
         name="image_generator",
-        description="Generates an image from a textual prompt. Use this tool if the user asks to 'generate an image of X', 'draw X', 'create a picture of X', or similar requests for visual content based on a description.",
-        # Ensure the HF_TOKEN is available to gradio-client if the space is private or requires auth
-        token=ACCESS_TOKEN if ACCESS_TOKEN and ACCESS_TOKEN.strip() != "" else None 
+        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.")
-    
-    # Initialize a model for the CodeAgent. This can be a simpler/faster model
-    # as it's mainly for orchestrating the tool call.
-    # Using a default InferenceClientModel from smolagents
-    smol_agent_model = SmolInferenceClientModel(token=ACCESS_TOKEN if ACCESS_TOKEN and ACCESS_TOKEN.strip() != "" else None)
-    print(f"Smolagent model initialized with: {smol_agent_model.model_id if hasattr(smol_agent_model, 'model_id') else 'default'}")
-    
-    image_agent = CodeAgent(
-        tools=[image_generation_tool], 
-        model=smol_agent_model,
-        verbosity_level=1 # Set to 0 for less verbose agent logging, 1 for info, 2 for debug
-    )
-    print("Image generation agent initialized successfully.")
 except Exception as e:
-    print(f"Error initializing smolagents components: {e}")
-    image_agent = None
-# --- End Smolagents Setup ---
+    print(f"Error loading image generation tool: {e}")
+    image_generation_tool = None
 
 # Function to encode image to base64
 def encode_image(image_path):
@@ -64,7 +47,7 @@ def encode_image(image_path):
         
         # Encode to base64
         buffered = io.BytesIO()
-        image.save(buffered, format="JPEG")
+        image.save(buffered, format="JPEG") 
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
         print("Image encoded successfully")
         return img_str
@@ -73,9 +56,9 @@ def encode_image(image_path):
         return None
 
 def respond(
-    message,
-    image_files,  # Changed parameter name and structure
-    history: list[tuple[str, str]],
+    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,
     max_tokens,
     temperature,
@@ -88,9 +71,9 @@ def respond(
     model_search_term,
     selected_model
 ):
-    print(f"Received message: {message}")
-    print(f"Received {len(image_files) if image_files else 0} images")
-    print(f"History: {history}")
+    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}")
@@ -100,136 +83,106 @@ def respond(
     print(f"Model search term: {model_search_term}")
     print(f"Selected model from radio: {selected_model}")
 
-    # --- Agent-based Image Generation ---
-    if message.startswith("/generate_image"):
-        if image_agent is None:
-            yield "Image generation agent is not initialized. Please check server logs."
-            return
+    # 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 ""
 
-        prompt_for_agent = message.replace("/generate_image", "").strip()
-        if not prompt_for_agent:
-            yield "Please provide a prompt for image generation. Usage: /generate_image <your prompt>"
+    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
 
-        print(f"Image generation requested with prompt: {prompt_for_agent}")
         try:
-            # Agent run is blocking and returns the final result
-            # Ensure the image_agent's model also has a token if needed for its operations (though it's for orchestration)
-            agent_response = image_agent.run(prompt_for_agent)
-            
-            if isinstance(agent_response, str) and agent_response.lower().startswith("error"): 
-                yield f"Agent error: {agent_response}"
-            elif hasattr(agent_response, 'to_string'): # Check if it's an AgentImage or similar
-                image_path = agent_response.to_string() # This is a local path to the generated image
-                print(f"Agent returned image path: {image_path}")
-                # Gradio's chatbot can display images if the content is a file path string
-                # or a tuple (filepath, alt_text)
-                yield image_path
-            else:
-                yield f"Agent returned an unexpected response: {str(agent_response)}"
+            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 running image agent: {e}")
-            yield f"Error generating image: {str(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
-    # --- End Agent-based Image Generation ---
+    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
 
-    # Determine which token to use for text generation
-    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")
-    
-    # Initialize the Inference Client with the provider and appropriate token
+    # 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 for text generation.")
+    print(f"Hugging Face Inference Client initialized with {provider} provider.")
 
-    # Convert seed to None if -1 (meaning random)
     if seed == -1:
         seed = None
 
-    # Create multimodal content if images are present
-    if image_files and len(image_files) > 0:
-        user_content = []
-        if message and message.strip():
-            user_content.append({
-                "type": "text",
-                "text": message
-            })
-        for img_path in image_files: # Assuming image_files contains paths from MultimodalTextbox
-            if img_path is not 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 from MultimodalTextbox
+                    encoded_image = encode_image(img_path) # img_path is already a path
                     if encoded_image:
-                        user_content.append({
+                        llm_user_content.append({
                             "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/jpeg;base64,{encoded_image}"
-                            }
+                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}
                         })
                 except Exception as e:
-                    print(f"Error encoding image: {e}")
-    else:
-        # Text-only message
-        user_content = message
+                    print(f"Error encoding image for LLM: {e}")
+    
+    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."}
+        return
 
-    # Prepare messages in the format expected by the API
-    messages = [{"role": "system", "content": system_message}]
-    print("Initial messages array constructed.")
+    messages_for_llm = [{"role": "system", "content": system_message}]
+    print("Initial messages array constructed for LLM.")
 
-    # Add conversation history to the context
-    for val in history:
-        user_part = val[0]
-        assistant_part = val[1]
-        
-        # Handle user messages (could be text or image markdown)
-        if user_part:
-            if isinstance(user_part, str) and user_part.startswith("![Image]("):
-                # This is an image path from a previous agent generation
-                # or a user upload represented as markdown
-                history_image_path = user_part.replace("![Image](", "").replace(")", "")
-                encoded_history_image = encode_image(history_image_path)
-                if encoded_history_image:
-                     messages.append({"role": "user", "content": [{
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/jpeg;base64,{encoded_history_image}"}
-                    }]})
-            elif isinstance(user_part, tuple) and len(user_part) == 2: # Multimodal input from user
-                history_content_list = []
-                if user_part[0]: # Text part
-                    history_content_list.append({"type": "text", "text": user_part[0]})
-                for img_hist_path in user_part[1]: # List of image paths
-                    encoded_img_hist = encode_image(img_hist_path)
-                    if encoded_img_hist:
-                        history_content_list.append({
-                            "type": "image_url",
-                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_img_hist}"}
-                        })
-                if history_content_list:
-                    messages.append({"role": "user", "content": history_content_list})
-            else: # Regular text message
-                messages.append({"role": "user", "content": user_part})
-            print(f"Added user message to context (type: {type(user_part)})")
+    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_part:
-            messages.append({"role": "assistant", "content": assistant_part})
-            print(f"Added assistant message to context: {assistant_part}")
+        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
 
-    # Append the latest user message
-    messages.append({"role": "user", "content": user_content})
-    print(f"Latest user message appended (content type: {type(user_content)})")
+    messages_for_llm.append({"role": "user", "content": llm_user_content})
+    # print(f"Full messages_for_llm: {messages_for_llm}") # Can be very verbose
 
-    # Determine which model to use, prioritizing custom_model if provided
     model_to_use = custom_model.strip() if custom_model.strip() != "" else selected_model
-    print(f"Model selected for inference: {model_to_use}")
+    print(f"Model selected for LLM inference: {model_to_use}")
 
-    # Start with an empty string to build the response as tokens stream in
-    response = ""
-    print(f"Sending request to {provider} provider.")
+    response_text = ""
+    print(f"Sending request to {provider} provider for LLM.")
 
-    # Prepare parameters for the chat completion request
     parameters = {
         "max_tokens": max_tokens,
         "temperature": temperature,
@@ -240,58 +193,50 @@ def respond(
     if seed is not None:
         parameters["seed"] = seed
 
-    # Use the InferenceClient for making the request
     try:
-        # Create a generator for the streaming response
         stream = client.chat_completion(
             model=model_to_use,
-            messages=messages,
+            messages=messages_for_llm,
             stream=True,
             **parameters
         )
         
-        print("Received tokens: ", end="", flush=True)
+        print("Received LLM tokens: ", end="", flush=True)
         
-        # Process the streaming response
         for chunk in stream:
             if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
-                # Extract the content from the response
                 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 += token_text
-                        yield response
+                        response_text += token_text
+                        yield {"type": "text", "content": response_text}
         
         print()
     except Exception as e:
-        print(f"Error during inference: {e}")
-        response += f"\nError: {str(e)}"
-        yield response
+        print(f"Error during LLM inference: {e}")
+        response_text += f"\nError: {str(e)}"
+        yield {"type": "text", "content": response_text}
 
-    print("Completed response generation.")
+    print("Completed LLM response generation.")
 
-# 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:
-    # Create the chatbot component
     chatbot = gr.Chatbot(
         height=600, 
         show_copy_button=True, 
-        placeholder="Select a model and begin chatting. Use '/generate_image your prompt' to create images.",
+        placeholder="Select a model and begin chatting. Now supports multiple inference providers and multimodal inputs. Try 'generate image of a cat playing chess'.",
         layout="panel",
-        show_share_button=True # Added for ease of sharing if deployed
+        bubble_full_width=False
     )
     print("Chatbot interface created.")
     
-    # Multimodal textbox for messages (combines text and file uploads)
     msg = gr.MultimodalTextbox(
-        placeholder="Type a message or upload images... (e.g., /generate_image a cat wearing a hat)",
+        placeholder="Type a message or upload images...",
         show_label=False,
         container=False,
         scale=12,
@@ -300,226 +245,138 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
         sources=["upload"]
     )
     
-    # Create accordion for settings
     with gr.Accordion("Settings", open=False):
-        # System message
         system_message_box = gr.Textbox(
-            value="You are a helpful AI assistant that can understand images and text. If asked to generate an image, use the image_generator tool.", 
+            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"
         )
         
-        # Generation parameters
         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=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")
             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)"
-                )
+                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)")
         
-        # Provider selection
-        providers_list = [
-            "hf-inference",  # Default Hugging Face Inference
-            "cerebras",      # Cerebras provider
-            "together",      # Together AI
-            "sambanova",     # SambaNova
-            "novita",        # Novita AI
-            "cohere",        # Cohere
-            "fireworks-ai",  # Fireworks AI
-            "hyperbolic",    # Hyperbolic
-            "nebius",        # Nebius
-        ]
-        
-        provider_radio = gr.Radio(
-            choices=providers_list,
-            value="hf-inference",
-            label="Inference Provider",
-        )
-        
-        # New BYOK textbox
-        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"  # Hide the API key for security
-        )
-        
-        # Custom model box
-        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"
-        )
+        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)
         
-        # Model search
-        model_search_box = gr.Textbox(
-            label="Filter Models",
-            placeholder="Search for a featured model...",
-            lines=1
-        )
-        
-        # Featured models list
-        # Updated to include multimodal models
         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.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",
+            "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.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",
         ]
-
-        featured_model_radio = gr.Radio(
-            label="Select a model below",
-            choices=models_list,
-            value="meta-llama/Llama-3.2-11B-Vision-Instruct",  # Default to a multimodal model
-            interactive=True
-        )
+        featured_model_radio = gr.Radio(label="Select a model below", choices=models_list, value="meta-llama/Llama-3.2-11B-Vision-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 state
     chat_history = gr.State([])
     
-    # Function to filter models
     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 if filtered else models_list, value=filtered[0] if filtered else models_list[0])
+        return gr.update(choices=filtered)
 
-
-    # Function to set custom model from radio
     def set_custom_model_from_radio(selected):
         print(f"Featured model selected: {selected}")
         return selected
 
-    # Function for the chat interface
-    def user(user_message_obj, history):
-        print(f"User message object received: {user_message_obj}")
-        
-        text_content = user_message_obj.get("text", "").strip()
-        files = user_message_obj.get("files", []) # files is a list of temp file paths
+    def user(user_multimodal_input, history):
+        print(f"User input (raw from gr.MultimodalTextbox): {user_multimodal_input}")
         
-        if not text_content and not files:
-            print("Empty message (no text, no files), skipping history update.")
-            return history # Or raise gr.Error("Please enter a message or upload an image.")
+        text_content = user_multimodal_input.get("text", "").strip()
+        files = user_multimodal_input.get("files", []) # These are temp file paths from Gradio
 
-        # Represent uploaded images in history using markdown syntax for local paths
-        # For multimodal models, the actual file path from 'files' will be used in 'respond'
-        display_message_parts = []
+        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
+
+        # Prepare content for history: a list of dicts for multimodal display
+        history_user_entry_content = []
         if text_content:
-            display_message_parts.append(text_content)
+            history_user_entry_content.append({"type": "text", "text": text_content})
+        
+        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}")
+            
+        if history_user_entry_content:
+            history.append([history_user_entry_content, None]) # User part, Bot part (initially None)
         
-        processed_files_for_history = []
-        if files:
-            for file_path_obj in files:
-                # Gradio's MultimodalTextbox provides file objects with a .name attribute for the path
-                file_path = file_path_obj.name if hasattr(file_path_obj, 'name') else str(file_path_obj)
-                display_message_parts.append(f"![Uploaded Image]({file_path})")
-                processed_files_for_history.append(file_path) # Store the actual path for 'respond'
-
-        # For history, we store the text and a list of file paths
-        # The 'respond' function will then re-encode these for the API
-        history_entry_user = (text_content, processed_files_for_history)
-        history.append([history_entry_user, None])
-        print(f"History updated with user input: {history_entry_user}")
         return history
     
-    # Define bot response function
     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 len(history) == 0 or history[-1][0] is None:
-            print("No user message in history to process for bot.")
-            yield history
+        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
             return
 
-        user_input_tuple = history[-1][0] # This is now (text, [file_paths])
-        text_message_from_history = user_input_tuple[0]
-        image_files_from_history = user_input_tuple[1]
-
-        print(f"Bot processing: text='{text_message_from_history}', images={image_files_from_history}")
+        user_content_list = history[-1][0] # This is now a list of content dicts
         
-        history[-1][1] = ""
+        # Extract text and image file paths from the user_content_list for the `respond` function
+        text_for_respond = ""
+        image_files_for_respond = []
+
+        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
         
-        # Pass text and image file paths to respond function
+        # Call the respond function which is now a generator
         for response_chunk in respond(
-            message=text_message_from_history,
-            image_files=image_files_from_history, 
-            history=history[:-1], # Pass history excluding the current user turn
-            system_message=system_msg,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            frequency_penalty=freq_penalty,
-            seed=seed,
-            provider=provider,
-            custom_api_key=api_key,
-            custom_model=custom_model,
-            model_search_term=search_term,
-            selected_model=selected_model
+            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
         ):
-            history[-1][1] = response_chunk
+            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
             yield history
-
-    # Event handlers
+            
     msg.submit(
         user,
-        [msg, chatbot], # msg is MultimodalTextboxOutput(text=str, files=List[FileData])
+        [msg, chatbot],
         [chatbot],
         queue=False
     ).then(
@@ -529,45 +386,25 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
          model_search_box, featured_model_radio],
         [chatbot]
     ).then(
-        lambda: gr.update(value={"text": "", "files": []}),  # Clear MultimodalTextbox
+        lambda: {"text": "", "files": []}, # Clear MultimodalTextbox
         None,
         [msg]
     )
     
-    # Connect the model filter to update the radio choices
-    model_search_box.change(
-        fn=filter_models,
-        inputs=model_search_box,
-        outputs=featured_model_radio
-    )
+    model_search_box.change(fn=filter_models, inputs=model_search_box, outputs=featured_model_radio)
     print("Model search box change event linked.")
 
-    # Connect the featured model radio to update the custom model box
-    featured_model_radio.change(
-        fn=set_custom_model_from_radio,
-        inputs=featured_model_radio,
-        outputs=custom_model_box
-    )
+    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.")
     
-    # Connect the BYOK textbox to validate provider selection
-    byok_textbox.change(
-        fn=validate_provider,
-        inputs=[byok_textbox, provider_radio],
-        outputs=provider_radio
-    )
+    byok_textbox.change(fn=validate_provider, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
     print("BYOK textbox change event linked.")
 
-    # Also validate provider when the radio changes to ensure consistency
-    provider_radio.change(
-        fn=validate_provider,
-        inputs=[byok_textbox, provider_radio],
-        outputs=provider_radio
-    )
+    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=False) # show_api=False for cleaner public interface, True for debugging
+    demo.launch(show_api=True)