app/main.py

import streamlit as st
from firebase_admin import credentials, firestore
from langchain_google_firestore import FirestoreChatMessageHistory
from langchain.memory import ConversationBufferWindowMemory
from langchain.schema import HumanMessage, AIMessage
import firebase_admin
import uuid
from agent.retrive_agent import run_llm
import os


OPENAI_API_KEY= os.getenv("OPENAI_API_KEY")


PROJECT_ID = "recommendation-system-mage"
COLLECTION_NAME = "data_user"

if not firebase_admin._apps:
    cred = credentials.Certificate("../recommendation-system-mage-firebase-adminsdk-ds1lw-1ac94ba6dd.json")
    firebase_admin.initialize_app(cred)

client = firestore.client()

# Streamlit UI
st.title(" Chat SITY")

# Login and User ID settings
st.sidebar.header("Login")
user_id = st.sidebar.text_input("Enter User ID", value="")

# Initialize session state variables if they are not present
if "user_prompt_history" not in st.session_state:
    st.session_state["user_prompt_history"] = []

if "chat_answers_history" not in st.session_state:
    st.session_state["chat_answers_history"] = []

if "session_id" not in st.session_state:
    st.session_state.session_id = ""

if "current_user" not in st.session_state:
    st.session_state.current_user = ""

# Check if the user exists in Firebase
if user_id:
    # If the user has changed, reset the session state
    if st.session_state.current_user != user_id:
        st.session_state.current_user = user_id
        st.session_state.user_prompt_history = []
        st.session_state.chat_answers_history = []
        st.session_state.session_id = ""

    user_doc_ref = client.collection(COLLECTION_NAME).document(user_id)
    user_doc = user_doc_ref.get()
    if not user_doc.exists:
        st.sidebar.error("Login failed: User ID not found in Firebase.")
    else:
        # Setup session and history if User ID is provided

        # Initialize Session ID and store in Firebase if new session
        if not st.session_state.session_id:
            session_id = str(uuid.uuid4())
            st.session_state.session_id = session_id
            user_doc_ref.set({"last_session_id": session_id}, merge=True)
            user_doc_ref.collection("message_history").document(session_id).set({})
            st.sidebar.success(f"New session created: {session_id}")

        if st.sidebar.button("New Chat"):
            session_id = str(uuid.uuid4())
            st.session_state.session_id = session_id
            user_doc_ref.set({"last_session_id": session_id}, merge=True)
            user_doc_ref.collection("message_history").document(session_id).set({})
            st.session_state.user_prompt_history = []
            st.session_state.chat_answers_history = []
            st.sidebar.success(f"New session created: {session_id}")

        sessions_collection_ref = user_doc_ref.collection("message_history").list_documents()
        session_ids = [session.id for session in sessions_collection_ref]

        st.sidebar.subheader("Session History")
        selected_session_id = st.sidebar.selectbox("Select a session", session_ids, index=session_ids.index(st.session_state.session_id) if session_ids else -1)

        # Load chat history when a new session is selected
        if selected_session_id and selected_session_id != st.session_state.session_id:
            st.session_state.session_id = selected_session_id

            # Load the chat history for the selected session
            chat_history = FirestoreChatMessageHistory(
                session_id=st.session_state.session_id,
                collection=f"{COLLECTION_NAME}/{user_id}/message_history",
                client=client,
            )

            # Load chat messages into session state variables
            st.session_state["user_prompt_history"] = []
            st.session_state["chat_answers_history"] = []

            for msg in chat_history.messages:
                if isinstance(msg, HumanMessage):
                    st.session_state["user_prompt_history"].append(msg.content)
                elif isinstance(msg, AIMessage):
                    st.session_state["chat_answers_history"].append(msg.content)

        # Display the loaded chat history
        if st.session_state.get("user_prompt_history") and st.session_state.get("chat_answers_history"):
            for user_prompt, answer in zip(st.session_state["user_prompt_history"], st.session_state["chat_answers_history"]):
                st.chat_message("user").write(user_prompt)
                st.chat_message("assistant").write(answer)
            



        # Initialize Firestore Chat Message History in user's subcollection
        chat_history = FirestoreChatMessageHistory(
            session_id=st.session_state.session_id,
            collection=f"{COLLECTION_NAME}/{user_id}/message_history",
            client=client,
        )

        # Initialize memory to store messages specific to the current user and session
        memory = ConversationBufferWindowMemory(
            k=5,
            chat_memory=chat_history
        )

        # Chatbot UI - Input and Chat Messages
        if prompt := st.chat_input("Enter your prompt here..."):
            with st.spinner("Generating response..."):
                # Use Firebase history in RAG query specific to the current session
                generated_response = run_llm(query=prompt, chat_history=memory.chat_memory.messages)
                generated_response = generated_response.replace("```", "").strip()
                memory.save_context(
                    {"input": prompt},
                    {"output": generated_response}
                )

                # Update session state
                st.session_state["user_prompt_history"].append(prompt)
                st.session_state["chat_answers_history"].append(generated_response)

                # Display the updated chat history
                st.chat_message("user").write(prompt)
                st.chat_message("assistant").write(generated_response)
        else:
            if len(st.session_state["user_prompt_history"]) == 0 and len(st.session_state["chat_answers_history"]) == 0:
                col1, col2, col3 = st.columns([1, 2, 1])  # Membagi halaman menjadi 3 bagian, dengan bagian tengah lebih besar
                with col2:
                    st.write("#")
                    st.write("#")
                    st.write("### Ada yang bisa saya bantu?")

else:
    st.sidebar.warning("Please enter your User ID to proceed.")