Update qlib/mcp_output/mcp_plugin/mcp_service.py

qlib/mcp_output/mcp_plugin/mcp_service.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Union, List, Dict, Any, Tuple, Generator
+from typing import Union, List, Dict, Any, Tuple, Generator, Optional
 import pandas as pd
 
 # Path settings
@@ -26,41 +26,43 @@ mcp = FastMCP("qlib_service")
 
 @mcp.tool(name="initialize_exchange", description="Initialize and return an Exchange object for backtesting or trading simulations.")
 def initialize_exchange(
-    exchange: Union[str, dict, object, os.PathLike] = None,
+    exchange: Optional[str] = None,
     freq: str = "day",
-    start_time: Union[str, pd.Timestamp] = None,
-    end_time: Union[str, pd.Timestamp] = None,
-    codes: Union[List[str], str] = "all",
-    subscribe_fields: List[str] = [],
+    start_time: Optional[str] = None,
+    end_time: Optional[str] = None,
+    codes: str = "all",
+    subscribe_fields: Optional[List[str]] = None,
     open_cost: float = 0.0015,
     close_cost: float = 0.0025,
     min_cost: float = 5.0,
-    limit_threshold: Union[Tuple[str, str], float, None] = None,
-    deal_price: Union[str, Tuple[str, str], List[str]] = None,
-    extra_kwargs: Dict[str, Any] = None,
+    limit_threshold: Optional[float] = None,
+    deal_price: Optional[str] = None,
+    extra_kwargs: Optional[Dict[str, Any]] = None,
 ) -> dict:
     """
     Initialize and return an Exchange object for backtesting or trading simulations.
 
     Parameters:
-        exchange (Union[str, dict, object, os.PathLike]): Existing exchange or configuration.
+        exchange (Optional[str]): Existing exchange name or configuration.
         freq (str): Frequency of data (e.g., 'day', 'minute').
-        start_time (Union[str, pd.Timestamp]): Start time for the exchange.
-        end_time (Union[str, pd.Timestamp]): End time for the exchange.
-        codes (Union[List[str], str]): List of instruments or a string (e.g., 'all', 'csi500').
-        subscribe_fields (List[str]): Data fields to subscribe to.
+        start_time (Optional[str]): Start time for the exchange (e.g., '2020-01-01').
+        end_time (Optional[str]): End time for the exchange (e.g., '2021-01-01').
+        codes (str): Instruments string (e.g., 'all', 'csi500').
+        subscribe_fields (Optional[List[str]]): Data fields to subscribe to.
         open_cost (float): Open transaction cost as a ratio.
         close_cost (float): Close transaction cost as a ratio.
         min_cost (float): Minimum transaction cost.
-        limit_threshold (Union[Tuple[str, str], float, None]): Price movement limits.
-        deal_price (Union[str, Tuple[str, str], List[str]]): Price configuration.
-        extra_kwargs (Dict[str, Any]): Additional keyword arguments as a dictionary.
+        limit_threshold (Optional[float]): Price movement limits.
+        deal_price (Optional[str]): Price configuration.
+        extra_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments as a dictionary.
 
     Returns:
         dict: A dictionary containing success, result, or error fields.
     """
     try:
-        # Handle extra_kwargs
+        # Handle None defaults
+        if subscribe_fields is None:
+            subscribe_fields = []
         if extra_kwargs is None:
             extra_kwargs = {}
         
@@ -78,27 +80,27 @@ def initialize_exchange(
             deal_price=deal_price,
             **extra_kwargs,
         )
