Update app.py

app.py

@@ -365,18 +365,22 @@ def extract_and_run_sympy_code(response_text):
         from contextlib import redirect_stdout
         output_buffer = io.StringIO()
         
-        logger.debug("Attempting to execute code")
-        # Execute with try/except to catch specific error
         try:
-            exec(sympy_code, globals_dict, local_vars)
-            
-            # Print all variables that were created
-            print("SymPy Calculation Results:")
-            print("-" * 25)
-            for var_name, value in local_vars.items():
-                if (not var_name.startswith('__') and 
-                    not hasattr(value, '__module__')):
-                    print(f"{var_name} = {value}")
+            # Execute the code and capture all output in one redirect_stdout block
+            with redirect_stdout(output_buffer):
+                # First print a header
+                print("SymPy Calculation Results:")
+                print("-" * 25)
+                
+                # Execute the code
+                exec(sympy_code, globals_dict, local_vars)
+                
+                # Print all variables created
+                for var_name, value in local_vars.items():
+                    if (not var_name.startswith('__') and 
+                        not hasattr(value, '__module__')):
+                        print(f"{var_name} = {value}")
+                
         except NameError as e:
             logger.error(f"NameError during execution: {str(e)}")
             logger.debug(f"Available names: {list(globals_dict.keys())}")
@@ -387,6 +391,7 @@ def extract_and_run_sympy_code(response_text):
             logger.error(f"Other error during execution: {str(e)}")
             return f"Error executing SymPy code: {str(e)}"
         
+        # Get the captured output
         result = output_buffer.getvalue()
         logger.debug(f"Final output:\n{result}")
         return result