Update pyfolio/mcp_output/mcp_plugin/mcp_service.py

pyfolio/mcp_output/mcp_plugin/mcp_service.py

@@ -38,34 +38,41 @@ from pyfolio.plotting import (
 mcp = FastMCP("pyfolio_service")
 
 def _save_plot_as_base64(plt, plot_name):
-    """Save plot as base64 string and to file"""
+    """Save plot as base64 string for HuggingFace deployment"""
     from io import BytesIO
     
     # Generate timestamp for unique filename
     timestamp = 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')
+    # For HuggingFace deployment, we primarily return base64 data
+    # Try to save to a temporary location if possible
+    try:
+        # Try to save to /tmp or current directory
+        import tempfile
+        temp_dir = tempfile.gettempdir()
+        filepath = f"{temp_dir}/{filename}"
+        plt.savefig(filepath, dpi=300, bbox_inches='tight')
+    except Exception as e:
+        # If file saving fails, we still have the base64 data
+        filepath = f"Unable to save file: {str(e)}"
     
-    # Convert to base64
+    # Convert to base64 (this is the main output for HuggingFace)
     buffer = BytesIO()
     plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
     buffer.seek(0)
-    image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    image_base64_bytes = base64.b64encode(buffer.getvalue())
+    image_base64 = image_base64_bytes.decode('utf-8')
     buffer.close()
     
-    # Create data URI
+    # Create data URI for direct display
     data_uri = f"data:image/png;base64,{image_base64}"
     
     return {
         "image_base64": image_base64,
         "image_path": filepath,
-        "data_uri": data_uri
+        "data_uri": data_uri,
+        "filename": filename
     }
 
 def _convert_to_series(data: List[float], name: str = "returns") -> pd.Series: