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 CV 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 CV content...'):
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="CV 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']
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()