Changing the code extract logic to preserve all code

app.py

@@ -83,22 +83,18 @@ def extract_and_format_code(text):
     else:
         code = text
 
-    # Remove any text before the function definition
-    code = re.sub(r'^.*?(?=def\s+\w+\s*\()', '', code, flags=re.DOTALL)
-
     # Dedent the code to remove any common leading whitespace
     code = textwrap.dedent(code)
 
     # Split the code into lines
     lines = code.split('\n')
 
-    # Find the function definition line
-    func_def_index = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), 0)
-
     # Ensure proper indentation
-    indented_lines = [lines[func_def_index]]  # Keep the function definition as is
-    for line in lines[func_def_index + 1:]:
-        if line.strip():  # If the line is not empty
+    indented_lines = []
+    for line in lines:
+        if line.strip().startswith('class') or line.strip().startswith('def'):
+            indented_lines.append(line)  # Keep class and function definitions as is
+        elif line.strip():  # If the line is not empty
             indented_lines.append('    ' + line)  # Add 4 spaces of indentation
         else:
             indented_lines.append(line)  # Keep empty lines as is
@@ -110,6 +106,7 @@ def extract_and_format_code(text):
     except:
         return formatted_code
 
+
 def select_random_problem():
     return random.choice(train_dataset)['instruction']