Update app.py

app.py

@@ -88,6 +88,7 @@ def respond(
         # Tags
         start_tag = "<start_working_out>"
         sol_start = "<SOLUTION>"
+        sol_end = "</SOLUTION>"
 
         thinking = ""
         final_answer = ""
@@ -105,7 +106,12 @@ def respond(
         # Extract "Final answer" section: everything after <SOLUTION>
         if sol_start in output:
             sol_start_idx = output.find(sol_start) + len(sol_start)
-            final_answer = output[sol_start_idx:].strip()
+            # If </SOLUTION> is present, stop at it
+            if sol_end in output:
+                sol_end_idx = output.find(sol_end)
+                final_answer = output[sol_start_idx:sol_end_idx].strip()
+            else:
+                final_answer = output[sol_start_idx:].strip()
 
         # Build formatted output
         formatted_output = ""
@@ -120,6 +126,10 @@ def respond(
 
         yield formatted_output
 
+        # If </SOLUTION> is found, end the response
+        if sol_end in output:
+            break
+
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface