import os
import google.generativeai as genai
import gradio as gr

# Configure the API
genai.configure(api_key=os.environ["API_KEY"])

# Create the model
generation_config = {
    "temperature": 0.9,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 1024,
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-flash-8b-exp-0924",
    generation_config=generation_config,
)

# Initialize the chat session
chat = model.start_chat(history=[])

def respond(message, history):
    response = chat.send_message(message)
    history.append((message, response.text))
    return "", history

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Gemini-1.5-flash-8b-exp-0924 chatbot")
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Your message")
    clear = gr.ClearButton([msg, chatbot])

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

if __name__ == "__main__":
    demo.launch()