app.py

# app.py

import streamlit as st
from config import APP_TITLE, LAST_N_MESSAGES_MEMORY
from utils.storage import load_users, save_users
from utils.chat_engine import get_response_stream, get_response
from utils.ui_components import login_ui, sidebar_ui

st.set_page_config(page_title=APP_TITLE, layout="wide")

# ---------- Initialize Session ----------
if "user" not in st.session_state:
    st.session_state.user = None
if "users" not in st.session_state:
    st.session_state.users = load_users()

# ---------- Login / Signup ----------
if st.session_state.user is None:
    login_ui(st.session_state.users)
else:
    username = st.session_state.user
    users = st.session_state.users
    sidebar_ui(username, users)

    if st.session_state.current_conv is None:
        st.info("Start a new conversation from the sidebar.")
    else:
        current_conv = users[username]["conversations"][st.session_state.current_conv]
        st.title(f"💬 {APP_TITLE}")

        # Buttons
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🗑️ Clear Chat", key=f"clear_{current_conv['id']}"):
                current_conv["messages"].clear()
                save_users(users)
                st.rerun()
        with col2:
            if st.button("🔥 Delete Chat", key=f"delete_{current_conv['id']}"):
                users[username]["conversations"] = [
                    c for c in users[username]["conversations"] if c["id"] != current_conv["id"]
                ]
                save_users(users)
                st.session_state.current_conv = None
                st.rerun()

        # Show messages
        for msg in current_conv["messages"]:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])
                
        # set chat title based on the first prompt       
        if current_conv["messages"]:
            current_conv["name"] = current_conv["messages"][0]["content"][:20]+"..."

        # Chat input
        if prompt := st.chat_input("Type your message..."):
            with st.chat_message("user"):
                st.markdown(prompt)

            current_conv["messages"].append({"role": "user", "content": prompt})
            model_messages = [
                {"role": m["role"], "content": m["content"]}
                for m in current_conv["messages"]
            ]

            # Assistant stream
            with st.chat_message("assistant"):
                response = st.write_stream(
                    get_response_stream(model_messages[-LAST_N_MESSAGES_MEMORY:])
                )

            current_conv["messages"].append({"role": "assistant", "content": response})
            save_users(users)