-        return {"success": True, "result": str(exchange_obj)}
+        return {"success": True, "result": str(exchange_obj), "message": "Exchange initialized successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
 
 @mcp.tool(name="create_account", description="Create and initialize an Account instance for trading simulations.")
 def create_account(
-    start_time: Union[str, pd.Timestamp],
-    end_time: Union[str, pd.Timestamp],
+    start_time: str,
+    end_time: str,
     benchmark: str,
-    account: Union[float, int, dict],
+    account: float,
     pos_type: str = "Position",
 ) -> dict:
     """
     Create and initialize an Account instance for trading simulations.
 
     Parameters:
-        start_time (Union[str, pd.Timestamp]): Start time of the benchmark.
-        end_time (Union[str, pd.Timestamp]): End time of the benchmark.
+        start_time (str): Start time of the benchmark (e.g., '2020-01-01').
+        end_time (str): End time of the benchmark (e.g., '2021-01-01').
         benchmark (str): Benchmark for reporting.
-        account (Union[float, int, dict]): Initial cash and/or stock positions.
+        account (float): Initial cash amount.
         pos_type (str): Type of position to use (default: "Position").
 
     Returns:
@@ -112,19 +114,19 @@ def create_account(
             account=account,
             pos_type=pos_type,
         )
-        return {"success": True, "result": str(account_obj)}
+        return {"success": True, "result": str(account_obj), "message": "Account created successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
 
 @mcp.tool(name="initialize_strategy_executor", description="Initialize and configure a trading strategy and its executor.")
 def initialize_strategy_executor(
-    start_time: Union[str, pd.Timestamp],
-    end_time: Union[str, pd.Timestamp],
+    start_time: str,
+    end_time: str,
     strategy: dict,
     executor: dict,
     benchmark: str,
-    account: Union[float, int, dict],
+    account: float,
     exchange_kwargs: dict,
     pos_type: str = "Position",
 ) -> dict:
@@ -132,12 +134,12 @@ def initialize_strategy_executor(
     Initialize and configure a trading strategy and its executor.
 
     Parameters:
-        start_time (Union[str, pd.Timestamp]): Start time for the strategy.
-        end_time (Union[str, pd.Timestamp]): End time for the strategy.
+        start_time (str): Start time for the strategy (e.g., '2020-01-01').
+        end_time (str): End time for the strategy (e.g., '2021-01-01').
         strategy (dict): Strategy configuration.
         executor (dict): Executor configuration.
         benchmark (str): Benchmark identifier.
-        account (Union[float, int, dict]): Account details.
+        account (float): Initial cash amount.
         exchange_kwargs (dict): Exchange-specific settings.
         pos_type (str): Type of position to use (default: "Position").
 
@@ -155,19 +157,19 @@ def initialize_strategy_executor(
             exchange_kwargs=exchange_kwargs,
             pos_type=pos_type,
         )
-        return {"success": True, "result": str(strategy_executor)}
+        return {"success": True, "result": str(strategy_executor), "message": "Strategy executor initialized successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
 
 @mcp.tool(name="run_backtest", description="Perform a backtest to evaluate a trading strategy.")
 def run_backtest(
-    start_time: Union[str, pd.Timestamp],
-    end_time: Union[str, pd.Timestamp],
+    start_time: str,
+    end_time: str,
     strategy: dict,
     executor: dict,
     benchmark: str,
-    account: Union[float, int, dict],
+    account: float,
     exchange_kwargs: dict,
     pos_type: str = "Position",
 ) -> dict:
@@ -175,12 +177,12 @@ def run_backtest(
     Perform a backtest to evaluate a trading strategy.
 
     Parameters:
-        start_time (Union[str, pd.Timestamp]): Start time for the backtest.
-        end_time (Union[str, pd.Timestamp]): End time for the backtest.
+        start_time (str): Start time for the backtest (e.g., '2020-01-01').
+        end_time (str): End time for the backtest (e.g., '2021-01-01').
         strategy (dict): Strategy configuration.
         executor (dict): Executor configuration.
         benchmark (str): Benchmark identifier.
-        account (Union[float, int, dict]): Account details.
+        account (float): Initial cash amount.
         exchange_kwargs (dict): Exchange-specific settings.
         pos_type (str): Type of position to use (default: "Position").
 
@@ -198,42 +200,41 @@ def run_backtest(
             exchange_kwargs=exchange_kwargs,
             pos_type=pos_type,
         )
-        return {
-            "success": True, 
-            "result": {
-                "portfolio_metrics": portfolio_metrics.to_dict() if hasattr(portfolio_metrics, 'to_dict') else str(portfolio_metrics),
-                "trading_indicators": trading_indicators.to_dict() if hasattr(trading_indicators, 'to_dict') else str(trading_indicators)
-            }
+        
+        # Convert to serializable format
+        result = {
+            "portfolio_metrics": portfolio_metrics.to_dict() if hasattr(portfolio_metrics, 'to_dict') else str(portfolio_metrics),
+            "trading_indicators": trading_indicators.to_dict() if hasattr(trading_indicators, 'to_dict') else str(trading_indicators)
         }
+        
+        return {"success": True, "result": result, "message": "Backtest completed successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
 
 @mcp.tool(name="collect_trade_data", description="Collect trade decision data for reinforcement learning training.")
 def collect_trade_data(
-    start_time: Union[str, pd.Timestamp],
-    end_time: Union[str, pd.Timestamp],
+    start_time: str,
+    end_time: str,
     strategy: dict,
     executor: dict,
     benchmark: str,
-    account: Union[float, int, dict],
+    account: float,
     exchange_kwargs: dict,
     pos_type: str = "Position",
-    return_value: Any = None,
 ) -> dict:
     """
     Collect trade decision data for reinforcement learning training.
 
     Parameters:
-        start_time (Union[str, pd.Timestamp]): Start time for data collection.
-        end_time (Union[str, pd.Timestamp]): End time for data collection.
+        start_time (str): Start time for data collection (e.g., '2020-01-01').
+        end_time (str): End time for data collection (e.g., '2021-01-01').
         strategy (dict): Strategy configuration.
         executor (dict): Executor configuration.
         benchmark (str): Benchmark identifier.
-        account (Union[float, int, dict]): Account details.
+        account (float): Initial cash amount.
         exchange_kwargs (dict): Exchange-specific settings.
         pos_type (str): Type of position to use (default: "Position").
-        return_value (Any): Optional container for return values.
 
     Returns:
         dict: A dictionary containing success, result, or error fields.
@@ -248,28 +249,28 @@ def collect_trade_data(
             account=account,
             exchange_kwargs=exchange_kwargs,
             pos_type=pos_type,
-            return_value=return_value,
+            return_value=None,
         )
         data = list(data_generator)
-        return {"success": True, "result": [str(item) for item in data]}
+        return {"success": True, "result": [str(item) for item in data], "message": "Trade data collected successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
 
 @mcp.tool(name="format_trade_decisions", description="Format trade decisions into a hierarchical structure.")
-def format_trade_decisions(decisions: List[Any]) -> dict:
+def format_trade_decisions(decisions: List[str]) -> dict:
     """
     Format trade decisions into a hierarchical structure.
 
     Parameters:
-        decisions (List[Any]): List of trade decisions.
+        decisions (List[str]): List of trade decisions as strings.
 
     Returns:
         dict: A dictionary containing success, result, or error fields.
     """
     try:
         formatted_decisions = format_decisions(decisions)
-        return {"success": True, "result": str(formatted_decisions)}
+        return {"success": True, "result": str(formatted_decisions), "message": "Decisions formatted successfully"}
     except Exception as e:
         return {"success": False, "error": str(e)}