Merge pull request #73 from ottlseo/master

Add a new streamlit UI for previous log

Author: dongjin-ml

genai/aws-gen-ai-kr/20_applications/02_qa_chatbot/04_web_ui/streamlit.py

@@ -2,6 +2,34 @@
 import bedrock as glib  # 로컬 라이브러리 스크립트에 대한 참조
 from langchain.callbacks import StreamlitCallbackHandler
 
+def context_showing_tab(contexts):
+    tab_titles = []
+    tab_contents = {}
+    for i, context in enumerate(contexts):
+        title = str(context[0])
+        tab_titles.append(title)
+        tab_contents[title] = context[1][0]
+    tabs = st.tabs(tab_titles)
+    for i, tab in enumerate(tabs):
+        with tab:
+            st.header(tab_titles[i])
+            st.write(tab_contents[tab_titles[i]])
+
+def multi_answer_column(answers):
+    col1, col2, col3, col4 = st.columns(4)
+    with col1:
+        st.markdown('''### option 1 ''')
+        st.write(answers[0])
+    with col2:
+        st.markdown('''### option 2 ''')
+        st.write(answers[1])
+    with col3:
+        st.markdown('''### option 3 ''')
+        st.write(answers[2])
+    with col4:
+        st.markdown('''### option 4 ''')
+        st.write(answers[3])
+
 st.set_page_config(layout="wide")
 st.title("AWS Q&A Bot with Advanced RAG!")  # page 제목
 
@@ -17,7 +45,6 @@
     st.session_state.showing_option = "Separately"
 
 with st.sidebar: # Sidebar 모델 옵션
-    # st.title("Set showing method 👇")
     with st.container(height=170):
         st.radio(
             "Set showing method 👇",
@@ -27,21 +54,37 @@
         )
 
     st.title("Set parameter for your Bot 👇")
-    parent = st.toggle("Parent_docs", disabled=st.session_state.showing_option=="All at once")
+
+    # semantic = st.toggle("Semantic", disabled=st.session_state.showing_option=="All at once")
+    # lexical = st.toggle("Lexical", disabled=st.session_state.showing_option=="All at once")
+    
+    # hybrid = st.slider('Alpha value of Hybrid Search: lexical(0.0) / semantic(1.0)', 0.0, 1.0, 0.5)
+    
     reranker = st.toggle("Reranker", disabled=st.session_state.showing_option=="All at once")
-    # hyde = st.toggle("HyDE")
-    # custom_model = st.toggle("Custom Model")
-    # alpha = st.slider('Alpha value of Hybrid Search: lexical(0.0) / semantic(1.0)', 0.0, 1.0, 0.5)
+    
+    # ragfusion = st.toggle("RAG-Fusion", disabled=st.session_state.showing_option=="All at once")
+    # hyde = st.toggle("HyDE", disabled=st.session_state.showing_option=="All at once")
+    
+    parent = st.toggle("Parent_docs", disabled=st.session_state.showing_option=="All at once")
 
 ### 1) 'Separately' 옵션 선택한 경우 ###
 if st.session_state.showing_option == "Separately":
     if "messages" not in st.session_state:
         st.session_state["messages"] = [
             {"role": "assistant", "content": "How can I help you?"}
         ]
+    # 지난 답변 출력
     for msg in st.session_state.messages:
-        st.chat_message(msg["role"]).write(msg["content"])
-
+        # 지난 답변에 대한 컨텍스트 출력
+        if msg["role"] == "assistant_context": 
+            with st.chat_message("assistant"):
+                with st.expander("정확도 별 답변 보기 ⬇️"):
+                    context_showing_tab(contexts=msg["content"])
+        if msg["role"] == "assistant_column": 
+            st.chat_message("assistant").write(msg["content"][0])
+        else:
+            st.chat_message(msg["role"]).write(msg["content"])
+    
     # 유저가 쓴 chat을 query라는 변수에 담음
     query = st.chat_input("Search documentation")
     if query:
@@ -74,53 +117,22 @@
         st.chat_message("assistant").write(answer)
 
         with st.chat_message("assistant"): 
-            with st.expander("정확도 별 답변 보기 (semantic) ⬇️"):
-                tab_titles = []
-                tab_contents = {}
-                for i, context in enumerate(contexts1):
-                    title = str(context[0])
-                    tab_titles.append(title)
-                    tab_contents[title] = context[1][0]
-                tabs = st.tabs(tab_titles)
-                for i, tab in enumerate(tabs):
-                    with tab:
-                        st.header(tab_titles[i])
-                        st.write(tab_contents[tab_titles[i]])
+            with st.expander("정확도 별 답변 보기 ⬇️"): # 정확도 별 답변 보기 (semantic)
+                context_showing_tab(contexts1)
+                
+        # with st.chat_message("assistant"): 
+        #     with st.expander("정확도 별 답변 보기 (lexical) ⬇️"):
+        #         context_showing_tab(contexts2)
 
-        with st.chat_message("assistant"): 
-            with st.expander("정확도 별 답변 보기 (lexical) ⬇️"):
-                tab_titles = []
-                tab_contents = {}
-                for i, context in enumerate(contexts2):
-                    title = str(context[0])
-                    tab_titles.append(title)
-                    tab_contents[title] = context[1][0]
-                tabs = st.tabs(tab_titles)
-                for i, tab in enumerate(tabs):
-                    with tab:
-                        st.header(tab_titles[i])
-                        st.write(tab_contents[tab_titles[i]])
-
-        with st.chat_message("assistant"): 
-            with st.expander("정확도 별 답변 보기 (reranker) ⬇️"):
-                tab_titles = []
-                tab_contents = {}
-                for i, context in enumerate(contexts3):
-                    title = str(context[0])
-                    tab_titles.append(title)
-                    tab_contents[title] = context[1][0]
-                tabs = st.tabs(tab_titles)
-                for i, tab in enumerate(tabs):
-                    with tab:
-                        st.header(tab_titles[i])
-                        st.write(tab_contents[tab_titles[i]])
+        # with st.chat_message("assistant"): 
+        #     with st.expander("정확도 별 답변 보기 (reranker) ⬇️"):
+        #         context_showing_tab(contexts3)
 
         # Session 메세지 저장
         st.session_state.messages.append({"role": "assistant", "content": answer})
