Fix: Better SBERT error handling + suppress invalid hand-off warnings

src/orchestrator.py

@@ -1316,7 +1316,7 @@ You receive quality reports from EDA agent and deliver clean data to modeling ag
             Dictionary with hand-off details
         """
         if target_agent not in self.specialist_agents:
-            print(f"⚠️ Invalid hand-off target: {target_agent}")
+            # Silently skip invalid hand-off targets (common during workflow transitions)
             return {"success": False, "error": "Invalid target agent"}
         
         # Update active agent
@@ -1999,6 +1999,11 @@ You receive quality reports from EDA agent and deliver clean data to modeling ag
                     models_val = arguments.pop("models")
                     print(f"   ✓ Stripped invalid parameter 'models': {models_val}")
                     print(f"   ℹ️ train_baseline_models trains all baseline models automatically")
+                # LLM often adds 'feature_columns' parameter that doesn't exist
+                if "feature_columns" in arguments:
+                    feature_cols = arguments.pop("feature_columns")
+                    print(f"   ✓ Stripped invalid parameter 'feature_columns': {feature_cols}")
+                    print(f"   ℹ️ train_baseline_models uses all numeric columns automatically")
             
             if tool_name == "generate_model_report":
                 # LLM uses 'file_path' instead of 'test_data_path'
@@ -3480,25 +3485,57 @@ You receive quality reports from EDA agent and deliver clean data to modeling ag
                             print(f"⚠️  INVALID TOOL NAME: '{tool_name}' (original: {tool_call.function.name})")
                             print(f"   Available tools: {', '.join(list(self.tool_functions.keys())[:10])}...")
                             
-                            # Try fuzzy matching to recover
-                            from difflib import get_close_matches
-                            close_matches = get_close_matches(tool_name, self.tool_functions.keys(), n=1, cutoff=0.6)
-                            if close_matches:
-                                tool_name = close_matches[0]
-                                print(f"   ✓ Recovered using fuzzy match: {tool_name}")
+                            # Explicit mappings for common LLM hallucinations
+                            tool_name_mappings = {
+                                "drop_columns": "execute_python_code",  # No drop_columns tool, use code
+                                "select_columns": "execute_python_code",  # No select_columns tool, use code
+                                "rename_columns": "execute_python_code",  # No rename_columns tool, use code
+                                "encode_categorical_variables": "encode_categorical",
+                                "train_model": "train_baseline_models",
+                                "train_models": "train_baseline_models",
+                                "baseline_models": "train_baseline_models",
+                                "tune_hyperparameters": "hyperparameter_tuning",
+                                "hyperparameter_search": "hyperparameter_tuning",
+                            }
+                            
+                            if tool_name in tool_name_mappings:
+                                mapped_tool = tool_name_mappings[tool_name]
+                                if mapped_tool == "execute_python_code":
+                                    print(f"   ✓ Tool '{tool_name}' not available - LLM should use execute_python_code instead")
+                                    # Skip and let LLM handle with code
+                                    messages.append({
+                                        "role": "tool",
+                                        "tool_call_id": tool_call_id,
+                                        "name": tool_name,
+                                        "content": json.dumps({
+                                            "error": f"Tool '{tool_name}' does not exist",
+                                            "hint": "Use execute_python_code with pandas to perform this operation. Example: df.drop(columns=['col1', 'col2'])"
+                                        })
+                                    })
+                                    continue
+                                else:
+                                    tool_name = mapped_tool
+                                    print(f"   ✓ Mapped to: {tool_name}")
                             else:
-                                print(f"   ❌ Cannot recover tool name, skipping")
-                                messages.append({
-                                    "role": "tool",
-                                    "tool_call_id": tool_call_id,
-                                    "name": "invalid_tool",
-                                    "content": json.dumps({
-                                        "error": f"Invalid tool: {tool_call.function.name}",
-                                        "message": "Tool does not exist in registry. Available tools can be found in the tools list.",
-                                        "hint": "Check spelling and use exact tool names from the tools registry."
+                                # Try fuzzy matching to recover
+                                from difflib import get_close_matches
+                                close_matches = get_close_matches(tool_name, self.tool_functions.keys(), n=1, cutoff=0.6)
+                                if close_matches:
+                                    tool_name = close_matches[0]
+                                    print(f"   ✓ Recovered using fuzzy match: {tool_name}")
+                                else:
+                                    print(f"   ❌ Cannot recover tool name, skipping")
+                                    messages.append({
+                                        "role": "tool",
+                                        "tool_call_id": tool_call_id,
+                                        "name": "invalid_tool",
+                                        "content": json.dumps({
+                                            "error": f"Invalid tool: {tool_call.function.name}",
+                                            "message": "Tool does not exist in registry. Available tools can be found in the tools list.",
+                                            "hint": "Check spelling and use exact tool names from the tools registry."
+                                        })
                                     })
-                                })
-                                continue
+                                    continue
                         
                         # CRITICAL FIX 3: Check for corrupted tool names (length check)
                         if len(str(tool_call.function.name)) > 100:

src/tools/model_training.py

@@ -128,6 +128,12 @@ def train_baseline_models(file_path: str, target_col: str,
     }
     
     # Train models based on task type
+    import sys
+    print(f"\n🚀 Training {6 if task_type == 'classification' else 6} baseline models...", flush=True)
+    print(f"   📊 Training set: {len(X_train):,} samples × {X_train.shape[1]} features", flush=True)
+    print(f"   📊 Test set: {len(X_test):,} samples", flush=True)
+    sys.stdout.flush()
+    
     if task_type == "classification":
         models = {
             "logistic_regression": LogisticRegression(max_iter=1000, random_state=random_state),
@@ -137,10 +143,17 @@ def train_baseline_models(file_path: str, target_col: str,
             "catboost": CatBoostClassifier(iterations=100, random_state=random_state, verbose=0, allow_writing_files=False)
         }
         
-        for model_name, model in models.items():
+        for idx, (model_name, model) in enumerate(models.items(), 1):
             try:
                 # Train
+                print(f"\n   [{idx}/{len(models)}] Training {model_name}...", flush=True)
+                sys.stdout.flush()
+                import time
+                start_time = time.time()
                 model.fit(X_train, y_train)
+                elapsed = time.time() - start_time
+                print(f"   ✓ {model_name} trained in {elapsed:.1f}s", flush=True)
+                sys.stdout.flush()
                 
                 # Predict
                 y_pred_train = model.predict(X_train)
@@ -206,10 +219,18 @@ def train_baseline_models(file_path: str, target_col: str,
             "catboost": CatBoostRegressor(iterations=100, random_state=random_state, verbose=0, allow_writing_files=False)
         }
         
-        for model_name, model in models.items():
+        for idx, (model_name, model) in enumerate(models.items(), 1):
             try:
                 # Train
+                import sys
+                print(f"\n   [{idx}/{len(models)}] Training {model_name}...", flush=True)
+                sys.stdout.flush()
+                import time
+                start_time = time.time()
                 model.fit(X_train, y_train)
+                elapsed = time.time() - start_time
+                print(f"   ✓ {model_name} trained in {elapsed:.1f}s", flush=True)
+                sys.stdout.flush()
                 
                 # Predict
                 y_pred_train = model.predict(X_train)
@@ -358,6 +379,20 @@ def train_baseline_models(file_path: str, target_col: str,
     else:
         results["visualization_generated"] = False
     
+    # Print final summary
+    print(f"\n{'='*60}")
+    print(f"✅ TRAINING COMPLETE")
+    print(f"{'='*60}")
+    print(f"📊 Best Model: {best_model_name}")
+    if task_type == "regression":
+        print(f"📈 Test R²: {best_score:.4f}")
+        print(f"📉 Test RMSE: {results['models'][best_model_name]['test_metrics']['rmse']:.4f}")
+    else:
+        print(f"📈 Test F1: {best_score:.4f}")
+        print(f"📉 Test Accuracy: {results['models'][best_model_name]['test_metrics']['accuracy']:.4f}")
+    print(f"💾 Model saved: {results['best_model']['model_path']}")
+    print(f"{'='*60}\\n")
+    
     return results
 
 

src/utils/semantic_layer.py

@@ -56,7 +56,8 @@ class SemanticLayer:
         if self.enabled:
             try:
                 print(f"🧠 Loading SBERT model: {model_name}...")
-                self.model = SentenceTransformer(model_name)
+                # Try loading with trust_remote_code for better compatibility
+                self.model = SentenceTransformer(model_name, trust_remote_code=True)
                 # Use GPU if available
                 if torch.cuda.is_available():
                     self.model = self.model.to('cuda')
@@ -65,6 +66,7 @@ class SemanticLayer:
                     print("✅ SBERT loaded on CPU")
             except Exception as e:
                 print(f"⚠️ Failed to load SBERT model: {e}")
+                print(f"   Falling back to keyword-based routing (semantic features disabled)")
                 self.enabled = False
         else:
             print("⚠️ SBERT semantic layer disabled (missing dependencies)")