Update app.py

app.py

@@ -345,55 +345,35 @@ def extract_and_run_sympy_code(response_text):
         sympy_code = '\n'.join(filtered_lines)
         logger.debug(f"Code after removing imports:\n{sympy_code}")
         
-        # Create globals dict
-        import sys
-        module = sys.modules[__name__]
+        # Create globals dict with all SymPy symbols
         globals_dict = {
             '__builtins__': __builtins__,
+            'sympy': sympy,  # Add the sympy module itself
         }
         
-        # Add all imported SymPy names to globals
-        for name in dir(module):
-            if not name.startswith('_'):  # Skip private names
-                globals_dict[name] = getattr(module, name)
+        # Add all sympy attributes to globals
+        for name in dir(sympy):
+            globals_dict[name] = getattr(sympy, name)
         
-        logger.debug(f"Global namespace keys: {list(globals_dict.keys())}")
-        
-        # Create locals dict and output buffer
+        # Create locals dict to store variables
         local_vars = {}
-        import io
-        from contextlib import redirect_stdout
-        output_buffer = io.StringIO()
         
-        try:
-            # Execute the code and capture all output in one redirect_stdout block
-            with redirect_stdout(output_buffer):
-                print("SymPy Calculation Results:")
-                print("-" * 25)
-                
-                # Execute the code
-                exec(sympy_code, globals_dict, local_vars)
-                
-                # Print all variables that were created
-                print("\nFinal Values:")
-                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())}")
-            import traceback
-            logger.error(f"Full traceback:\n{traceback.format_exc()}")
-            return f"Error executing SymPy code: {str(e)}"
-        except Exception as e:
-            logger.error(f"Other error during execution: {str(e)}")
-            return f"Error executing SymPy code: {str(e)}"
+        # Execute the code
+        exec(sympy_code, globals_dict, local_vars)
+        
+        # Format output
+        output = ["SymPy Calculation Results:", "-" * 25, ""]
+        
+        # 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}")
         
-        result = output_buffer.getvalue().strip()
-        logger.debug(f"Final output:\n{result}")
-        return result
+        # 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:
         logger.error(f"Top-level error in extract_and_run_sympy_code: {str(e)}")