Update app.py

app.py

@@ -346,33 +346,60 @@ def extract_and_run_sympy_code(response_text):
         logger.debug(f"Code after removing imports:\n{sympy_code}")
         
         # Create globals dict with all SymPy symbols
+        import io
         import sys
-        globals_dict = {}
+        from contextlib import redirect_stdout
         
-        # Add built-ins
-        globals_dict['__builtins__'] = __builtins__
+        # Import SymPy first at the function level
+        import sympy
         
-        # Add sympy itself by loading it from sys.modules
-        globals_dict.update({name: value for name, value in globals().items() if name in dir(sys.modules['sympy'])})
+        # Create output buffer
+        output_buffer = io.StringIO()
+        
+        # Create globals dict with SymPy module
+        globals_dict = {
+            '__builtins__': __builtins__,
+            'symbols': sympy.symbols,
+            'diff': sympy.diff,
+            'solve': sympy.solve,
+            'integrate': sympy.integrate,
+            'oo': sympy.oo,
+            'Sum': sympy.Sum
+        }
+        
+        logger.debug(f"Globals dictionary contains: {list(globals_dict.keys())}")
         
         # Create locals dict to store variables
         local_vars = {}
         
-        # Execute the code
-        exec(sympy_code, globals_dict, local_vars)
+        # Execute the code and capture its output
+        try:
+            with redirect_stdout(output_buffer):
+                exec(sympy_code, globals_dict, local_vars)
+            stdout_output = output_buffer.getvalue()
+            logger.debug(f"Code execution stdout: {stdout_output}")
+        except Exception as e:
+            logger.error(f"Error during code execution: {str(e)}")
+            return f"Error executing code: {str(e)}"
+        
+        # Format the final output
+        output_lines = ["SymPy Calculation Results:", "-" * 25]
+        
+        # Add stdout if there was any
+        if stdout_output.strip():
+            output_lines.append(f"Output:")
+            output_lines.append(stdout_output.strip())
         
-        # Format output
-        output = ["SymPy Calculation Results:", "-" * 25, ""]
+        # Add variables
+        output_lines.append("\nVariables:")
+        for name, value in local_vars.items():
+            if not name.startswith('__') and not hasattr(value, '__module__'):
+                output_lines.append(f"{name} = {value}")
         
-        # Add all created variables and their values
-        for var_name, value in local_vars.items():
-            if (not var_name.startswith('__') and 
-                not hasattr(value, '__module__')):
-                output.append(f"{var_name} = {value}")
+        # Join all lines
+        final_output = "\n".join(output_lines)
+        logger.debug(f"Final formatted output:\n{final_output}")
         
-        # Join all output lines with newlines
-        final_output = "\n".join(output)
-        logger.debug(f"Final output:\n{final_output}")
         return final_output
         
     except Exception as e: