Update app.py

app.py

@@ -79,16 +79,54 @@ retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
     llm=chat_openai
 )
 
-@cl.on_chat_start # marks a function that will be executed at the start of a user session
-async def start_chat():
-    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(vector_db_retriever=vector_db, llm=chat_openai, wandb_project=wandb_project)
-    cl.user_session.set("pipeline", retrieval_augmented_qa_pipeline)
-
-@cl.on_message # marks a function that should be run each time the chatbot receives a message from a user
-async def main(message: str):
-    retrieval_augmented_qa_pipeline = cl.user_session.get("pipeline")
-    completion = retrieval_augmented_qa_pipeline.run_pipeline(message)
+@cl.on_message
+async def main(message: cl.Message):
+    settings = cl.user_session.get("settings")
+
+    client = AsyncOpenAI()
+
+    context_list = vector_db.search_by_text(message.content, 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=message.content)
+    print(formatted_system_prompt)
+    print(formatted_user_prompt)
+    prompt = Prompt(
+        provider=ChatOpenAI.id,
+        messages=[
+            PromptMessage(
+                role="system",
+                template=RAQA_PROMPT_TEMPLATE,
+                formatted=formatted_system_prompt['content'],
+            ),
+            PromptMessage(
+                role="user",
+                template=USER_PROMPT_TEMPLATE,
+                formatted=formatted_user_prompt['content'],
+            ),
+        ],
+        inputs={"context": context_prompt,
+            "user_query": message.content},
+        settings=settings,
+    )
+
+    msg = cl.Message(content="")
+    
+    async for stream_resp in await client.chat.completions.create(
+        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
+    ):
+        token = stream_resp.choices[0].delta.content
+        if not token:
+            token = ""
+        await msg.stream_token(token)
+        
+    # Update the prompt object with the completion
+    prompt.completion = msg.content
+    msg.prompt = prompt
 
-    await cl.Message(content=completion).send()
+    # Send and close the message stream
+    await msg.send()