Remove api key from app.py

app.py

@@ -1,245 +1,243 @@
-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()
+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')
+
+    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()