Update app.py

app.py

@@ -31,58 +31,28 @@ def extract_code(response):
         # Find all content between ```lean4 and ``` tags
         matches = re.findall(r'```lean4(.*?)```', response, re.DOTALL)
         if not matches:
-            # If no matches found in code blocks, return the full response
             return response.strip()
         
-        # Get the last complete response including the model output
+        # Get the last complete response
         full_content = matches[-1].strip()
         
-        # Clean up any duplicate headers or content
+        # Split into lines and clean up
         lines = full_content.split('\n')
         cleaned_lines = []
-        seen_headers = set()
-        header_done = False
         
         for line in lines:
             # Skip empty lines at the start
             if not cleaned_lines and not line.strip():
                 continue
-                
-            # Process headers only once
-            if not header_done and any(header in line for header in ["import Mathlib", "import Aesop", "set_option", "open BigOperators"]):
-                if line not in seen_headers:
-                    seen_headers.add(line)
-                    cleaned_lines.append(line)
+            # Skip lines containing "Complete the following"
+            if "Complete the following" in line:
                 continue
-            
-            # Mark header section as complete after processing imports and settings
-            if not header_done and cleaned_lines:
-                header_done = True
-                cleaned_lines.append("")  # Add a blank line after headers
-            
-            # Include the line if it's not a duplicate goal statement
-            if "Goal:" in line and line in cleaned_lines:
-                continue
-            if "Complete the following" not in line:
-                cleaned_lines.append(line)
-        
-        # Remove any trailing "Complete the following..." text
-        while cleaned_lines and "Complete the following" in cleaned_lines[-1]:
-            cleaned_lines.pop()
-            
-        # Ensure the model's output is included
-        if "===========================" in full_content:
-            output_start = full_content.find("============================")
-            if output_start != -1:
-                output_text = full_content[output_start:].strip()
-                if ":= by" in output_text:
-                    proof_part = output_text[output_text.find(":= by"):]
-                    cleaned_lines.append(proof_part)
+            cleaned_lines.append(line)
         
         return '\n'.join(cleaned_lines)
     except Exception as e:
         print(f"Error in extract_code: {str(e)}")
-        return "Error processing code"
+        return "Error processing code""
         
 
 # Example problems