更新 README 和代码，数据源切换为 Tushare，优化请求重试机制

README.md

@@ -11,7 +11,7 @@ pinned: false
 
 基于[清华大学 Kronos 金融 K 线基础大模型](https://arxiv.org/abs/2508.02739)的 A 股概率预测 REST API。
 
-- **数据源**：AkShare 东方财富 A 股日线，前复权（qfq）
+- **数据源**：Tushare Pro A 股日线，前复权（qfq）
 - **推理方式**：蒙特卡洛分批采样，输出预测方向、置信度及 95% 交易区间
 - **异步任务**：POST 提交 → 返回 `task_id` → GET 轮询结果
 
@@ -20,7 +20,7 @@ pinned: false
 ## 更新说明（2026-03-16）
 
 - 请求/响应主字段由 `ts_code` 统一为 `symbol`
-- 数据源切换为 AkShare 东财日线接口（`stock_zh_a_hist`，前复权 `qfq`）
+- 数据源切换回 Tushare Pro（`pro_bar`，前复权 `qfq`）
 - 方向概率 `direction.probability` 定义为“预测区间内看涨概率（0–1）”
 - 推理路径改为分批采样（可通过 `MC_BATCH_SIZE` 调整批大小）
 - 新增阶段耗时日志（`fetch/calendar/infer/build/cache/total`）
@@ -311,8 +311,11 @@ curl "https://yingfeng64-kronos-api.hf.space/api/v1/cache"
 
 | 环境变量 | 默认值 | 说明 |
 |---|---|---|
+| `TUSHARE_TOKEN` | — | Tushare 访问令牌（必填） |
 | `KRONOS_DIR` | `/app/Kronos` | Kronos 源码目录 |
 | `MC_BATCH_SIZE` | `8` | 蒙特卡洛分批采样大小（越大通常越快，但占用显存/内存更高） |
+| `TS_RETRY_COUNT` | `3` | Tushare 请求重试次数 |
+| `TS_RETRY_BASE_SLEEP` | `0.8` | Tushare 重试退避基准秒数（指数退避） |
 
 ---
 

data_fetcher.py

@@ -3,88 +3,80 @@ import os
 import threading
 import time
 
-import akshare as ak
 import pandas as pd
+import tushare as ts
 
-_TRADE_CALENDAR_CACHE: pd.DatetimeIndex | None = None
-_TRADE_CALENDAR_CACHED_AT: datetime | None = None
-_TRADE_CALENDAR_CACHE_TTL = timedelta(hours=12)
-_TRADE_CALENDAR_LOCK = threading.Lock()
-_AK_RETRY_COUNT = max(1, int(os.environ.get("AK_RETRY_COUNT", "3")))
-_AK_RETRY_BASE_SLEEP = float(os.environ.get("AK_RETRY_BASE_SLEEP", "0.8"))
-_AK_TIMEOUT = float(os.environ.get("AK_TIMEOUT", "15"))
+_TUSHARE_TOKEN = os.environ.get("TUSHARE_TOKEN", "").strip()
+_TS_RETRY_COUNT = max(1, int(os.environ.get("TS_RETRY_COUNT", "3")))
+_TS_RETRY_BASE_SLEEP = float(os.environ.get("TS_RETRY_BASE_SLEEP", "0.8"))
 
+_PRO = None
+_PRO_LOCK = threading.Lock()
 
-def _normalize_symbol(raw_symbol: str) -> str:
+
+def _get_pro():
+    global _PRO
+    with _PRO_LOCK:
+        if _PRO is not None:
+            return _PRO
+        if not _TUSHARE_TOKEN:
+            raise RuntimeError("TUSHARE_TOKEN is not set")
+        ts.set_token(_TUSHARE_TOKEN)
+        _PRO = ts.pro_api()
+        return _PRO
+
+
+def _normalize_to_ts_code(raw_symbol: str) -> str:
     """
-    Convert user input into the 6-digit stock code expected by
-    `ak.stock_zh_a_hist`.
+    Normalize user `symbol` input to Tushare `ts_code`.
 
     Accepted examples:
-      - "603777"
-      - "600900.SH"
-      - "000063.SZ"
+      - "603777"     -> "603777.SH"
+      - "300065"     -> "300065.SZ"
+      - "430047"     -> "430047.BJ"
+      - "000063.SZ"  -> "000063.SZ"
     """
     symbol = raw_symbol.strip().upper()
     if "." in symbol:
-        symbol = symbol.split(".", 1)[0]
+        code, market = symbol.split(".", 1)
+        if len(code) != 6 or not code.isdigit() or market not in {"SH", "SZ", "BJ"}:
+            raise ValueError(
+                f"Invalid symbol {raw_symbol!r}; expected e.g. '603777' or '000063.SZ'."
+            )
+        return f"{code}.{market}"
+
     if len(symbol) != 6 or not symbol.isdigit():
         raise ValueError(
-            f"Invalid stock code {raw_symbol!r}; expected 6 digits like '603777' "
-            "or Tushare-style code like '600900.SH'."
+            f"Invalid symbol {raw_symbol!r}; expected 6 digits or code with suffix."
         )
-    return symbol
 
+    if symbol.startswith("6"):
+        market = "SH"
+    elif symbol.startswith(("0", "3")):
+        market = "SZ"
+    elif symbol.startswith(("4", "8")):
+        market = "BJ"
+    else:
+        raise ValueError(f"Cannot infer exchange suffix for symbol={raw_symbol!r}")
 
-def _retry_ak_call(fn, *, call_name: str):
-    """
-    Retry wrapper for AkShare calls to handle transient network disconnects.
-    """
+    return f"{symbol}.{market}"
+
+
+def _retry_call(fn, *, call_name: str):
     last_exc: Exception | None = None
-    for attempt in range(1, _AK_RETRY_COUNT + 1):
+    for attempt in range(1, _TS_RETRY_COUNT + 1):
         try:
             return fn()
         except Exception as exc:  # pragma: no cover - external IO
             last_exc = exc
-            if attempt >= _AK_RETRY_COUNT:
+            if attempt >= _TS_RETRY_COUNT:
                 break
-            time.sleep(_AK_RETRY_BASE_SLEEP * (2 ** (attempt - 1)))
+            time.sleep(_TS_RETRY_BASE_SLEEP * (2 ** (attempt - 1)))
     raise RuntimeError(
-        f"AkShare call failed after {_AK_RETRY_COUNT} attempts ({call_name}): {last_exc}"
+        f"Tushare call failed after {_TS_RETRY_COUNT} attempts ({call_name}): {last_exc}"
     ) from last_exc
 
 
-def _get_trade_calendar_cached() -> pd.DatetimeIndex:
-    """
-    Fetch and cache exchange trading dates in-process to avoid repeated
-    network calls on each request.
-    """
-    global _TRADE_CALENDAR_CACHE, _TRADE_CALENDAR_CACHED_AT
-
-    now = datetime.now()
-    with _TRADE_CALENDAR_LOCK:
-        if (
-            _TRADE_CALENDAR_CACHE is not None
-            and _TRADE_CALENDAR_CACHED_AT is not None
-            and (now - _TRADE_CALENDAR_CACHED_AT) < _TRADE_CALENDAR_CACHE_TTL
-        ):
-            return _TRADE_CALENDAR_CACHE
-
-    cal = _retry_ak_call(
-        lambda: ak.tool_trade_date_hist_sina(),
-        call_name="tool_trade_date_hist_sina",
-    )
-    cal_col = "trade_date" if "trade_date" in cal.columns else "日期"
-    all_dates = pd.to_datetime(cal[cal_col]).sort_values().drop_duplicates()
-    cached = pd.DatetimeIndex(all_dates)
-
-    with _TRADE_CALENDAR_LOCK:
-        _TRADE_CALENDAR_CACHE = cached
-        _TRADE_CALENDAR_CACHED_AT = now
-
-    return cached
-
-
 def fetch_stock_data(
     symbol: str, lookback: int
 ) -> tuple[pd.DataFrame, pd.Series, str]:
@@ -94,55 +86,36 @@ def fetch_stock_data(
         x_timestamp  : pd.Series[datetime], aligned to x_df
         last_trade_date: str "YYYYMMDD", the most recent bar date
     """
-    normalized_symbol = _normalize_symbol(symbol)
+    ts_code = _normalize_to_ts_code(symbol)
     end_date = datetime.today().strftime("%Y%m%d")
-    # 4x buffer to account for weekends/long holidays.
-    start_date = (datetime.today() - timedelta(days=lookback * 4)).strftime("%Y%m%d")
-
-    df = _retry_ak_call(
-        lambda: ak.stock_zh_a_hist(
-            symbol=normalized_symbol,
-            period="daily",
+    # 2x buffer to account for weekends/holidays.
+    start_date = (datetime.today() - timedelta(days=lookback * 2)).strftime("%Y%m%d")
+
+    pro = _get_pro()
+    df = _retry_call(
+        lambda: ts.pro_bar(
+            ts_code=ts_code,
+            adj="qfq",
             start_date=start_date,
             end_date=end_date,
-            adjust="qfq",
-            timeout=_AK_TIMEOUT,
+            asset="E",
         ),
-        call_name=f"stock_zh_a_hist(symbol={normalized_symbol})",
+        call_name=f"pro_bar(ts_code={ts_code})",
     )
 
     if df is None or df.empty:
-        raise ValueError(f"No data returned for symbol={symbol!r}")
-
-    df = df.rename(
-        columns={
-            "日期": "trade_date",
-            "开盘": "open",
-            "最高": "high",
-            "最低": "low",
-            "收盘": "close",
-            "成交量": "volume",
-            "成交额": "amount",
-        }
-    )
-    required_cols = ["trade_date", "open", "high", "low", "close", "volume", "amount"]
-    missing = [c for c in required_cols if c not in df.columns]
-    if missing:
-        raise ValueError(f"AkShare response missing columns: {missing}")
-
-    df["trade_date"] = pd.to_datetime(df["trade_date"])
-    for col in ["open", "high", "low", "close", "volume", "amount"]:
-        df[col] = pd.to_numeric(df[col], errors="coerce")
-    df = df.dropna(subset=["trade_date", "open", "high", "low", "close", "volume", "amount"])
+        raise ValueError(f"No data returned for symbol={symbol!r} (ts_code={ts_code})")
+
     df = df.sort_values("trade_date").reset_index(drop=True)
-    df["timestamps"] = df["trade_date"]
+    df = df.rename(columns={"vol": "volume"})
+    df["timestamps"] = pd.to_datetime(df["trade_date"], format="%Y%m%d")
 
     # Keep the most recent `lookback` bars
     df = df.tail(lookback).reset_index(drop=True)
 
     x_df = df[["open", "high", "low", "close", "volume", "amount"]].copy()
     x_timestamp = df["timestamps"].copy()
-    last_trade_date = df["trade_date"].iloc[-1].strftime("%Y%m%d")
+    last_trade_date = str(df["trade_date"].iloc[-1])
 
     return x_df, x_timestamp, last_trade_date
 
@@ -153,20 +126,27 @@ def get_future_trading_dates(last_trade_date: str, pred_len: int) -> pd.Series:
     follow `last_trade_date` (format: YYYYMMDD).
     """
     last_dt = datetime.strptime(last_trade_date, "%Y%m%d")
-    dates: list[pd.Timestamp] = []
-
-    # Prefer real exchange trade dates from AkShare.
-    try:
-        all_dates = _get_trade_calendar_cached()
-        dates.extend([d for d in all_dates if d > pd.Timestamp(last_dt)][:pred_len])
-    except Exception:
-        # If calendar fetch fails, fall back to weekday-based dates.
-        pass
-
-    candidate = last_dt + timedelta(days=1)
-    while len(dates) < pred_len:
-        if candidate.weekday() < 5:
-            dates.append(pd.Timestamp(candidate))
-        candidate += timedelta(days=1)
-
-    return pd.Series(pd.DatetimeIndex(dates[:pred_len]))
+    # 3x buffer so we always have enough dates even over a long holiday
+    end_dt = last_dt + timedelta(days=pred_len * 3)
+
+    pro = _get_pro()
+    cal = _retry_call(
+        lambda: pro.trade_cal(
+            exchange="SSE",
+            start_date=(last_dt + timedelta(days=1)).strftime("%Y%m%d"),
+            end_date=end_dt.strftime("%Y%m%d"),
+            is_open="1",
+        ),
+        call_name="trade_cal(exchange=SSE)",
+    )
+
+    cal = cal.sort_values("cal_date")
+    dates = pd.to_datetime(cal["cal_date"].values[:pred_len], format="%Y%m%d")
+
+    if len(dates) < pred_len:
+        raise ValueError(
+            f"Could only obtain {len(dates)} future trading dates; "
+            f"increase buffer or check Tushare calendar coverage."
+        )
+
+    return pd.Series(dates)

requirements.txt

@@ -9,4 +9,4 @@ huggingface_hub==0.33.1
 matplotlib==3.9.3
 tqdm==4.67.1
 safetensors==0.6.2
-akshare
+tushare