Update pyfolio/mcp_output/mcp_plugin/mcp_service.py

pyfolio/mcp_output/mcp_plugin/mcp_service.py

@@ -4,6 +4,11 @@ import pandas as pd
 from typing import Optional, Dict, List, Any
 import io
 import contextlib
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend
+import matplotlib.pyplot as plt
+import base64
+from datetime import datetime
 
 source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
 sys.path.insert(0, source_path)
@@ -32,6 +37,35 @@ from pyfolio.plotting import (
 
 mcp = FastMCP("pyfolio_service")
 
+def _save_plot_as_base64(fig=None, filename_prefix="plot") -> str:
+    """Save matplotlib plot as base64 encoded image."""
+    try:
+        if fig is None:
+            fig = plt.gcf()
+        
+        # Save to BytesIO
+        buffer = io.BytesIO()
+        fig.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
+        buffer.seek(0)
+        
+        # Convert to base64
+        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+        
+        # Also save to file
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filepath = f"/tmp/{filename_prefix}_{timestamp}.png"
+        fig.savefig(filepath, format='png', dpi=150, bbox_inches='tight')
+        
+        plt.close(fig)  # Close to free memory
+        
+        return {
+            "base64": image_base64,
+            "filepath": filepath,
+            "data_uri": f"data:image/png;base64,{image_base64}"
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
 def _convert_to_series(data: List[float], name: str = "returns") -> pd.Series:
     """Convert list to pandas Series with date index."""
     if isinstance(data, list):
@@ -389,12 +423,32 @@ def plot_annual_returns_tool(returns: list) -> dict:
         returns: List of daily returns
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status, image data, and result/error message
     """
     try:
         returns_series = _convert_to_series(returns)
-        ax = plot_annual_returns(returns_series)
-        return {"success": True, "result": "Annual returns plot generated successfully.", "error": None}
+        
+        # Create the plot
+        fig, ax = plt.subplots(figsize=(10, 6))
+        plot_annual_returns(returns_series, ax=ax)
+        plt.title('Annual Returns')
+        
+        # Save plot as base64
+        image_data = _save_plot_as_base64(fig, "annual_returns")
+        
+        if "error" in image_data:
+            return {"success": False, "result": None, "error": image_data["error"]}
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Annual returns plot generated successfully.",
+                "image_base64": image_data["base64"],
+                "image_path": image_data["filepath"],
+                "data_uri": image_data["data_uri"]
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -407,12 +461,32 @@ def plot_monthly_returns_heatmap_tool(returns: list) -> dict:
         returns: List of daily returns
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status, image data, and result/error message
     """
     try:
         returns_series = _convert_to_series(returns)
-        ax = plot_monthly_returns_heatmap(returns_series)
-        return {"success": True, "result": "Monthly returns heatmap generated successfully.", "error": None}
+        
+        # Create the plot
+        fig, ax = plt.subplots(figsize=(12, 8))
+        plot_monthly_returns_heatmap(returns_series, ax=ax)
+        plt.title('Monthly Returns Heatmap')
+        
+        # Save plot as base64
+        image_data = _save_plot_as_base64(fig, "monthly_heatmap")
+        
+        if "error" in image_data:
+            return {"success": False, "result": None, "error": image_data["error"]}
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Monthly returns heatmap generated successfully.",
+                "image_base64": image_data["base64"],
+                "image_path": image_data["filepath"],
+                "data_uri": image_data["data_uri"]
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -426,12 +500,32 @@ def plot_drawdown_periods_tool(returns: list, top: int = 10) -> dict:
         top: Number of top drawdown periods to highlight
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status, image data, and result/error message
     """
     try:
         returns_series = _convert_to_series(returns)
-        ax = plot_drawdown_periods(returns_series, top)
-        return {"success": True, "result": "Drawdown periods plot generated successfully.", "error": None}
+        
+        # Create the plot
+        fig, ax = plt.subplots(figsize=(12, 6))
+        plot_drawdown_periods(returns_series, top=top, ax=ax)
+        plt.title(f'Top {top} Drawdown Periods')
+        
+        # Save plot as base64
+        image_data = _save_plot_as_base64(fig, "drawdown_periods")
+        
+        if "error" in image_data:
+            return {"success": False, "result": None, "error": image_data["error"]}
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Drawdown periods plot generated successfully.",
+                "image_base64": image_data["base64"],
+                "image_path": image_data["filepath"],
+                "data_uri": image_data["data_uri"]
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -446,13 +540,34 @@ def plot_rolling_returns_tool(returns: list, factor_returns: list = None, live_s
         live_start_date: Start date for live trading period (optional)
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status, image data, and result/error message
     """
     try:
         returns_series = _convert_to_series(returns)
         factor_series = _convert_to_series(factor_returns, "factor") if factor_returns else None
-        ax = plot_rolling_returns(returns_series, factor_series, live_start_date)
-        return {"success": True, "result": "Rolling returns plot generated successfully.", "error": None}
+        
+        # Create the plot
+        fig, ax = plt.subplots(figsize=(12, 8))
+        plot_rolling_returns(returns_series, factor_returns=factor_series, 
+                           live_start_date=live_start_date, ax=ax)
+        plt.title('Rolling Returns Analysis')
+        
+        # Save plot as base64
+        image_data = _save_plot_as_base64(fig, "rolling_returns")
+        
+        if "error" in image_data:
+            return {"success": False, "result": None, "error": image_data["error"]}
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Rolling returns plot generated successfully.",
+                "image_base64": image_data["base64"],
+                "image_path": image_data["filepath"],
+                "data_uri": image_data["data_uri"]
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}