Update causalml/mcp_output/mcp_plugin/mcp_service.py

causalml/mcp_output/mcp_plugin/mcp_service.py

@@ -11,15 +11,16 @@ import numpy as np
 import pandas as pd
 from typing import Optional
 
-# Import causalml components
+# Import causalml components (avoiding tree modules due to sklearn compatibility issues)
 from causalml.inference.meta import (
     BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, 
     BaseDRLearner, TMLELearner
 )
-from causalml.inference.tree import (
-    UpliftTreeClassifier, UpliftRandomForestClassifier,
-    CausalTreeRegressor, CausalRandomForestRegressor
-)
+# Note: Tree models temporarily disabled due to sklearn.utils.validation compatibility
+# from causalml.inference.tree import (
+#     UpliftTreeClassifier, UpliftRandomForestClassifier,
+#     CausalTreeRegressor, CausalRandomForestRegressor
+# )
 from causalml.propensity import (
     LogisticRegressionPropensityModel,
     GradientBoostedPropensityModel
@@ -34,7 +35,6 @@ mcp = FastMCP("causalml_service")
 
 # Session storage for models
 _meta_learners = {}
-_tree_models = {}
 _propensity_models = {}
 _matches = {}
 
@@ -188,88 +188,6 @@ def estimate_ate(learner_id: str, X: list, treatment: list, y: list) -> dict:
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-# ==================== Uplift Tree Tools ====================
-
-@mcp.tool()
-def create_uplift_tree(model_id: str, max_depth: int = 3, min_samples_leaf: int = 100) -> dict:
-    """Create an Uplift Tree classifier."""
-    try:
-        model = UpliftTreeClassifier(
-            max_depth=max_depth,
-            min_samples_leaf=min_samples_leaf
-        )
-        _tree_models[model_id] = model
-        
-        return {
-            "success": True,
-            "model_id": model_id,
-            "type": "UpliftTree"
-        }
-    except Exception as e:
-        return {"success": False, "error": str(e)}
-
-@mcp.tool()
-def create_causal_tree(model_id: str, max_depth: int = 3, min_samples_leaf: int = 100) -> dict:
-    """Create a Causal Tree regressor."""
-    try:
-        model = CausalTreeRegressor(
-            max_depth=max_depth,
-            min_samples_leaf=min_samples_leaf
-        )
-        _tree_models[model_id] = model
-        
-        return {
-            "success": True,
-            "model_id": model_id,
-            "type": "CausalTree"
-        }
-    except Exception as e:
-        return {"success": False, "error": str(e)}
-
-@mcp.tool()
-def fit_tree_model(model_id: str, X: list, treatment: list, y: list) -> dict:
-    """Fit an uplift/causal tree model."""
-    try:
-        if model_id not in _tree_models:
-            return {"success": False, "error": f"Model {model_id} not found"}
-        
-        model = _tree_models[model_id]
-        X_arr = np.array(X)
-        treatment_arr = np.array(treatment)
-        y_arr = np.array(y)
-        
-        model.fit(X_arr, treatment_arr, y_arr)
-        
-        return {
-            "success": True,
-            "model_id": model_id,
-            "n_samples": len(y),
-            "fitted": True
-        }
-    except Exception as e:
-        return {"success": False, "error": str(e)}
-
-@mcp.tool()
-def predict_uplift(model_id: str, X: list) -> dict:
-    """Predict uplift scores."""
-    try:
-        if model_id not in _tree_models:
-            return {"success": False, "error": f"Model {model_id} not found"}
-        
-        model = _tree_models[model_id]
-        X_arr = np.array(X)
-        
-        uplift = model.predict(X_arr)
-        
-        return {
-            "success": True,
-            "model_id": model_id,
-            "uplift_scores": uplift.tolist(),
-            "n_samples": len(uplift)
-        }
-    except Exception as e:
-        return {"success": False, "error": str(e)}
-
 # ==================== Propensity Score Tools ====================
 
 @mcp.tool()
@@ -465,22 +383,6 @@ def list_meta_learners() -> dict:
         "count": len(learners)
     }
 
-@mcp.tool()
-def list_tree_models() -> dict:
-    """List all created tree models."""
-    models = []
-    for model_id, model in _tree_models.items():
-        models.append({
-            "id": model_id,
-            "type": model.__class__.__name__
-        })
-    
-    return {
-        "success": True,
-        "models": models,
-        "count": len(models)
-    }
-
 @mcp.tool()
 def get_causalml_info() -> dict:
     """Get information about causalml library."""
@@ -494,10 +396,7 @@ def get_causalml_info() -> dict:
                 "BaseSLearner", "BaseTLearner", "BaseXLearner", 
                 "BaseRLearner", "BaseDRLearner", "TMLELearner"
             ],
-            "available_tree_models": [
-                "UpliftTreeClassifier", "UpliftRandomForestClassifier",
-                "CausalTreeRegressor", "CausalRandomForestRegressor"
-            ]
+            "note": "Tree models (UpliftTree, CausalTree) temporarily disabled due to sklearn compatibility"
         }
     except Exception as e:
         return {"success": False, "error": str(e)}