Update pyfolio/mcp_output/mcp_plugin/mcp_service.py

pyfolio/mcp_output/mcp_plugin/mcp_service.py

@@ -2,11 +2,14 @@ import os
 import sys
 import pandas as pd
 from typing import Optional, Dict, List, Any
+import io
+import contextlib
 
 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)
 
 from fastmcp import FastMCP
+from pyfolio import timeseries
 from pyfolio.tears import (
     create_full_tear_sheet,
     create_simple_tear_sheet,
@@ -44,46 +47,122 @@ def _convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
         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:
+def _calculate_performance_stats(returns_series: pd.Series) -> Dict[str, float]:
+    """Calculate key performance statistics."""
+    try:
+        stats = {}
+        stats['total_return'] = timeseries.cum_returns_final(returns_series)
+        stats['annual_return'] = timeseries.annual_return(returns_series)
+        stats['annual_volatility'] = timeseries.annual_volatility(returns_series)
+        stats['sharpe_ratio'] = timeseries.sharpe_ratio(returns_series)
+        stats['max_drawdown'] = timeseries.max_drawdown(returns_series)
+        stats['calmar_ratio'] = timeseries.calmar_ratio(returns_series)
+        stats['stability'] = timeseries.stability_of_timeseries(returns_series)
+        stats['omega_ratio'] = timeseries.omega_ratio(returns_series)
+        stats['sortino_ratio'] = timeseries.sortino_ratio(returns_series)
+        stats['skew'] = timeseries.stats.skew(returns_series)
+        stats['kurtosis'] = timeseries.stats.kurtosis(returns_series)
+        stats['tail_ratio'] = timeseries.tail_ratio(returns_series)
+        
+        # Convert numpy types to Python types for JSON serialization
+        return {k: float(v) if pd.notna(v) else None for k, v in stats.items()}
+    except Exception as e:
+        return {"error": str(e)}
+
+@mcp.tool(name="calculate_statistics", description="Calculate key performance statistics for a returns series.")
+def calculate_statistics(returns: list) -> dict:
     """
-    Generate a comprehensive tear sheet for portfolio analysis.
+    Calculate key performance statistics for a returns series.
     
     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
+        Dictionary with performance statistics including total return, Sharpe ratio, max drawdown, etc.
     """
     try:
-        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
+        returns_series = _convert_to_series(returns)
+        stats = _calculate_performance_stats(returns_series)
         
-        create_full_tear_sheet(returns_series, positions_df, transactions_df, benchmark_series)
-        return {"success": True, "result": "Full tear sheet generated successfully.", "error": None}
+        return {
+            "success": True,
+            "result": stats,
+            "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.")
+@mcp.tool(name="generate_simple_tear_sheet", description="Generate a basic tear sheet with performance statistics.")
 def generate_simple_tear_sheet(returns: list) -> dict:
     """
-    Generate a basic tear sheet for portfolio analysis.
+    Generate a basic tear sheet for portfolio analysis with performance statistics.
     
     Args:
         returns: List of daily returns
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status and performance statistics
     """
     try:
         returns_series = _convert_to_series(returns)
-        create_simple_tear_sheet(returns_series)
-        return {"success": True, "result": "Simple tear sheet generated successfully.", "error": None}
+        
+        # Calculate performance statistics
+        stats = _calculate_performance_stats(returns_series)
+        
+        # Capture tear sheet output
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_simple_tear_sheet(returns_series)
+        
+        return {
+            "success": True, 
+            "result": {
+                "statistics": stats,
+                "message": "Simple tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@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 performance statistics
+    """
+    try:
+        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
+        
+        # Calculate performance statistics
+        stats = _calculate_performance_stats(returns_series)
+        
+        # Capture tear sheet output
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_full_tear_sheet(returns_series, positions_df, transactions_df, benchmark_series)
+        
+        return {
+            "success": True, 
+            "result": {
+                "statistics": stats,
+                "message": "Full tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None  # Limit output
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -96,12 +175,28 @@ def generate_returns_tear_sheet(returns: list) -> dict:
         returns: List of daily returns
     
     Returns:
-        Dictionary with success status and result/error message
+        Dictionary with success status and performance statistics
     """
     try:
         returns_series = _convert_to_series(returns)
-        create_returns_tear_sheet(returns_series)
-        return {"success": True, "result": "Returns tear sheet generated successfully.", "error": None}
+        
+        # Calculate performance statistics
+        stats = _calculate_performance_stats(returns_series)
+        
+        # Capture tear sheet output
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_returns_tear_sheet(returns_series)
+        
+        return {
+            "success": True, 
+            "result": {
+                "statistics": stats,
+                "message": "Returns tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -118,8 +213,19 @@ def generate_position_tear_sheet(positions: dict) -> dict:
     """
     try:
         positions_df = _convert_to_dataframe(positions)
-        create_position_tear_sheet(positions_df)
-        return {"success": True, "result": "Position tear sheet generated successfully.", "error": None}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_position_tear_sheet(positions_df)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Position tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -136,8 +242,19 @@ def generate_transaction_tear_sheet(transactions: dict) -> dict:
     """
     try:
         transactions_df = _convert_to_dataframe(transactions)
-        create_txn_tear_sheet(transactions_df)
-        return {"success": True, "result": "Transaction tear sheet generated successfully.", "error": None}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_txn_tear_sheet(transactions_df)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Transaction tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -154,8 +271,19 @@ def generate_round_trip_tear_sheet(round_trips: dict) -> dict:
     """
     try:
         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}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_round_trip_tear_sheet(round_trips_df)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Round trip tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -173,8 +301,19 @@ def generate_interesting_times_tear_sheet(returns: list, events: list) -> dict:
     """
     try:
         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}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_interesting_times_tear_sheet(returns_series, events)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Interesting times tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -193,8 +332,19 @@ def generate_capacity_tear_sheet(returns: list, positions: dict) -> dict:
     try:
         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}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_capacity_tear_sheet(returns_series, positions_df)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Capacity tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
@@ -213,8 +363,19 @@ def generate_performance_attribution_tear_sheet(returns: list, factors: dict) ->
     try:
         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}
+        
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            create_perf_attrib_tear_sheet(returns_series, factors_df)
+        
+        return {
+            "success": True, 
+            "result": {
+                "message": "Performance attribution tear sheet generated successfully.",
+                "tear_sheet_output": output.getvalue()[:1000] if output.getvalue() else None
+            }, 
+            "error": None
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}