Update app.py

app.py

@@ -317,7 +317,6 @@ def extract_and_run_sympy_code(response_text):
     """
     Extract SymPy code from the response and execute it.
     Returns only the essential calculation results as a string.
-    Filters out debugging info and module imports.
     """
     try:
         # Find the SymPy code block
@@ -341,7 +340,7 @@ def extract_and_run_sympy_code(response_text):
         # Create a string buffer to capture print output
         output_buffer = io.StringIO()
         
-        # Create a clean globals dict with only essential SymPy components
+        # Create globals dict with minimal SymPy elements
         sympy_globals = {
             "Symbol": sympy.Symbol,
             "symbols": sympy.symbols,
@@ -350,8 +349,10 @@ def extract_and_run_sympy_code(response_text):
             "diff": sympy.diff,
             "limit": sympy.limit,
             "sqrt": sympy.sqrt,
+            "exp": sympy.exp,
             "simplify": sympy.simplify,
-            "expand": sympy.expand
+            "expand": sympy.expand,
+            "print": print  # Allow print statements
         }
         
         # Create locals dict to capture new variables
@@ -364,37 +365,39 @@ def extract_and_run_sympy_code(response_text):
         # Get the printed output
         printed_output = output_buffer.getvalue().strip()
         
-        # Filter the local variables to only include calculation results
-        # Exclude built-in variables, modules, and functions
-        calculation_results = {
-            name: value for name, value in local_vars.items()
-            if (
-                not name.startswith('__') and
-                not hasattr(value, '__call__') and  # Exclude functions
-                not hasattr(value, '__module__')    # Exclude modules
-            )
-        }
-        
-        # Build the final output string
-        result_parts = []
+        # Initialize result string
+        result = []
         
-        # Add any printed output first
-        if printed_output:
-            result_parts.append(printed_output)
+        # Add numerical and symbolic calculation results
+        for name, value in local_vars.items():
+            # Skip private variables and functions
+            if name.startswith('__') or callable(value) or isinstance(value, type):
+                continue
+                
+            # Skip module references
+            if hasattr(value, '__module__'):
+                continue
+                
+            # Format the output
+            try:
+                # If it's a sympy expression or numerical result
+                if isinstance(value, (sympy.Basic, int, float)) or hasattr(value, 'free_symbols'):
+                    result.append(f"{name} = {value}")
+                # If it's a list or tuple of sympy expressions
+                elif isinstance(value, (list, tuple)) and all(isinstance(x, (sympy.Basic, int, float)) for x in value):
+                    result.append(f"{name} = {value}")
+            except:
+                continue
         
-        # Add calculation results if there are any
-        if calculation_results:
-            if printed_output:  # Add a separator if we already have printed output
-                result_parts.append("\nCalculation Results:")
-                result_parts.append("-" * 20)
-            for var_name, value in calculation_results.items():
-                result_parts.append(f"{var_name} = {value}")
+        # Add printed output if it exists and isn't just whitespace
+        if printed_output and not printed_output.isspace():
+            result.insert(0, printed_output)
         
-        # If we have no results at all, provide a message
-        if not result_parts:
+        # If we have no results, provide a message
+        if not result:
             return "No calculation results produced."
         
-        return "\n".join(result_parts)
+        return "\n".join(result)
     
     except Exception as e:
         return f"Error executing SymPy code: {str(e)}"