Enable streaming responses in Chainlit chatbot

app.py

@@ -116,7 +116,7 @@ Question:
 prompt = ChatPromptTemplate.from_template(template)
 
 # Define the primary LLM
-primary_llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
+primary_llm = ChatOpenAI(model_name="gpt-4o", temperature=0, streaming=True)
 
 #-----Creating a Retrieval Augmented Generation (RAG) Chain-----#
 # The RAG chain:
@@ -161,11 +161,28 @@ async def start_chat():
 async def handle_message(message: cl.Message):
     settings = cl.user_session.get("settings")
 
-    response = retrieval_augmented_qa_chain.invoke({"question": message.content})
 
+    # Initialize the stream message in Chainlit
+    stream_msg = cl.Message(content="")
+    await stream_msg.send()  # Send initial empty message to start the stream
 
-    # Extracting and sending just the content
-    content = response["response"].content
-    pretty_content = content.strip()  # Remove any leading/trailing whitespace
+    # Create a generator from the RAG chain
+    response_generator = retrieval_augmented_qa_chain.stream({"question": message.content})
 
-    await cl.Message(content=pretty_content).send()
+    async for response_chunk in response_generator:
+        # Extract the content from the chunk
+        chunk_content = response_chunk.get("response", {}).get("content", "")
+        if chunk_content:
+            # Append the chunk to the streaming message content
+            stream_msg.content += chunk_content
+            await stream_msg.update()  # Update the message in Chainlit
+
+    ## Remove to stream the response
+    # response = retrieval_augmented_qa_chain.invoke({"question": message.content})
+
+
+##    # Extracting and sending just the content
+##    content = response["response"].content
+##    pretty_content = content.strip()  # Remove any leading/trailing whitespace
+
+##   await cl.Message(content=pretty_content).send()