Streaming on/off switch

code/main.py

@@ -114,6 +114,9 @@ class Chatbot:
                 cl.input_widget.Switch(
                     id="view_sources", label="View Sources", initial=False
                 ),
+                cl.input_widget.Switch(
+                    id="stream_response", label="Stream response", initial=True
+                ),
                 cl.input_widget.Select(
                     id="llm_style",
                     label="Type of Conversation (Default Normal)",
@@ -233,26 +236,18 @@ class Chatbot:
         """
         await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
 
-    async def main(self, message):
+    async def stream_response(self, response):
         """
-        Process and Display the Conversation.
+        Stream the response from the LLM.
 
         Args:
-            message: The incoming chat message.
+            response: The response from the LLM.
         """
-        chain = cl.user_session.get("chain")
-        llm_settings = cl.user_session.get("llm_settings", {})
-        view_sources = llm_settings.get("view_sources", False)
-
-        processor = cl.user_session.get("chat_processor")
-        res = await processor.rag(message.content, chain)
-
-        # TODO: STREAM MESSAGE
         msg = cl.Message(content="")
         await msg.send()
 
         output = {}
-        for chunk in res:
+        for chunk in response:
             if 'answer' in chunk:
                 await msg.stream_token(chunk['answer'])
 
@@ -261,11 +256,32 @@ class Chatbot:
                     output[key] = chunk[key]
                 else:
                     output[key] += chunk[key]
+        return output
+
+    async def main(self, message):
+        """
+        Process and Display the Conversation.
+
+        Args:
+            message: The incoming chat message.
+        """
+        chain = cl.user_session.get("chain")
+        llm_settings = cl.user_session.get("llm_settings", {})
+        view_sources = llm_settings.get("view_sources", False)
+        stream = (llm_settings.get("stream_response", True)) or (not self.config["llm_params"]["stream"])
+        print("Streaming", stream)
+
+        processor = cl.user_session.get("chat_processor")
+        res = await processor.rag(message.content, chain, stream)
+
+        if stream:
+            res = await self.stream_response(res)
 
-        answer = output.get("answer", output.get("result"))
+        print(res)
+        answer = res.get("answer", res.get("result"))
 
         answer_with_sources, source_elements, sources_dict = get_sources(
-            output, answer, view_sources=view_sources
+            res, answer, stream=stream, view_sources=view_sources
         )
         processor._process(message.content, answer, sources_dict)
 

code/modules/chat/helpers.py

@@ -2,7 +2,7 @@ from modules.config.prompts import prompts
 import chainlit as cl
 
 
-def get_sources(res, answer, view_sources=False):
+def get_sources(res, answer, stream=True, view_sources=False):
     source_elements = []
     source_dict = {}  # Dictionary to store URL elements
 
@@ -35,10 +35,11 @@ def get_sources(res, answer, view_sources=False):
         else:
             source_dict[url_name]["text"] += f"\n\n{source.page_content}"
 
-    # First, display the answer
-    #full_answer = "**Answer:**\n"
-    #full_answer += answer
-    full_answer = ""  # Not to include the answer again
+    full_answer = ""  # Not to include the answer again if streaming
+
+    if not stream:  # First, display the answer if not streaming
+        full_answer = "**Answer:**\n"
+        full_answer += answer
 
     if view_sources:
 

code/modules/chat_processor/chat_processor.py

@@ -33,7 +33,7 @@ class ChatProcessor:
         else:
             pass
 
-    async def rag(self, user_query: str, chain):
+    async def rag(self, user_query: str, chain, stream):
         user_query_dict = {"input": user_query}
         # Define the base configuration
         config = {
@@ -50,4 +50,6 @@ class ChatProcessor:
                 user_query=user_query_dict, config=config, chain=chain
             )
         else:
-            return chain.stream(user_query=user_query_dict, config=config)
+            if stream:
+                return chain.stream(user_query=user_query_dict, config=config)
+            return chain.invoke(user_query=user_query_dict, config=config)

code/modules/config/config.yml

@@ -35,6 +35,7 @@ llm_params:
     temperature: 0.7 # float
   local_llm_params:
     temperature: 0.7 # float
+  stream: True # bool
 
 chat_logging:
   log_chat: False # bool