Update pymc/mcp_output/mcp_plugin/mcp_service.py

pymc/mcp_output/mcp_plugin/mcp_service.py

@@ -1,445 +1,226 @@
-import os
-import sys
-from typing import Dict, Any, List, Optional
-
-# Add the local source directory to sys.path
-source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
-if source_path not in sys.path:
-    sys.path.insert(0, source_path)
-
 from fastmcp import FastMCP
-import numpy as np
-
-# Import core modules from the local source directory
-import pymc as pm
-from pymc import Model, sample
-from pymc.distributions import Normal, Beta, Uniform, Exponential, Gamma, HalfNormal
-import arviz as az
 
 # Create the FastMCP service application
 mcp = FastMCP("pymc_service")
 
-# Storage for models and traces
-_models: Dict[str, Model] = {}
-_traces: Dict[str, Any] = {}
-
-
-@mcp.tool(name="create_simple_model", description="Create a simple regression model with normal priors")
-def create_simple_model(
-    model_id: str,
-    n_predictors: int = 2
-) -> Dict[str, Any]:
+@mcp.tool(name="model_specification", description="Define a probabilistic model using PyMC's syntax")
+def model_specification(model_code: str) -> dict:
     """
-    Create a simple Bayesian regression model.
+    Define a probabilistic model using PyMC's syntax.
 
-    Args:
-        model_id: Unique identifier for the model
-        n_predictors: Number of predictor variables
+    Parameters:
+    - model_code (str): The PyMC model definition code.
 
     Returns:
-        Dictionary with success status and model information
+    - dict: A dictionary with success status and model object.
     """
     try:
-        with Model() as model:
-            # Priors
-            alpha = Normal('alpha', mu=0, sigma=10)
-            beta = Normal('beta', mu=0, sigma=10, shape=n_predictors)
-            sigma = HalfNormal('sigma', sigma=1)
-            
-            # Store model structure info
-            model_info = {
-                "model_id": model_id,
-                "type": "simple_regression",
-                "n_predictors": n_predictors,
-                "variables": [var.name for var in model.free_RVs]
-            }
-        
-        _models[model_id] = model
-        
-        return {
-            "success": True,
-            "result": model_info,
-            "error": None
-        }
+        # Execute the model code
+        local_vars = {}
+        exec(model_code, {}, local_vars)
+        model = local_vars.get("model")
+        if model is None:
+            raise ValueError("Model definition is invalid or missing.")
+        return {"success": True, "model": model}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
-
-@mcp.tool(name="create_hierarchical_model", description="Create a hierarchical model")
-def create_hierarchical_model(
-    model_id: str,
-    n_groups: int = 3
-) -> Dict[str, Any]:
+@mcp.tool(name="mcmc_sampling", description="Perform Markov Chain Monte Carlo (MCMC) sampling")
+def mcmc_sampling(model: object, num_samples: int, tune: int) -> dict:
     """
-    Create a simple hierarchical model with group-level effects.
+    Perform Markov Chain Monte Carlo (MCMC) sampling.
 
-    Args:
-        model_id: Unique identifier for the model
-        n_groups: Number of groups
+    Parameters:
+    - model (object): The PyMC model object.
+    - num_samples (int): Number of samples to draw.
+    - tune (int): Number of tuning steps.
 
     Returns:
-        Dictionary with success status and model information
+    - dict: A dictionary with success status and sampling results.
     """
     try:
