app.py

# -*- coding: utf-8 -*-
import os
import gradio as gr
# Corrected import: Import the main module and use an alias
import google.generativeai as genai
# Types will be accessed via genai.types

import requests
import markdownify
from urllib.robotparser import RobotFileParser
from urllib.parse import urlparse
import traceback
import json # Useful for debugging args

# --- Browser/Web Tool Functions ---

def can_crawl_url(url: str, user_agent: str = "PythonGoogleGenAIAgent/1.0") -> bool:
    """Check robots.txt permissions for a URL"""
    if not url:
        print("No URL provided to can_crawl_url")
        return False
    try:
        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
        robots_url = f"{parsed_url.scheme}://{parsed_url.netloc}/robots.txt"
        print(f"Checking robots.txt at: {robots_url} for URL: {url}")
        rp = RobotFileParser()
        rp.set_url(robots_url)
        rp.read()
        can_fetch = rp.can_fetch(user_agent, url)
        print(f"Can fetch {url} with agent '{user_agent}': {can_fetch}")
        return can_fetch
    except Exception as e:
        print(f"Error checking robots.txt for {url}: {e}")
        return False

def load_page(url: str) -> str:
    """
    Load webpage content as markdown. Designed to be used as a Gemini Function.
    Args:
        url: The URL of the webpage to load.
    Returns:
        Markdown content of the page or an error message.
    """
    print(f"Attempting to load page: {url}")
    if not url:
        return "Error: No URL provided."
    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)"
    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()
        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 f"Content at {url} is of type '{content_type}'. Size: {len(response.content)} bytes. Cannot convert to Markdown."

        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.")
             try:
                 html_content = response.content[:MAX_CONTENT_SIZE].decode(response.apparent_encoding or 'utf-8', errors='ignore')
             except Exception as decode_err:
                 print(f"Decoding error after truncation: {decode_err}. Falling back to utf-8 ignore.")
                 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
            truncated_msg = ""

        markdown_content = markdownify.markdownify(html_content, heading_style="ATX", strip=['script', 'style'], escape_underscores=False)
        markdown_content = '\n'.join([line.strip() for line in markdown_content.splitlines() if line.strip()])
        print(f"Successfully loaded and converted {url} to markdown.")
        return f"Content from {url}:\n\n" + markdown_content + truncated_msg

    except requests.exceptions.Timeout:
        print(f"Timeout error loading page: {url}")
        return f"Error: Timeout while trying to load {url}"
    except requests.exceptions.RequestException as e:
        print(f"Request error loading page {url}: {str(e)}")
        return f"Error loading page {url}: {str(e)}"
    except Exception as e:
        print(f"General error loading page {url}: {str(e)}")
        traceback.print_exc()
        return f"Error loading page {url}: An unexpected error occurred ({type(e).__name__})."


# --- Gemini Client Initialization and Configuration ---
try:
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY environment variable not set.")
    genai.configure(api_key=api_key)

    MODEL_NAME = "gemini-2.5-pro-exp-03-25"
    print(f"Attempting to use EXPERIMENTAL model: {MODEL_NAME}")

    # Define tools using genai.types
    browse_tool = genai.types.Tool(
        function_declarations=[
            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.',
                # Use basic string types for schema definition
                parameters={
                    'type': 'object',
                    'properties': {
                        'url': {
                             'type': 'string',
                             'description': "The *full* URL of the webpage to load (must start with http:// or https://)."
                        }
                    },
                    'required': ['url']
                }
            )
        ]
    )
    # --- CORRECTED Code Execution Tool Definition ---
    # Pass an empty dictionary to enable code execution capability
    code_execution_tool = genai.types.Tool(code_execution={})
    # --- END OF CORRECTION ---

    tools = [browse_tool, code_execution_tool]

    model = genai.GenerativeModel(
        model_name=MODEL_NAME,
        tools=tools,
        safety_settings={
             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="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.")

except Exception as e:
    print(f"CRITICAL ERROR: Error initializing Gemini client: {e}")
    traceback.print_exc()
    # Set model to None so the app can report initialization failure
    model = None
    tools = []


# --- 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 # Dict-like object

    print(f"Executing Function Call: {function_name} with args: {dict(args)}")

    try:
        if function_name == 'load_page':
            url = args.get('url')
            if url:
                function_response_content = load_page(url=url)
                MAX_RESPONSE_LEN = 50000 # Limit response size
                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:
            print(f"Error: Received call for unknown function '{function_name}'")
            function_response_content = f"Error: Unknown function '{function_name}' called by the model."

        # 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} # Response must be a dict
            )
        )
        print(f"Function Response generated for {function_name}")
        return function_response_part

    except Exception as e:
         print(f"Error during execution of function '{function_name}': {e}")
         traceback.print_exc()
         # Use genai.types here too, return error details in the response dict
         return genai.types.Part(
             function_response=genai.types.FunctionResponse(
                 name=function_name,
                 response={'error': f"Failed to execute function {function_name}: {str(e)}"}
             )
         )

