app.py

# HRBuddy
# app.py

# Imported libraries
import streamlit as st
import hashlib
import datetime
from pymongo import MongoClient
from core.config_loader import cfg
from core.logger import log
from core.ingestion import process_pdf
from core.rag_engine import HRBuddyEngine

# Streamlit UI Configuration
st.set_page_config(page_title=cfg.get("app_name", "HR Buddy"), layout="wide")
st.title(cfg.get("app_name", "HR Policy Assistant"))

# Backend Initialization
@st.cache_resource
def initialize_system():
    log.info(f"System boot sequence initiated for {cfg['app_name']}...")
    
    # Connect to MongoDB
    client = MongoClient(cfg["vector_store"]["uri"])
    db = client[cfg.get("db_name", "hr_buddy_db")]
    
    # Process PDF
    PDF_PATH = cfg["ingestion"]["pdf_path"]
    chunks = process_pdf(PDF_PATH)
    
    # Initialize RAG Engine
    engine = HRBuddyEngine(chunks)
    
    return db["chat_history"], db["users"], engine

chat_collection, users_collection, engine = initialize_system()

# Login System
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

# Login Page
if not st.session_state.logged_in:
    st.sidebar.subheader("Login or Register")
    user = st.sidebar.text_input("Username")
    pwd = st.sidebar.text_input("Password", type="password")
    
    # Login or Register
    if st.sidebar.button("Login/Register"):
        pwd_hash = hashlib.sha256(pwd.encode()).hexdigest()
        user_record = users_collection.find_one({"username": user})
        
        # Check if user exists
        if user_record:
            if user_record["password_hash"] == pwd_hash:
                st.session_state.logged_in = True
                st.session_state.user_id = user
                st.rerun()
            else:
                st.sidebar.error("Invalid password")
        
        # Create new user
        else:
            users_collection.insert_one({"username": user, "password_hash": pwd_hash})
            st.session_state.logged_in = True
            st.session_state.user_id = user
            st.rerun()

    # Stop the app if not logged in
    st.stop()

# Chat Interface
SESSION_ID = st.session_state.user_id
st.sidebar.success(f"Active Session: {SESSION_ID}")

# Initialize Chat History
if "messages" not in st.session_state:
    st.session_state.messages = []

    # Sync with MongoDB
    past_msgs = chat_collection.find({"session_id": SESSION_ID}).sort("timestamp", 1)

    for m in past_msgs:
        role = "user" if m["role"] == "User" else "assistant"
        st.session_state.messages.append({"role": role, "content": m["content"]})

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

# Chat Input
if user_input := st.chat_input("Ask about company policy..."):
    log.info(f"User Query: {user_input}")
    
    # Display user input
    with st.chat_message("user"):
        st.markdown(user_input)
    
    # Add user input to chat history
    st.session_state.messages.append({"role": "user", "content": user_input})

    # Build conversation history (last 3 exchanges)
    recent = st.session_state.messages[-6:]
    history_text = "\n".join([f"{m['role']}: {m['content']}" for m in recent])

    # Execute RAG Engine
    with st.chat_message("assistant"):
        stream = engine.generate_response(user_input, history_text, SESSION_ID)
        
        # Stream response
        def stream_parser(s):
            for chunk in s:
                yield chunk['message']['content']
        
        # Display response
        ai_response = st.write_stream(stream_parser(stream))
    
    st.session_state.messages.append({"role": "assistant", "content": ai_response})

    # Save to MongoDB
    chat_collection.insert_many([
        {"session_id": SESSION_ID, "role": "User", "content": user_input, "timestamp": datetime.datetime.now()},
        {"session_id": SESSION_ID, "role": "AI", "content": ai_response, "timestamp": datetime.datetime.now()}
    ])