Update app.py

app.py

@@ -24,27 +24,32 @@ def format_prompt(formal_statement, informal_prefix=""):
         f"{informal_prefix}\n"
         f"{formal_statement}"
     )
-    
+
 def extract_code(response):
     """Extract code between lean4 code blocks and the model's output"""
     try:
-        # Find all content between ```lean4 and ``` tags
-        matches = re.findall(r'```lean4(.*?)```', response, re.DOTALL)
-        if not matches:
+        # Find the last occurrence of ```lean4 and extract everything until the last ```
+        start_idx = response.rfind("```lean4")
+        if start_idx == -1:
             return response.strip()
         
-        # Get the last complete response
-        full_content = matches[-1].strip()
+        # Get content after ```lean4
+        content = response[start_idx + 7:]
+        
+        # Find the last closing ```
+        end_idx = content.rfind("```")
+        if end_idx != -1:
+            content = content[:end_idx]
         
-        # Split into lines and clean up
-        lines = full_content.split('\n')
+        # Clean up the content
+        lines = content.split('\n')
         cleaned_lines = []
         
         for line in lines:
-            # Skip empty lines at the start
+            # Skip empty lines at start
             if not cleaned_lines and not line.strip():
                 continue
-            # Skip lines containing "Complete the following"
+            # Skip "Complete the following" lines
             if "Complete the following" in line:
                 continue
             cleaned_lines.append(line)