Update app.py

adding streaming

app.py

@@ -10,14 +10,33 @@ from langchain.schema import (
     AIMessage
 )
 from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.callbacks.base import BaseCallbackHandler
 
 import streamlit as st
 
 load_dotenv()
 
 
-os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
-llm = ChatOpenAI(temperature=0.3, model="gpt-3.5-turbo-16k")
+class StreamHandler(BaseCallbackHandler):
+    def __init__(self, container, initial_text=""):
+        self.container = container
+        self.text = initial_text
+
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        self.text += token
+        self.container.markdown(self.text)
+
+
+def init_gpt(gpt_model, stream_handler):
+    global llm
+    llm = ChatOpenAI(
+        temperature=0.3,
+        model=gpt_model,
+        streaming=True,
+        callbacks=[stream_handler]
+    )
+
+
 embeddings = OpenAIEmbeddings()
 
 
@@ -70,14 +89,16 @@ def main():
             st.markdown(prompt)
 
         with st.chat_message("assistant"):
-            message_placeholder = st.empty()
+            # message_placeholder = st.empty()
+            stream_handler = StreamHandler(st.empty())
+            init_gpt(gpt_model, stream_handler)
             content = generate_content(
                 st.session_state.messages, knowledge_base
             )
         st.session_state.messages.append(
             {"role": "assistant", "content": content}
         )
-        message_placeholder.markdown(content)
+        # message_placeholder.markdown(content)
 
     with st.sidebar:
         # remove last 2 messages
@@ -87,4 +108,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()