ds-agent-cli 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (67) hide show
  1. package/bin/ds-agent.js +451 -0
  2. package/ds_agent/__init__.py +8 -0
  3. package/package.json +28 -0
  4. package/requirements.txt +126 -0
  5. package/setup.py +35 -0
  6. package/src/__init__.py +7 -0
  7. package/src/_compress_tool_result.py +118 -0
  8. package/src/api/__init__.py +4 -0
  9. package/src/api/app.py +1626 -0
  10. package/src/cache/__init__.py +5 -0
  11. package/src/cache/cache_manager.py +561 -0
  12. package/src/cli.py +2886 -0
  13. package/src/dynamic_prompts.py +281 -0
  14. package/src/orchestrator.py +4799 -0
  15. package/src/progress_manager.py +139 -0
  16. package/src/reasoning/__init__.py +332 -0
  17. package/src/reasoning/business_summary.py +431 -0
  18. package/src/reasoning/data_understanding.py +356 -0
  19. package/src/reasoning/model_explanation.py +383 -0
  20. package/src/reasoning/reasoning_trace.py +239 -0
  21. package/src/registry/__init__.py +3 -0
  22. package/src/registry/tools_registry.py +3 -0
  23. package/src/session_memory.py +448 -0
  24. package/src/session_store.py +370 -0
  25. package/src/storage/__init__.py +19 -0
  26. package/src/storage/artifact_store.py +620 -0
  27. package/src/storage/helpers.py +116 -0
  28. package/src/storage/huggingface_storage.py +694 -0
  29. package/src/storage/r2_storage.py +0 -0
  30. package/src/storage/user_files_service.py +288 -0
  31. package/src/tools/__init__.py +335 -0
  32. package/src/tools/advanced_analysis.py +823 -0
  33. package/src/tools/advanced_feature_engineering.py +708 -0
  34. package/src/tools/advanced_insights.py +578 -0
  35. package/src/tools/advanced_preprocessing.py +549 -0
  36. package/src/tools/advanced_training.py +906 -0
  37. package/src/tools/agent_tool_mapping.py +326 -0
  38. package/src/tools/auto_pipeline.py +420 -0
  39. package/src/tools/autogluon_training.py +1480 -0
  40. package/src/tools/business_intelligence.py +860 -0
  41. package/src/tools/cloud_data_sources.py +581 -0
  42. package/src/tools/code_interpreter.py +390 -0
  43. package/src/tools/computer_vision.py +614 -0
  44. package/src/tools/data_cleaning.py +614 -0
  45. package/src/tools/data_profiling.py +593 -0
  46. package/src/tools/data_type_conversion.py +268 -0
  47. package/src/tools/data_wrangling.py +433 -0
  48. package/src/tools/eda_reports.py +284 -0
  49. package/src/tools/enhanced_feature_engineering.py +241 -0
  50. package/src/tools/feature_engineering.py +302 -0
  51. package/src/tools/matplotlib_visualizations.py +1327 -0
  52. package/src/tools/model_training.py +520 -0
  53. package/src/tools/nlp_text_analytics.py +761 -0
  54. package/src/tools/plotly_visualizations.py +497 -0
  55. package/src/tools/production_mlops.py +852 -0
  56. package/src/tools/time_series.py +507 -0
  57. package/src/tools/tools_registry.py +2133 -0
  58. package/src/tools/visualization_engine.py +559 -0
  59. package/src/utils/__init__.py +42 -0
  60. package/src/utils/error_recovery.py +313 -0
  61. package/src/utils/parallel_executor.py +402 -0
  62. package/src/utils/polars_helpers.py +248 -0
  63. package/src/utils/schema_extraction.py +132 -0
  64. package/src/utils/semantic_layer.py +392 -0
  65. package/src/utils/token_budget.py +411 -0
  66. package/src/utils/validation.py +377 -0
  67. package/src/workflow_state.py +154 -0
@@ -0,0 +1,708 @@
1
+ """
2
+ Advanced Feature Engineering Tools
3
+ Tools for creating interaction features, aggregations, text features, and auto feature engineering.
4
+ """
5
+
6
+ import polars as pl
7
+ import numpy as np
8
+ from typing import Dict, Any, List, Optional, Tuple
9
+ from pathlib import Path
10
+ import sys
11
+ import os
12
+ import json
13
+ import warnings
14
+ from itertools import combinations
15
+
16
+ warnings.filterwarnings('ignore')
17
+
18
+ # Add parent directory to path for imports
19
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
20
+
21
+ from sklearn.preprocessing import PolynomialFeatures
22
+ from sklearn.decomposition import PCA, TruncatedSVD
23
+ from sklearn.feature_selection import mutual_info_classif, mutual_info_regression, SelectKBest, f_classif, f_regression
24
+ from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
25
+ from textblob import TextBlob
26
+ import re
27
+
28
+ from ds_agent.utils.polars_helpers import (
29
+ load_dataframe, save_dataframe, get_numeric_columns,
30
+ get_categorical_columns, get_datetime_columns
31
+ )
32
+ from ds_agent.utils.validation import (
33
+ validate_file_exists, validate_file_format, validate_dataframe,
34
+ validate_column_exists
35
+ )
36
+
37
+
38
+ def create_interaction_features(
39
+ file_path: str,
40
+ method: str = "polynomial",
41
+ degree: int = 2,
42
+ n_components: Optional[int] = None,
43
+ columns: Optional[List[str]] = None,
44
+ max_features: int = 50,
45
+ output_path: Optional[str] = None
46
+ ) -> Dict[str, Any]:
47
+ """
48
+ Create interaction features using polynomial features, PCA, or feature crossing.
49
+
50
+ Args:
51
+ file_path: Path to dataset
52
+ method: Feature interaction method:
53
+ - 'polynomial': Polynomial features (degree 2 or 3)
54
+ - 'pca': Principal Component Analysis
55
+ - 'cross': Manual feature crossing (multiply pairs)
56
+ - 'mutual_info': Select best features by mutual information
57
+ degree: Polynomial degree (for polynomial method)
58
+ n_components: Number of components (for PCA, None = auto)
59
+ columns: Columns to use (None = all numeric)
60
+ max_features: Maximum number of new features to create
61
+ output_path: Path to save dataset with new features
62
+
63
+ Returns:
64
+ Dictionary with feature engineering results
65
+ """
66
+ # Validation
67
+ validate_file_exists(file_path)
68
+ validate_file_format(file_path)
69
+
70
+ # Load data
71
+ df = load_dataframe(file_path)
72
+ validate_dataframe(df)
73
+
74
+ # Get numeric columns if not specified
75
+ if columns is None:
76
+ columns = get_numeric_columns(df)
77
+ print(f"🔢 Auto-detected {len(columns)} numeric columns")
78
+ else:
79
+ for col in columns:
80
+ validate_column_exists(df, col)
81
+
82
+ if not columns:
83
+ return {
84
+ 'status': 'skipped',
85
+ 'message': 'No numeric columns found for interaction features'
86
+ }
87
+
88
+ # Limit columns if too many
89
+ if len(columns) > 20:
90
+ print(f"⚠️ Too many columns ({len(columns)}). Using top 20 by variance.")
91
+ variances = df[columns].select([
92
+ (pl.col(col).var().alias(col)) for col in columns
93
+ ]).to_dicts()[0]
94
+ columns = sorted(variances.keys(), key=lambda x: variances[x], reverse=True)[:20]
95
+
96
+ # Handle NaN values before transformation
97
+ print(f"🧬 Checking for NaN values...")
98
+ df_subset = df.select(columns)
99
+ has_nulls = df_subset.null_count().sum_horizontal()[0] > 0
100
+
101
+ if has_nulls:
102
+ print(f"⚠️ Found NaN values, imputing with column medians...")
103
+ # Impute NaN with median for each column
104
+ impute_exprs = []
105
+ for col in columns:
106
+ median_val = df_subset[col].median()
107
+ if median_val is None: # All NaN
108
+ median_val = 0.0
109
+ impute_exprs.append(pl.col(col).fill_null(median_val).alias(col))
110
+ df_subset = df_subset.select(impute_exprs)
111
+
112
+ X = df_subset.to_numpy()
113
+ original_features = len(columns)
114
+
115
+ # Create interaction features based on method
116
+ if method == "polynomial":
117
+ print(f"🔄 Creating polynomial features (degree={degree})...")
118
+ poly = PolynomialFeatures(degree=degree, include_bias=False, interaction_only=False)
119
+ X_poly = poly.fit_transform(X)
120
+
121
+ # Get feature names
122
+ feature_names = poly.get_feature_names_out(columns)
123
+
124
+ # Limit features
125
+ if X_poly.shape[1] > max_features + original_features:
126
+ # Keep original + top max_features new ones by variance
127
+ variances = np.var(X_poly[:, original_features:], axis=0)
128
+ top_indices = np.argsort(variances)[::-1][:max_features]
129
+ X_new = np.hstack([X, X_poly[:, original_features:][:, top_indices]])
130
+ new_feature_names = [feature_names[i + original_features] for i in top_indices]
131
+ else:
132
+ X_new = X_poly
133
+ new_feature_names = feature_names[original_features:].tolist()
134
+
135
+ # Create new dataframe
136
+ df_new = df.clone()
137
+ for i, name in enumerate(new_feature_names):
138
+ clean_name = name.replace(' ', '_').replace('^', '_pow_')
139
+ df_new = df_new.with_columns(
140
+ pl.Series(f"poly_{clean_name}", X_new[:, original_features + i])
141
+ )
142
+
143
+ created_features = new_feature_names
144
+
145
+ elif method == "pca":
146
+ print(f"🔄 Creating PCA features...")
147
+ if n_components is None:
148
+ n_components = min(len(columns), max_features)
149
+
150
+ pca = PCA(n_components=n_components)
151
+ X_pca = pca.fit_transform(X)
152
+
153
+ # Create new dataframe
154
+ df_new = df.clone()
155
+ for i in range(n_components):
156
+ df_new = df_new.with_columns(
157
+ pl.Series(f"pca_{i+1}", X_pca[:, i])
158
+ )
159
+
160
+ created_features = [f"pca_{i+1}" for i in range(n_components)]
161
+
162
+ explained_variance = pca.explained_variance_ratio_
163
+ cumulative_variance = np.cumsum(explained_variance)
164
+
165
+ elif method == "cross":
166
+ print(f"🔄 Creating feature crosses...")
167
+ # Create pairwise interactions
168
+ pairs = list(combinations(columns, 2))
169
+
170
+ # Limit number of pairs
171
+ if len(pairs) > max_features:
172
+ pairs = pairs[:max_features]
173
+
174
+ df_new = df.clone()
175
+ created_features = []
176
+
177
+ for col1, col2 in pairs:
178
+ new_name = f"{col1}_x_{col2}"
179
+ df_new = df_new.with_columns(
180
+ (pl.col(col1) * pl.col(col2)).alias(new_name)
181
+ )
182
+ created_features.append(new_name)
183
+
184
+ elif method == "mutual_info":
185
+ print(f"🔄 Selecting features by mutual information...")
186
+ # This requires a target column - for now, create interaction features
187
+ # and let the user select based on their target
188
+ return {
189
+ 'status': 'error',
190
+ 'message': 'mutual_info method requires a target column. Use polynomial or cross instead, then use feature selection.'
191
+ }
192
+
193
+ else:
194
+ raise ValueError(f"Unsupported method: {method}")
195
+
196
+ # Save if output path provided
197
+ if output_path:
198
+ save_dataframe(df_new, output_path)
199
+ print(f"💾 Dataset with interaction features saved to: {output_path}")
200
+
201
+ result = {
202
+ 'status': 'success',
203
+ 'method': method,
204
+ 'original_features': original_features,
205
+ 'new_features_created': len(created_features),
206
+ 'total_features': len(df_new.columns),
207
+ 'feature_names': created_features[:20], # Show first 20
208
+ 'output_path': output_path
209
+ }
210
+
211
+ if method == "pca":
212
+ result['explained_variance_ratio'] = explained_variance.tolist()
213
+ result['cumulative_variance'] = cumulative_variance.tolist()
214
+ result['variance_explained_by_top_5'] = float(cumulative_variance[min(4, len(cumulative_variance)-1)])
215
+
216
+ return result
217
+
218
+
219
+ def create_aggregation_features(
220
+ file_path: str,
221
+ group_col: str,
222
+ agg_columns: Optional[List[str]] = None,
223
+ agg_functions: Optional[List[str]] = None,
224
+ rolling_window: Optional[int] = None,
225
+ time_col: Optional[str] = None,
226
+ lag_periods: Optional[List[int]] = None,
227
+ output_path: Optional[str] = None
228
+ ) -> Dict[str, Any]:
229
+ """
230
+ Create aggregation features including group-by aggregations, rolling windows, and lags.
231
+
232
+ Args:
233
+ file_path: Path to dataset
234
+ group_col: Column to group by (e.g., 'customer_id', 'category')
235
+ agg_columns: Columns to aggregate (None = all numeric)
236
+ agg_functions: Aggregation functions ('mean', 'sum', 'std', 'min', 'max', 'count')
237
+ rolling_window: Window size for rolling aggregations (requires sorted data)
238
+ time_col: Time column for sorting (required for rolling/lag features)
239
+ lag_periods: Lag periods to create (e.g., [1, 7, 30] for 1-day, 7-day, 30-day lags)
240
+ output_path: Path to save dataset with new features
241
+
242
+ Returns:
243
+ Dictionary with aggregation results
244
+ """
245
+ # Validation
246
+ validate_file_exists(file_path)
247
+ validate_file_format(file_path)
248
+
249
+ # Load data
250
+ df = load_dataframe(file_path)
251
+ validate_dataframe(df)
252
+ validate_column_exists(df, group_col)
253
+
254
+ if time_col:
255
+ validate_column_exists(df, time_col)
256
+ df = df.sort(time_col)
257
+
258
+ # Get numeric columns if not specified
259
+ if agg_columns is None:
260
+ agg_columns = [col for col in get_numeric_columns(df) if col != group_col]
261
+ print(f"🔢 Auto-detected {len(agg_columns)} numeric columns for aggregation")
262
+ else:
263
+ for col in agg_columns:
264
+ validate_column_exists(df, col)
265
+
266
+ if not agg_columns:
267
+ return {
268
+ 'status': 'skipped',
269
+ 'message': 'No numeric columns found for aggregation'
270
+ }
271
+
272
+ # Default aggregation functions
273
+ if agg_functions is None:
274
+ agg_functions = ['mean', 'sum', 'std', 'min', 'max', 'count']
275
+
276
+ df_new = df.clone()
277
+ created_features = []
278
+
279
+ # Group-by aggregations
280
+ print(f"📊 Creating group-by aggregations for {group_col}...")
281
+
282
+ for agg_col in agg_columns:
283
+ for agg_func in agg_functions:
284
+ try:
285
+ if agg_func == 'mean':
286
+ agg_df = df.group_by(group_col).agg(
287
+ pl.col(agg_col).mean().alias(f"{agg_col}_{group_col}_mean")
288
+ )
289
+ elif agg_func == 'sum':
290
+ agg_df = df.group_by(group_col).agg(
291
+ pl.col(agg_col).sum().alias(f"{agg_col}_{group_col}_sum")
292
+ )
293
+ elif agg_func == 'std':
294
+ agg_df = df.group_by(group_col).agg(
295
+ pl.col(agg_col).std().alias(f"{agg_col}_{group_col}_std")
296
+ )
297
+ elif agg_func == 'min':
298
+ agg_df = df.group_by(group_col).agg(
299
+ pl.col(agg_col).min().alias(f"{agg_col}_{group_col}_min")
300
+ )
301
+ elif agg_func == 'max':
302
+ agg_df = df.group_by(group_col).agg(
303
+ pl.col(agg_col).max().alias(f"{agg_col}_{group_col}_max")
304
+ )
305
+ elif agg_func == 'count':
306
+ agg_df = df.group_by(group_col).agg(
307
+ pl.col(agg_col).count().alias(f"{agg_col}_{group_col}_count")
308
+ )
309
+ else:
310
+ continue
311
+
312
+ # Join back to original dataframe
313
+ df_new = df_new.join(agg_df, on=group_col, how='left')
314
+ created_features.append(f"{agg_col}_{group_col}_{agg_func}")
315
+ except Exception as e:
316
+ print(f"⚠️ Skipping {agg_col}_{agg_func}: {str(e)}")
317
+
318
+ # Rolling window features
319
+ if rolling_window and time_col:
320
+ print(f"📈 Creating rolling window features (window={rolling_window})...")
321
+
322
+ for agg_col in agg_columns[:5]: # Limit to first 5 columns to avoid explosion
323
+ try:
324
+ # Rolling mean
325
+ df_new = df_new.with_columns(
326
+ pl.col(agg_col).rolling_mean(window_size=rolling_window)
327
+ .over(group_col)
328
+ .alias(f"{agg_col}_rolling_{rolling_window}_mean")
329
+ )
330
+ created_features.append(f"{agg_col}_rolling_{rolling_window}_mean")
331
+
332
+ # Rolling std
333
+ df_new = df_new.with_columns(
334
+ pl.col(agg_col).rolling_std(window_size=rolling_window)
335
+ .over(group_col)
336
+ .alias(f"{agg_col}_rolling_{rolling_window}_std")
337
+ )
338
+ created_features.append(f"{agg_col}_rolling_{rolling_window}_std")
339
+ except Exception as e:
340
+ print(f"⚠️ Skipping rolling for {agg_col}: {str(e)}")
341
+
342
+ # Lag features
343
+ if lag_periods and time_col:
344
+ print(f"⏰ Creating lag features (periods={lag_periods})...")
345
+
346
+ for agg_col in agg_columns[:5]: # Limit to avoid explosion
347
+ for lag in lag_periods:
348
+ try:
349
+ df_new = df_new.with_columns(
350
+ pl.col(agg_col).shift(lag)
351
+ .over(group_col)
352
+ .alias(f"{agg_col}_lag_{lag}")
353
+ )
354
+ created_features.append(f"{agg_col}_lag_{lag}")
355
+ except Exception as e:
356
+ print(f"⚠️ Skipping lag {lag} for {agg_col}: {str(e)}")
357
+
358
+ # Save if output path provided
359
+ if output_path:
360
+ save_dataframe(df_new, output_path)
361
+ print(f"💾 Dataset with aggregation features saved to: {output_path}")
362
+
363
+ return {
364
+ 'status': 'success',
365
+ 'group_column': group_col,
366
+ 'aggregated_columns': agg_columns,
367
+ 'aggregation_functions': agg_functions,
368
+ 'new_features_created': len(created_features),
369
+ 'total_features': len(df_new.columns),
370
+ 'feature_names': created_features[:30], # Show first 30
371
+ 'rolling_window': rolling_window,
372
+ 'lag_periods': lag_periods,
373
+ 'output_path': output_path
374
+ }
375
+
376
+
377
+ def engineer_text_features(
378
+ file_path: str,
379
+ text_column: str,
380
+ methods: Optional[List[str]] = None,
381
+ max_features: int = 100,
382
+ ngram_range: Tuple[int, int] = (1, 2),
383
+ output_path: Optional[str] = None
384
+ ) -> Dict[str, Any]:
385
+ """
386
+ Extract features from text columns using TF-IDF, n-grams, and text statistics.
387
+
388
+ Args:
389
+ file_path: Path to dataset
390
+ text_column: Name of text column
391
+ methods: List of methods to apply:
392
+ - 'tfidf': TF-IDF vectorization
393
+ - 'count': Count vectorization (bag of words)
394
+ - 'sentiment': Sentiment analysis
395
+ - 'stats': Text statistics (length, word count, etc.)
396
+ - 'ngrams': N-gram features
397
+ max_features: Maximum number of TF-IDF/count features
398
+ ngram_range: N-gram range (e.g., (1, 2) for unigrams and bigrams)
399
+ output_path: Path to save dataset with new features
400
+
401
+ Returns:
402
+ Dictionary with text feature engineering results
403
+ """
404
+ # Validation
405
+ validate_file_exists(file_path)
406
+ validate_file_format(file_path)
407
+
408
+ # Load data
409
+ df = load_dataframe(file_path)
410
+ validate_dataframe(df)
411
+ validate_column_exists(df, text_column)
412
+
413
+ # Default methods
414
+ if methods is None:
415
+ methods = ['stats', 'sentiment', 'tfidf']
416
+
417
+ df_new = df.clone()
418
+ created_features = []
419
+
420
+ # Get text data
421
+ texts = df[text_column].fill_null("").to_list()
422
+
423
+ # Text statistics
424
+ if 'stats' in methods:
425
+ print("📝 Extracting text statistics...")
426
+
427
+ char_counts = [len(str(text)) for text in texts]
428
+ word_counts = [len(str(text).split()) for text in texts]
429
+ avg_word_lengths = [np.mean([len(word) for word in str(text).split()]) if text else 0 for text in texts]
430
+ special_char_counts = [len(re.findall(r'[^a-zA-Z0-9\s]', str(text))) for text in texts]
431
+ digit_counts = [len(re.findall(r'\d', str(text))) for text in texts]
432
+ uppercase_counts = [len(re.findall(r'[A-Z]', str(text))) for text in texts]
433
+
434
+ df_new = df_new.with_columns([
435
+ pl.Series(f"{text_column}_char_count", char_counts),
436
+ pl.Series(f"{text_column}_word_count", word_counts),
437
+ pl.Series(f"{text_column}_avg_word_length", avg_word_lengths),
438
+ pl.Series(f"{text_column}_special_char_count", special_char_counts),
439
+ pl.Series(f"{text_column}_digit_count", digit_counts),
440
+ pl.Series(f"{text_column}_uppercase_count", uppercase_counts)
441
+ ])
442
+
443
+ created_features.extend([
444
+ f"{text_column}_char_count",
445
+ f"{text_column}_word_count",
446
+ f"{text_column}_avg_word_length",
447
+ f"{text_column}_special_char_count",
448
+ f"{text_column}_digit_count",
449
+ f"{text_column}_uppercase_count"
450
+ ])
451
+
452
+ # Sentiment analysis
453
+ if 'sentiment' in methods:
454
+ print("💭 Performing sentiment analysis...")
455
+
456
+ sentiments = []
457
+ subjectivities = []
458
+
459
+ for text in texts:
460
+ try:
461
+ blob = TextBlob(str(text))
462
+ sentiments.append(blob.sentiment.polarity)
463
+ subjectivities.append(blob.sentiment.subjectivity)
464
+ except:
465
+ sentiments.append(0.0)
466
+ subjectivities.append(0.0)
467
+
468
+ df_new = df_new.with_columns([
469
+ pl.Series(f"{text_column}_sentiment", sentiments),
470
+ pl.Series(f"{text_column}_subjectivity", subjectivities)
471
+ ])
472
+
473
+ created_features.extend([
474
+ f"{text_column}_sentiment",
475
+ f"{text_column}_subjectivity"
476
+ ])
477
+
478
+ # TF-IDF features
479
+ if 'tfidf' in methods:
480
+ print(f"🔤 Creating TF-IDF features (max_features={max_features})...")
481
+
482
+ tfidf = TfidfVectorizer(
483
+ max_features=max_features,
484
+ ngram_range=ngram_range,
485
+ stop_words='english',
486
+ min_df=2
487
+ )
488
+
489
+ try:
490
+ tfidf_matrix = tfidf.fit_transform([str(text) for text in texts])
491
+ feature_names = tfidf.get_feature_names_out()
492
+
493
+ # Add TF-IDF features to dataframe
494
+ for i, feature_name in enumerate(feature_names):
495
+ clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', feature_name)[:30]
496
+ df_new = df_new.with_columns(
497
+ pl.Series(f"tfidf_{clean_name}", tfidf_matrix[:, i].toarray().flatten())
498
+ )
499
+ created_features.append(f"tfidf_{clean_name}")
500
+ except Exception as e:
501
+ print(f"⚠️ TF-IDF failed: {str(e)}")
502
+
503
+ # Count vectorization
504
+ if 'count' in methods:
505
+ print(f"🔢 Creating count features (max_features={max_features})...")
506
+
507
+ count_vec = CountVectorizer(
508
+ max_features=max_features,
509
+ ngram_range=ngram_range,
510
+ stop_words='english',
511
+ min_df=2
512
+ )
513
+
514
+ try:
515
+ count_matrix = count_vec.fit_transform([str(text) for text in texts])
516
+ feature_names = count_vec.get_feature_names_out()
517
+
518
+ # Add count features to dataframe
519
+ for i, feature_name in enumerate(feature_names[:50]): # Limit to 50
520
+ clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', feature_name)[:30]
521
+ df_new = df_new.with_columns(
522
+ pl.Series(f"count_{clean_name}", count_matrix[:, i].toarray().flatten())
523
+ )
524
+ created_features.append(f"count_{clean_name}")
525
+ except Exception as e:
526
+ print(f"⚠️ Count vectorization failed: {str(e)}")
527
+
528
+ # Save if output path provided
529
+ if output_path:
530
+ save_dataframe(df_new, output_path)
531
+ print(f"💾 Dataset with text features saved to: {output_path}")
532
+
533
+ return {
534
+ 'status': 'success',
535
+ 'text_column': text_column,
536
+ 'methods_applied': methods,
537
+ 'new_features_created': len(created_features),
538
+ 'total_features': len(df_new.columns),
539
+ 'feature_names': created_features[:30], # Show first 30
540
+ 'output_path': output_path
541
+ }
542
+
543
+
544
+ def auto_feature_engineering(
545
+ file_path: str,
546
+ target_col: str,
547
+ groq_api_key: Optional[str] = None,
548
+ max_suggestions: int = 10,
549
+ implement_top_k: int = 5,
550
+ output_path: Optional[str] = None
551
+ ) -> Dict[str, Any]:
552
+ """
553
+ Use LLM (Groq or Gemini) to automatically generate and implement feature engineering ideas.
554
+
555
+ Args:
556
+ file_path: Path to dataset
557
+ target_col: Target column name
558
+ groq_api_key: Groq API key (optional - will try to use environment variable or Gemini)
559
+ max_suggestions: Maximum number of feature suggestions to generate
560
+ implement_top_k: Number of top suggestions to implement
561
+ output_path: Path to save dataset with new features
562
+
563
+ Returns:
564
+ Dictionary with feature suggestions and implementation results
565
+ """
566
+ import os
567
+
568
+ # Validation
569
+ validate_file_exists(file_path)
570
+ validate_file_format(file_path)
571
+
572
+ # Load data
573
+ df = load_dataframe(file_path)
574
+ validate_dataframe(df)
575
+ validate_column_exists(df, target_col)
576
+
577
+ # Get dataset summary
578
+ numeric_cols = get_numeric_columns(df)
579
+ categorical_cols = get_categorical_columns(df)
580
+
581
+ # Sample data for analysis
582
+ sample_df = df.head(5)
583
+
584
+ # Create prompt for LLM
585
+ prompt = f"""You are a data science expert. Analyze this dataset and suggest {max_suggestions} creative feature engineering ideas.
586
+
587
+ Dataset Overview:
588
+ - Target column: {target_col}
589
+ - Numeric columns ({len(numeric_cols)}): {', '.join(numeric_cols[:10])}
590
+ - Categorical columns ({len(categorical_cols)}): {', '.join(categorical_cols[:5])}
591
+ - Rows: {len(df)}
592
+
593
+ Sample data (first 5 rows):
594
+ {sample_df.head(5)}
595
+
596
+ Suggest {max_suggestions} feature engineering ideas that could improve model performance. For each idea:
597
+ 1. Describe the feature clearly
598
+ 2. Provide Python code using Polars to create it
599
+ 3. Explain why it might be valuable
600
+
601
+ Format your response as JSON:
602
+ {{
603
+ "suggestions": [
604
+ {{
605
+ "name": "feature_name",
606
+ "description": "what it does",
607
+ "code": "pl.col('a') * pl.col('b')",
608
+ "reasoning": "why it helps"
609
+ }}
610
+ ]
611
+ }}
612
+ """
613
+
614
+ print("🤖 Asking LLM for feature engineering suggestions...")
615
+
616
+ # Try multiple LLM providers in order of preference
617
+ llm_response = None
618
+
619
+ # Try Groq first if API key provided
620
+ if groq_api_key or os.getenv("GROQ_API_KEY"):
621
+ try:
622
+ from groq import Groq
623
+ api_key = groq_api_key or os.getenv("GROQ_API_KEY")
624
+ client = Groq(api_key=api_key)
625
+ response = client.chat.completions.create(
626
+ model="llama-3.3-70b-versatile",
627
+ messages=[{"role": "user", "content": prompt}],
628
+ temperature=0.7,
629
+ max_tokens=2000
630
+ )
631
+ llm_response = response.choices[0].message.content
632
+ print(" ✓ Using Groq LLM")
633
+ except Exception as e:
634
+ print(f" ⚠️ Groq failed: {str(e)}, trying Gemini...")
635
+
636
+ # Try Gemini if Groq failed or not available
637
+ if not llm_response and os.getenv("GEMINI_API_KEY"):
638
+ try:
639
+ import google.generativeai as genai
640
+ genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
641
+ model = genai.GenerativeModel('gemini-2.0-flash-exp')
642
+ response = model.generate_content(prompt)
643
+ llm_response = response.text
644
+ print(" ✓ Using Gemini LLM")
645
+ except Exception as e:
646
+ print(f" ⚠️ Gemini failed: {str(e)}")
647
+
648
+ if not llm_response:
649
+ return {
650
+ "status": "error",
651
+ "message": "No LLM API key available. Set GROQ_API_KEY or GEMINI_API_KEY environment variable."
652
+ }
653
+
654
+ try:
655
+ # Parse JSON response
656
+ import json
657
+ # Extract JSON from response (might be wrapped in markdown code blocks)
658
+ if "```json" in llm_response:
659
+ llm_response = llm_response.split("```json")[1].split("```")[0].strip()
660
+ elif "```" in llm_response:
661
+ llm_response = llm_response.split("```")[1].split("```")[0].strip()
662
+
663
+ suggestions = json.loads(llm_response)
664
+
665
+ # Implement top K suggestions
666
+ df_new = df.clone()
667
+ implemented = []
668
+
669
+ for i, suggestion in enumerate(suggestions['suggestions'][:implement_top_k]):
670
+ try:
671
+ # Execute feature creation code
672
+ feature_name = suggestion['name']
673
+ code = suggestion['code']
674
+
675
+ # Create new column using eval (be careful in production!)
676
+ df_new = df_new.with_columns(
677
+ eval(code).alias(feature_name)
678
+ )
679
+
680
+ implemented.append({
681
+ 'name': feature_name,
682
+ 'description': suggestion['description'],
683
+ 'reasoning': suggestion['reasoning']
684
+ })
685
+
686
+ print(f"✅ Implemented: {feature_name}")
687
+ except Exception as e:
688
+ print(f"⚠️ Failed to implement {suggestion.get('name', 'unknown')}: {str(e)}")
689
+
690
+ # Save if output path provided
691
+ if output_path:
692
+ save_dataframe(df_new, output_path)
693
+ print(f"💾 Dataset with auto-generated features saved to: {output_path}")
694
+
695
+ return {
696
+ 'status': 'success',
697
+ 'total_suggestions': len(suggestions['suggestions']),
698
+ 'suggestions': suggestions['suggestions'],
699
+ 'implemented': implemented,
700
+ 'new_features_created': len(implemented),
701
+ 'output_path': output_path
702
+ }
703
+
704
+ except Exception as e:
705
+ return {
706
+ 'status': 'error',
707
+ 'message': f"Auto feature engineering failed: {str(e)}"
708
+ }