Update causalml/mcp_output/mcp_plugin/mcp_service.py

causalml/mcp_output/mcp_plugin/mcp_service.py

@@ -7,158 +7,501 @@ if source_path not in sys.path:
     sys.path.insert(0, source_path)
 
 from fastmcp import FastMCP
-from causalml.inference.meta import BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, BaseDRLearner, TMLELearner
-from causalml.inference.tree import UpliftTreeClassifier, CausalTreeRegressor
-from causalml.metrics import AUUC, Qini
-from causalml.optimize import CounterfactualUnitSelector
+import numpy as np
+import pandas as pd
+from typing import Optional
+
+# Import causalml components
+from causalml.inference.meta import (
+    BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, 
+    BaseDRLearner, TMLELearner
+)
+from causalml.inference.tree import (
+    UpliftTreeClassifier, UpliftRandomForestClassifier,
+    CausalTreeRegressor, CausalRandomForestRegressor
+)
+from causalml.propensity import (
+    LogisticRegressionPropensityModel,
+    GradientBoostedPropensityModel
+)
+from causalml.match import NearestNeighborMatch, smd
+from causalml.metrics import (
+    auuc_score, qini_score, 
+    get_cumgain, get_cumlift, get_qini
+)
 
 mcp = FastMCP("causalml_service")
 
-@mcp.tool(name="s_learner", description="Estimate treatment effect using S-Learner")
-def s_learner(X: list, treatment: list, y: list) -> dict:
-    """
-    Estimate treatment effect using S-Learner.
+# Session storage for models
+_meta_learners = {}
+_tree_models = {}
+_propensity_models = {}
+_matches = {}
 
-    Parameters:
-    - X: list of features
-    - treatment: list of treatment indicators
-    - y: list of outcomes
+# ==================== Meta-Learner Tools ====================
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def create_s_learner(learner_id: str, learner_type: str = "dummy", control_name: int = 0) -> dict:
+    """Create an S-Learner for treatment effect estimation."""
     try:
-        model = BaseSLearner()
-        model.fit(X, treatment, y)
-        result = model.estimate_ate(X, treatment, y)
-        return {"success": True, "result": result}
+        from sklearn.ensemble import RandomForestRegressor
+        from sklearn.linear_model import Ridge
+        
+        if learner_type == "lr":
+            base_learner = Ridge()
+        elif learner_type == "rf":
+            base_learner = RandomForestRegressor(n_estimators=100, random_state=42)
+        else:
+            base_learner = None
+        
+        learner = BaseSLearner(learner=base_learner, control_name=control_name)
+        _meta_learners[learner_id] = learner
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "type": "S-Learner",
+            "base_learner": learner_type
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="t_learner", description="Estimate treatment effect using T-Learner")
-def t_learner(X: list, treatment: list, y: list) -> dict:
-    """
-    Estimate treatment effect using T-Learner.
+@mcp.tool()
+def create_t_learner(learner_id: str, learner_type: str = "dummy", control_name: int = 0) -> dict:
+    """Create a T-Learner for treatment effect estimation."""
+    try:
+        from sklearn.ensemble import RandomForestRegressor
+        from sklearn.linear_model import Ridge
+        
+        if learner_type == "lr":
+            base_learner = Ridge()
+        elif learner_type == "rf":
+            base_learner = RandomForestRegressor(n_estimators=100, random_state=42)
+        else:
+            base_learner = None
+        
+        learner = BaseTLearner(learner=base_learner, control_name=control_name)
+        _meta_learners[learner_id] = learner
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "type": "T-Learner"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Parameters:
-    - X: list of features
-    - treatment: list of treatment indicators
-    - y: list of outcomes
+@mcp.tool()
+def create_x_learner(learner_id: str, learner_type: str = "dummy", control_name: int = 0) -> dict:
+    """Create an X-Learner for treatment effect estimation."""
+    try:
+        from sklearn.ensemble import RandomForestRegressor
+        from sklearn.linear_model import Ridge
+        
+        if learner_type == "lr":
+            base_learner = Ridge()
+        elif learner_type == "rf":
+            base_learner = RandomForestRegressor(n_estimators=100, random_state=42)
+        else:
+            base_learner = None
+        
+        learner = BaseXLearner(learner=base_learner, control_name=control_name)
+        _meta_learners[learner_id] = learner
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "type": "X-Learner"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def fit_meta_learner(learner_id: str, X: list, treatment: list, y: list, p: Optional[list] = None) -> dict:
+    """Fit a meta-learner on training data."""
     try:
-        model = BaseTLearner()
-        model.fit(X, treatment, y)
-        result = model.estimate_ate(X, treatment, y)
-        return {"success": True, "result": result}
+        if learner_id not in _meta_learners:
+            return {"success": False, "error": f"Learner {learner_id} not found"}
+        
+        learner = _meta_learners[learner_id]
+        X_arr = np.array(X)
+        treatment_arr = np.array(treatment)
+        y_arr = np.array(y)
+        p_arr = np.array(p) if p is not None else None
+        
+        learner.fit(X_arr, treatment_arr, y_arr, p=p_arr)
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "n_samples": len(y),
+            "fitted": True
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="uplift_tree", description="Classify using Uplift Tree")
-def uplift_tree(X: list, treatment: list, y: list) -> dict:
-    """
-    Classify using Uplift Tree.
+@mcp.tool()
+def predict_treatment_effects(learner_id: str, X: list) -> dict:
+    """Predict treatment effects for new samples."""
+    try:
+        if learner_id not in _meta_learners:
+            return {"success": False, "error": f"Learner {learner_id} not found"}
+        
+        learner = _meta_learners[learner_id]
+        X_arr = np.array(X)
+        
+        te = learner.predict(X_arr, verbose=False)
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "treatment_effects": te.tolist(),
+            "n_samples": len(te)
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Parameters:
-    - X: list of features
-    - treatment: list of treatment indicators
-    - y: list of outcomes
+@mcp.tool()
+def estimate_ate(learner_id: str, X: list, treatment: list, y: list) -> dict:
+    """Estimate Average Treatment Effect (ATE)."""
+    try:
+        if learner_id not in _meta_learners:
+            return {"success": False, "error": f"Learner {learner_id} not found"}
+        
+        learner = _meta_learners[learner_id]
+        X_arr = np.array(X)
+        treatment_arr = np.array(treatment)
+        y_arr = np.array(y)
+        
+        ate = learner.estimate_ate(X_arr, treatment_arr, y_arr)
+        
+        if isinstance(ate, np.ndarray):
+            ate = ate.tolist()
+        elif isinstance(ate, (np.float32, np.float64, np.int32, np.int64)):
+            ate = float(ate)
+        
+        return {
+            "success": True,
+            "learner_id": learner_id,
+            "ate": ate
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+# ==================== 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()
-        model.fit(X, treatment, y)
-        result = model.predict(X)
-        return {"success": True, "result": result}
+        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(name="causal_tree", description="Regress using Causal Tree")
-def causal_tree(X: list, treatment: list, y: list) -> dict:
-    """
-    Regress using Causal Tree.
+@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)}
 
-    Parameters:
-    - X: list of features
-    - treatment: list of treatment indicators
-    - y: list of outcomes
+@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)}
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def predict_uplift(model_id: str, X: list) -> dict:
+    """Predict uplift scores."""
     try:
-        model = CausalTreeRegressor()
-        model.fit(X, treatment, y)
-        result = model.predict(X)
-        return {"success": True, "result": result}
+        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)}
 
