| import streamlit as st |
| from streamlit_option_menu import option_menu |
| |
| from modules import nav_about, nav_query_docs, nav_researcher, nav_summarizer, nav_file_manager |
| from modules import app_constants, app_logger, common_utils |
| from modules.message_store import MessageStore |
|
|
| app_logger = app_logger.app_logger |
| WORKSPACE_DIRECTORY = app_constants.WORKSPACE_DIRECTORY |
|
|
| |
| st.set_page_config(page_title="ZySec AI", page_icon=":sparkles:", layout="wide") |
|
|
| |
| if 'message_store' not in st.session_state: |
| st.session_state['message_store'] = MessageStore() |
| def request_username(): |
| st.title("Welcome to ZySec AI") |
| app_logger.info("App started") |
| username = st.text_input("How do you want me to call you?",value="Security Ninja" ,placeholder="Enter your name") |
| submit_button = st.button('Submit') |
|
|
| if submit_button and username: |
| st.session_state['username'] = username |
| return True |
| return False |
|
|
| def main(): |
| common_utils.setup_initial_folders() |
| if 'messages' not in st.session_state: |
| st.session_state['messages'] = [] |
|
|
| if 'username' not in st.session_state or not st.session_state['username']: |
| if request_username(): |
| st.rerun() |
| return |
|
|
| |
| with st.sidebar: |
| selected = option_menu( |
| "ZySec AI", |
| ["Private AI", "Playbooks", "Standards", "Policies", "Researcher", "Summarizer","Files", "About"], |
| icons=["shield-lock", "book-half", "file-earmark-ruled", "journal-bookmark", "search", "file-text","files", "info-circle"], |
| default_index=0, |
| menu_icon="cast", |
| styles={} |
| ) |
| st.markdown("---") |
|
|
|
|
| try: |
| message_store = st.session_state['message_store'] |
|
|
|
|
| if selected == "Private AI": |
| |
| nav_query_docs.app(message_store,current_page="nav_private_ai") |
| elif selected == "Playbooks": |
| nav_query_docs.app(message_store,current_page="nav_playbooks",use_retrieval_chain=True) |
| elif selected == "Standards": |
| nav_query_docs.app(message_store,current_page="nav_standards",use_retrieval_chain=True) |
| elif selected == "Policies": |
| nav_query_docs.app(message_store,current_page="nav_policies",use_retrieval_chain=True) |
| elif selected == "Researcher": |
| nav_researcher.app(message_store) |
| elif selected == "Summarizer": |
| nav_summarizer.app() |
| elif selected == "Files": |
| nav_file_manager.app() |
| elif selected == "About": |
| nav_about.app() |
| else: |
| pass |
|
|
| except Exception as e: |
| st.error(f"Looks like there's a gap here; maybe we forgot to add some data!: {e}") |
|
|
| if __name__ == "__main__": |
| main() |