-        # st.session_state.messages.append({"role": "assistant", "content": contexts1})
-        # st.session_state.messages.append({"role": "assistant", "content": contexts2})
-        # st.session_state.messages.append({"role": "assistant", "content": contexts3})
-        # st.session_state.messages.append({"role": "assistant", "content": contexts4})
+        st.session_state.messages.append({"role": "assistant_context", "content": contexts1})
+        # st.session_state.messages.append({"role": "assistant_context", "content": contexts2})
+        # st.session_state.messages.append({"role": "assistant_context", "content": contexts3})
 
         # Thinking을 complete로 수동으로 바꾸어 줌
         st_cb._complete_current_thought()
@@ -131,9 +143,14 @@
         st.session_state["messages"] = [
             {"role": "assistant", "content": "How can I help you?"}
         ]
+    # 지난 답변 출력
     for msg in st.session_state.messages:
-        st.chat_message(msg["role"]).write(msg["content"])
-
+        if msg["role"] == "assistant_column":
+            answers = msg["content"]
+            multi_answer_column(answers)
+        else:
+            st.chat_message(msg["role"]).write(msg["content"])
+    
     # 유저가 쓴 chat을 query라는 변수에 담음
     query = st.chat_input("Search documentation")
     if query:
@@ -145,80 +162,69 @@
 
         col1, col2, col3, col4 = st.columns(4)
         with col1:
-            st.markdown('''#### parent=:red[False], reranker=:red[False]''')
+            st.markdown('''### option 1 ''')
         with col2:
-            st.markdown('''#### parent=:green[True], reranker=:red[False]''')
+            st.markdown('''### option 2 ''')
         with col3:
-            st.markdown('''#### parent=:red[False], reranker=:green[True]''')
+            st.markdown('''### option 3 ''')
         with col4:
-            st.markdown('''#### parent=:green[True], reranker=:green[True]''')
+            st.markdown('''### option 4 ''')
 
         with col1:
             # Streamlit callback handler로 bedrock streaming 받아오는 컨테이너 설정
             st_cb = StreamlitCallbackHandler(
                 st.chat_message("assistant"), 
                 collapse_completed_thoughts=True
                 )
-            answer = glib.invoke(
+            answer1 = glib.invoke(
                 query=query, 
                 streaming_callback=st_cb, 
                 parent=False, 
                 reranker=False
                 )[0]
-            # st.subheader("parent=False, reranker=False ⬇️")
-            st.write(answer)
+            st.write(answer1)
             st_cb._complete_current_thought() # Thinking을 complete로 수동으로 바꾸어 줌
         with col2:
-            # Streamlit callback handler로 bedrock streaming 받아오는 컨테이너 설정
             st_cb = StreamlitCallbackHandler(
                 st.chat_message("assistant"), 
                 collapse_completed_thoughts=True
                 )
-            answer = glib.invoke(
+            answer2 = glib.invoke(
                 query=query, 
                 streaming_callback=st_cb, 
                 parent=True, 
                 reranker=False
                 )[0]
-            # st.subheader("parent=True, reranker=False ⬇️")
-            st.write(answer)
-            st_cb._complete_current_thought() # Thinking을 complete로 수동으로 바꾸어 줌
+            st.write(answer2)
+            st_cb._complete_current_thought() 
         with col3:
-            # Streamlit callback handler로 bedrock streaming 받아오는 컨테이너 설정
             st_cb = StreamlitCallbackHandler(
                 st.chat_message("assistant"), 
                 collapse_completed_thoughts=True
                 )
-            answer = glib.invoke(
+            answer3 = glib.invoke(
                 query=query, 
                 streaming_callback=st_cb, 
                 parent=False, 
                 reranker=True
                 )[0]
-            # st.subheader("parent=False, reranker=True ⬇️")
-            st.write(answer)
-            st_cb._complete_current_thought() # Thinking을 complete로 수동으로 바꾸어 줌
+            st.write(answer3)
+            st_cb._complete_current_thought() 
         with col4:
-            # Streamlit callback handler로 bedrock streaming 받아오는 컨테이너 설정
             st_cb = StreamlitCallbackHandler(
                 st.chat_message("assistant"), 
                 collapse_completed_thoughts=True
             )
-            answer = glib.invoke(
+            answer4 = glib.invoke(
                 query=query, 
                 streaming_callback=st_cb, 
                 parent=True, 
                 reranker=True
                 )[0]
-            # st.subheader("parent=True, reranker=True ⬇️")
-            st.write(answer)
-            st_cb._complete_current_thought() # Thinking을 complete로 수동으로 바꾸어 줌
+            st.write(answer4)
+            st_cb._complete_current_thought()
 
         # Session 메세지 저장
-        st.session_state.messages.append({"role": "assistant", "content": answer})
-        # st.session_state.messages.append({"role": "assistant", "content": contexts})
-        
-        # UI 출력
-        # st.chat_message("assistant").write(answer)
-        # st.chat_message("assistant").write(contexts)
-        # st.chat_message("assistant").write(ref)
+        answer = [answer1, answer2, answer3, answer4]
+        st.session_state.messages.append({"role": "assistant_column", "content": answer})
+