utils.py

import logging
import pathlib
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import TextLoader
from langchain.memory import ConversationBufferMemory
from langchain.schema import Document
import hmac
import streamlit as st

def init_memory(key):
    """
    Initialize the memory for contextual conversation.

    We are caching this, so it won't be deleted every time, we restart the server.
    """
    return ConversationBufferMemory(
        memory_key=key,
        return_messages=True,
        output_key='answer'
    )
MEMORY = init_memory('chat_history')

class DocumentLoaderException(Exception):
    pass

class DocumentLoader(object):
    supported_extensions = {
        ".pdf": PyPDFLoader,
        ".txt": TextLoader
    }

def load_document(temp_filepath: str) -> list[Document]:
    ext = pathlib.Path(temp_filepath).suffix
    loader = DocumentLoader.supported_extensions.get(ext)
    if not loader:
        raise DocumentLoaderException(
            f"Invalid file extension: <{ext}>"
        )

    loaded = loader(temp_filepath)
    docs = loaded.load()
    logging.info(docs)
    return docs


def check_password():
    st.header("")
    def password_entered():
        if hmac.compare_digest(st.session_state["password"], st.secrets["adminpassword"]):
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password.
        else:
            st.session_state["password_correct"] = False

    if st.session_state.get("password_correct", False):
        return True
    
    st.text_input(
        "Enter Password 🚀", type="password", on_change=password_entered, key="password"
    )
    if "password_correct" in st.session_state:
        st.error("Password incorrect 😕")
    return False