app.py

import streamlit as st
import requests
import firebase_admin
from firebase_admin import credentials, firestore, auth
from streamlit_authenticator import Authenticate


API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
from google.cloud import secretmanager

# Initialize Firebase app
def access_secret_version(project_id, secret_id, version_id="latest"):
    client = secretmanager.SecretManagerServiceClient()
    name = f"projects/380679199682/secrets/serviceAccountKey/versions/1"
    response = client.access_secret_version(request={"name": name})
    return response.payload.data.decode("UTF-8")

if not firebase_admin._apps:
    project_id = "allyarc-edu"
    secret_id = "serviceAccountKey"
    service_account_key = access_secret_version(project_id, secret_id)
    cred = credentials.Certificate(json.loads(service_account_key))
    firebase_admin.initialize_app(cred)
    

db = firestore.client()

# User authentication
users = {
    "user1@example.com": {"password": "hashed_password"},
    "user2@example.com": {"password": "hashed_password"},
}

authenticator = Authenticate(users, "chat_app", "auth", cookie_expiry_days=30)

# Adding a unique separator "||" to easily identify the start of the actual response.
prompt = """
As a concise and supportive chatbot designed for classroom interaction with autistic kids, respond to greetings with a simple and direct offer of help.

Here is the text:

"""
separator = "||"  # Unique separator to help split the prompt from the model's response.

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

def main():
    st.title("Mistral Chatbot")

    # User authentication
    name, authentication_status, username = authenticator.login("Login", "main")

    if authentication_status == False:
        st.error("Username/password is incorrect")
    elif authentication_status == None:
        st.warning("Please enter your username and password")
    elif authentication_status:
        # Get the current user's ID
        user_id = auth.get_user_by_email(username).uid

        # Initialize chat history for the current user
        if "messages" not in st.session_state:
            st.session_state.messages = []

        # Load chat history from Firebase for the current user
        user_ref = db.collection('users').document(user_id)
        chat_history_ref = user_ref.collection('chat_history').document('current_chat')
        chat_history_doc = chat_history_ref.get()
        if chat_history_doc.exists:
            st.session_state.messages = chat_history_doc.to_dict().get('messages', [])

        # Display chat messages from history on app rerun
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # React to user input
        if user_input := st.chat_input("Enter your message:"):
            # Display user message in chat message container
            st.chat_message("user").markdown(user_input)

            # Add user message to chat history
            st.session_state.messages.append({"role": "user", "content": user_input, "timestamp": firestore.SERVER_TIMESTAMP})

            # Make API request
            output = query({"inputs": prompt + user_input + separator})
            generated_text = output[0]['generated_text']

            # Splitting the generated text at the separator and taking the second part as the actual response.
            actual_response_parts = generated_text.split(separator)
            if len(actual_response_parts) > 1:
                actual_response = actual_response_parts[1].strip()  # The actual response comes after the separator.
            else:
                actual_response = "I'm sorry, I couldn't generate a response. Please try again."

            # Display assistant response in chat message container
            with st.chat_message("assistant"):
                st.markdown(actual_response)

            # Add assistant response to chat history
            st.session_state.messages.append({"role": "assistant", "content": actual_response, "timestamp": firestore.SERVER_TIMESTAMP})

            # Save updated chat history to Firebase for the current user
            chat_history_ref.set({'messages': st.session_state.messages})

        authenticator.logout("Logout", "sidebar")

if __name__ == "__main__":
    main()