from openai import OpenAI
import streamlit as st
import os 
from datetime import datetime 

client = OpenAI(
  base_url = "https://integrate.api.nvidia.com/v1",
  api_key = "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"
)



st.title("Nemotron 4 340B")

with st.sidebar:
    st.markdown("This is a basic chatbot. Ask anything. The app is supported by Nazmul Hasan Nihal")
    if st.button("Clear Session"):
        st.session_state.clear()
    st.write(f"Copyright 2023-{datetime.now().year} Present Nazmul Hasan Nihal")

if "openai_model" not in st.session_state:
    st.session_state['openai_model'] = "invidia/nemotron_4"

if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content":"you are a helpful assistant"}]

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

if prompt := st.chat_input("What is up"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("The assistant is thinking... Please wait."):
            stream = client.chat.completions.create(
                model=st.session_state["openai_model"],
                messages = st.session_state.messages,
                temperature = 0.5,
                top_p = 0.7,
                max_tokens = 1024,
                stream = True,
            )
            response_chunks = []
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    response_chunks.append(chunk.choices[0].delta.content)
            respose = "".join(response_chunks)
            st.markdown(response)

    st.session_state.messages.append({"role":"assistant", "content": response})