server.py

#!/usr/bin/env python3
"""
Qlib MCP Server
Microsoft Qlib 量化研究平台的 MCP 接口
提供数据查询、因子计算、策略回测等工具
"""

import os
import sys
import json
import warnings
from typing import Optional, Any
from datetime import datetime, timedelta

# 压制 gym 警告
warnings.filterwarnings("ignore")
os.environ["GYM_DISABLE_WARNINGS"] = "1"

from mcp.server.fastmcp import FastMCP

mcp = FastMCP(
    "qlib",
    instructions="Microsoft Qlib 量化研究平台。提供 A 股/美股历史数据查询、因子计算、策略回测和绩效评估。",
)

# ─────────────────────────────────────────────
# 全局状态
# ─────────────────────────────────────────────

_qlib_initialized = False
_qlib_region = None
_qlib_data_path = None


def _ensure_initialized():
    """确保 Qlib 已初始化，否则返回错误提示"""
    if not _qlib_initialized:
        raise RuntimeError(
            "Qlib 尚未初始化，请先调用 qlib_init 工具，指定数据路径和区域（cn/us）。"
        )


def _safe_df_to_dict(df, max_rows: int = 200) -> dict:
    """将 DataFrame 安全转换为可序列化字典"""
    if df is None:
        return {}
    if hasattr(df, "reset_index"):
        df = df.reset_index()
    if len(df) > max_rows:
        df = df.tail(max_rows)
    # 处理 MultiIndex columns
    if hasattr(df.columns, "levels"):
        df.columns = ["_".join(str(v) for v in col).strip("_") for col in df.columns]
    # 处理 Timestamp 列
    for col in df.columns:
        if hasattr(df[col], "dt"):
            df[col] = df[col].astype(str)
    return {
        "shape": list(df.shape),
        "columns": list(df.columns.astype(str)),
        "data": df.to_dict(orient="records"),
        "truncated": len(df) == max_rows,
    }


# ─────────────────────────────────────────────
# 工具 1：初始化
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_init(
    data_path: str,
    region: str = "cn",
    provider_uri: Optional[str] = None,
) -> str:
    """
    初始化 Qlib 数据环境。必须在使用其他工具前调用。

    Args:
        data_path: 本地 Qlib 数据目录路径，例如 ~/.qlib/qlib_data/cn_data
        region: 市场区域，"cn"（A 股）或 "us"（美股），默认 "cn"
        provider_uri: 可选，覆盖 provider_uri（通常与 data_path 相同）

    Returns:
        初始化结果信息
    """
    global _qlib_initialized, _qlib_region, _qlib_data_path

    import qlib
    from qlib.config import REG_CN, REG_US

    data_path = os.path.expanduser(data_path)
    if not os.path.exists(data_path):
        return f"❌ 数据路径不存在：{data_path}\n提示：可以用 qlib_download_data 工具下载数据，或先创建目录。"

    uri = provider_uri or data_path
    reg = REG_CN if region.lower() == "cn" else REG_US

    try:
        qlib.init(provider_uri=uri, region=reg)
        _qlib_initialized = True
        _qlib_region = region.lower()
        _qlib_data_path = data_path

        return f"✅ Qlib 初始化成功\n- 区域：{region.upper()}\n- 数据路径：{data_path}\n- 版本：{qlib.__version__}"
    except Exception as e:
        return f"❌ 初始化失败：{e}"


# ─────────────────────────────────────────────
# 工具 2：下载数据
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_download_data(
    target_dir: str = "~/.qlib/qlib_data/cn_data",
    region: str = "cn",
    interval: str = "1d",
) -> str:
    """
    下载 Qlib 示例数据集到本地（约 500MB~2GB）。

    Args:
        target_dir: 数据保存目录，默认 ~/.qlib/qlib_data/cn_data
        region: 区域 "cn"（A 股）或 "us"（美股），默认 "cn"
        interval: 数据频率，"1d"（日线）或 "1min"（分钟线），默认 "1d"

    Returns:
        下载命令和状态
    """
    target_dir = os.path.expanduser(target_dir)
    os.makedirs(target_dir, exist_ok=True)

    python_bin = sys.executable
    cmd = (
        f"{python_bin} -m qlib.run.get_data qlib_data "
        f"--target_dir {target_dir} "
        f"--region {region} "
        f"--interval {interval}"
    )

    return (
        f"📦 数据下载命令（在终端执行）：\n```bash\n{cmd}\n```\n"
        f"- 目标目录：{target_dir}\n"
        f"- 区域：{region.upper()}，频率：{interval}\n"
        f"- 下载完成后，调用 qlib_init(data_path='{target_dir}') 初始化\n"
        f"⚠️ 下载需要较长时间（取决于网速），建议在终端后台运行。"
    )


# ─────────────────────────────────────────────
# 工具 3：股票列表 & 基本信息
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_list_instruments(
    market: str = "csi300",
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
) -> str:
    """
    获取指定市场/指数的成分股列表。

    Args:
        market: 市场/指数名称。A 股常用：all（全部）、csi300（沪深 300）、csi500、
                csi100、sz50（上证 50）。美股：sp500、nasdaq100 等。
        start_date: 开始日期 YYYY-MM-DD，默认今天
        end_date: 结束日期 YYYY-MM-DD，默认今天

    Returns:
        成分股代码列表
    """
    _ensure_initialized()

    from qlib.data import D

    today = datetime.now().strftime("%Y-%m-%d")
    start = start_date or today
    end = end_date or today

    try:
        instruments = D.list_instruments(
            instruments=D.instruments(market),
            start_time=start,
            end_time=end,
            as_list=True,
        )

        if not instruments:
            return f"⚠️ 未找到 {market} 在 {start}~{end} 的成分股数据"

        sample = instruments[:20]
        return (
            f"📊 {market.upper()} 成分股（共 {len(instruments)} 只）\n"
            f"前 20 只：{', '.join(sample)}\n"
            f"{'...' if len(instruments) > 20 else ''}"
        )
    except Exception as e:
        return f"❌ 获取成分股失败：{e}"


# ─────────────────────────────────────────────
# 工具 4：价格/因子数据查询
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_get_data(
    symbols: str,
    fields: str = "$close,$open,$high,$low,$volume,$factor",
    start_date: str = "2024-01-01",
    end_date: Optional[str] = None,
    freq: str = "day",
    max_rows: int = 100,
) -> str:
    """
    查询股票历史价格/因子数据。

    Args:
        symbols: 股票代码，逗号分隔，A 股格式如 "SH600000,SZ000001"，
                 或市场名 "csi300"。
        fields: 数据字段，逗号分隔。常用：$close（收盘价）、$open（��盘价）、
                $high（最高价）、$low（最低价）、$volume（成交量）、
                $factor（复权因子）、$change（涨跌幅）、$vwap（均价）。
                也支持表达式如 "Ref($close,1)/Div($close,1)-1"。
        start_date: 开始日期 YYYY-MM-DD
        end_date: 结束日期 YYYY-MM-DD，默认今天
        freq: 频率，"day" 或 "1min"
        max_rows: 返回最大行数，默认 100

    Returns:
        价格数据（JSON 格式）
    """
    _ensure_initialized()

    from qlib.data import D

    today = datetime.now().strftime("%Y-%m-%d")
    end = end_date or today

    # 解析标的
    syms = [s.strip() for s in symbols.split(",")]
    if len(syms) == 1 and not any(s in syms[0] for s in ["SH", "SZ", "."]):
        # 可能是市场名
        try:
            instrument = D.instruments(syms[0])
        except Exception:
            instrument = syms
    else:
        instrument = syms

    # 解析字段
    field_list = [f.strip() for f in fields.split(",")]

    try:
        df = D.features(
            instruments=instrument,
            fields=field_list,
            start_time=start_date,
            end_time=end,
            freq=freq,
        )

        if df is None or df.empty:
            return f"⚠️ ���找到数据：{symbols}（{start_date} ~ {end}）"

        result = _safe_df_to_dict(df, max_rows=max_rows)
        return json.dumps(result, ensure_ascii=False, default=str)

    except Exception as e:
        return f"❌ 数据查询失败：{e}"


# ─────────────────────────────────────────────
# 工具 5：简单策略回测
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_backtest_topk(
    market: str = "csi300",
    start_date: str = "2020-01-01",
    end_date: str = "2023-12-31",
    topk: int = 50,
    n_drop: int = 5,
    hold_thresh: float = 1.0,
    signal_field: str = "$close / Ref($close, 20) - 1",
    benchmark: str = "SH000300",
    initial_cash: float = 1_000_000,
) -> str:
    """
    基于动量或自定义信号的 TopK 持仓策略回测。

    Args:
        market: 股票池，如 "csi300"、"csi500"
        start_date: 回测开始日期 YYYY-MM-DD
        end_date: 回测结束日期 YYYY-MM-DD
        topk: 持仓只数（信号最高的 K 只）
        n_drop: 每次调仓最多换出的只数
        hold_thresh: 持仓阈值（0~1），信号超过该分位数才���入
        signal_field: 选股信号表达式（Qlib 表达式语法）。
                      示例：
                      - "$close / Ref($close, 20) - 1"（20 日动量）
                      - "Mean($volume, 5) / Mean($volume, 20)"（短期放量比）
                      - "-$pe"（低 PE 价值策略，需有 PE 数据）
        benchmark: 基准指数代码，默认沪深 300（SH000300）
        initial_cash: 初始资金，默认 100 万

    Returns:
        回测绩效摘要（年化收益、夏普、最大回撤等）
    """
    _ensure_initialized()

    try:
        from qlib.data import D
        from qlib.backtest import backtest
        from qlib.contrib.strategy import TopkDropoutStrategy
        from qlib.contrib.evaluate import (
            risk_analysis,
            backtest_daily,
        )

        # 生成信号分数
        instruments = D.instruments(market)
        score_df = D.features(
            instruments=instruments,
            fields=[signal_field],
            start_time=start_date,
            end_time=end_date,
        )
        score_df.columns = ["score"]
        score_df = score_df.dropna()

        # 策略配置
        strategy = TopkDropoutStrategy(
            signal=score_df["score"],
            topk=topk,
            n_drop=n_drop,
        )

        # 执行回测
        portfolio_metrics, positions = backtest_daily(
            start_time=start_date,
            end_time=end_date,
            strategy=strategy,
        )

        # 绩效分析
        analysis = risk_analysis(portfolio_metrics["return"] - portfolio_metrics["bench"])

        result = {
            "策略": f"TopK={topk} 信号={signal_field[:40]}...",
            "股票池": market,
            "回测区间": f"{start_date} ~ {end_date}",
            "年化收益": f"{portfolio_metrics.get('annualized_return', 0):.2%}",
            "基准收益": f"{portfolio_metrics.get('bench', {}).get('annualized_return', 0):.2%}",
            "超额收益": f"{analysis.get('mean', 0) * 252:.2%}",
            "夏普比率": f"{analysis.get('information_ratio', 0):.3f}",
            "最大回撤": f"{portfolio_metrics.get('max_drawdown', 0):.2%}",
            "信息比率": f"{analysis.get('information_ratio', 0):.3f}",
        }

        return json.dumps(result, ensure_ascii=False)

    except ImportError as e:
        return f"❌ 缺少依赖：{e}\n提示：完整回测需要 Qlib 数据和相关依赖。"
    except Exception as e:
        return f"❌ 回测失败：{e}\n\n提示：确保已初始化 Qlib 并下载了对应数据集。"


# ─────────────────────────────────────────────
# 工具 6：因子分析（IC / 收益率分析）
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_factor_analysis(
    factor_expr: str,
    market: str = "csi300",
    start_date: str = "2020-01-01",
    end_date: str = "2023-12-31",
    forward_periods: str = "1,5,20",
    quantiles: int = 5,
) -> str:
    """
    对自定义因子做 IC 分析和分位数收益分析。

    Args:
        factor_expr: 因子表达式（Qlib 语法）。
                     示例：
                     - "$volume / Mean($volume, 20)"（相对成交量）
                     - "($close - $open) / $open"（日内振幅）
                     - "EMA($close, 12) - EMA($close, 26)"（MACD 信号）
                     - "Corr($close, $volume, 10)"（价量相关性）
        market: 股票池，如 "csi300"、"csi500"、"all"
        start_date: 分析开始日期 YYYY-MM-DD
        end_date: 分析结束日期 YYYY-MM-DD
        forward_periods: 预测期（天数），逗号分隔，默认 "1,5,20"
        quantiles: 分位数，默认 5 分组

    Returns:
        因子 IC 统计和分位数收益摘要
    """
    _ensure_initialized()

    try:
        import numpy as np
        import pandas as pd
        from qlib.data import D
        from scipy import stats as scipy_stats

        periods = [int(p.strip()) for p in forward_periods.split(",")]
        instruments = D.instruments(market)

        # 获取因子值和未来收益率
        fields = [factor_expr] + [f"Ref($close, -{p}) / $close - 1" for p in periods]
        df = D.features(
            instruments=instruments,
            fields=fields,
            start_time=start_date,
            end_time=end_date,
        )

        if df is None or df.empty:
            return "⚠️ 未获取到数据，请检查数据路径和日期范围"

        df.columns = ["factor"] + [f"ret_{p}d" for p in periods]
        df = df.dropna()

        result = {
            "因子": factor_expr[:60],
            "股票池": market,
            "样本量": len(df),
            "因子统计": {
                "均值": float(df["factor"].mean()),
                "标准差": float(df["factor"].std()),
                "偏度": float(df["factor"].skew()),
            },
            "IC 分析": {},
        }

        for p in periods:
            col = f"ret_{p}d"
            ic, p_value = scipy_stats.spearmanr(df["factor"], df[col])
            result["IC 分析"][f"{p}日 IC"] = f"{ic:.4f}（p={p_value:.4f}）"

            # 分位数均值收益
            df[f"q_{p}"] = pd.qcut(df["factor"], q=quantiles, labels=False, duplicates="drop")
            q_ret = df.groupby(f"q_{p}")[col].mean()
            result[f"{p}日分位收益"] = {
                f"Q{int(q)+1}": f"{v:.2%}" for q, v in q_ret.items()
            }

        return json.dumps(result, ensure_ascii=False, default=str)

    except ImportError as e:
        return f"❌ 缺少依赖 scipy，请安装：pip install scipy\n详情：{e}"
    except Exception as e:
        return f"❌ 因子分析失败：{e}"


# ─────────────────────────────────────────────
# 工具 7：Qlib 表达式帮助
# ─────────────────────────────────────────────


