from openai import OpenAI
import streamlit as st

with st.sidebar:
    IFC_API_KEY = st.text_input("HF access token", key="chat_bot_api_key", type="password")
    selected_model = st.selectbox(
        "Choose a model to chat with",
        ("meta-llama/Llama-3.3-70B-Instruct", "meta-llama/Llama-3.1-8B-Instruct", "Qwen/QwQ-32B-Preview", "Qwen/Qwen2.5-72B-Instruct", "microsoft/Phi-3.5-mini-instruct"),
    )
    temperature = st.slider("Temperature", 0.01, 0.99, 0.5)
    top_p = st.slider("Top_p", 0.01, 0.99, 0.7)
    max_tokens = st.slider("Max Tokens", 64, 131072, 2048)

st.title("💬 Chatbot")
st.caption(" A HF chatbot powered by HF")
if "messages" not in st.session_state:
    st.session_state.messages = []

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

if prompt := st.chat_input():
    if not IFC_API_KEY:
        st.info("Please add your access token to continue.")
        st.stop()

    client = OpenAI(api_key=IFC_API_KEY, base_url="https://api-inference.huggingface.co/v1/")
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        stream = client.chat.completions.create(
            model = selected_model,
            messages=[
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages
            ],
            temperature = temperature,
            top_p = top_p,
            max_tokens = max_tokens,
            stream = True,
        )
        response = st.write_stream(stream)
    st.session_state.messages.append({"role": "assistant", "content": response})