import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
import os
load_dotenv()

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

st.title("Chat - Gemini Bot")

model = genai.GenerativeModel(
    model_name="gemini-pro"
)

if "messages" not in st.session_state:
    st.session_state.messages = [
    { 
        "role": "assistant",
        "content": "Ask me anything"
    }
    ]

for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.markdown(message['content'])

def llm_function(query):
    response = model.generate_content(query)
    with st.chat_message("assistant"):
        st.markdown(response.text)
    
    st.session_state.messages.append(
        {
            "role": "assistant",
            "content": response.text
        }
    )

query = st.chat_input("Type your query here...")

if query:
    with st.chat_message("user"):
        st.markdown(query)
    
    llm_function(query=query)