from dotenv import load_dotenv import pandas as pd import streamlit as st import streamlit_authenticator as stauth from streamlit_modal import Modal from utils import new_file, clear_memory, append_documentation_to_sidebar, load_authenticator_config, init_qa, \ append_header from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack import Document load_dotenv() OPENAI_MODELS = ['gpt-3.5-turbo', "gpt-4", "gpt-4-1106-preview"] OPEN_MODELS = [ 'mistralai/Mistral-7B-Instruct-v0.1', 'HuggingFaceH4/zephyr-7b-beta' ] def reset_chat_memory(): st.button( 'Reset chat memory', key="reset-memory-button", on_click=clear_memory, help="Clear the conversational memory. Currently implemented to retain the 4 most recent messages.", disabled=False) def manage_files(modal, document_store): open_modal = st.sidebar.button("Manage Files", use_container_width=True) if open_modal: modal.open() if modal.is_open(): with modal.container(): uploaded_file = st.file_uploader( "Upload a document in PDF format", type=("pdf",), on_change=new_file(), disabled=st.session_state['document_qa_model'] is None, label_visibility="collapsed", help="The document is used to answer your questions. The system will process the document and store it in a RAG to answer your questions.", ) edited_df = st.data_editor(use_container_width=True, data=st.session_state['files'], num_rows='dynamic', column_order=['name', 'size', 'is_active'], column_config={'name': {'editable': False}, 'size': {'editable': False}, 'is_active': {'editable': True, 'type': 'checkbox', 'width': 100}} ) st.session_state['files'] = pd.DataFrame(columns=['name', 'content', 'size', 'is_active']) if uploaded_file: st.session_state['file_uploaded'] = True st.session_state['files'] = pd.concat([st.session_state['files'], edited_df]) with st.spinner('Processing the document...'): store_file_in_table(document_store, uploaded_file) ingest_document(uploaded_file) def ingest_document(uploaded_file): if not st.session_state['document_qa_model']: st.warning('Please select a model to start asking questions') else: try: st.session_state['document_qa_model'].ingest_pdf(uploaded_file) st.success('Document processed successfully') except Exception as e: st.error(f"Error processing the document: {e}") st.session_state['file_uploaded'] = False def store_file_in_table(document_store, uploaded_file): pdf_content = uploaded_file.getvalue() st.session_state['pdf_content'] = pdf_content st.session_state.messages = [] document = Document(content=pdf_content, meta={"name": uploaded_file.name}) df = pd.DataFrame(st.session_state['files']) df['is_active'] = False st.session_state['files'] = pd.concat([df, pd.DataFrame( [{"name": uploaded_file.name, "content": pdf_content, "size": len(pdf_content), "is_active": True}])]) document_store.write_documents([document]) def init_session_state(): st.session_state.setdefault('files', pd.DataFrame(columns=['name', 'content', 'size', 'is_active'])) st.session_state.setdefault('models', []) st.session_state.setdefault('api_keys', {}) st.session_state.setdefault('current_selected_model', 'gpt-3.5-turbo') st.session_state.setdefault('current_api_key', '') st.session_state.setdefault('messages', []) st.session_state.setdefault('pdf_content', None) st.session_state.setdefault('memory', None) st.session_state.setdefault('pdf', None) st.session_state.setdefault('document_qa_model', None) st.session_state.setdefault('file_uploaded', False) def set_page_config(): st.set_page_config( page_title="Document Insights AI Assistant", page_icon=":shark:", initial_sidebar_state="expanded", layout="wide", menu_items={ 'Get Help': 'https://www.extremelycoolapp.com/help', 'Report a bug': "https://www.extremelycoolapp.com/bug", 'About': "# This is a header. This is an *extremely* cool app!" } ) def update_running_model(api_key, model): st.session_state['api_keys'][model] = api_key st.session_state['document_qa_model'] = init_qa(model, api_key) def init_api_key_dict(): # st.session_state['models'] = OPENAI_MODELS + list(OPEN_MODELS) + ['local LLM'] st.session_state['models'] = OPENAI_MODELS for model_name in OPENAI_MODELS: st.session_state['api_keys'][model_name] = None def display_chat_messages(chat_box, chat_input): with chat_box: if chat_input: for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"], unsafe_allow_html=True) st.chat_message("user").markdown(chat_input) with st.chat_message("assistant"): # process user input and generate response response = st.session_state['document_qa_model'].inference(chat_input, st.session_state.messages) st.markdown(response) st.session_state.messages.append({"role": "user", "content": chat_input}) st.session_state.messages.append({"role": "assistant", "content": response}) def setup_model_selection(): model = st.selectbox( "Model:", options=st.session_state['models'], index=0, # default to the first model in the list gpt-3.5-turbo placeholder="Select model", help="Select an LLM:" ) if model: if model != st.session_state['current_selected_model']: st.session_state['current_selected_model'] = model if model == 'local LLM': st.session_state['document_qa_model'] = init_qa(model) # api_key = st.sidebar.text_input("Enter LLM-authorization Key:", type="password", # disabled=st.session_state['current_selected_model'] == 'local LLM') api_key = "sk-proj-vQgkXQKYjy8m3waKtDFQT3BlbkFJ7uuMeDinKxql7J0Q161N" if api_key and api_key != st.session_state['current_api_key']: update_running_model(api_key, model) st.session_state['current_api_key'] = api_key return model def setup_task_selection(model): # enable extractive and generative tasks if we're using a local LLM or an OpenAI model with an API key if model == 'local LLM' or st.session_state['api_keys'].get(model): task_options = ['Extractive', 'Generative'] else: task_options = ['Extractive'] task_selection = st.sidebar.radio('Select the task:', task_options) # TODO: Add the task selection logic here (initializing the model based on the task) def setup_page_body(): chat_box = st.container(height=350, border=False) chat_input = st.chat_input( placeholder="Upload a document to start asking questions...", disabled=not st.session_state['file_uploaded'], ) if st.session_state['file_uploaded']: display_chat_messages(chat_box, chat_input) class StreamlitApp: def __init__(self): self.authenticator_config = load_authenticator_config() self.document_store = InMemoryDocumentStore() set_page_config() self.authenticator = self.init_authenticator() init_session_state() init_api_key_dict() def init_authenticator(self): return stauth.Authenticate( self.authenticator_config['credentials'], self.authenticator_config['cookie']['name'], self.authenticator_config['cookie']['key'], self.authenticator_config['cookie']['expiry_days'] ) def setup_sidebar(self): with st.sidebar: st.sidebar.image("resources/puma.png", use_column_width=True) # Sidebar for Task Selection st.sidebar.header('Options:') model = setup_model_selection() # setup_task_selection(model) st.divider() self.authenticator.logout() reset_chat_memory() modal = Modal("Manage Files", key="demo-modal") manage_files(modal, self.document_store) st.divider() append_documentation_to_sidebar() def run(self): name, authentication_status, username = self.authenticator.login() if authentication_status: self.run_authenticated_app() elif st.session_state["authentication_status"] is False: st.error('Username/password is incorrect') elif st.session_state["authentication_status"] is None: st.warning('Please enter your username and password') def run_authenticated_app(self): self.setup_sidebar() append_header() setup_page_body() app = StreamlitApp() app.run()