Update pyfolio/mcp_output/mcp_plugin/mcp_service.py

pyfolio/mcp_output/mcp_plugin/mcp_service.py

@@ -1,5 +1,7 @@
 import os
 import sys
+import pandas as pd
+from typing import Optional, Dict, List, Any
 
 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)
@@ -17,155 +19,301 @@ from pyfolio.tears import (
     create_perf_attrib_tear_sheet,
 )
 from pyfolio.plotting import (
-    axes_style,
-    customize,
     plot_annual_returns,
-    plot_capacity_sweep,
-    plot_cones,
-    plot_daily_turnover_hist,
-    plot_daily_volume,
-    plot_drawdown_periods,
-    plot_drawdown_underwater,
-    plot_exposures,
-    plot_gross_leverage,
-    plot_holdings,
-    plot_long_short_holdings,
-    plot_max_median_position_concentration,
-    plot_monthly_returns_dist,
     plot_monthly_returns_heatmap,
-    plot_monthly_returns_timeseries,
-    plot_perf_stats,
-    plot_prob_profit_trade,
-    plot_return_quantiles,
-    plot_returns,
-    plot_rolling_beta,
+    plot_drawdown_periods,
     plot_rolling_returns,
-    plot_rolling_sharpe,
-    plot_rolling_volatility,
-    plot_round_trip_lifetimes,
-    plot_sector_allocations,
-    plot_slippage_sensitivity,
-    plot_slippage_sweep,
     plot_turnover,
-    plot_txn_time_hist,
-    plotting_context,
-    show_and_plot_top_positions,
-    show_perf_stats,
-    show_profit_attribution,
-    show_worst_drawdown_periods,
 )
 
 mcp = FastMCP("pyfolio_service")
 
+def _convert_to_series(data: List[float], name: str = "returns") -> pd.Series:
+    """Convert list to pandas Series with date index."""
+    if isinstance(data, list):
+        return pd.Series(
+            data, 
+            index=pd.date_range("2020-01-01", periods=len(data), freq="D"),
+            name=name
+        )
+    return data
+
+def _convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
+    """Convert dict to pandas DataFrame."""
+    if isinstance(data, dict):
+        return pd.DataFrame(data)
+    return data
+
 @mcp.tool(name="generate_full_tear_sheet", description="Generate a comprehensive tear sheet for portfolio analysis.")
 def generate_full_tear_sheet(returns: list, positions: dict = None, transactions: dict = None, benchmark_rets: list = None) -> dict:
+    """
+    Generate a comprehensive tear sheet for portfolio analysis.
+    
+    Args:
+        returns: List of daily returns
+        positions: Dictionary of positions over time (optional)
+        transactions: Dictionary of transactions (optional)
+        benchmark_rets: List of benchmark returns (optional)
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_full_tear_sheet(returns, positions, transactions, benchmark_rets)
+        returns_series = _convert_to_series(returns, "returns")
+        positions_df = _convert_to_dataframe(positions) if positions else None
+        transactions_df = _convert_to_dataframe(transactions) if transactions else None
+        benchmark_series = _convert_to_series(benchmark_rets, "benchmark") if benchmark_rets else None
+        
+        create_full_tear_sheet(returns_series, positions_df, transactions_df, benchmark_series)
         return {"success": True, "result": "Full tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_simple_tear_sheet", description="Generate a basic tear sheet for portfolio analysis.")
 def generate_simple_tear_sheet(returns: list) -> dict:
+    """
+    Generate a basic tear sheet for portfolio analysis.
+    
+    Args:
+        returns: List of daily returns
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_simple_tear_sheet(returns)
+        returns_series = _convert_to_series(returns)
+        create_simple_tear_sheet(returns_series)
         return {"success": True, "result": "Simple tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_returns_tear_sheet", description="Generate a tear sheet for returns analysis.")
 def generate_returns_tear_sheet(returns: list) -> dict:
+    """
+    Generate a tear sheet for returns analysis.
+    
+    Args:
+        returns: List of daily returns
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_returns_tear_sheet(returns)
+        returns_series = _convert_to_series(returns)
+        create_returns_tear_sheet(returns_series)
         return {"success": True, "result": "Returns tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_position_tear_sheet", description="Generate a tear sheet for position analysis.")
 def generate_position_tear_sheet(positions: dict) -> dict:
+    """
+    Generate a tear sheet for position analysis.
+    
+    Args:
+        positions: Dictionary of positions over time
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_position_tear_sheet(positions)
+        positions_df = _convert_to_dataframe(positions)
+        create_position_tear_sheet(positions_df)
         return {"success": True, "result": "Position tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_transaction_tear_sheet", description="Generate a tear sheet for transaction analysis.")
 def generate_transaction_tear_sheet(transactions: dict) -> dict:
+    """
+    Generate a tear sheet for transaction analysis.
+    
+    Args:
+        transactions: Dictionary of transactions
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_txn_tear_sheet(transactions)
+        transactions_df = _convert_to_dataframe(transactions)
+        create_txn_tear_sheet(transactions_df)
         return {"success": True, "result": "Transaction tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_round_trip_tear_sheet", description="Generate a tear sheet for round trip analysis.")
 def generate_round_trip_tear_sheet(round_trips: dict) -> dict:
+    """
+    Generate a tear sheet for round trip analysis.
+    
+    Args:
+        round_trips: Dictionary of round trip trades
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_round_trip_tear_sheet(round_trips)
+        round_trips_df = _convert_to_dataframe(round_trips)
+        create_round_trip_tear_sheet(round_trips_df)
         return {"success": True, "result": "Round trip tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_interesting_times_tear_sheet", description="Generate a tear sheet for performance during key events.")
 def generate_interesting_times_tear_sheet(returns: list, events: list) -> dict:
