Update qlib/mcp_output/mcp_plugin/mcp_service.py

qlib/mcp_output/mcp_plugin/mcp_service.py

@@ -7,33 +7,9 @@ import qlib
 data_dir = os.getenv("QLIB_DATA_DIR", "/app/qlib_data")
 print(f"Using Qlib data path: {data_dir}")
 
+# 初始化 qlib，使用 day_future 作为默认频率
 qlib.init(provider_uri=data_dir, region="cn")
 
-# 根据你的数据包实际目录结构调整
-instruments_path = os.path.join(data_dir, "instruments")  # 或 "cn_data/instruments"
-calendars_path = os.path.join(data_dir, "calendars")
-features_path = os.path.join(data_dir, "features")
-
-# 打印 instruments
-if os.path.exists(instruments_path):
-    print("\n=== Instruments directory ===")
-    for freq in os.listdir(instruments_path):
-        freq_path = os.path.join(instruments_path, freq)
-        if os.path.isdir(freq_path):
-            files = os.listdir(freq_path)
-            print(f"{freq}/ : {len(files)} files -> {files[:10]}{'...' if len(files) > 10 else ''}")
-else:
-    print("Instruments directory not found!")
-
-# 尝试列出 day 级标的
-try:
-    instruments = D.list_instruments(freq="day", market="cn")
-    print(f"\nNumber of day-frequency instruments: {len(instruments)}")
-    print(f"Sample instruments: {instruments[:10]}")
-except Exception as e:
-    print("Error listing day-frequency instruments:", e)
-
-
 # Path settings
 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)
@@ -55,13 +31,13 @@ from qlib.workflow.task.manage import TaskManager
 mcp = FastMCP("qlib_service")
 
 
-@mcp.tool(name="initialize_exchange", description="Initialize and return an Exchange object for backtesting or trading simulations.")
+@mcp.tool(name="initialize_exchange", description="Initialize and return an Exchange object for backtesting or trading simulations. Uses day_future frequency by default to match available data.")
 def initialize_exchange(
     exchange: Optional[str] = None,
-    freq: str = "day",
+    freq: str = "day_future",  # 修改默认值为 day_future
     start_time: Optional[str] = None,
     end_time: Optional[str] = None,
-    codes: str = "all",
+    codes: str = "csi300",  # 修改默认值为 csi300（更常用）
     subscribe_fields: Optional[List[str]] = None,
     open_cost: float = 0.0015,
     close_cost: float = 0.0025,
@@ -72,23 +48,32 @@ def initialize_exchange(
 ) -> dict:
     """
     Initialize and return an Exchange object for backtesting or trading simulations.
-
+    
     Parameters:
         exchange (Optional[str]): Existing exchange name or configuration.
-        freq (str): Frequency of data (e.g., 'day', 'minute').
+        freq (str): Frequency of data. Default is 'day_future' (available: 'day_future').
         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 (Optional[float]): Price movement limits.
-        deal_price (Optional[str]): Price configuration.
+        codes (str): Instruments string. Available options: 'all', 'csi300', 'csi500', 'csi800', 'csi1000', 'csiall'. Default: 'csi300'.
+        subscribe_fields (Optional[List[str]]): Data fields to subscribe to (e.g., ['$close', '$volume', '$open', '$high', '$low']).
+        open_cost (float): Open transaction cost as a ratio (default: 0.0015 = 0.15%).
+        close_cost (float): Close transaction cost as a ratio (default: 0.0025 = 0.25%).
+        min_cost (float): Minimum transaction cost in currency units (default: 5.0).
+        limit_threshold (Optional[float]): Price movement limits (e.g., 0.095 for 9.5% limit).
+        deal_price (Optional[str]): Price configuration for order execution.
         extra_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments as a dictionary.
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, result, or error information.
+        
+    Example:
+        initialize_exchange(
+            freq="day_future",
+            start_time="2020-01-01",
+            end_time="2021-12-31",
+            codes="csi300",
+            subscribe_fields=["$close", "$volume"]
+        )
     """
     try:
         # Handle None defaults
@@ -97,6 +82,14 @@ def initialize_exchange(
         if extra_kwargs is None:
             extra_kwargs = {}
         
+        # 验证 codes 参数
+        valid_codes = ['all', 'csi300', 'csi500', 'csi800', 'csi1000', 'csiall']
+        if codes not in valid_codes:
+            return {
+                "success": False, 
+                "error": f"Invalid codes parameter. Available options: {', '.join(valid_codes)}"
+            }
+        
         exchange_obj = get_exchange(
             exchange=exchange,
             freq=freq,
@@ -111,7 +104,11 @@ def initialize_exchange(
             deal_price=deal_price,
             **extra_kwargs,
         )
-        return {"success": True, "result": str(exchange_obj), "message": "Exchange initialized successfully"}
+        return {
+            "success": True, 
+            "result": str(exchange_obj), 
+            "message": f"Exchange initialized successfully with freq={freq}, codes={codes}"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -126,16 +123,24 @@ def create_account(
 ) -> dict:
     """
     Create and initialize an Account instance for trading simulations.
-
+    
     Parameters:
         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 (float): Initial cash amount.
+        benchmark (str): Benchmark for performance comparison (e.g., 'SH000300' for CSI300, 'SH000905' for CSI500).
+        account (float): Initial cash amount (e.g., 1000000 for 1 million).
         pos_type (str): Type of position to use (default: "Position").
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, result, or error information.
+        
+    Example:
+        create_account(
+            start_time="2020-01-01",
+            end_time="2021-12-31",
+            benchmark="SH000300",
+            account=1000000.0
+        )
     """
     try:
         account_obj = create_account_instance(
@@ -145,7 +150,11 @@ def create_account(
             account=account,
             pos_type=pos_type,
         )
-        return {"success": True, "result": str(account_obj), "message": "Account created successfully"}
+        return {
+            "success": True, 
+            "result": str(account_obj), 
+            "message": f"Account created successfully with initial capital: {account}"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -163,19 +172,22 @@ def initialize_strategy_executor(
 ) -> dict:
     """
     Initialize and configure a trading strategy and its executor.
-
+    
     Parameters:
         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.
+        strategy (dict): Strategy configuration with class and module info.
+            Example: {"class": "TopkDropoutStrategy", "module_path": "qlib.contrib.strategy", "kwargs": {"topk": 30, "n_drop": 5}}
         executor (dict): Executor configuration.
-        benchmark (str): Benchmark identifier.
-        account (float): Initial cash amount.
-        exchange_kwargs (dict): Exchange-specific settings.
+            Example: {"class": "SimulatorExecutor", "module_path": "qlib.backtest.executor", "kwargs": {"time_per_step": "day"}}
+        benchmark (str): Benchmark identifier (e.g., 'SH000300' for CSI300).
+        account (float): Initial cash amount (e.g., 1000000).
+        exchange_kwargs (dict): Exchange-specific settings. Should include 'freq' and 'codes'.
+            Example: {"freq": "day_future", "codes": "csi300", "start_time": "2020-01-01", "end_time": "2021-12-31"}
         pos_type (str): Type of position to use (default: "Position").
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, result, or error information.
     """
     try:
         strategy_executor = get_strategy_executor(
@@ -188,7 +200,11 @@ def initialize_strategy_executor(
             exchange_kwargs=exchange_kwargs,
             pos_type=pos_type,
         )
-        return {"success": True, "result": str(strategy_executor), "message": "Strategy executor initialized successfully"}
+        return {
+            "success": True, 
+            "result": str(strategy_executor), 
+            "message": "Strategy executor initialized successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -206,21 +222,39 @@ def run_backtest(
 ) -> dict:
     """
     Perform a backtest to evaluate a trading strategy.
-
+    
     Parameters:
         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.
+            Example: {"class": "TopkDropoutStrategy", "module_path": "qlib.contrib.strategy", "kwargs": {"topk": 30, "n_drop": 5}}
         executor (dict): Executor configuration.
-        benchmark (str): Benchmark identifier.
-        account (float): Initial cash amount.
-        exchange_kwargs (dict): Exchange-specific settings.
+            Example: {"class": "SimulatorExecutor", "module_path": "qlib.backtest.executor", "kwargs": {"time_per_step": "day"}}
+        benchmark (str): Benchmark identifier (e.g., 'SH000300' for CSI300, 'SH000905' for CSI500).
+        account (float): Initial cash amount (e.g., 1000000).
+        exchange_kwargs (dict): Exchange-specific settings. Must include 'freq' as 'day_future' and valid 'codes'.
+            Example: {"freq": "day_future", "codes": "csi300", "open_cost": 0.0015, "close_cost": 0.0025, "min_cost": 5}
         pos_type (str): Type of position to use (default: "Position").
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, portfolio metrics, trading indicators, or error information.
+        
+    Example:
+        run_backtest(
+            start_time="2020-01-01",
+            end_time="2021-12-31",
+            strategy={"class": "TopkDropoutStrategy", "module_path": "qlib.contrib.strategy", "kwargs": {"topk": 30}},
+            executor={"class": "SimulatorExecutor", "module_path": "qlib.backtest.executor"},
+            benchmark="SH000300",
+            account=1000000,
+            exchange_kwargs={"freq": "day_future", "codes": "csi300"}
+        )
     """
     try:
+        # 确保 exchange_kwargs 中包含正确的 freq
+        if 'freq' not in exchange_kwargs:
+            exchange_kwargs['freq'] = 'day_future'
+        
         portfolio_metrics, trading_indicators = backtest(
             start_time=start_time,
             end_time=end_time,
@@ -238,7 +272,11 @@ def run_backtest(
             "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"}
+        return {
+            "success": True, 
+            "result": result, 
+            "message": "Backtest completed successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -256,7 +294,7 @@ def collect_trade_data(
 ) -> dict:
     """
     Collect trade decision data for reinforcement learning training.
-
+    
     Parameters:
         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').
@@ -264,13 +302,17 @@ def collect_trade_data(
         executor (dict): Executor configuration.
         benchmark (str): Benchmark identifier.
         account (float): Initial cash amount.
-        exchange_kwargs (dict): Exchange-specific settings.
+        exchange_kwargs (dict): Exchange-specific settings. Should use 'freq': 'day_future'.
         pos_type (str): Type of position to use (default: "Position").
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, collected data, or error information.
     """
     try:
+        # 确保使用正确的 freq
+        if 'freq' not in exchange_kwargs:
+            exchange_kwargs['freq'] = 'day_future'
+            
         data_generator = collect_data(
             start_time=start_time,
             end_time=end_time,
@@ -283,7 +325,11 @@ def collect_trade_data(
             return_value=None,
         )
         data = list(data_generator)
-        return {"success": True, "result": [str(item) for item in data], "message": "Trade data collected successfully"}
+        return {
+            "success": True, 
+            "result": [str(item) for item in data], 
+            "message": f"Trade data collected successfully. Total records: {len(data)}"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -292,16 +338,64 @@ def collect_trade_data(
 def format_trade_decisions(decisions: List[str]) -> dict:
     """
     Format trade decisions into a hierarchical structure.
-
+    
     Parameters:
         decisions (List[str]): List of trade decisions as strings.
-
+    
     Returns:
-        dict: A dictionary containing success, result, or error fields.
+        dict: A dictionary containing success status, formatted decisions, or error information.
     """
     try:
         formatted_decisions = format_decisions(decisions)
-        return {"success": True, "result": str(formatted_decisions), "message": "Decisions formatted successfully"}
+        return {
+            "success": True, 
+            "result": str(formatted_decisions), 
+            "message": "Decisions formatted successfully"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="get_available_instruments", description="Get information about available instrument pools in the dataset.")
+def get_available_instruments() -> dict:
+    """
+    Get information about available instrument pools in the dataset.
+    
+    Returns:
+        dict: A dictionary containing available instrument pools and their descriptions.
+    """
+    try:
+        instruments_info = {
+            "available_pools": [
+                "all",
+                "csi300",  # 沪深300指数成分股
+                "csi500",  # 中证500指数成分股
+                "csi800",  # 中证800指数成分股
+                "csi1000", # 中证1000指数成分股
+                "csiall"   # 所有中证指数成分股
+            ],
+            "descriptions": {
+                "all": "All available stocks in the dataset",
+                "csi300": "CSI 300 Index constituents (top 300 stocks by market cap)",
+                "csi500": "CSI 500 Index constituents (mid-cap stocks)",
+                "csi800": "CSI 800 Index constituents (CSI 300 + CSI 500)",
+                "csi1000": "CSI 1000 Index constituents (small-cap stocks)",
+                "csiall": "All CSI Index constituents"
+            },
+            "benchmarks": {
+                "csi300": "SH000300",
+                "csi500": "SH000905",
+                "csi800": "SH000906",
+                "csi1000": "SH000852"
+            },
+            "frequency": "day_future",
+            "data_directory": data_dir
+        }
+        return {
+            "success": True, 
+            "result": instruments_info, 
+            "message": "Available instruments information retrieved successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
@@ -309,7 +403,7 @@ def format_trade_decisions(decisions: List[str]) -> dict:
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.
-
+    
     Returns:
         FastMCP: The FastMCP application instance.
     """