import base64
import os
from google import genai
from google.genai import types
import gradio as gr
from IPython.display import display
from IPython.display import Markdown

client = genai.Client(
    api_key=os.environ.get("GEMINI_API_KEY"),
)

def generate(prompt, *args):

    #model = "gemini-2.5-pro-exp-03-25"
    model = "gemini-2.0-flash"
    #model = "gemini-2.0-pro-exp-02-05"
    #model = "gemini-2.0-flash-thinking-exp-01-21"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=f"{prompt}"),
            ],
        ),
    ]
    tools = [
        types.Tool(google_search=types.GoogleSearch())
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.35,
        top_p=0.95,
        top_k=40,
        #max_output_tokens=65536,
        max_output_tokens=8192,
        tools=tools,
        response_mime_type="text/plain",
    )

    response = ""
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        response += chunk.text
        
    display(Markdown(response))
    return response


if __name__ == "__main__":
    iface = gr.ChatInterface(
        fn=generate,
        title=gr.Markdown("# Chatbot with Google Search"),
        description="Chatbot powered by Gemini-2.5 Pro and Google Search.",
    )
    iface.launch()