Update app.py

app.py

@@ -332,34 +332,35 @@ def extract_and_run_sympy_code(response_text):
         
         sympy_code = response_text[code_start:code_end].strip()
         
-        # Create a string buffer to capture print output
+        # Import necessary modules
         import io
         import sys
         from contextlib import redirect_stdout
-        
-        # Create string buffer for output
+        import sympy
+
+        # Create a string buffer to capture print output
         output_buffer = io.StringIO()
         
-        # Create globals dict with all SymPy symbols
-        sympy_globals = {}
-        sympy_globals.update({name: getattr(sympy, name) for name in dir(sympy)})
+        # Create globals dict with SymPy explicitly included
+        sympy_globals = {"sympy": sympy}
+        sympy_globals.update(vars(sympy))
         
         # Create locals dict to capture new variables
         local_vars = {}
-        
-        # Execute the original code
-        exec(sympy_code, sympy_globals, local_vars)
-        
-        # Print all variables that were created (excluding built-ins and modules)
+
+        # Redirect stdout and execute the code
         with redirect_stdout(output_buffer):
-            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}")
+            exec(sympy_code, sympy_globals, local_vars)
         
-        return output_buffer.getvalue()
+        # Append the calculated variables
+        output_buffer.write("SymPy Calculation Results:\n")
+        output_buffer.write("-" * 25 + "\n")
+        for var_name, value in local_vars.items():
+            if not var_name.startswith('__') and not hasattr(value, '__module__'):
+                output_buffer.write(f"{var_name}: {value}\n")
         
+        return output_buffer.getvalue()
+    
     except Exception as e:
         return f"Error executing SymPy code: {str(e)}"