Update app.py

app.py

@@ -1,11 +1,14 @@
 # -*- coding: utf-8 -*-
 import os
 import gradio as gr
+# Corrected import: Import the main module and use an alias
 import google.generativeai as genai
-from google import genai
-# Make sure to import necessary types from the SDK
-from google.generative_ai import types
-from google.generative_ai.types import HarmCategory, HarmBlockThreshold # For safety settings
+# Types will be accessed via genai.types
+
+# Removed direct type imports, will use genai.types instead
+# from google.generativeai import types # No longer needed
+# from google.generativeai.types import HarmCategory, HarmBlockThreshold # No longer needed
+
 import requests
 import markdownify
 from urllib.robotparser import RobotFileParser
@@ -17,7 +20,6 @@ import json # Although not directly used in the final code, useful for debugging
 
 def can_crawl_url(url: str, user_agent: str = "PythonGoogleGenAIAgent/1.0") -> bool:
     """Check robots.txt permissions for a URL"""
-    # Use a more specific user agent, but '*' is a fallback
     if not url:
         print("No URL provided to can_crawl_url")
         return False
@@ -25,12 +27,9 @@ def can_crawl_url(url: str, user_agent: str = "PythonGoogleGenAIAgent/1.0") -> b
         parsed_url = urlparse(url)
         if not parsed_url.scheme or not parsed_url.netloc:
             print(f"Invalid URL format for robots.txt check: {url}")
-            return False # Cannot determine robots.txt location
-
+            return False
         robots_url = f"{parsed_url.scheme}://{parsed_url.netloc}/robots.txt"
         print(f"Checking robots.txt at: {robots_url} for URL: {url}")
-
-        # Using RobotFileParser's default opener which handles redirects
         rp = RobotFileParser()
         rp.set_url(robots_url)
         rp.read()
@@ -39,7 +38,6 @@ def can_crawl_url(url: str, user_agent: str = "PythonGoogleGenAIAgent/1.0") -> b
         return can_fetch
     except Exception as e:
         print(f"Error checking robots.txt for {url}: {e}")
-        # Default to false if unsure, to be polite to servers
         return False
 
 def load_page(url: str) -> str:
@@ -56,27 +54,22 @@ def load_page(url: str) -> str:
     if not url.startswith(('http://', 'https://')):
          return f"Error: Invalid URL scheme. Please provide http or https URL. Got: {url}"
 
-    USER_AGENT = "PythonGoogleGenAIAgent/1.0 (Function Calling)" # Be identifiable
+    USER_AGENT = "PythonGoogleGenAIAgent/1.0 (Function Calling)"
     if not can_crawl_url(url, user_agent=USER_AGENT):
         print(f"URL {url} failed robots.txt check for agent {USER_AGENT}")
         return f"Error: Access denied by robots.txt for URL {url}"
     try:
         headers = {'User-Agent': USER_AGENT}
         response = requests.get(url, timeout=15, headers=headers, allow_redirects=True)
-        response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx)
-
-        # Check content type - try to only process HTML
+        response.raise_for_status()
         content_type = response.headers.get('content-type', '').lower()
         if 'html' not in content_type:
             print(f"Non-HTML content type '{content_type}' at {url}. Returning summary.")
-            # Return limited info for non-html types
             return f"Content at {url} is of type '{content_type}'. Size: {len(response.content)} bytes. Cannot convert to Markdown."
 
-        # Limit content size before markdown conversion to avoid excessive memory/CPU
-        MAX_CONTENT_SIZE = 1_000_000 # 1MB limit
+        MAX_CONTENT_SIZE = 1_000_000
         if len(response.content) > MAX_CONTENT_SIZE:
              print(f"Content size {len(response.content)} exceeds limit {MAX_CONTENT_SIZE}. Truncating.")
-             # Decode potentially large content carefully
              try:
                  html_content = response.content[:MAX_CONTENT_SIZE].decode(response.apparent_encoding or 'utf-8', errors='ignore')
              except Exception as decode_err:
@@ -84,18 +77,12 @@ def load_page(url: str) -> str:
                  html_content = response.content[:MAX_CONTENT_SIZE].decode('utf-8', errors='ignore')
              truncated_msg = "\n\n[Content truncated due to size limit]"
         else:
-            html_content = response.text # Use response.text which handles encoding better for smaller content
+            html_content = response.text
             truncated_msg = ""
 
-        # Convert to Markdown
-        # Added heading_style for potentially better formatting
         markdown_content = markdownify.markdownify(html_content, heading_style="ATX", strip=['script', 'style'], escape_underscores=False)
