Update app.py

app.py

@@ -35,7 +35,8 @@ def generate_response(message, history, enable_history=False, max_new_tokens=102
     # Add history only if enabled
     if enable_history and history:
         prev = "\n".join(
-            [f"User: {h[0]}\nGamunu: {h[1]}" for h in history if h[1] is not None]
+            [f"User: {h['content']}\nGamunu: {h.get('content', '')}" 
+             for h in history if h.get('role') == 'assistant']
         )
         context = f"{prev}\n\n{message}"
     else:
@@ -68,22 +69,31 @@ def generate_response(message, history, enable_history=False, max_new_tokens=102
     thread.start()
     
     # Stream the response
-    partial_text = ""
-    response_started = False
+    full_text = ""
     
     for new_text in streamer:
-        partial_text += new_text
+        full_text += new_text
         
-        # Check if we've reached the response section
-        if "### ප්‍රතිචාරය:" in partial_text and not response_started:
-            partial_text = partial_text.split("### ප්‍රතිචාරය:")[-1].strip()
-            response_started = True
-        
-        if response_started:
-            yield partial_text
+        # Check if we've reached the response section and extract it
+        if "### ප්‍රතිචාරය:" in full_text:
+            response_text = full_text.split("### ප්‍රතිචාරය:")[-1].strip()
+            yield response_text
+        else:
+            # Still building up to the response marker, yield what we have
+            yield full_text.strip()
     
     # Make sure thread completes
     thread.join()
+    
+    # Final yield with cleaned response
+    if "### ප්‍රතිචාරය:" in full_text:
+        final_response = full_text.split("### ප්‍රතිචාරය:")[-1].strip()
+    else:
+        final_response = full_text.strip()
+    
+    # Ensure we yield at least once
+    if final_response:
+        yield final_response
 
 # Custom CSS for styling with copy button
 custom_css = """