import gradio as gr

import os
import google.generativeai as genai

genai.configure(
    api_key=os.getenv("api_key")
)

models = "models/gemini-pro"

def generate_text(prompt, examples, model=models, temperature=0.25):
    return genai.chat(model=model, temperature=temperature, examples=examples, messages=prompt).last

def generate(prompt, history, temperature=0.25):
    examples = [(item[0], item[1]) for item in history]
    output = generate_text(prompt, examples, models, temperature)
    return output

additional_inputs=[
    gr.Slider(
        label="Temperature",
        value=0.25,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Higher values produce more diverse outputs",
    )
]

css = """
  #mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
  }
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("<h1><center>LLaVa Chat by <a href='https://huggingface.co/Satyam-Singh'>Satyam Singh</a></center></h1>")
    gr.ChatInterface(
        generate,
        additional_inputs=additional_inputs,
        examples=[["What is the secret to life?"], ["Write me a recipe for pancakes."]]
    )

demo.launch(debug=True)