-
-        # Simple cleaning (optional, can be expanded)
         markdown_content = '\n'.join([line.strip() for line in markdown_content.splitlines() if line.strip()])
-
         print(f"Successfully loaded and converted {url} to markdown.")
-        # Add URL source attribution
         return f"Content from {url}:\n\n" + markdown_content + truncated_msg
 
     except requests.exceptions.Timeout:
@@ -106,7 +93,7 @@ def load_page(url: str) -> str:
         return f"Error loading page {url}: {str(e)}"
     except Exception as e:
         print(f"General error loading page {url}: {str(e)}")
-        traceback.print_exc() # Print full traceback for debugging
+        traceback.print_exc()
         return f"Error loading page {url}: An unexpected error occurred ({type(e).__name__})."
 
 
@@ -115,50 +102,43 @@ try:
     api_key = os.environ.get("GEMINI_API_KEY")
     if not api_key:
         raise ValueError("GEMINI_API_KEY environment variable not set.")
+    # Use genai (the alias) to configure
     genai.configure(api_key=api_key)
 
-    # *** Use the requested experimental model ***
     MODEL_NAME = "gemini-2.5-pro-exp-03-25"
     print(f"Attempting to use EXPERIMENTAL model: {MODEL_NAME}")
 
-    # Define the browse tool using FunctionDeclaration
-    browse_tool = types.Tool(
+    # Define tools using genai.types
+    browse_tool = genai.types.Tool(
         function_declarations=[
-            types.FunctionDeclaration(
+            genai.types.FunctionDeclaration(
                 name='load_page',
                 description='Fetches the content of a specific web page URL as Markdown text. Use this when the user asks for information from a specific URL they provide, or when you need to look up live information mentioned alongside a specific source URL.',
-                parameters=types.Schema(
-                    type=types.Type.OBJECT,
+                parameters=genai.types.Schema(
+                    type=genai.types.Type.OBJECT,
                     properties={
-                        'url': types.Schema(type=types.Type.STRING, description="The *full* URL of the webpage to load (must start with http:// or https://).")
+                        'url': genai.types.Schema(type=genai.types.Type.STRING, description="The *full* URL of the webpage to load (must start with http:// or https://).")
                     },
                     required=['url']
                 )
             )
         ]
     )
-    # Define the code execution tool
-    # Enables the model to suggest and potentially execute Python code.
-    code_execution_tool = types.Tool(code_execution=types.ToolCodeExecution())
+    code_execution_tool = genai.types.Tool(code_execution=genai.types.ToolCodeExecution()) # Note: Simplified access
 
-    # Combine tools that the model can use
     tools = [browse_tool, code_execution_tool]
 
-    # Create the model instance
+    # Create the model instance using genai alias
     model = genai.GenerativeModel(
         model_name=MODEL_NAME,
         tools=tools,
-        # Relax safety settings slightly *if needed* for code/complex generation,
-        # but be aware of the implications. BLOCK_NONE is risky. Use with caution.
-        # Consider BLOCK_LOW_AND_ABOVE or MEDIUM as safer alternatives.
         safety_settings={
-             HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-             HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-             HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-             HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-             # Adjust specific categories if you face frequent blocking for safe content.
+             # Access HarmCategory and HarmBlockThreshold via genai.types
+             genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+             genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+             genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+             genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
         },
-        # System instruction (optional but recommended for setting context)
         system_instruction="You are a helpful AI assistant called Gemini-Toolkit. You can browse specific web pages provided by the user via the 'load_page' tool. You can also execute Python code using the 'code_execution' tool to perform calculations, analyze data, or demonstrate programming concepts. Explain your reasoning and the steps you take. If asked to browse, confirm the URL you are accessing. If providing code, explain what it does.",
     )
     print(f"Gemini client initialized with model: {MODEL_NAME} and tools.")
@@ -166,43 +146,37 @@ try:
 except Exception as e:
     print(f"CRITICAL ERROR: Error initializing Gemini client: {e}")
     traceback.print_exc()
-    # Provide a fallback model or handle the error gracefully in the UI
     model = None
     tools = []
-    # Consider exiting if the core functionality is unavailable
-    # raise SystemExit("Failed to initialize core Gemini model.") from e
+
 
 # --- Gradio App Logic ---
 
 def handle_function_call(function_call):
     """Executes the function call requested by the model."""
     function_name = function_call.name
-    args = function_call.args # This is now a dict-like object
+    args = function_call.args # Dict-like object
 
-    print(f"Executing Function Call: {function_name} with args: {dict(args)}") # Log args
+    print(f"Executing Function Call: {function_name} with args: {dict(args)}")
 
     try:
         if function_name == 'load_page':
             url = args.get('url')
             if url:
-                # Execute the actual function
                 function_response_content = load_page(url=url)
-                # Limit response size to send back to Gemini
-                MAX_RESPONSE_LEN = 50000 # Limit characters sent back
+                MAX_RESPONSE_LEN = 50000
                 if len(function_response_content) > MAX_RESPONSE_LEN:
                     print(f"Tool Response truncated from {len(function_response_content)} to {MAX_RESPONSE_LEN} chars.")
                     function_response_content = function_response_content[:MAX_RESPONSE_LEN] + "\n\n[... Tool Response Truncated Due to Size Limit ...]"
             else:
                 function_response_content = "Error: URL parameter was missing in the function call. Please ensure the 'url' argument is provided."
         else:
-            # Should not happen if tools are defined correctly and model uses them
             print(f"Error: Received call for unknown function '{function_name}'")
             function_response_content = f"Error: Unknown function '{function_name}' called by the model."
 
-        # Create the FunctionResponse part to send back to the model
-        # API expects the response arg to be a dict, typically {'content': <result>}
-        function_response_part = types.Part(
-            function_response=types.FunctionResponse(
+        # Use genai.types for FunctionResponse and Part
+        function_response_part = genai.types.Part(
+            function_response=genai.types.FunctionResponse(
                 name=function_name,
                 response={'content': function_response_content}
             )
@@ -213,9 +187,9 @@ def handle_function_call(function_call):
     except Exception as e:
          print(f"Error during execution of function '{function_name}': {e}")
          traceback.print_exc()
-         # Return an error message back to the model
-         return types.Part(
-             function_response=types.FunctionResponse(
+         # Use genai.types here too
+         return genai.types.Part(
+             function_response=genai.types.FunctionResponse(
                  name=function_name,
                  response={'error': f"Failed to execute function {function_name}: {str(e)}"}
              )
@@ -224,140 +198,142 @@ def handle_function_call(function_call):
 def generate_response_with_tools(user_input, history_state):
     """Handles user input, interacts with Gemini (incl. tools), and manages history."""
     if not model:
-         # Handle case where model initialization failed
-         return "Error: The AI model (Gemini) could not be initialized. Please check the logs or API key configuration.", history_state or []
+         return [[None, "Error: The AI model (Gemini) could not be initialized. Please check the logs or API key configuration."]], history_state or []
 
     if not user_input.strip():
-        # Return immediately if input is empty, don't update history
-        # Let the UI handle showing this message without clearing history state
-        # For chatbot, we might just not send anything or return a specific tuple
-        # Returning just a message for the chatbot display:
         return [[None, "Please enter a valid query."]], history_state or []
 
-
     # --- History Management ---
-    # Load history from state (should be list of Content objects)
-    # Initialize if state is None or empty
     conversation_history = history_state if isinstance(history_state, list) else []
-
-    # Append the user's new message to the history
-    conversation_history.append(types.Content(role="user", parts=[types.Part.from_text(user_input)]))
+    # Use genai.types for Content and Part
+    conversation_history.append(genai.types.Content(role="user", parts=[genai.types.Part.from_text(user_input)]))
     print(f"\n--- Sending to Gemini (History length: {len(conversation_history)}) ---")
 
-    # Limit history length *before* sending to API to avoid excessive token usage/cost
-    # Keep the system instruction + last N turns. A turn = user msg + model response (potentially with tool calls/responses)
     MAX_HISTORY_TURNS = 10
-    max_history_items = MAX_HISTORY_TURNS * 2 + (1 if conversation_history and conversation_history[0].role == "system" else 0) # Approx items to keep
-
+    max_history_items = MAX_HISTORY_TURNS * 2 + (1 if conversation_history and conversation_history[0].role == "system" else 0)
     if len(conversation_history) > max_history_items:
         print(f"Trimming conversation history from {len(conversation_history)} items to ~{max_history_items}")
         if conversation_history[0].role == "system":
-             # Keep system instruction and the latest items
              conversation_history = [conversation_history[0]] + conversation_history[-(max_history_items-1):]
         else:
-             # Just keep the latest items
              conversation_history = conversation_history[-max_history_items:]
 
-
-    # --- Interaction Loop (for potential tool calls) ---
-    MAX_TOOL_LOOPS = 5 # Prevent infinite loops if the model keeps calling tools without finishing
+    # --- Interaction Loop ---
+    MAX_TOOL_LOOPS = 5
     loop_count = 0
-    current_history_for_api = list(conversation_history) # Work with a copy in the loop
+    current_history_for_api = list(conversation_history)
+    final_bot_message = "" # Initialize variable to hold the final message text
 
     try:
         while loop_count < MAX_TOOL_LOOPS:
             loop_count += 1
             print(f"Generation loop {loop_count}/{MAX_TOOL_LOOPS}...")
 
-            # Send context and query to Gemini
-            # Use the potentially trimmed history for this API call
             response = model.generate_content(
                 current_history_for_api,
-                request_options={"timeout": 120}, # Increase timeout for complex/tool calls
-                # generation_config=genai.types.GenerationConfig( # If you need temperature etc.
-                #     temperature=0.7
-                # )
+                request_options={"timeout": 120},
             )
 
-            # --- Process Response Candidate ---
             if not response.candidates:
                  print("Warning: No candidates received from Gemini.")
-                 # Append a message indicating no response
                  final_bot_message = "[No response generated by the model.]"
-                 current_history_for_api.append(types.Content(role="model", parts=[types.Part.from_text(final_bot_message)]))
-                 break # Exit loop
+                 # Use genai.types here
+                 current_history_for_api.append(genai.types.Content(role="model", parts=[genai.types.Part.from_text(final_bot_message)]))
+                 break
 
             candidate = response.candidates[0]
+            # Access FinishReason via genai.types
+            finish_reason = candidate.finish_reason
 
-            # Check for safety blocks or finish reasons other than STOP or TOOL use
-            if candidate.finish_reason not in (types.Candidate.FinishReason.STOP, types.Candidate.FinishReason.TOOL_CALL):
-                 print(f"Warning: Generation stopped unexpectedly. Reason: {candidate.finish_reason.name}")
-                 # Append the reason to the conversation for context, if desired
-                 stop_reason_msg = f"[Model stopped generating. Reason: {candidate.finish_reason.name}]"
-                 # Check if there's any text content before adding the stop reason
-                 if candidate.content and candidate.content.parts and any(p.text for p in candidate.content.parts):
-                     current_history_for_api.append(candidate.content) # Add what content there was
-                     # Extract text to display if needed, before adding stop reason
-                     final_bot_message = "".join([p.text for p in candidate.content.parts if p.text]) + f"\n{stop_reason_msg}"
-                 else:
-                      # No text, just add the stop reason message as the model turn
-                      final_bot_message = stop_reason_msg
-                      current_history_for_api.append(types.Content(role="model", parts=[types.Part.from_text(final_bot_message)]))
+            # Append model's turn to history *before* potentially executing tools
+            # This includes text parts and potential function_call parts
+            if candidate.content:
+                current_history_for_api.append(candidate.content)
+            else:
+                 print("Warning: Candidate content is empty.")
+                 # Decide how to handle this - perhaps break or log and continue?
+                 # If finish_reason indicates a stop, maybe just break.
+                 # If it indicates TOOL_CALL without content, that's an error state.
+
+            # Check for safety or unexpected stops first
+            # Use genai.types for FinishReason comparison
+            if finish_reason not in (genai.types.Candidate.FinishReason.STOP, genai.types.Candidate.FinishReason.TOOL_CALL):
+                 print(f"Warning: Generation stopped unexpectedly. Reason: {finish_reason.name}")
+                 stop_reason_msg = f"[Model stopped generating. Reason: {finish_reason.name}]"
+                 # Extract any partial text response
+                 partial_text = ""
+                 if candidate.content and candidate.content.parts:
+                     partial_text = "".join([p.text for p in candidate.content.parts if p.text])
+                 final_bot_message = (partial_text + "\n" if partial_text else "") + stop_reason_msg
+                 # We already appended the content, so the history is up-to-date with the partial model turn.
                  break # Exit loop
 
-            # --- Handle Potential Tool Call ---
-            has_tool_call = candidate.finish_reason == types.Candidate.FinishReason.TOOL_CALL
-
-            # Append the model's response (which might contain text and/or tool calls) to history *before* execution
-            # The API expects the model's turn asking for the tool first.
-            current_history_for_api.append(candidate.content)
+            # Check for Tool Call
+            # Use genai.types for FinishReason comparison
+            has_tool_call = finish_reason == genai.types.Candidate.FinishReason.TOOL_CALL
 
             if has_tool_call:
                 print("Tool call requested by model.")
-                tool_calls_to_process = [part.function_call for part in candidate.content.parts if part.function_call]
-
-                if not tool_calls_to_process:
-                     print("Warning: Model indicated TOOL_CALL finish reason but no function_call part found.")
-                     # Maybe append an error message? Or just break?
-                     # Let's try to continue, maybe there's text output.
-                     final_bot_message = "".join([p.text for p in candidate.content.parts if p.text])
+                if not candidate.content or not candidate.content.parts:
+                     print("Error: TOOL_CALL indicated but candidate content is empty.")
+                     final_bot_message = "[Model indicated tool use but provided no details.]"
+                     # Append error message as model turn?
+                     # current_history_for_api.append(genai.types.Content(role="model", parts=[genai.types.Part.from_text(final_bot_message)]))
+                     break # Exit loop
+
+                function_calls = [part.function_call for part in candidate.content.parts if hasattr(part, 'function_call')]
+
+                if not function_calls:
+                     print("Warning: TOOL_CALL finish reason but no function_call part found in content.")
+                     final_bot_message = "".join([p.text for p in candidate.content.parts if p.text]) # Capture any text
                      if not final_bot_message:
-                          final_bot_message = "[Model indicated tool use but provided no details or text.]"
-                     break # Exit loop as we can't proceed with tool call
+                          final_bot_message = "[Model indicated tool use but provided no callable function.]"
+                     # Model turn with text (if any) is already in history
+                     break # Exit loop
 
-                # Execute the function(s) and get responses
                 tool_responses = []
-                for function_call in tool_calls_to_process:
-                     function_response_part = handle_function_call(function_call)
-                     tool_responses.append(function_response_part)
-
-                # Add the tool execution results to history for the *next* API call
-                current_history_for_api.append(types.Content(role="tool", parts=tool_responses)) # Use role="tool"
+                for func_call in function_calls:
+                     if func_call: # Ensure it's not None
+                         function_response_part = handle_function_call(func_call)
+                         tool_responses.append(function_response_part)
+                     else:
+                          print("Warning: Encountered None value where function_call was expected.")
+
+                if not tool_responses:
+                     print("Warning: No valid tool responses generated despite TOOL_CALL.")
+                     # Decide how to proceed. Maybe break?
+                     final_bot_message = "[Failed to process tool call request.]"
+                     break
+
+                # Add the tool execution results to history
+                # Use genai.types for Content
+                current_history_for_api.append(genai.types.Content(role="tool", parts=tool_responses))
                 print("Added tool response(s) to history. Continuing loop...")
-                continue # Go back to the start of the while loop to call the API again
+                continue # Go back to the start of the while loop
 
-            else:
-                # No tool call, this is the final response from the model
+            else: # FinishReason == STOP
                 print("No tool call requested. Final response received.")
-                final_bot_message = "".join([part.text for part in candidate.content.parts if part.text])
-
-                # Also check for code execution *suggestions* or *results* in the final turn
+                # Extract final text and any code suggestions/results
+                final_bot_message = ""
                 code_parts_display = []
-                for part in candidate.content.parts:
-                     if part.executable_code:
-                         lang = part.executable_code.language.name.lower() if part.executable_code.language else "python"
-                         code = part.executable_code.code
-                         code_parts_display.append(f"Suggested Code ({lang}):\n```{'python' if lang == 'unknown_language' else lang}\n{code}\n```")
-                     elif part.code_execution_result:
-                         outcome_str = "Success" if part.code_execution_result.outcome == part.code_execution_result.Outcome.OK else "Failure"
-                         code_parts_display.append(f"Code Execution Result ({outcome_str}):\n```\n{part.code_execution_result.output}\n```")
+                if candidate.content and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, 'text'):
+                             final_bot_message += part.text
+                        if hasattr(part, 'executable_code') and part.executable_code:
+                            lang = part.executable_code.language.name.lower() if part.executable_code.language else "python"
+                            code = part.executable_code.code
+                            code_parts_display.append(f"Suggested Code ({lang}):\n```{'python' if lang == 'unknown_language' else lang}\n{code}\n```")
+                        elif hasattr(part, 'code_execution_result') and part.code_execution_result:
+                            outcome_str = "Success" if part.code_execution_result.outcome == genai.types.ExecutableCodeResponse.Outcome.OK else "Failure" # Adjusted reference
+                            code_parts_display.append(f"Code Execution Result ({outcome_str}):\n```\n{part.code_execution_result.output}\n```")
 
                 if code_parts_display:
                     final_bot_message += "\n\n" + "\n\n".join(code_parts_display)
 
-                # Handle empty final message case
                 if not final_bot_message.strip():
                      final_bot_message = "[Assistant completed its turn without generating text output.]"
+                     # The empty model turn is already in history
 
                 break # Exit the while loop
 
@@ -365,83 +341,90 @@ def generate_response_with_tools(user_input, history_state):
         if loop_count >= MAX_TOOL_LOOPS:
              print(f"Warning: Reached maximum tool execution loops ({MAX_TOOL_LOOPS}).")
              final_bot_message = (final_bot_message + "\n\n" if final_bot_message else "") + f"[Warning: Reached maximum tool execution loops ({MAX_TOOL_LOOPS}). The final response might be incomplete.]"
-             # Ensure the last model message is added even if loop limit reached
-             if current_history_for_api[-1].role != "model":
-                  current_history_for_api.append(types.Content(role="model", parts=[types.Part.from_text(final_bot_message)]))
+             # Ensure the loop warning is part of the last model message if needed
+             if current_history_for_api[-1].role == "model":
+                 # Append warning to the existing last message parts (simplistic)
+                 # Use genai.types for Part
+                 current_history_for_api[-1].parts.append(genai.types.Part.from_text(f"\n[Warning: Max loops reached]"))
+             else:
+                  # Use genai.types for Content and Part
+                  current_history_for_api.append(genai.types.Content(role="model", parts=[genai.types.Part.from_text(final_bot_message)]))
 
 
         print("--- Response Generation Complete ---")
-        # Update the main history state with the final state of the conversation
-        # We return the *final* bot message text for display, and the *full* history state
-        # The chatbot UI needs [[user, bot], [user, bot], ...] format
-        # Create the Gradio chatbot display format from our history
+        # --- Format final output for Gradio Chatbot ---
         chatbot_display_list = []
-        user_msg = None
+        user_msg_buffer = None # To hold user message until bot reply comes
         for i, content in enumerate(current_history_for_api):
-             # Skip system instruction for display
-             if content.role == "system": continue
-             # Combine multi-part messages for display
-             msg_text = ""
-             for part in content.parts:
-                 if part.text:
-                      msg_text += part.text + "\n"
-                 # Display code suggestions nicely
-                 elif part.executable_code:
-                      lang = part.executable_code.language.name.lower() if part.executable_code.language else "python"
-                      code = part.executable_code.code
-                      msg_text += f"\nSuggested Code ({lang}):\n```{'python' if lang == 'unknown_language' else lang}\n{code}\n```\n"
-                 # We don't display tool calls/responses directly in chat bubbles usually
-                 # elif part.function_call: msg_text += f"[Requesting tool: {part.function_call.name}]\n"
-                 # elif part.function_response: msg_text += f"[Tool response received for {part.function_response.name}]\n"
-                 elif part.code_execution_result:
-                      outcome_str = "Success" if part.code_execution_result.outcome == part.code_execution_result.Outcome.OK else "Failure"
-                      msg_text += f"\nCode Execution Result ({outcome_str}):\n```\n{part.code_execution_result.output}\n```\n"
-
-             msg_text = msg_text.strip()
-             if not msg_text: continue # Skip empty parts/turns
-
-             if content.role == "user":
-                  # If there was a pending user message, start a new pair
-                  user_msg = msg_text
-                  # Append None temporarily for the bot response, it will be filled if available
-                  chatbot_display_list.append([user_msg, None])
-             elif content.role == "model":
-                  if chatbot_display_list and chatbot_display_list[-1][1] is None:
-                      # Fill in the bot response for the last user message
-                      chatbot_display_list[-1][1] = msg_text
-                  else:
-                      # Model message without a preceding user message (unlikely here, but handle)
-                      # Or potentially consecutive model messages after tool use. Append as separate bot message.
-                      chatbot_display_list.append([None, msg_text])
-                  user_msg = None # Reset pending user message
-
-        # Ensure the very last bot message is captured if the loop ended correctly
-        # This logic might be redundant if the history appending handles it correctly
-        # Let's rely on history build up and the formatting loop above.
-
-        return chatbot_display_list, current_history_for_api # Return display list and history state
+            if content.role == "system": continue # Skip system prompt in display
+
+            # Combine parts into a single message string for display
+            display_text = ""
+            if content.parts: # Check if parts exist
+                for part in content.parts:
+                    # Check attributes safely before accessing
+                    if hasattr(part, 'text'):
+                        display_text += part.text + "\n"
+                    elif hasattr(part, 'executable_code') and part.executable_code:
+                        lang = part.executable_code.language.name.lower() if hasattr(part.executable_code, 'language') and part.executable_code.language else "python"
+                        code = part.executable_code.code if hasattr(part.executable_code, 'code') else ""
+                        display_text += f"\nSuggested Code ({lang}):\n```{'python' if lang == 'unknown_language' else lang}\n{code}\n```\n"
+                    elif hasattr(part, 'code_execution_result') and part.code_execution_result:
+                        # Use genai.types.ExecutableCodeResponse.Outcome
+                        outcome_ok = genai.types.ExecutableCodeResponse.Outcome.OK if hasattr(genai.types, 'ExecutableCodeResponse') else 1 # Fallback if type not found? Be careful
+                        outcome_str = "Success" if part.code_execution_result.outcome == outcome_ok else "Failure"
+                        output = part.code_execution_result.output if hasattr(part.code_execution_result, 'output') else ""
+                        display_text += f"\nCode Execution Result ({outcome_str}):\n```\n{output}\n```\n"
+                    # Optional: Display indications of tool use (can make chat noisy)
+                    # elif hasattr(part, 'function_call') and part.function_call: display_text += f"[Requesting tool: {part.function_call.name}...]\n"
+                    # elif hasattr(part, 'function_response') and part.function_response: display_text += f"[Tool '{part.function_response.name}' response processed.]\n"
+
+            display_text = display_text.strip()
+
+            if not display_text and content.role != 'tool': continue # Skip empty non-tool turns
+
+            if content.role == "user":
+                user_msg_buffer = display_text # Store user message
+                # Don't append to display list yet, wait for model response
+            elif content.role == "model":
+                if user_msg_buffer is not None:
+                    # We have a user message and now the model's response
+                    chatbot_display_list.append([user_msg_buffer, display_text])
+                    user_msg_buffer = None # Clear buffer
+                else:
+                    # Model message without preceding user message (e.g., initial greeting or consecutive model turns)
+                    chatbot_display_list.append([None, display_text])
+            # Ignore 'tool' role messages in the chatbot display list
+
+        # If the loop ended with a user message still in the buffer (e.g., error before model reply)
+        if user_msg_buffer is not None:
+             chatbot_display_list.append([user_msg_buffer, None]) # Show user msg, no bot reply yet
+
+
+        return chatbot_display_list, current_history_for_api
 
     except Exception as e:
         print(f"ERROR during Gemini generation or tool processing: {str(e)}")
         traceback.print_exc()
-        error_message = f"An error occurred while processing your request: {str(e)}"
-        # Return error in chatbot format and the history state *before* the error
-        chatbot_error_display = [[None, error_message]]
-        # Try to get the display history before error if possible
-        if 'current_history_for_api' in locals():
-             # Rebuild display list up to the point before error for continuity
-             # (This is simplified, full rebuild might be complex)
-             existing_display = []
-             for c in current_history_for_api[:-1]: # Exclude potentially problematic last addition
-                  if c.role == "user": existing_display.append([c.parts[0].text, None])
-                  elif c.role == "model" and existing_display and existing_display[-1][1] is None:
-                      existing_display[-1][1] = "".join([p.text for p in c.parts if p.text])
-             existing_display.append([None, error_message]) # Add error message at end
-             chatbot_error_display = existing_display
-
-
-        # Return the history *before* this failed turn started
-        return chatbot_error_display, conversation_history # Revert state to before this turn
+        error_message = f"An error occurred: {str(e)}"
+        # Return error in chatbot format, maintain previous history state
+        # Build display history from existing state + error
+        error_display_list = []
+        if isinstance(history_state, list):
+             # Simplified history-to-display conversion for error case
+             temp_user_msg = None
+             for content in history_state:
+                  if content.role == "user": temp_user_msg = content.parts[0].text
+                  elif content.role == "model" and temp_user_msg:
+                       model_text = "".join([p.text for p in content.parts if hasattr(p, 'text')])
+                       error_display_list.append([temp_user_msg, model_text])
+                       temp_user_msg = None
+             if temp_user_msg: error_display_list.append([temp_user_msg, None]) # Append dangling user message
+
+        error_display_list.append([None, error_message]) # Add the error message
+
+        # Return the state *before* the error occurred
+        return error_display_list, conversation_history[:-1] # Exclude the failed user turn
 
 
 # --- Gradio Interface ---
@@ -450,72 +433,62 @@ with gr.Blocks(title="Gemini AI Assistant w/ Tools", theme=gr.themes.Soft()) as
     gr.Markdown(f"# 🚀 Gemini AI Assistant ({MODEL_NAME})")
     gr.Markdown("Ask questions, request info from specific URLs, or ask for code/calculations. Uses function calling and code execution.")
 
-    # Chatbot component to display conversation
     chatbot_display = gr.Chatbot(
         label="Conversation",
         bubble_full_width=False,
-        height=600, # Increased height
+        height=600,
         show_copy_button=True,
-        render_markdown=True # Ensure markdown inc code blocks is rendered
+        render_markdown=True
     )
 
-    # Textbox for user input
-    msg_input = gr.Textbox(
-        label="Your Query",
-        placeholder="Ask anything... (e.g., 'Summarize example.com', 'Calculate 2^64', 'Write python code to list files')",
-        lines=3, # Start with more lines
-        scale=4 # Take more horizontal space
-    )
-
-    # Use ClearButton which handles multiple components
-    clear_btn = gr.ClearButton(value="🗑️ Clear Chat")
-
-    # Submit button (using default value seems fine)
-    send_btn = gr.Button("➡️ Send", variant="primary", scale=1)
-
+    with gr.Row(): # Arrange input and buttons horizontally
+        msg_input = gr.Textbox(
+            label="Your Query",
+            placeholder="Ask anything...",
+            lines=3,
+            scale=4 # Input takes more space
+        )
+        with gr.Column(scale=1, min_width=150): # Column for buttons
+            send_btn = gr.Button("➡️ Send", variant="primary")
+            clear_btn = gr.ClearButton(value="🗑️ Clear Chat")
 
     # Hidden state to store the raw conversation history (list of genai.types.Content)
     chat_history_state = gr.State([])
 
     def user_message_update(user_message, history_display_list):
         """Appends the user's message to the display list and clears the input."""
-        if not user_message.strip(): # Avoid adding empty messages
-             return gr.update(value=""), history_display_list # Clear input, return unchanged history display
-        # Append user message with None placeholder for bot response
-        return gr.update(value=""), history_display_list + [[user_message, None]]
+        if not user_message.strip():
+             return gr.update(value=""), history_display_list
+        return gr.update(value=""), history_display_list + [[user_message, None]] # Add placeholder for bot response
 
     def bot_response_update(history_display_list, history_state):
         """Calls the backend Gemini function and updates display/state."""
-        if not history_display_list or history_display_list[-1][0] is None:
-             # Should not happen if user_message_update ran first, but safeguard
-             print("Warning: bot_response_update called without preceding user message in display.")
-             # Return unchanged display, maybe signal error? For now, just return current state.
-             return history_display_list, history_state
+        if not history_display_list or history_display_list[-1][1] is not None:
+            # Only proceed if there is a pending user message (placeholder is None)
+            print("Bot update called without pending user message.")
+            # Should return current state if called incorrectly
+            return history_display_list, history_state
 
-        user_message = history_display_list[-1][0] # Get the last user message from display list
+        user_message = history_display_list[-1][0]
         print(f"User message being sent to backend: {user_message}")
 
         # Call the main Gemini interaction function
-        # It now returns the *entire* chat history for display, and the updated state
         updated_display_list, updated_history_state = generate_response_with_tools(user_message, history_state)
 
-        # The backend function now returns the full display list
-        # Update the state variable directly
         return updated_display_list, updated_history_state
 
-    # Define the action for sending a message (Enter key in Textbox)
+    # --- Event Listeners ---
     msg_input.submit(
-        user_message_update,    # 1. Update display with user msg, clear input
+        user_message_update,
         [msg_input, chatbot_display],
         [msg_input, chatbot_display],
-        queue=False,            # Run immediately UI update
+        queue=False,
     ).then(
-        bot_response_update,    # 2. Call backend, get full display list & new state
-        [chatbot_display, chat_history_state], # Pass current display (for last msg) & state
-        [chatbot_display, chat_history_state]  # Update display & state from backend return
+        bot_response_update,
+        [chatbot_display, chat_history_state],
+        [chatbot_display, chat_history_state] # Update display and state
     )
 
-    # Define the action for clicking the Send button
     send_btn.click(
         user_message_update,
         [msg_input, chatbot_display],
@@ -528,14 +501,17 @@ with gr.Blocks(title="Gemini AI Assistant w/ Tools", theme=gr.themes.Soft()) as
     )
 
     # Setup the ClearButton to target the necessary components, including the state
-    clear_btn.add(components=[msg_input, chatbot_display, chat_history_state])
-    # The ClearButton itself doesn't need a custom function when using .add()
-    # It will set components to their default/initial values (Textbox="", Chatbot=None, State=[])
+    # Use list comprehension/lambda if add doesn't accept state directly, or define custom clear fn
+    #clear_btn.add(components=[msg_input, chatbot_display, chat_history_state]) # May not work with state
+
+    # Custom clear function is safer for state
+    def clear_all():
+         return ["", None, []] # Clears Textbox, Chatbot display, State
+
+    clear_btn.click(clear_all, [], [msg_input, chatbot_display, chat_history_state], queue=False)
 
 
 if __name__ == "__main__":
     print("Starting Gradio App...")
-    # Enable queue for handling potentially long API calls/tool executions
-    # Set share=True to get a public link (remove if only running locally)
     demo.queue().launch(server_name="0.0.0.0", server_port=7860)
     print("Gradio App Stopped.")