v115

app.py

@@ -187,7 +187,7 @@ def run_crewai_process(user_query, model, temperature):
                 with open(temp_script_path, "w") as f:
                     f.write(generated_code)
 
-                # Add auto_adjust=True to yf.download() calls in the generated script
+                # Read the generated script
                 with open(temp_script_path, 'r') as f:
                     script_content = f.read()
                 
@@ -198,6 +198,63 @@ def run_crewai_process(user_query, model, temperature):
                     script_content
                 )
                 
+                # Add helper functions at the beginning of the script
+                helpers = """
+# Helper functions for data processing
+def safe_get_column(df, column):
+    # Try exact match first
+    if column in df.columns:
+        return df[column]
+    # Try case-insensitive match
+    col_lower = column.lower()
+    for col in df.columns:
+        if col.lower() == col_lower:
+            return df[col]
+    # If not found, try common variations
+    variations = {
+        'close': ['Close', 'Adj Close', 'close', 'adj close', 'CLOSE'],
+        'adj close': ['Adj Close', 'adj close', 'ADJ CLOSE', 'Close', 'close', 'CLOSE']
+    }
+    for var in variations.get(column.lower(), []):
+        if var in df.columns:
+            return df[var]
+    # If still not found, raise a helpful error
+    raise KeyError(f"Column '{column}' not found in DataFrame. Available columns: {list(df.columns)}")
+
+def plot_and_save(plt, filename='plot.png'):
+    import os
+    # Ensure the directory exists
+    os.makedirs(os.path.dirname(os.path.abspath(filename)) or '.', exist_ok=True)
+    plt.savefig(filename)
+    plt.close()
+    print(f"Plot saved as {os.path.abspath(filename)}")
+
+# Monkey patch DataFrame to add safe column access
+import pandas as pd
+pd.DataFrame.safe_get = safe_get_column
+"""
+                # Insert the helper functions after imports
+                if 'import ' in script_content:
+                    # Insert after the last import
+                    last_import = script_content.rfind('import ')
+                    insert_pos = script_content.find('\n', last_import) + 1
+                    script_content = script_content[:insert_pos] + '\n' + helpers + script_content[insert_pos:]
+                else:
+                    # Insert at the beginning if no imports found
+                    script_content = helpers + '\n' + script_content
+                
+                # Replace common column access patterns with our safe version
+                script_content = script_content.replace("['Adj Close']", ".safe_get('close')")
+                script_content = script_content.replace("['Close']", ".safe_get('close')")
+                script_content = script_content.replace("['close']", ".safe_get('close')")
+                
+                # Replace plt.savefig() calls with our helper
+                script_content = re.sub(
+                    r'plt\.savefig\(([^)]+)\)',
+                    r'plot_and_save(plt, \1)',
+                    script_content
+                )
+                
                 # Write the updated script back
                 with open(temp_script_path, 'w') as f:
                     f.write(script_content)