+    """
+    Generate a tear sheet for performance during key events.
+    
+    Args:
+        returns: List of daily returns
+        events: List of event periods
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_interesting_times_tear_sheet(returns, events)
+        returns_series = _convert_to_series(returns)
+        create_interesting_times_tear_sheet(returns_series, events)
         return {"success": True, "result": "Interesting times tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_capacity_tear_sheet", description="Generate a tear sheet for strategy capacity analysis.")
 def generate_capacity_tear_sheet(returns: list, positions: dict) -> dict:
+    """
+    Generate a tear sheet for strategy capacity analysis.
+    
+    Args:
+        returns: List of daily returns
+        positions: Dictionary of positions over time
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_capacity_tear_sheet(returns, positions)
+        returns_series = _convert_to_series(returns)
+        positions_df = _convert_to_dataframe(positions)
+        create_capacity_tear_sheet(returns_series, positions_df)
         return {"success": True, "result": "Capacity tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_performance_attribution_tear_sheet", description="Generate a tear sheet for performance attribution analysis.")
 def generate_performance_attribution_tear_sheet(returns: list, factors: dict) -> dict:
+    """
+    Generate a tear sheet for performance attribution analysis.
+    
+    Args:
+        returns: List of daily returns
+        factors: Dictionary of factor returns
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        create_perf_attrib_tear_sheet(returns, factors)
+        returns_series = _convert_to_series(returns)
+        factors_df = _convert_to_dataframe(factors)
+        create_perf_attrib_tear_sheet(returns_series, factors_df)
         return {"success": True, "result": "Performance attribution tear sheet generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="plot_annual_returns", description="Plot annual returns as a bar chart.")
 def plot_annual_returns_tool(returns: list) -> dict:
+    """
+    Plot annual returns as a bar chart.
+    
+    Args:
+        returns: List of daily returns
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        ax = plot_annual_returns(returns)
-        return {"success": True, "result": ax, "error": None}
+        returns_series = _convert_to_series(returns)
+        ax = plot_annual_returns(returns_series)
+        return {"success": True, "result": "Annual returns plot generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="plot_monthly_returns_heatmap", description="Plot a heatmap of monthly returns.")
 def plot_monthly_returns_heatmap_tool(returns: list) -> dict:
+    """
+    Plot a heatmap of monthly returns.
+    
+    Args:
+        returns: List of daily returns
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        ax = plot_monthly_returns_heatmap(returns)
-        return {"success": True, "result": ax, "error": None}
+        returns_series = _convert_to_series(returns)
+        ax = plot_monthly_returns_heatmap(returns_series)
+        return {"success": True, "result": "Monthly returns heatmap generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="plot_drawdown_periods", description="Plot cumulative returns highlighting top drawdown periods.")
 def plot_drawdown_periods_tool(returns: list, top: int = 10) -> dict:
+    """
+    Plot cumulative returns highlighting top drawdown periods.
+    
+    Args:
+        returns: List of daily returns
+        top: Number of top drawdown periods to highlight
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        ax = plot_drawdown_periods(returns, top)
-        return {"success": True, "result": ax, "error": None}
+        returns_series = _convert_to_series(returns)
+        ax = plot_drawdown_periods(returns_series, top)
+        return {"success": True, "result": "Drawdown periods plot generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="plot_rolling_returns", description="Plot cumulative rolling returns versus benchmarks.")
 def plot_rolling_returns_tool(returns: list, factor_returns: list = None, live_start_date: str = None) -> dict:
+    """
+    Plot cumulative rolling returns versus benchmarks.
+    
+    Args:
+        returns: List of daily returns
+        factor_returns: List of factor returns (optional)
+        live_start_date: Start date for live trading period (optional)
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        ax = plot_rolling_returns(returns, factor_returns, live_start_date)
-        return {"success": True, "result": ax, "error": None}
+        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}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="plot_turnover", description="Plot turnover over time.")
 def plot_turnover_tool(returns: list, transactions: dict, positions: dict, turnover_denom: str = "AGB") -> dict:
+    """
+    Plot turnover over time.
+    
+    Args:
+        returns: List of daily returns
+        transactions: Dictionary of transactions
+        positions: Dictionary of positions over time
+        turnover_denom: Turnover denominator calculation method
+    
+    Returns:
+        Dictionary with success status and result/error message
+    """
     try:
-        ax = plot_turnover(returns, transactions, positions, turnover_denom)
-        return {"success": True, "result": ax, "error": None}
+        returns_series = _convert_to_series(returns)
+        transactions_df = _convert_to_dataframe(transactions)
+        positions_df = _convert_to_dataframe(positions)
+        ax = plot_turnover(returns_series, transactions_df, positions_df, turnover_denom)
+        return {"success": True, "result": "Turnover plot generated successfully.", "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}