@mcp.tool()
def qlib_expression_help(topic: str = "overview") -> str:
    """
    获取 Qlib 表达式语法帮助，包括常用算子��示例。

    Args:
        topic: 帮助主题。可选：
               - "overview"（概览）
               - "operators"（所有算子列表）
               - "factors"（常用技术因子示例）
               - "alpha158"（Qlib 内置 Alpha158 因子说明）

    Returns:
        帮助文档
    """
    docs = {
        "overview": """
# Qlib 表达式语法概览

## 基础字段（需初始化数据后可用）
- `$close` 收盘价（复权）
- `$open` 开盘价
- `$high` 最高价
- `$low` 最低价
- `$volume` 成交量
- `$amount` 成交额
- `$factor` 复权因子
- `$vwap` 均价
- `$change` 涨跌幅

## 常用算子
- `Ref($close, N)` — N 日前的值
- `Mean($close, N)` — N 日均值（=MA）
- `Std($close, N)` — N 日标准差
- `Max($high, N)` / `Min($low, N)` — N 日最高/最低
- `EMA($close, N)` — 指数移动平均
- `Corr($close, $volume, N)` — N 日价量相关系数
- `Sum($volume, N)` — N 日成交量之和
- `Rank($close, N)` — N 日内的排名（横截面）

## 表达式示例
```
# 20 日动量
$close / Ref($close, 20) - 1

# 布林带上轨距离
($close - Mean($close, 20)) / Std($close, 20)

# MACD
EMA($close, 12) - EMA($close, 26)

# 价量相关（聪明钱）
Corr($close, $volume, 10)
```
""",
        "operators": """
# Qlib 表达式算子完整列表

## 时序算子（对单股票的时间序列操作）
Ref(x, N)      — N 日前的值
Mean(x, N)     — 滚动均值
Std(x, N)      — 滚动标准差
Var(x, N)      — 滚动方差
Skew(x, N)     — 滚动偏度
Kurt(x, N)     — 滚动峰度
Max(x, N)      — 滚动最大值
Min(x, N)      — 滚动最小值
Quantile(x, N, qscore) — 滚动分位数
Med(x, N)      — 滚动中位数
Mad(x, N)      — 滚动绝对中位差
EMA(x, N)      — 指数移动平均
WMA(x, N)      — 加权移动平均
DEMA(x, N)     — 双重指数移动平均
Sum(x, N)      — 滚动求和
Prod(x, N)     — 滚动乘积
Slope(x, N)    — 滚动线性回归斜率
Rsquare(x, N)  — 滚动 R² 值
Resi(x, N)     — 滚动线性回归残差
Corr(x, y, N)  — 滚动相关系数
Cov(x, y, N)   — 滚动协方差
IdxMax(x, N)   — 最大值所在位置（距今天数）
IdxMin(x, N)   — 最小值所在位置
ArgMax(x, N)   — 同 IdxMax
ArgMin(x, N)   — 同 IdxMin
Rank(x, N)     — 时间序列排名

## 横截面算子（在同一天对所有股票排名/归一化）
CSRank(x)      — 横截面排名（归一化到 0~1）
CSZScore(x)    — 横截面 Z-Score 标准化
CSMean(x)      — 横截面均值
CSStd(x)       — 横截面标准差

## 数学算子
Abs(x)         — 绝对值
Sign(x)        — 符号函数（-1/0/1）
Log(x)         — 自然对数
Power(x, e)    — 幂次
Div(x, y)      — 除法（处理 0 除）
Greater(x, y)  — max(x, y)
Less(x, y)     — min(x, y)
Not(x)         — 布尔取反
And(x, y)      — 布尔与
Or(x, y)       — 布尔或
If(cond, x, y) — 条件选择
""",
        "factors": """
# 常用技术因子示例

## 动量类
20 日动量：  $close / Ref($close, 20) - 1
5 日动量：   $close / Ref($close, 5) - 1
60 日动量：  $close / Ref($close, 60) - 1

## 均值回复类
布林带位置：  ($close - Mean($close,20)) / Std($close,20)
RSI 类：     Mean(Greater($close-Ref($close,1),0),14) / Mean(Abs($close-Ref($close,1)),14)

## 成交量类
相对成交量：  $volume / Mean($volume, 20)
放量比：      Mean($volume,5) / Mean($volume,20)
成交额加速：  $amount / Mean($amount, 20)

## 波动率类
20 日历史波动率：  Std($close/Ref($close,1)-1, 20)
ATR（平均真实波幅）：Mean(Greater($high-$low, Abs($high-Ref($close,1)), Abs($low-Ref($close,1))), 14)

## 价量相关
聪明钱指标：  Corr($close, Log($volume+1), 10)
日内强度：    ($close - $open) / ($high - $low + 1e-6)

## 估值类（需要财务数据）
PE 的倒数（盈利收益率）：1 / $pe_ttm
PB 的倒数（净资产收益率比）：1 / $pb
""",
        "alpha158": """
# Alpha158 因子集（Qlib 内置）

Alpha158 是 Qlib 官方提供的 158 个技术因子集合，涵盖：

## 使用方式
```python
from qlib.contrib.data.handler import Alpha158

handler = Alpha158(
    instruments="csi300",
    start_time="2020-01-01",
    end_time="2023-12-31",
    fit_start_time="2020-01-01",
    fit_end_time="2021-12-31",
)
df = handler.fetch()
```

## 因子分组（共 158 个）
- **价格类（30+）**：各周期动量、均线交叉、价格位置
- **成交量类（20+）**：相对成交量、换手率、量价关系
- **波动率类（20+）**：历史波动率、ATR、振幅
- **相关性类（20+）**：价量相关、与市场相关
- **技术指标（60+）**：RSI、MACD、布林带、KDJ 变体

## 完整文档
GitHub: https://github.com/microsoft/qlib/tree/main/qlib/contrib/data
""",
    }

    return docs.get(topic, f"❌ 未知主题 '{topic}'，可选：overview / operators / factors / alpha158")


# ─────────────────────────────────────────────
# 入口
# ─────────────────────────────────────────────

if __name__ == "__main__":
    mcp.run(transport="stdio")