Add password & streaming

app.py

@@ -1,5 +1,13 @@
 ## Import Modules
 
+from langchain.document_loaders.csv_loader import CSVLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import FAISS
+from langchain import OpenAI
+from langchain.chains import RetrievalQA
+## Import Modules
+
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings.openai import OpenAIEmbeddings
@@ -11,12 +19,43 @@ from langchain.prompts import PromptTemplate
 from langchain.prompts.few_shot import FewShotPromptTemplate
 import os
 
+# Web App
+import hmac
 import streamlit as st
 from streamlit_chat import message
 from PIL import Image
 
 os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"]
 
+def check_password():
+    """Returns `True` if the user had the correct password."""
+
+    def password_entered():
+        """Checks whether a password entered by the user is correct."""
+        # Add at secret
+        if hmac.compare_digest(st.session_state["password"], st.secrets["password"]):
+            st.session_state["password_correct"] = True
+            del st.session_state["password"]  # Don't store the password.
+        else:
+            st.session_state["password_correct"] = False
+
+    # Return True if the passward is validated.
+    if st.session_state.get("password_correct", False):
+        return True
+
+    # Show input for password.
+    st.text_input(
+        "Password", type="password", on_change=password_entered, key="password"
+    )
+    if "password_correct" in st.session_state:
+        st.error("😕 Password incorrect")
+    return False
+
+
+if not check_password():
+    st.stop()  # Do not continue if check_password is not True.
+
+
 # 주석 부분 자동으로 되는 영역 같음
 loader = CSVLoader(file_path='./parsed_texts.csv',
                    encoding='utf-8',
@@ -148,9 +187,31 @@ fewshot_prompt = FewShotPromptTemplate(
 
 # print(fewshot_prompt.format(context="Hello, world!", question="Who was the father of Mary Ball Washington?"))
 
+
+# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.callbacks.base import BaseCallbackHandler
+
+class MyCustomHandler(BaseCallbackHandler):  
+    def __init__(self):
+        super().__init__()
+        self.tokens = []
+
+    def on_llm_new_token(self, token: str, **kwargs) -> None:  
+        # print(f"My custom handler, token: {token}")
+        global full_response
+        global message_placeholder
+        self.tokens.append(token)
+        # print(self.tokens)
+        
+        full_response += token
+        message_placeholder.markdown(full_response + "▌")
+
+
+
 ## Load up your LLM
 # llm = OpenAI() # 'text-davinci-003', model_name="gpt-4"
-chat = ChatOpenAI(model_name="gpt-3.5-turbo-16k")  # gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4-32k : 정제된 Prompt 7장을 넣으려면 Prompt만 4k 이상이어야 함
+# , streaming=True, callbacks=[StreamingStdOutCallbackHandler()]
+chat = ChatOpenAI(model_name="gpt-3.5-turbo-16k", streaming=True, callbacks=[MyCustomHandler()])  # gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4-32k : 정제된 Prompt 7장을 넣으려면 Prompt만 4k 이상이어야 함
 
 chain_type_kwargs = {"prompt": fewshot_prompt}
 qa = RetrievalQA.from_chain_type(llm=chat,
@@ -163,22 +224,22 @@ qa = RetrievalQA.from_chain_type(llm=chat,
 
 img = Image.open('resource/hume.jpg')
  
-def generate_response(prompt):
-    # query = "How do you define the notion of a cause in his A Treatise of Human Nature? And how is it different from the traditional definition that you reject?"
-    result = qa({"query": prompt})
-    message = result['result']
-    sources = []
-    for src in result['source_documents']:
-        if src.page_content.startswith('Paragraph:'):
-            sources.append(src.metadata['source'])
+# def generate_response(prompt):
+#     # query = "How do you define the notion of a cause in his A Treatise of Human Nature? And how is it different from the traditional definition that you reject?"
+#     result = qa({"query": prompt})
+#     message = result['result']
+#     sources = []
+#     for src in result['source_documents']:
+#         if src.page_content.startswith('Paragraph:'):
+#             sources.append(src.metadata['source'])
     
-    if len(sources)==0:
-        message = message + "\n\n[No sources]"
-    else:
-        message = message + "\n\n[" + ", ".join(sources) + "]"
+#     if len(sources)==0:
+#         message = message + "\n\n[No sources]"
+#     else:
+#         message = message + "\n\n[" + ", ".join(sources) + "]"
         
-    return message
- 
+#     return message
+
 
 col1, col2, col3 = st.columns(3)
 
@@ -194,22 +255,33 @@ with col3:
 
 st.header("Chat with Hume (Demo)")
  
-if 'generated' not in st.session_state:
-    st.session_state['generated'] = []
- 
-if 'past' not in st.session_state:
-    st.session_state['past'] = []
- 
-with st.form('form', clear_on_submit=True):
-    user_input = st.text_input('You: ', '', key='input')
-    submitted = st.form_submit_button('Send')
- 
-if submitted and user_input:
-    output = generate_response(user_input)
-    st.session_state.past.append(user_input)
-    st.session_state.generated.append(output)
- 
-if st.session_state['generated']:
-    for i in range(len(st.session_state['generated'])-1, -1, -1):
-        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
-        message(st.session_state["generated"][i], key=str(i))
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if prompt := st.chat_input("What is up?"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+
+        full_response = ""
+        result = qa({"query": prompt})
+        
+        sources = []
+        for src in result['source_documents']:
+            if src.page_content.startswith('Paragraph:'):
+                sources.append(src.metadata['source'])
+        
+        if len(sources)==0:
+            full_response = full_response + "\n\n[No sources]"
+        else:
+            full_response = full_response + "\n\n[" + ", ".join(sources) + "]"
+            
+        message_placeholder.markdown(full_response)
+    st.session_state.messages.append({"role": "assistant", "content": full_response})