-        with Model() as model:
-            # Hyperpriors
-            mu_alpha = Normal('mu_alpha', mu=0, sigma=10)
-            sigma_alpha = HalfNormal('sigma_alpha', sigma=1)
-            
-            # Group-level parameters
-            alpha = Normal('alpha', mu=mu_alpha, sigma=sigma_alpha, shape=n_groups)
-            sigma = HalfNormal('sigma', sigma=1)
-            
-            model_info = {
-                "model_id": model_id,
-                "type": "hierarchical",
-                "n_groups": n_groups,
-                "variables": [var.name for var in model.free_RVs]
-            }
-        
-        _models[model_id] = model
-        
-        return {
-            "success": True,
-            "result": model_info,
-            "error": None
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="add_observed_variable", description="Add observed data to a model")
-def add_observed_variable(
-    model_id: str,
-    var_name: str,
-    distribution: str = "Normal",
-    observed_data: Optional[List[float]] = None
-) -> Dict[str, Any]:
-    """
-    Add an observed variable to an existing model.
-
-    Args:
-        model_id: ID of the model to modify
-        var_name: Name of the observed variable
-        distribution: Distribution type (Normal, Beta, etc.)
-        observed_data: Optional observed data values
-
-    Returns:
-        Dictionary with success status
-    """
-    try:
-        if model_id not in _models:
-            return {"success": False, "result": None, "error": f"Model '{model_id}' not found"}
-        
-        model = _models[model_id]
-        
+        import pymc as pm
         with model:
-            if distribution == "Normal":
-                if observed_data:
-                    Normal(var_name, mu=0, sigma=1, observed=np.array(observed_data))
-                else:
-                    Normal(var_name, mu=0, sigma=1)
-            elif distribution == "Beta":
-                if observed_data:
-                    Beta(var_name, alpha=1, beta=1, observed=np.array(observed_data))
-                else:
-                    Beta(var_name, alpha=1, beta=1)
-        
-        return {
-            "success": True,
-            "result": {
-                "model_id": model_id,
-                "variable_added": var_name,
-                "distribution": distribution
-            },
-            "error": None
-        }
+            trace = pm.sample(draws=num_samples, tune=tune)
+        return {"success": True, "trace": trace}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
-
-@mcp.tool(name="sample_model", description="Perform MCMC sampling on a model")
-def sample_model(
-    model_id: str,
-    draws: int = 1000,
-    tune: int = 1000,
-    chains: int = 2,
-    cores: int = 1
-) -> Dict[str, Any]:
+@mcp.tool(name="variational_inference", description="Perform Variational Inference (VI)")
+def variational_inference(model: object, num_samples: int) -> dict:
     """
-    Perform MCMC sampling on a stored model.
+    Perform Variational Inference (VI).
 
-    Args:
-        model_id: ID of the model to sample
-        draws: Number of samples to draw
-        tune: Number of tuning steps
-        chains: Number of chains
-        cores: Number of cores to use
+    Parameters:
+    - model (object): The PyMC model object.
+    - num_samples (int): Number of samples to draw.
 
     Returns:
-        Dictionary with success status and trace ID
+    - dict: A dictionary with success status and VI results.
     """
     try:
-        if model_id not in _models:
-            return {"success": False, "result": None, "error": f"Model '{model_id}' not found"}
-        
-        model = _models[model_id]
-        
+        import pymc as pm
         with model:
-            trace = sample(draws=draws, tune=tune, chains=chains, cores=cores)
-        
-        trace_id = f"{model_id}_trace"
-        _traces[trace_id] = trace
-        
-        # Get summary statistics
-        summary = az.summary(trace)
-        
-        return {
-            "success": True,
-            "result": {
-                "trace_id": trace_id,
-                "model_id": model_id,
-                "draws": draws,
-                "chains": chains,
-                "summary": summary.to_dict()
-            },
-            "error": None
-        }
+            approx = pm.fit(n=num_samples)
+        return {"success": True, "approximation": approx}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="get_trace_summary", description="Get summary statistics from a trace")
-def get_trace_summary(trace_id: str) -> Dict[str, Any]:
+@mcp.tool(name="posterior_predictive_sampling", description="Generate posterior predictive samples")
+def posterior_predictive_sampling(model: object, trace: object) -> dict:
     """
-    Get summary statistics from a stored trace.
+    Generate posterior predictive samples.
 
-    Args:
-        trace_id: ID of the trace
+    Parameters:
+    - model (object): The PyMC model object.
+    - trace (object): The trace object from MCMC or VI.
 
     Returns:
-        Dictionary with summary statistics
+    - dict: A dictionary with success status and posterior predictive samples.
     """
     try:
