add stream pipeline

aimakerspace/openai_utils/chatmodel.py

@@ -25,3 +25,10 @@ class ChatOpenAI:
             return response.choices[0].message.content
 
         return response
+
+    def run_stream(self, messages, settings, chainlit_msg, text_only: bool = True):
+        async for stream_resp in await openai.Completion.acreate(
+            model=self.model_name, prompt=messages, stream=True, **settings
+        ):
+            token = stream_resp.get("choices")[0].get("text")
+        await chainlit_msg.stream_token(token)

app.py

@@ -76,6 +76,19 @@ class RetrievalAugmentedQAPipeline:
 
         return self.llm.run([formatted_system_prompt, formatted_user_prompt])
 
+    def stream_pipeline(self, user_query: str, msg: cl.Message) -> str:
+        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
+
+        context_prompt = ""
+        for context in context_list:
+            context_prompt += context[0] + "\n"
+
+        formatted_system_prompt = raqa_prompt.create_message(context=context_prompt)
+
+        formatted_user_prompt = user_prompt.create_message(user_query=user_query)
+
+        self.llm.stream([formatted_system_prompt, formatted_user_prompt])
+
 
 @cl.on_chat_start # marks a function that will be executed at the start of a user session
 def start_chat():
@@ -97,7 +110,10 @@ def start_chat():
 async def main(message: str):
 
     qaPipeline =  RetrievalAugmentedQAPipeline(vector_db_retriever=vector_db, llm=chat_openai)
-    qaPipeline.run_pipeline(user_query=message)
+    msg = cl.Message(content="")
+
+    qaPipeline.stream_pipeline(user_query=message, msg=msg)
+    await msg.send()