def generate_response_with_tools(user_input, history_state):
    """Handles user input, interacts with Gemini (incl. tools), and manages history."""
    if not model:
         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 [[None, "Please enter a valid query."]], history_state or []

    # --- History Management ---
    conversation_history = history_state if isinstance(history_state, list) else []
    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)}) ---")

    MAX_HISTORY_TURNS = 10
    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":
             conversation_history = [conversation_history[0]] + conversation_history[-(max_history_items-1):]
        else:
             conversation_history = conversation_history[-max_history_items:]

    # --- Interaction Loop ---
    MAX_TOOL_LOOPS = 5
    loop_count = 0
    current_history_for_api = list(conversation_history)
    final_bot_message = ""

    try:
        while loop_count < MAX_TOOL_LOOPS:
            loop_count += 1
            print(f"Generation loop {loop_count}/{MAX_TOOL_LOOPS}...")

            response = model.generate_content(
                current_history_for_api,
                request_options={"timeout": 120},
            )

            if not response.candidates:
                 print("Warning: No candidates received from Gemini.")
                 final_bot_message = "[No response generated by the model.]"
                 current_history_for_api.append(genai.types.Content(role="model", parts=[genai.types.Part.from_text(final_bot_message)]))
                 break

            candidate = response.candidates[0]
            finish_reason = candidate.finish_reason

            if candidate.content:
                current_history_for_api.append(candidate.content)

            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}]"
                 partial_text = ""
                 if candidate.content and candidate.content.parts:
                     partial_text = "".join([p.text for p in candidate.content.parts if hasattr(p, 'text') and p.text])
                 final_bot_message = (partial_text + "\n" if partial_text else "") + stop_reason_msg
                 break

            has_tool_call = finish_reason == genai.types.Candidate.FinishReason.TOOL_CALL

            if has_tool_call:
                print("Tool call requested by model.")
                if not candidate.content or not candidate.content.parts:
                     print("Error: TOOL_CALL indicated but candidate content/parts is missing.")
                     final_bot_message = "[Model indicated tool use but provided no details.]"
                     break

                function_calls = [part.function_call for part in candidate.content.parts if hasattr(part, 'function_call') and part.function_call]

                if not function_calls:
                     print("Warning: TOOL_CALL finish reason but no valid function_call part found.")
                     final_bot_message = "".join([p.text for p in candidate.content.parts if hasattr(p, 'text') and p.text])
                     if not final_bot_message:
                          final_bot_message = "[Model indicated tool use but provided no callable function or text.]"
                     break

                tool_responses = []
                for func_call in function_calls:
                     function_response_part = handle_function_call(func_call)
                     tool_responses.append(function_response_part)

                if not tool_responses:
                     print("Warning: No valid tool responses generated despite TOOL_CALL.")
                     final_bot_message = "[Failed to process tool call request.]"
                     break

                current_history_for_api.append(genai.types.Content(role="tool", parts=tool_responses))
                print("Added tool response(s) to history. Continuing loop...")
                final_bot_message = ""
                continue

            else: # FinishReason == STOP
                print("No tool call requested. Final response received.")
                final_bot_message = ""
                code_parts_display = []
                if current_history_for_api and current_history_for_api[-1].role == "model":
                    last_model_content = current_history_for_api[-1]
                    if last_model_content.parts:
                        for part in last_model_content.parts:
                            if hasattr(part, 'text') and part.text:
                                final_bot_message += part.text
                            if hasattr(part, 'executable_code') and part.executable_code:
                                lang = getattr(getattr(part.executable_code, 'language', None), 'name', 'python').lower()
                                code = getattr(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_enum = getattr(genai.types, 'ExecutableCodeResponse', None)
                                outcome_ok_val = getattr(outcome_enum.Outcome, 'OK', None) if outcome_enum and hasattr(outcome_enum, 'Outcome') else 1
                                outcome_val = getattr(part.code_execution_result, 'outcome', None)
                                outcome_str = "Success" if outcome_val == outcome_ok_val else "Failure"
                                output = getattr(part.code_execution_result, 'output', '')
                                code_parts_display.append(f"Code Execution Result ({outcome_str}):\n```\n{output}\n```")

                if code_parts_display:
                    final_bot_message += "\n\n" + "\n\n".join(code_parts_display)

                if not final_bot_message.strip():
                     final_bot_message = "[Assistant completed its turn without generating text output.]"
                     if current_history_for_api[-1].role == "model" and not any(hasattr(p,'text') and p.text for p in current_history_for_api[-1].parts):
                          if not hasattr(current_history_for_api[-1], 'parts') or not current_history_for_api[-1].parts:
                              current_history_for_api[-1].parts = [] # Initialize if needed
                          current_history_for_api[-1].parts.append(genai.types.Part.from_text(final_bot_message))

                break

        # End of while loop
        if loop_count >= MAX_TOOL_LOOPS:
             print(f"Warning: Reached maximum tool execution loops ({MAX_TOOL_LOOPS}).")
             warning_msg = f"\n\n[Warning: Reached maximum tool execution loops ({MAX_TOOL_LOOPS}). The final response might be incomplete.]"
             final_bot_message += warning_msg
             if current_history_for_api and current_history_for_api[-1].role == "model":
                 if not hasattr(current_history_for_api[-1], 'parts') or not current_history_for_api[-1].parts:
                     current_history_for_api[-1].parts = []
                 current_history_for_api[-1].parts.append(genai.types.Part.from_text(warning_msg))

        print("--- Response Generation Complete ---")

        # --- Format final output for Gradio Chatbot ---
        chatbot_display_list = []
        user_msg_buffer = None
        for content in current_history_for_api:
            if content.role == "system": continue

            display_text = ""
            if hasattr(content, 'parts') and content.parts:
                for part in content.parts:
                    if hasattr(part, 'text') and part.text:
                        display_text += part.text + "\n"
                    elif hasattr(part, 'executable_code') and part.executable_code:
                        lang = getattr(getattr(part.executable_code, 'language', None), 'name', 'python').lower()
                        code = getattr(part.executable_code, 'code', '')
                        display_text += f"\n```{'python' if lang == 'unknown_language' else lang}\n{code}\n```\n"
                    elif hasattr(part, 'code_execution_result') and part.code_execution_result:
                        outcome_enum = getattr(genai.types, 'ExecutableCodeResponse', None)
                        outcome_ok_val = getattr(outcome_enum.Outcome, 'OK', None) if outcome_enum and hasattr(outcome_enum, 'Outcome') else 1
                        outcome_val = getattr(part.code_execution_result, 'outcome', None)
                        outcome_str = "Success" if outcome_val == outcome_ok_val else "Failure"
                        output = getattr(part.code_execution_result, 'output', '')
                        display_text += f"\nCode Execution Result ({outcome_str}):\n```\n{output}\n```\n"

            display_text = display_text.strip()

            if not display_text and content.role not in ["tool"]: continue

            if content.role == "user":
                user_msg_buffer = display_text
            elif content.role == "model":
                if user_msg_buffer is not None:
                    chatbot_display_list.append([user_msg_buffer, display_text or "[Processing...]"])
                    user_msg_buffer = None
                else:
                    chatbot_display_list.append([None, display_text])

        if user_msg_buffer is not None:
             chatbot_display_list.append([user_msg_buffer, None])

        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: {str(e)}"
        error_display_list = []
        if isinstance(history_state, list):
            temp_user_msg = None
            for content in history_state:
                 if content.role == "system": continue
                 text = ""
                 if hasattr(content, 'parts') and content.parts:
                     text = "".join([p.text for p in content.parts if hasattr(p, 'text')])
                 if content.role == "user": temp_user_msg = text
                 elif content.role == "model" and temp_user_msg:
                      error_display_list.append([temp_user_msg, text])
                      temp_user_msg = None
                 elif content.role == "model": error_display_list.append([None, text])
            if temp_user_msg: error_display_list.append([temp_user_msg, None])

        error_display_list.append([None, error_message])

        previous_history = conversation_history[:-1] if isinstance(conversation_history, list) and conversation_history else []
        return error_display_list, previous_history


# --- Gradio Interface ---

with gr.Blocks(title="Gemini AI Assistant w/ Tools", theme=gr.themes.Soft()) as demo:
    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_display = gr.Chatbot(
        label="Conversation",
        bubble_full_width=False,
        height=600,
        show_copy_button=True,
        render_markdown=True,
    )

    with gr.Row():
        msg_input = gr.Textbox(
            label="Your Query",
            placeholder="Ask anything...",
            lines=3,
            scale=4
        )
        with gr.Column(scale=1, min_width=150):
            send_btn = gr.Button("➡️ Send", variant="primary")
            clear_btn = gr.ClearButton(value="🗑️ Clear Chat") # Uses custom clear_all below

    chat_history_state = gr.State([])

    def user_message_update(user_message, history_display_list):
        """Appends user message to display list and clears input."""
        if not user_message.strip():
             return gr.update(value=""), history_display_list
        return gr.update(value=""), history_display_list + [[user_message, None]]

    def bot_response_update(history_display_list, history_state):
        """Calls backend Gemini function and updates display/state."""
        # Check if the last item in display list is a user message awaiting response
        if not history_display_list or (len(history_display_list[-1]) > 1 and history_display_list[-1][1] is not None):
            print("Bot update called without pending user message in display list.")
            return history_display_list, history_state # Return current state

        user_message = history_display_list[-1][0]
        print(f"User message being sent to backend: {user_message}")

        updated_display_list, updated_history_state = generate_response_with_tools(user_message, history_state)

        return updated_display_list, updated_history_state

    # --- Event Listeners ---
    msg_input.submit(
        user_message_update,
        [msg_input, chatbot_display],
        [msg_input, chatbot_display],
        queue=False,
    ).then(
        bot_response_update,
        [chatbot_display, chat_history_state],
        [chatbot_display, chat_history_state]
    )

    send_btn.click(
        user_message_update,
        [msg_input, chatbot_display],
        [msg_input, chatbot_display],
        queue=False,
    ).then(
        bot_response_update,
        [chatbot_display, chat_history_state],
        [chatbot_display, chat_history_state]
    )

    # Custom clear function to reset state as well
    def clear_all():
         # Returns values for outputs=[msg_input, chatbot_display, chat_history_state]
         return ["", None, []]

    # Wire clear button to the custom function
    clear_btn.click(clear_all, [], [msg_input, chatbot_display, chat_history_state], queue=False)


if __name__ == "__main__":
    print("Starting Gradio App...")
    # show_error=True can be helpful for seeing Gradio errors in the browser console
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)
    print("Gradio App Stopped.")