-        if trace_id not in _traces:
-            return {"success": False, "result": None, "error": f"Trace '{trace_id}' not found"}
-        
-        trace = _traces[trace_id]
-        summary = az.summary(trace)
-        
-        return {
-            "success": True,
-            "result": {
-                "trace_id": trace_id,
-                "summary": summary.to_dict()
-            },
-            "error": None
-        }
+        import pymc as pm
+        with model:
+            ppc = pm.sample_posterior_predictive(trace)
+        return {"success": True, "posterior_predictive": ppc}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
-
-@mcp.tool(name="compute_loo", description="Compute leave-one-out cross-validation")
-def compute_loo(trace_id: str) -> Dict[str, Any]:
+@mcp.tool(name="model_visualization", description="Visualize the model using Graphviz or Mermaid")
+def model_visualization(model: object, format: str) -> dict:
     """
-    Compute LOO-CV for model comparison.
+    Visualize the model using Graphviz or Mermaid.
 
-    Args:
-        trace_id: ID of the trace
+    Parameters:
+    - model (object): The PyMC model object.
+    - format (str): Visualization format (e.g., "graphviz", "mermaid").
 
     Returns:
-        Dictionary with LOO results
+    - dict: A dictionary with success status and visualization data.
     """
     try:
-        if trace_id not in _traces:
-            return {"success": False, "result": None, "error": f"Trace '{trace_id}' not found"}
-        
-        trace = _traces[trace_id]
-        loo = az.loo(trace)
-        
-        return {
-            "success": True,
-            "result": {
-                "trace_id": trace_id,
-                "loo": float(loo.loo),
-                "loo_se": float(loo.loo_se),
-                "p_loo": float(loo.p_loo)
-            },
-            "error": None
-        }
+        from pymc.model_graph import model_to_graphviz, model_to_mermaid
+        if format == "graphviz":
+            visualization = model_to_graphviz(model)
+        elif format == "mermaid":
+            visualization = model_to_mermaid(model)
+        else:
+            raise ValueError("Unsupported visualization format.")
+        return {"success": True, "visualization": visualization}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
-
-@mcp.tool(name="compute_waic", description="Compute Watanabe-Akaike information criterion")
-def compute_waic(trace_id: str) -> Dict[str, Any]:
+@mcp.tool(name="model_diagnostics", description="Perform diagnostics on the model")
+def model_diagnostics(trace: object) -> dict:
     """
-    Compute WAIC for model comparison.
+    Perform diagnostics on the model.
 
-    Args:
-        trace_id: ID of the trace
+    Parameters:
+    - trace (object): The trace object from MCMC or VI.
 
     Returns:
-        Dictionary with WAIC results
+    - dict: A dictionary with diagnostic results (e.g., Rhat, ESS).
     """
     try:
-        if trace_id not in _traces:
-            return {"success": False, "result": None, "error": f"Trace '{trace_id}' not found"}
-        
-        trace = _traces[trace_id]
-        waic = az.waic(trace)
-        
-        return {
-            "success": True,
-            "result": {
-                "trace_id": trace_id,
-                "waic": float(waic.waic),
-                "waic_se": float(waic.waic_se),
-                "p_waic": float(waic.p_waic)
-            },
-            "error": None
+        import arviz as az
+        diagnostics = {
+            "rhat": az.rhat(trace).to_dict(),
+            "ess": az.ess(trace).to_dict()
         }
+        return {"success": True, "diagnostics": diagnostics}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="get_posterior_samples", description="Get posterior samples for a variable")
-def get_posterior_samples(
-    trace_id: str,
-    var_name: str,
-    n_samples: int = 100
-) -> Dict[str, Any]:
+@mcp.tool(name="load_data", description="Load data for modeling")
+def load_data(data_path: str) -> dict:
     """
-    Get posterior samples for a specific variable.
+    Load data for modeling.
 
-    Args:
-        trace_id: ID of the trace
-        var_name: Name of the variable
-        n_samples: Number of samples to return
+    Parameters:
+    - data_path (str): Path to the data file.
 
     Returns:
-        Dictionary with posterior samples
+    - dict: A dictionary with success status and loaded data.
     """
     try:
-        if trace_id not in _traces:
-            return {"success": False, "result": None, "error": f"Trace '{trace_id}' not found"}
-        
-        trace = _traces[trace_id]
-        
-        # Convert to InferenceData if needed
-        if not isinstance(trace, az.InferenceData):
-            idata = az.from_pymc3(trace)
-        else:
-            idata = trace
-        
-        # Get samples
-        samples = idata.posterior[var_name].values.flatten()[:n_samples].tolist()
-        
-        return {
-            "success": True,
-            "result": {
-                "trace_id": trace_id,
-                "var_name": var_name,
-                "samples": samples,
-                "n_samples": len(samples)
-            },
-            "error": None
-        }
+        import pandas as pd
+        data = pd.read_csv(data_path)
+        return {"success": True, "data": data}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="list_models", description="List all stored models")
-def list_models() -> Dict[str, Any]:
+@mcp.tool(name="save_model", description="Save a PyMC model to a file")
+def save_model(model: object, file_path: str) -> dict:
     """
-    List all stored models.
+    Save a PyMC model to a file.
 
-    Returns:
-        Dictionary with list of model IDs
-    """
-    try:
-        return {
-            "success": True,
-            "result": {"model_ids": list(_models.keys())},
-            "error": None
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="list_traces", description="List all stored traces")
-def list_traces() -> Dict[str, Any]:
-    """
-    List all stored traces.
+    Parameters:
+    - model (object): The PyMC model object.
+    - file_path (str): Path to save the model.
 
     Returns:
-        Dictionary with list of trace IDs
+    - dict: A dictionary with success status.
     """
     try:
-        return {
-            "success": True,
-            "result": {"trace_ids": list(_traces.keys())},
-            "error": None
-        }
+        import pickle
+        with open(file_path, "wb") as f:
+            pickle.dump(model, f)
+        return {"success": True}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="delete_model", description="Delete a stored model")
-def delete_model(model_id: str) -> Dict[str, Any]:
+@mcp.tool(name="load_model", description="Load a PyMC model from a file")
+def load_model(file_path: str) -> dict:
     """
-    Delete a stored model.
+    Load a PyMC model from a file.
 
-    Args:
-        model_id: ID of the model to delete
+    Parameters:
+    - file_path (str): Path to the model file.
 
     Returns:
-        Dictionary with success status
+    - dict: A dictionary with success status and loaded model.
     """
     try:
-        if model_id not in _models:
-            return {"success": False, "result": None, "error": f"Model '{model_id}' not found"}
-        
-        del _models[model_id]
-        return {
-            "success": True,
-            "result": {"deleted": model_id},
-            "error": None
-        }
+        import pickle
+        with open(file_path, "rb") as f:
+            model = pickle.load(f)
+        return {"success": True, "model": model}
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
-@mcp.tool(name="delete_trace", description="Delete a stored trace")
-def delete_trace(trace_id: str) -> Dict[str, Any]:
+@mcp.tool(name="export_results", description="Export results to a specified format")
+def export_results(results: object, output_format: str) -> dict:
     """
-    Delete a stored trace.
+    Export results to a specified format.
 
-    Args:
-        trace_id: ID of the trace to delete
+    Parameters:
+    - results (object): The results object (e.g., trace, posterior predictive samples).
+    - output_format (str): Desired output format (e.g., "csv", "json").
 
     Returns:
-        Dictionary with success status
+    - dict: A dictionary with success status and export details.
     """
     try:
-        if trace_id not in _traces:
-            return {"success": False, "result": None, "error": f"Trace '{trace_id}' not found"}
-        
-        del _traces[trace_id]
-        return {
-            "success": True,
-            "result": {"deleted": trace_id},
-            "error": None
-        }
+        if output_format == "csv":
+            import pandas as pd
+            df = pd.DataFrame(results)
+            df.to_csv("results.csv", index=False)
+            return {"success": True, "file": "results.csv"}
+        elif output_format == "json":
+            import json
+            with open("results.json", "w") as f:
+                json.dump(results, f)
+            return {"success": True, "file": "results.json"}
+        else:
+            raise ValueError("Unsupported output format.")
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        return {"success": False, "error": str(e)}
 
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.
 
     Returns:
-        FastMCP: The FastMCP application instance.
+    - FastMCP: The FastMCP application instance.
     """
     return mcp