-@mcp.tool(name="auuc_metric", description="Calculate AUUC metric")
-def auuc_metric(y_true: list, uplift: list) -> dict:
-    """
-    Calculate AUUC metric.
+# ==================== Propensity Score Tools ====================
 
-    Parameters:
-    - y_true: list of true outcomes
-    - uplift: list of uplift predictions
+@mcp.tool()
+def create_propensity_model(model_id: str, model_type: str = "lr", calibrate: bool = True) -> dict:
+    """Create a propensity score model."""
+    try:
+        if model_type == "gbm":
+            model = GradientBoostedPropensityModel(calibrate=calibrate)
+        else:
+            model = LogisticRegressionPropensityModel(calibrate=calibrate)
+        
+        _propensity_models[model_id] = model
+        
+        return {
+            "success": True,
+            "model_id": model_id,
+            "type": model_type
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def fit_propensity_model(model_id: str, X: list, treatment: list) -> dict:
+    """Fit a propensity score model."""
     try:
-        result = AUUC(y_true, uplift)
-        return {"success": True, "result": result}
+        if model_id not in _propensity_models:
+            return {"success": False, "error": f"Model {model_id} not found"}
+        
+        model = _propensity_models[model_id]
+        X_arr = np.array(X)
+        treatment_arr = np.array(treatment)
+        
+        model.fit(X_arr, treatment_arr)
+        
+        return {
+            "success": True,
+            "model_id": model_id,
+            "n_samples": len(treatment),
+            "fitted": True
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="qini_metric", description="Calculate Qini metric")
-def qini_metric(y_true: list, uplift: list) -> dict:
-    """
-    Calculate Qini metric.
+@mcp.tool()
+def predict_propensity(model_id: str, X: list) -> dict:
+    """Predict propensity scores."""
+    try:
+        if model_id not in _propensity_models:
+            return {"success": False, "error": f"Model {model_id} not found"}
+        
+        model = _propensity_models[model_id]
+        X_arr = np.array(X)
+        
+        p_scores = model.predict(X_arr)
+        
+        return {
+            "success": True,
+            "model_id": model_id,
+            "propensity_scores": p_scores.tolist()
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# ==================== Matching Tools ====================
 
-    Parameters:
-    - y_true: list of true outcomes
-    - uplift: list of uplift predictions
+@mcp.tool()
+def create_nearest_neighbor_match(match_id: str, caliper: float = 0.2, replace: bool = False, ratio: int = 1) -> dict:
+    """Create a nearest neighbor matching object."""
+    try:
+        matcher = NearestNeighborMatch(
+            caliper=caliper,
+            replace=replace,
+            ratio=ratio
+        )
+        _matches[match_id] = {"matcher": matcher, "matched_data": None}
+        
+        return {
+            "success": True,
+            "match_id": match_id
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def perform_matching(match_id: str, data_dict: dict, treatment_col: str, score_cols: list) -> dict:
+    """Perform propensity score matching."""
     try:
-        result = Qini(y_true, uplift)
-        return {"success": True, "result": result}
+        if match_id not in _matches:
+            return {"success": False, "error": f"Matcher {match_id} not found"}
+        
+        matcher = _matches[match_id]["matcher"]
+        data = pd.DataFrame(data_dict)
+        
+        matched_data = matcher.match(data, treatment_col, score_cols)
+        _matches[match_id]["matched_data"] = matched_data
+        
+        n_treated = len(matched_data[matched_data[treatment_col] == 1])
+        n_control = len(matched_data[matched_data[treatment_col] == 0])
+        
+        return {
+            "success": True,
+            "match_id": match_id,
+            "n_matched_treated": int(n_treated),
+            "n_matched_control": int(n_control)
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="counterfactual_selector", description="Select counterfactual units")
-def counterfactual_selector(X: list, treatment: list, y: list) -> dict:
-    """
-    Select counterfactual units.
+@mcp.tool()
+def calculate_smd(feature: list, treatment: list) -> dict:
+    """Calculate Standardized Mean Difference (SMD)."""
+    try:
+        feature_series = pd.Series(feature)
+        treatment_series = pd.Series(treatment)
+        
+        smd_value = smd(feature_series, treatment_series)
+        
+        return {
+            "success": True,
+            "smd": float(smd_value)
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Parameters:
-    - X: list of features
-    - treatment: list of treatment indicators
-    - y: list of outcomes
+# ==================== Evaluation Metrics ====================
 
-    Returns:
-    - dict: containing success, result, or error
-    """
+@mcp.tool()
+def calculate_auuc(y_true: list, uplift: list, treatment: list) -> dict:
+    """Calculate Area Under the Uplift Curve (AUUC)."""
     try:
-        selector = CounterfactualUnitSelector()
-        result = selector.select(X, treatment, y)
-        return {"success": True, "result": result}
+        y_arr = np.array(y_true)
+        uplift_arr = np.array(uplift)
+        treatment_arr = np.array(treatment)
+        
+        auuc = auuc_score(y_arr, uplift_arr, treatment_arr)
+        
+        return {
+            "success": True,
+            "auuc": float(auuc)
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-def create_app() -> FastMCP:
-    """
-    Create and return the FastMCP application instance.
+@mcp.tool()
+def calculate_qini_score(y_true: list, uplift: list, treatment: list) -> dict:
+    """Calculate Qini coefficient."""
+    try:
+        y_arr = np.array(y_true)
+        uplift_arr = np.array(uplift)
+        treatment_arr = np.array(treatment)
+        
+        qini = qini_score(y_arr, uplift_arr, treatment_arr)
+        
+        return {
+            "success": True,
+            "qini_score": float(qini)
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-    Returns:
-    - FastMCP: the application instance
-    """
-    return mcp
+@mcp.tool()
+def calculate_cumulative_gain(y_true: list, uplift: list, treatment: list) -> dict:
+    """Calculate cumulative gain curve data."""
+    try:
+        y_arr = np.array(y_true)
+        uplift_arr = np.array(uplift)
+        treatment_arr = np.array(treatment)
+        
+        cum_gain = get_cumgain(y_arr, uplift_arr, treatment_arr)
+        
+        return {
+            "success": True,
+            "cumulative_gain": cum_gain.tolist() if isinstance(cum_gain, np.ndarray) else cum_gain
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# ==================== Utility Tools ====================
+
+@mcp.tool()
+def list_meta_learners() -> dict:
+    """List all created meta-learners."""
+    learners = []
+    for learner_id, learner in _meta_learners.items():
+        learners.append({
+            "id": learner_id,
+            "type": learner.__class__.__name__
+        })
+    
+    return {
+        "success": True,
+        "learners": learners,
+        "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."""
+    try:
+        import causalml
+        
+        return {
+            "success": True,
+            "version": causalml.__version__ if hasattr(causalml, '__version__') else "unknown",
+            "available_meta_learners": [
+                "BaseSLearner", "BaseTLearner", "BaseXLearner", 
+                "BaseRLearner", "BaseDRLearner", "TMLELearner"
+            ],
+            "available_tree_models": [
+                "UpliftTreeClassifier", "UpliftRandomForestClassifier",
+                "CausalTreeRegressor", "CausalRandomForestRegressor"
+            ]
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+def create_app() -> FastMCP:
+    """Create and return the FastMCP application instance."""
+    return mcp