Update pyfolio/mcp_output/mcp_plugin/mcp_service.py

pyfolio/mcp_output/mcp_plugin/mcp_service.py

@@ -37,34 +37,39 @@ 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 _save_plot_as_base64(plt, plot_name):
+    """Save plot as base64 string and to file"""
+    import base64
+    from io import BytesIO
+    import datetime
+    import os
+    
+    # Generate timestamp for unique filename
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"{plot_name}_{timestamp}.png"
+    
+    # Save to current working directory instead of /tmp
+    current_dir = os.getcwd()
+    filepath = os.path.join(current_dir, filename)
+    
+    # Save to file
+    plt.savefig(filepath, dpi=300, bbox_inches='tight')
+    
+    # Convert to base64
+    buffer = BytesIO()
+    plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
+    buffer.seek(0)
+    image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    buffer.close()
+    
+    # Create data URI
+    data_uri = f"data:image/png;base64,{image_base64}"
+    
+    return {
+        "image_base64": image_base64,
+        "image_path": filepath,
+        "data_uri": data_uri
+    }
 
 def _convert_to_series(data: List[float], name: str = "returns") -> pd.Series:
     """Convert list to pandas Series with date index."""