Update pyro/mcp_output/mcp_plugin/mcp_service.py

pyro/mcp_output/mcp_plugin/mcp_service.py

@@ -1,97 +1,466 @@
+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", description="List all available distributions in Pyro")
+@mcp.tool(name="list_distributions")
 def list_distributions() -> dict:
     """
     List all available distributions in Pyro.
 
     Returns:
-    - dict: A dictionary with success status and list of distributions.
+        dict: A list of available distribution names organized by category.
     """
     try:
-        from pyro.distributions import __all__ as distributions
+        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,
-            "distributions": distributions,
-            "count": len(distributions)
+            "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, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
 
-@mcp.tool(name="sample_from_distribution", description="Sample from a specific Pyro distribution")
-def sample_from_distribution(distribution_name: str, *args, **kwargs) -> dict:
+@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 specific Pyro distribution.
+    Sample from a Normal distribution.
 
     Parameters:
-    - distribution_name: Name of the distribution (e.g., 'Normal', 'Beta')
-    - *args, **kwargs: Parameters for the distribution
+        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: Sampled values or error message.
+        dict: Samples from the Normal distribution.
     """
     try:
-        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."
-            }
+        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)}
+
 
-        distribution = dist_dict[distribution_name](*args, **kwargs)
-        samples = distribution.sample()
+@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,
-            "samples": samples.tolist()
+            "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, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
 
-@mcp.tool(name="run_inference", description="Run inference using a specific Pyro model")
-def run_inference(model_name: str, data: dict) -> dict:
+@mcp.tool(name="sample_categorical")
+def sample_categorical(probs: List[float], sample_shape: Optional[List[int]] = None) -> dict:
     """
-    Run inference using a specific Pyro model.
+    Sample from a Categorical distribution.
 
     Parameters:
-    - model_name: Name of the model to use
-    - data: Input data for the model
+        probs (List[float]): Probabilities for each category (must sum to 1).
+        sample_shape (Optional[List[int]]): Shape of samples to draw.
 
     Returns:
-    - dict: Inference results or error message.
+        dict: Samples from the Categorical distribution.
     """
     try:
-        from pyro.infer import SVI, Trace_ELBO
-        from pyro.optim import Adam
-        from pyro import __dict__ as pyro_dict
+        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)}
 
-        if model_name not in pyro_dict:
-            return {
-                "success": False,
-                "error": f"Model '{model_name}' not found."
-            }
 
-        model = pyro_dict[model_name]
-        guide = pyro_dict.get(f"{model_name}_guide")
-        if not guide:
-            return {
-                "success": False,
-                "error": f"Guide for model '{model_name}' not found."
-            }
+@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.
+    """
+    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
+        
+        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,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "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:
+    """
+    Run Stochastic Variational Inference.
+
+    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.
+
+    Returns:
+        dict: Training results including loss history.
+    """
+    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)}
+
 
-        svi = SVI(model, guide, Adam({"lr": 0.01}), Trace_ELBO())
-        loss = svi.step(data)
+@mcp.tool(name="list_models")
+def list_models() -> dict:
+    """
+    List all stored models.
+
+    Returns:
+        dict: List of model IDs.
+    """
+    try:
         return {
             "success": True,
-            "loss": loss
+            "result": {
+                "models": list(_models.keys()),
+                "guides": list(_guides.keys()),
+                "svi_instances": list(_svi_instances.keys()),
+                "mcmc_instances": list(_mcmc_instances.keys()),
+            },
+            "error": None,
         }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="delete_model")
+def delete_model(model_id: str) -> dict:
+    """
+    Delete a stored model.
+
+    Parameters:
+        model_id (str): ID of the model to delete.
+
+    Returns:
+        dict: Confirmation of deletion.
+    """
+    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)}
+
+
+@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:
+            return {
+                "success": True,
+                "result": {
+                    "distribution": distribution_name,
+                    **dist_info[distribution_name],
+                },
+                "error": None,
+            }
+        else:
+            return {
+                "success": False,
+                "result": None,
+                "error": f"Distribution '{distribution_name}' not found in info database",
+            }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
 
 
 def create_app() -> FastMCP:
@@ -99,6 +468,6 @@ def create_app() -> FastMCP:
     Create and return the FastMCP application instance.
 
     Returns:
-    - FastMCP: The FastMCP application instance.
+        The FastMCP application instance.
     """
     return mcp