Update app.py

app.py

@@ -33,11 +33,10 @@ def extract_code(response):
         if not matches:
             return "None"
         
-        # Get the last match which should contain the complete code including model output
+        # Get the last complete response including the model output
         full_content = matches[-1].strip()
         
         # Clean up any duplicate headers or content
-        # Split the content into lines
         lines = full_content.split('\n')
         cleaned_lines = []
         seen_headers = set()
@@ -46,14 +45,20 @@ def extract_code(response):
             # Skip empty lines at the start
             if not cleaned_lines and not line.strip():
                 continue
-            # Skip duplicate headers
-            if "import Mathlib" in line or "import Aesop" in line or "set_option" in line:
+            # Skip duplicate headers and "Complete the following..." text
+            if any(header in line for header in ["import Mathlib", "import Aesop", "set_option", "Complete the following"]):
                 if line in seen_headers:
                     continue
                 seen_headers.add(line)
+            # Skip duplicate goal statements
+            if "Goal:" in line and line in cleaned_lines:
+                continue
             cleaned_lines.append(line)
         
-        # Join the lines back together
+        # Remove any trailing "Complete the following..." text
+        while cleaned_lines and "Complete the following" in cleaned_lines[-1]:
+            cleaned_lines.pop()
+            
         return '\n'.join(cleaned_lines)
     except:
         return "None"