src/streamlit_app.py

import os
import urllib
import uuid

import requests
import streamlit as st
from huggingface_hub import InferenceClient

oauth_parameters = {
    "client_id": "",
    "client_secret": "",
    "redirect_uri": "http://localhost:8501",
    "scope": "openid profile inference-api",
    "auth_uri": "https://huggingface.co/oauth/authorize",
    "token_uri": "https://huggingface.co/oauth/token",
}

chat_parameters = {
    "system_prompt": "You are a friendly Chatbot.",
    "max_tokens": 512,
    "temperature": 0.7,
    "top_p": 0.95,
}


def switch_page():
    if "message" in st.session_state:
        st.session_state.message.empty()
    st.rerun()


def get_oauth_parameters():
    client_id = os.getenv("OAUTH_CLIENT_ID")
    client_secret = os.getenv("OAUTH_CLIENT_SECRET")

    if not client_id or not client_secret:
        st.warning("OAuth credentials missing.")
        st.markdown(
            "### Missing OAuth Configuration\n\n"
            "To use this app, you need to "
            "set up OAuth credentials from Hugging Face.\n\n"
            "1. Visit [Hugging Face OAuth documentation]"
            "(https://huggingface.co/docs/hub/oauth).\n"
            "2. Create a new OAuth App.\n"
            "3. Copy the **Client ID** and **Client Secret**.\n"
            "4. Set them as environment variables on your system:\n\n"
            "```bash\n"
            "export OAUTH_CLIENT_ID=\"your-client-id\"\n"
            "export OAUTH_CLIENT_SECRET=\"your-client-secret\"\n"
            "```\n\n"
            "Restart the app after setting the variables."
        )
        st.stop()

    oauth_parameters["client_id"] = client_id
    oauth_parameters["client_secret"] = client_secret
    space_host = os.getenv("SPACE_HOST")
    if space_host:
        oauth_parameters["redirect_uri"] = f"https://{space_host}"


def get_auth_url():
    params = {
        "response_type": "code",
        "client_id": oauth_parameters["client_id"],
        "redirect_uri": oauth_parameters["redirect_uri"],
        "scope": oauth_parameters["scope"],
        "state": uuid.uuid4(),
    }
    return f"{oauth_parameters['auth_uri']}?{urllib.parse.urlencode(params)}"


def get_token(code):
    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": oauth_parameters["redirect_uri"],
        "client_id": oauth_parameters["client_id"],
        "client_secret": oauth_parameters["client_secret"],
    }
    response = requests.post(oauth_parameters["token_uri"], data=data)
    return response.json()


def get_chat_parameters():
    with st.sidebar:
        chat_parameters["system_prompt"] = st.text_area(
            "System Prompt", value=chat_parameters["system_prompt"])
        chat_parameters["max_tokens"] = st.slider(
            "Max Tokens", min_value=64, max_value=2048,
            value=chat_parameters["max_tokens"], step=64)
        chat_parameters["temperature"] = st.slider(
            "Temperature", min_value=0.1, max_value=2.0,
            value=chat_parameters["temperature"], step=0.1)
        chat_parameters["top_p"] = st.slider(
            "Top-p", min_value=0.1, max_value=1.0,
            value=chat_parameters["top_p"], step=0.05)


def get_reponse(client, user_input):
    # Construct message list with system prompt and history
    messages = [{"role": "system",
                 "content": chat_parameters["system_prompt"]}]
    messages.extend(st.session_state.chat_history)
    messages.append({"role": "user", "content": user_input})

    with st.spinner("Generating response..."):
        response_text = ""
        placeholder = st.empty()
        for chunk in client.chat_completion(
            messages=messages,
            max_tokens=chat_parameters["max_tokens"],
            stream=True,
            temperature=chat_parameters["temperature"],
            top_p=chat_parameters["top_p"],
        ):
            if chunk.choices and chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                response_text += token
                placeholder.markdown(f"**Assistant:** {response_text}")

        # Update chat history
        st.session_state.chat_history.append(
            {"role": "user", "content": user_input})
        st.session_state.chat_history.append(
            {"role": "assistant", "content": response_text})


def login():
    st.title("🔐 Hugging Face OAuth Login")
    get_oauth_parameters()

    query_params = st.query_params
    code = query_params.get("code")

    if code:
        token_data = get_token(code)
        access_token = token_data.get("access_token")

        if access_token:
            st.success("🎉 Login successful!")
            st.session_state.hf_token = access_token
            switch_page()
        else:
            st.error("❌ Failed to get access token")
    else:
        auth_link = get_auth_url()
        st.markdown(f"""
            <a href="{auth_link}" target="_blank">🔐 Login with Hugging Face</a>
            """, unsafe_allow_html=True)


def chatbot():
    model_id = "openai/gpt-oss-20b"

    # Initialize InferenceClient
    client = InferenceClient(token=st.session_state.hf_token, model=model_id)

    # Initialize chat history
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    # Page title
    st.title("🤖 Chatbot")

    # Sidebar: system prompt and generation parameters
    get_chat_parameters()

    # User input field
    user_input = st.text_input("Type your message...")

    # Handle user input
    if st.button("✅ Send"):
        if user_input.strip():
            get_reponse(client, user_input)
        else:
            st.warning("Please enter a message.")

    # Display chat history
    if st.session_state.chat_history:
        st.markdown("---")
        st.subheader("Chat History")
        for msg in st.session_state.chat_history:
            role = "🧑 You" if msg["role"] == "user" else "🤖 Assistant"
            st.markdown(f"**{role}:** {msg['content']}")


def main():
    if "hf_token" not in st.session_state:
        st.session_state.hf_token = None

    if st.session_state.hf_token:
        chatbot()
    else:
        login()


if __name__ == "__main__":
    main()