Update app.py

app.py

@@ -18,27 +18,46 @@ def process_files_into_docs(files, progress=gr.Progress()):
 def rag(history, question):
     if history is None:
         history = []
-        
-    messages = [system_message, user_message]
+
+    # Run Haystack pipeline
     res = conversational_rag.run(
-        data = {
-            'query_rephrase_prompt_builder' : {'query': question},
-            'prompt_builder': {'template': messages, 'query': question},
-            'memory_joiner': {'values': [ChatMessage.from_user(question)]}
+        data={
+            "query_rephrase_prompt_builder": {"query": question},
+            "prompt_builder": {
+                "template": [system_message, user_message],
+                "query": question,
+            },
+            "memory_joiner": {
+                "values": [ChatMessage.from_user(question)]
+            }
         },
-        include_outputs_from=['llm', 'query_rephrase_llm']
+        include_outputs_from=["llm"]
     )
-    
-    bot_message = res['llm']['replies'][0].content
-    streamed_message = ""
-    
+
+    bot_message = res["llm"]["replies"][0].content
+
+    # Add user message
+    history = history + [
+        {"role": "user", "content": question}
+    ]
+
+    # Stream assistant message
+    streamed = ""
     for token in bot_message.split():
-        streamed_message += f"{token} "
-        yield history + [(question, streamed_message.strip())], " "
+        streamed += token + " "
+        yield (
+            history + [{"role": "assistant", "content": streamed.strip()}],
+            ""
+        )
         time.sleep(0.05)
-        
-    history.append((question, bot_message))
-    yield history, " "
+
+    # Final assistant message
+    history = history + [
+        {"role": "assistant", "content": bot_message}
+    ]
+
+    yield history, ""
+
 
 EXAMPLE_FILE = "RAG Survey.pdf"