Merge pull request #175 from aws-samples/spy_dev

add force_set_cookie

Author: supinyu

application/Dockerfile

@@ -8,7 +8,7 @@ WORKDIR /app
 
 COPY requirements.txt /app/
 
-ARG AWS_REGION=us-east-1
+#ARG AWS_REGION=us-east-1
 ENV AWS_REGION=${AWS_REGION}
 
 # Print the AWS_REGION for verification

application/Dockerfile-api

@@ -4,7 +4,7 @@ WORKDIR /app
 
 COPY . /app/
 
-ARG AWS_REGION=us-east-1
+#ARG AWS_REGION=us-east-1
 ENV AWS_REGION=${AWS_REGION}
 
 # Print the AWS_REGION for verification

application/Index.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from utils.navigation import get_authenticator
+from utils.navigation import get_authenticator, force_set_cookie
 
 st.set_page_config(
     page_title="Intelligent BI",
@@ -10,6 +10,7 @@
 name, authentication_status, username = authenticator.login('main')
 
 if authentication_status:
+    force_set_cookie(authenticator)
     st.switch_page("pages/mainpage.py")
 elif authentication_status == False:
     st.error('Username/password is incorrect')

application/pages/5_🪙_Prompt_Management.py

@@ -36,18 +36,18 @@ def main():
             prompt_type_selected_table = st.selectbox("Prompt Type", prompt_map.keys(), index=None,
                                                       format_func=lambda x: prompt_map[x].get('title'),
                                                       placeholder="Please select a prompt type")
-
-            profile_detail = ProfileManagement.get_profile_by_name(current_profile)
-            prompt_map = profile_detail.prompt_map
-
             if prompt_type_selected_table is not None:
                 single_type_prompt_map = prompt_map.get(prompt_type_selected_table)
                 system_prompt = single_type_prompt_map.get('system_prompt')
-                user_prompt = single_type_prompt_map.get('user_prompt')
                 model_selected_table = st.selectbox("LLM Model", system_prompt.keys(), index=None,
                                                     placeholder="Please select a model")
 
                 if model_selected_table is not None:
+                    profile_detail = ProfileManagement.get_profile_by_name(current_profile)
+                    prompt_map = profile_detail.prompt_map
+                    single_type_prompt_map = prompt_map.get(prompt_type_selected_table)
+                    system_prompt = single_type_prompt_map.get('system_prompt')
+                    user_prompt = single_type_prompt_map.get('user_prompt')
                     system_prompt_input = st.text_area('System Prompt', system_prompt[model_selected_table], height=300)
                     user_prompt_input = st.text_area('User Prompt', user_prompt[model_selected_table], height=500)
 

application/utils/navigation.py

@@ -20,6 +20,18 @@ def get_authenticator():
     )
 
 
+def force_set_cookie(authenticator):
+    """
+    Force the cookie
+    :param authenticator:
+    :return:
+    """
+    try:
+        authenticator.cookie_handler.set_cookie()
+    except:
+        pass
+
+
 def get_current_page_name():
     ctx = get_script_run_ctx()
     if ctx is None:
@@ -35,12 +47,14 @@ def make_sidebar():
         if st.session_state.get('authentication_status'):
             st.page_link("pages/mainpage.py", label="Index")
             st.page_link("pages/1_🌍_Generative_BI_Playground.py", label="Generative BI Playground", icon="🌍")
-            st.markdown(":gray[Data Customization Management]", help='Add your own datasources and customize description for LLM to better understand them')
+            st.markdown(":gray[Data Customization Management]",
+                        help='Add your own datasources and customize description for LLM to better understand them')
             st.page_link("pages/2_🪙_Data_Connection_Management.py", label="Data Connection Management", icon="🪙")
             st.page_link("pages/3_🪙_Data_Profile_Management.py", label="Data Profile Management", icon="🪙")
             st.page_link("pages/4_🪙_Schema_Description_Management.py", label="Schema Description Management", icon="🪙")
             st.page_link("pages/5_🪙_Prompt_Management.py", label="Prompt Management", icon="🪙")
-            st.markdown(":gray[Performance Enhancement]", help='Optimize your LLM for better performance by adding RAG or agent')
+            st.markdown(":gray[Performance Enhancement]",
+                        help='Optimize your LLM for better performance by adding RAG or agent')
             st.page_link("pages/6_📚_Index_Management.py", label="Index Management", icon="📚")
             st.page_link("pages/7_📚_Entity_Management.py", label="Entity Management", icon="📚")
             st.page_link("pages/8_📚_Agent_Cot_Management.py", label="Agent Cot Management", icon="📚")