Update pyro/mcp_output/mcp_plugin/mcp_service.py

pyro/mcp_output/mcp_plugin/mcp_service.py

@@ -1,466 +1,97 @@
-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 torch
-import pyro
-import pyro.distributions as dist
-from pyro.infer import SVI, Trace_ELBO, MCMC, NUTS, Predictive
-from pyro.optim import Adam
 
 # Create the FastMCP service application
 mcp = FastMCP("pyro_service")
 
-# Store models and guides as string IDs
-_models = {}
-_guides = {}
-_svi_instances = {}
-_mcmc_instances = {}
-
-@mcp.tool(name="get_pyro_info")
-def get_pyro_info() -> dict:
-    """
-    Get information about the Pyro library.
-
-    Returns:
-        dict: Version and configuration information about Pyro.
-    """
-    try:
-        return {
-            "success": True,
-            "result": {
-                "version": pyro.__version__,
-                "torch_version": torch.__version__,
-                "cuda_available": torch.cuda.is_available(),
-                "backend": "torch",
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
 
-@mcp.tool(name="list_distributions")
+@mcp.tool(name="list_distributions", description="List all available distributions in Pyro")
 def list_distributions() -> dict:
     """
     List all available distributions in Pyro.
 
     Returns:
-        dict: A list of available distribution names organized by category.
-    """
-    try:
-        basic_dists = [
-            "Normal", "Bernoulli", "Beta", "Binomial", "Categorical", "Cauchy",
-            "Dirichlet", "Exponential", "Gamma", "Geometric", "LogNormal",
-            "Multinomial", "MultivariateNormal", "Poisson", "Uniform"
-        ]
-        
-        hmm_dists = [
-            "DiscreteHMM", "GaussianHMM", "GammaGaussianHMM", "LinearHMM",
-            "GaussianMRF", "IndependentHMM"
-        ]
-        
-        advanced_dists = [
-            "Delta", "Empirical", "MixtureOfDiagNormals", "TransformedDistribution",
-            "ConditionalDistribution", "ZeroInflatedPoisson", "ZeroInflatedNegativeBinomial"
-        ]
-        
-        return {
-            "success": True,
-            "result": {
-                "basic": basic_dists,
-                "hidden_markov_models": hmm_dists,
-                "advanced": advanced_dists,
-                "total_count": len(basic_dists) + len(hmm_dists) + len(advanced_dists),
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="sample_normal")
-def sample_normal(loc: float = 0.0, scale: float = 1.0, sample_shape: Optional[List[int]] = None) -> dict:
-    """
-    Sample from a Normal distribution.
-
-    Parameters:
-        loc (float): Mean of the distribution.
-        scale (float): Standard deviation of the distribution.
-        sample_shape (Optional[List[int]]): Shape of samples to draw.
-
-    Returns:
-        dict: Samples from the Normal distribution.
-    """
-    try:
-        normal = dist.Normal(loc, scale)
-        if sample_shape:
-            samples = normal.sample(torch.Size(sample_shape))
-        else:
-            samples = normal.sample()
-        
-        return {
-            "success": True,
-            "result": {
-                "samples": samples.tolist() if isinstance(samples, torch.Tensor) else float(samples),
-                "distribution": "Normal",
-                "parameters": {"loc": loc, "scale": scale},
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="sample_bernoulli")
-def sample_bernoulli(probs: float = 0.5, sample_shape: Optional[List[int]] = None) -> dict:
-    """
-    Sample from a Bernoulli distribution.
-
-    Parameters:
-        probs (float): Probability of success (between 0 and 1).
-        sample_shape (Optional[List[int]]): Shape of samples to draw.
-
-    Returns:
-        dict: Samples from the Bernoulli distribution.
-    """
-    try:
-        bernoulli = dist.Bernoulli(probs)
-        if sample_shape:
-            samples = bernoulli.sample(torch.Size(sample_shape))
-        else:
-            samples = bernoulli.sample()
-        
-        return {
-            "success": True,
-            "result": {
-                "samples": samples.tolist() if isinstance(samples, torch.Tensor) else int(samples),
-                "distribution": "Bernoulli",
-                "parameters": {"probs": probs},
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="sample_categorical")
-def sample_categorical(probs: List[float], sample_shape: Optional[List[int]] = None) -> dict:
-    """
-    Sample from a Categorical distribution.
-
-    Parameters:
-        probs (List[float]): Probabilities for each category (must sum to 1).
-        sample_shape (Optional[List[int]]): Shape of samples to draw.
-
-    Returns:
-        dict: Samples from the Categorical distribution.
-    """
-    try:
-        probs_tensor = torch.tensor(probs)
-        categorical = dist.Categorical(probs_tensor)
-        if sample_shape:
-            samples = categorical.sample(torch.Size(sample_shape))
-        else:
-            samples = categorical.sample()
-        
-        return {
-            "success": True,
-            "result": {
-                "samples": samples.tolist() if isinstance(samples, torch.Tensor) else int(samples),
-                "distribution": "Categorical",
-                "parameters": {"probs": probs},
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="create_simple_model")
-def create_simple_model(model_id: str, model_type: str = "normal_normal") -> dict:
-    """
-    Create a simple probabilistic model.
-
-    Parameters:
-        model_id (str): Unique identifier for the model.
-        model_type (str): Type of model ('normal_normal', 'coin_flip', 'linear_regression').
-
-    Returns:
-        dict: Information about the created model.
-    """
-    try:
-        if model_type == "normal_normal":
-            def model(data=None):
-                loc = pyro.sample("loc", dist.Normal(0.0, 1.0))
-                scale = pyro.sample("scale", dist.LogNormal(0.0, 1.0))
-                with pyro.plate("data", len(data) if data is not None else 1):
-                    return pyro.sample("obs", dist.Normal(loc, scale), obs=data)
-            
-        elif model_type == "coin_flip":
-            def model(data=None):
-                p = pyro.sample("p", dist.Beta(2.0, 2.0))
-                with pyro.plate("data", len(data) if data is not None else 1):
-                    return pyro.sample("obs", dist.Bernoulli(p), obs=data)
-                    
-        elif model_type == "linear_regression":
-            def model(x=None, y=None):
-                a = pyro.sample("a", dist.Normal(0.0, 10.0))
-                b = pyro.sample("b", dist.Normal(0.0, 10.0))
-                sigma = pyro.sample("sigma", dist.LogNormal(0.0, 1.0))
-                if x is not None:
-                    mean = a + b * x
-                    with pyro.plate("data", len(x)):
-                        return pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
-        else:
-            return {"success": False, "result": None, "error": f"Unknown model type: {model_type}"}
-        
-        _models[model_id] = model
-        
-        return {
-            "success": True,
-            "result": {
-                "model_id": model_id,
-                "model_type": model_type,
-                "message": f"Model '{model_id}' created successfully",
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="create_guide")
-def create_guide(guide_id: str, model_id: str, guide_type: str = "auto_normal") -> dict:
-    """
-    Create a variational guide for a model.
-
-    Parameters:
-        guide_id (str): Unique identifier for the guide.
-        model_id (str): ID of the model to create a guide for.
-        guide_type (str): Type of guide ('auto_normal', 'auto_delta').
-
-    Returns:
-        dict: Information about the created guide.
+    - dict: A dictionary with success status and list of distributions.
     """
     try:
-        if model_id not in _models:
-            return {"success": False, "result": None, "error": f"Model '{model_id}' not found"}
-        
-        model = _models[model_id]
-        
-        if guide_type == "auto_normal":
-            from pyro.infer.autoguide import AutoNormal
-            guide = AutoNormal(model)
-        elif guide_type == "auto_delta":
-            from pyro.infer.autoguide import AutoDelta
-            guide = AutoDelta(model)
-        else:
-            return {"success": False, "result": None, "error": f"Unknown guide type: {guide_type}"}
-        
-        _guides[guide_id] = guide
-        
+        from pyro.distributions import __all__ as distributions
         return {
             "success": True,
-            "result": {
-                "guide_id": guide_id,
-                "model_id": model_id,
-                "guide_type": guide_type,
-                "message": f"Guide '{guide_id}' created successfully",
-            },
-            "error": None,
+            "distributions": distributions,
+            "count": len(distributions)
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
 
-@mcp.tool(name="run_svi")
-def run_svi(
-    svi_id: str,
-    model_id: str,
-    guide_id: str,
-    num_steps: int = 1000,
-    learning_rate: float = 0.01
-) -> dict:
+@mcp.tool(name="sample_from_distribution", description="Sample from a specific Pyro distribution")
+def sample_from_distribution(distribution_name: str, *args, **kwargs) -> dict:
     """
-    Run Stochastic Variational Inference.
+    Sample from a specific Pyro distribution.
 
     Parameters:
-        svi_id (str): Unique identifier for this SVI instance.
-        model_id (str): ID of the model.
-        guide_id (str): ID of the guide.
-        num_steps (int): Number of optimization steps.
-        learning_rate (float): Learning rate for optimization.
+    - distribution_name: Name of the distribution (e.g., 'Normal', 'Beta')
+    - *args, **kwargs: Parameters for the distribution
 
     Returns:
-        dict: Training results including loss history.
+    - dict: Sampled values or error message.
     """
     try:
-        if model_id not in _models:
-            return {"success": False, "result": None, "error": f"Model '{model_id}' not found"}
-        if guide_id not in _guides:
-            return {"success": False, "result": None, "error": f"Guide '{guide_id}' not found"}
-        
-        model = _models[model_id]
-        guide = _guides[guide_id]
-        
-        # Clear parameter store
-        pyro.clear_param_store()
-        
-        # Create SVI
-        optimizer = Adam({"lr": learning_rate})
-        svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
-        
-        # Training loop with synthetic data
-        data = torch.randn(100)
-        losses = []
-        for step in range(num_steps):
-            loss = svi.step(data)
-            losses.append(loss)
-            if step % 100 == 0:
-                print(f"Step {step}, Loss: {loss}")
-        
-        _svi_instances[svi_id] = svi
-        
-        return {
-            "success": True,
-            "result": {
-                "svi_id": svi_id,
-                "model_id": model_id,
-                "guide_id": guide_id,
-                "num_steps": num_steps,
-                "final_loss": losses[-1],
-                "loss_history": losses[::max(1, num_steps // 20)],  # Return ~20 points
-                "message": "SVI completed successfully",
-            },
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
-
-@mcp.tool(name="list_models")
-def list_models() -> dict:
-    """
-    List all stored models.
+        from pyro.distributions import __dict__ as dist_dict
+        if distribution_name not in dist_dict:
+            return {
+                "success": False,
+                "error": f"Distribution '{distribution_name}' not found. Use list_distributions to see available options."
+            }
 
-    Returns:
-        dict: List of model IDs.
-    """
-    try:
+        distribution = dist_dict[distribution_name](*args, **kwargs)
+        samples = distribution.sample()
         return {
             "success": True,
-            "result": {
-                "models": list(_models.keys()),
-                "guides": list(_guides.keys()),
-                "svi_instances": list(_svi_instances.keys()),
-                "mcmc_instances": list(_mcmc_instances.keys()),
-            },
-            "error": None,
+            "samples": samples.tolist()
         }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
 
-@mcp.tool(name="delete_model")
-def delete_model(model_id: str) -> dict:
+@mcp.tool(name="run_inference", description="Run inference using a specific Pyro model")
+def run_inference(model_name: str, data: dict) -> dict:
     """
-    Delete a stored model.
+    Run inference using a specific Pyro model.
 
     Parameters:
-        model_id (str): ID of the model to delete.
+    - model_name: Name of the model to use
+    - data: Input data for the model
 
     Returns:
-        dict: Confirmation of deletion.
+    - dict: Inference results or error message.
     """
     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": {"message": f"Model '{model_id}' deleted successfully"},
-            "error": None,
-        }
-    except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
-
+        from pyro.infer import SVI, Trace_ELBO
+        from pyro.optim import Adam
+        from pyro import __dict__ as pyro_dict
 
-@mcp.tool(name="get_distribution_info")
-def get_distribution_info(distribution_name: str) -> dict:
-    """
-    Get information about a specific distribution.
-
-    Parameters:
-        distribution_name (str): Name of the distribution (e.g., 'Normal', 'Bernoulli').
-
-    Returns:
-        dict: Information about the distribution including parameters.
-    """
-    try:
-        dist_info = {
-            "Normal": {
-                "parameters": ["loc (mean)", "scale (std dev)"],
-                "support": "real numbers",
-                "description": "Gaussian/Normal distribution",
-            },
-            "Bernoulli": {
-                "parameters": ["probs (probability of 1)"],
-                "support": "{0, 1}",
-                "description": "Binary distribution",
-            },
-            "Beta": {
-                "parameters": ["concentration1 (alpha)", "concentration0 (beta)"],
-                "support": "[0, 1]",
-                "description": "Beta distribution for probabilities",
-            },
-            "Categorical": {
-                "parameters": ["probs (probability vector)"],
-                "support": "{0, 1, ..., K-1}",
-                "description": "Categorical distribution over K categories",
-            },
-            "Dirichlet": {
-                "parameters": ["concentration (alpha vector)"],
-                "support": "probability simplex",
-                "description": "Dirichlet distribution for probability vectors",
-            },
-            "DiscreteHMM": {
-                "parameters": ["initial_logits", "transition_logits", "observation_dist"],
-                "support": "sequences of discrete states",
-                "description": "Discrete Hidden Markov Model",
-            },
-            "GaussianHMM": {
-                "parameters": ["initial_dist", "transition_matrix", "observation_matrix"],
-                "support": "sequences of continuous observations",
-                "description": "Gaussian Hidden Markov Model",
-            },
-        }
-        
-        if distribution_name in dist_info:
+        if model_name not in pyro_dict:
             return {
-                "success": True,
-                "result": {
-                    "distribution": distribution_name,
-                    **dist_info[distribution_name],
-                },
-                "error": None,
+                "success": False,
+                "error": f"Model '{model_name}' not found."
             }
-        else:
+
+        model = pyro_dict[model_name]
+        guide = pyro_dict.get(f"{model_name}_guide")
+        if not guide:
             return {
                 "success": False,
-                "result": None,
-                "error": f"Distribution '{distribution_name}' not found in info database",
+                "error": f"Guide for model '{model_name}' not found."
             }
+
+        svi = SVI(model, guide, Adam({"lr": 0.01}), Trace_ELBO())
+        loss = svi.step(data)
+        return {
+            "success": True,
+            "loss": loss
+        }
     except Exception as e:
-        return {"success": False, "result": None, "error": str(e)}
+        return {"success": False, "error": str(e)}
 
 
 def create_app() -> FastMCP:
@@ -468,6 +99,6 @@ def create_app() -> FastMCP:
     Create and return the FastMCP application instance.
 
     Returns:
-        The FastMCP application instance.
+    - FastMCP: The FastMCP application instance.
     """
     return mcp