import os
import gradio as gr
from google import genai
from google.genai import types

# Load API key from environment variable
API_KEY = os.getenv("API_KEY")
if not API_KEY:
    raise ValueError("API_KEY environment variable is not set")

# Initialize the genai client with API key
client = genai.Client(api_key=API_KEY)

def generate_grounded_response(prompt):
    grounding_tool = types.Tool(
        google_search=types.GoogleSearch()
    )

    config = types.GenerateContentConfig(
        tools=[grounding_tool]
    )

    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config=config,
    )

    return response.text

with gr.Blocks() as demo:
    gr.Markdown("# Gemini 2.5 Flash Grounded Content Generator")

    prompt_input = gr.Textbox(label="Enter your question or prompt", lines=3, value="Who won the euro 2024?")
    generate_btn = gr.Button("Generate Response")
    output_text = gr.Textbox(label="Grounded Response", lines=10)

    generate_btn.click(fn=generate_grounded_response, inputs=prompt_input, outputs=output_text)

demo.launch()