初始化工程

This commit is contained in:
2026-07-01 22:07:55 +08:00
parent 9bb6d7a654
commit fc2b0944d2
248 changed files with 65020 additions and 316 deletions
+1
View File
@@ -0,0 +1 @@
"""数据源适配层 — 能力探测 / 调度 / Repository。"""
+76
View File
@@ -0,0 +1,76 @@
"""Capability 定义(§5.1)。
业务代码只依赖 CapabilitySet,不读 tiers.yaml,不感知"档位"
"""
from __future__ import annotations
from dataclasses import dataclass
from enum import StrEnum
class Cap(StrEnum):
"""所有 capability 的命名常量。新增能力时只在这里加一行。"""
QUOTE_BY_SYMBOL = "quote.by_symbol"
QUOTE_BATCH = "quote.batch"
QUOTE_POOL = "quote.pool"
KLINE_DAILY_BY_SYMBOL = "kline.daily.by_symbol"
KLINE_DAILY_BATCH = "kline.daily.batch"
KLINE_MINUTE_BY_SYMBOL = "kline.minute.by_symbol"
KLINE_MINUTE_BATCH = "kline.minute.batch"
INTRADAY = "intraday"
INTRADAY_BATCH = "intraday.batch"
DEPTH5 = "depth5"
DEPTH5_BATCH = "depth5.batch"
WEBSOCKET = "websocket"
FINANCIAL = "financial"
ADJ_FACTOR = "adj_factor"
@dataclass(slots=True, frozen=True)
class CapabilityLimits:
"""单个 capability 的运行时限制。"""
rpm: int | None = None # 次/分钟,None 表示未知或不限
batch: int | None = None # 标的/次
subscribe: int | None = None # WS 订阅上限
class CapabilitySet:
"""探测得到的"用户当前可用能力"。业务代码的唯一真理源。"""
def __init__(self, caps: dict[Cap, CapabilityLimits] | None = None) -> None:
self._caps: dict[Cap, CapabilityLimits] = dict(caps or {})
def has(self, cap: Cap) -> bool:
return cap in self._caps
def limits(self, cap: Cap) -> CapabilityLimits | None:
return self._caps.get(cap)
def require(self, cap: Cap) -> CapabilityLimits:
"""断言可用,否则抛 CapabilityDenied。"""
if cap not in self._caps:
raise CapabilityDenied(cap)
return self._caps[cap]
def all(self) -> dict[Cap, CapabilityLimits]:
return dict(self._caps)
def to_dict(self) -> dict[str, dict]:
return {
str(cap): {
"rpm": lim.rpm,
"batch": lim.batch,
"subscribe": lim.subscribe,
}
for cap, lim in self._caps.items()
}
class CapabilityDenied(Exception):
"""请求的 capability 当前不可用。"""
def __init__(self, cap: Cap, suggestion: str | None = None) -> None:
self.cap = cap
self.suggestion = suggestion or f"加购『{cap}』能力可解锁"
super().__init__(f"capability not available: {cap}; {self.suggestion}")
+109
View File
@@ -0,0 +1,109 @@
"""数据源 SDK 封装(§5)。
进程内单例;Key 来源(优先级):secrets.json > .env。
用户改 Key 后需要 `reset_clients()`,然后 `get_client()` 会拿新的。
5 档体系下服务器归属:
- none 档(无 key / 无效 key) → 数据源 SDK .free()(free-api 服务器)
- free 档(免费有效 key) → 数据源 SDK .free()(key 被 SDK 忽略,运行时走 free-api)
- starter/pro/expert(付费 key) → 数据源 SDK (api_key=key, base_url)
"""
from __future__ import annotations
import os
from tickflow import AsyncTickFlow, TickFlow
from app import secrets_store
_sync_client: TickFlow | None = None
_async_client: AsyncTickFlow | None = None
# ===== 服务器归属判定 =====
# free-api 服务器默认节点(SDK 默认值),none/free 档运行时走这里。
FREE_ENDPOINT = "https://free-api.tickflow.org"
# 付费端点默认节点(starter+ 运行时走这里,也是端点切换的默认值)。
PAID_ENDPOINT = "https://api.tickflow.org"
def _should_use_free_server() -> bool:
"""是否应走 free-api 服务器。
判定依据:无 key,或当前档位为 none/free。
付费档(starter+)走付费端点。
"""
if not secrets_store.get_tickflow_key():
return True
# 有 key 时按探测出的档位判定(避免读 capabilities.json 在首次启动前未生成的边界)
from app.tickflow.policy import base_tier_name
return base_tier_name() in ("none", "free")
def _base_url() -> str | None:
"""从 secrets.json 读取用户自定义端点,没有则返回 None(用 SDK 默认)。"""
return secrets_store.load().get("tickflow_base_url") or None
def get_client() -> TickFlow:
"""同步客户端。能力探测、盘后管道用。"""
global _sync_client
if _sync_client is None:
key = secrets_store.get_tickflow_key()
if _should_use_free_server():
# none/free 档:走 free-api 服务器(无 key 或免费 key 被 SDK 忽略)
_sync_client = TickFlow.free()
else:
_sync_client = TickFlow(api_key=key, base_url=_base_url())
return _sync_client
def get_async_client() -> AsyncTickFlow:
"""异步客户端。FastAPI 请求路径上用。"""
global _async_client
if _async_client is None:
key = secrets_store.get_tickflow_key()
if _should_use_free_server():
_async_client = AsyncTickFlow.free()
else:
_async_client = AsyncTickFlow(api_key=key, base_url=_base_url())
return _async_client
def reset_clients() -> None:
"""Key 变化后调用 — 让下一次 get_client() 拿新实例。"""
global _sync_client, _async_client
_sync_client = None
_async_client = None
def current_mode() -> str:
"""供 UI 显示当前模式。三态:
- "none" : 无 key / 无效 key(走 free-api,仅历史日K)
- "free" : 免费有效 key(走 free-api,仅历史日K)
- "api_key" : 付费 key(starter+,走付费端点,有实时行情)
"""
if not secrets_store.get_tickflow_key():
return "none"
from app.tickflow.policy import base_tier_name
tier = base_tier_name()
if tier in ("none", "free"):
return "free" if tier == "free" else "none"
return "api_key"
def current_endpoint() -> str:
"""返回当前显示用的端点 URL(对应 endpoints.json 列表项)。
- none/free 档:显示 free-api 服务器节点
- 付费档:显示用户自定义端点(测速切换后)或默认付费节点 api.tickflow.org
"""
if _should_use_free_server():
return FREE_ENDPOINT
# 自定义端点(付费模式测速切换后):优先返回
base = _base_url()
if base:
return base.rstrip("/")
return PAID_ENDPOINT
+550
View File
@@ -0,0 +1,550 @@
"""能力探测 + CapabilitySet 持久化(§5.3)。
探测策略:逐 capability 用最小代价请求试探。
- 成功 → 记录可用,优先取响应头 X-RateLimit-* 否则用 tiers.yaml 默认
- 抛权限错 → 不可用
- 抛其他错 → 不可用(谨慎,保留日志)
Tier Label 算法见 §5.3:基线档 + 补丁能力。
"""
from __future__ import annotations
import json
import logging
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import yaml
from app.config import settings
from app import secrets_store
from .capabilities import Cap, CapabilityLimits, CapabilitySet
logger = logging.getLogger(__name__)
_CAPSET_CACHE_FILE = "capabilities.json"
# 缓存 schema 版本。capabilities 模型有结构性变更时 bump(如新增/拆分 Cap),
# 旧缓存(无此字段或版本更低)会被判定过期,触发重新探测。
# v2: 拆分 depth5 → depth5(单只) + depth5.batch(批量)
# v3: 探测补全 quote.batch(此前 tiers.yaml 声明了但 _probe_real 漏探测)
# v4: 5 档重构 —— 新增 none 档(无key/无效key),free 档重定义(走 free-api 服务器,
# 仅历史日K)。判定改为复权因子分水岭:_classify_tier 接管档位判定。
_CACHE_SCHEMA_VERSION = 4
# 探测用最小代价请求:挑流通性最好的 1 只标的试
_PROBE_SYMBOL = "600000.SH" # 浦发银行,长期不会退市
def _load_tiers_yaml() -> dict[str, dict[str, dict[str, Any]]]:
for path in [settings.tiers_yaml, Path("/app/tiers.yaml"), Path("../tiers.yaml")]:
if path.exists():
with path.open(encoding="utf-8") as f:
return yaml.safe_load(f)
raise FileNotFoundError("tiers.yaml not found")
def _tier_to_capset(tier_def: dict[str, dict[str, Any]]) -> CapabilitySet:
caps: dict[Cap, CapabilityLimits] = {}
for cap_name, limits_dict in tier_def.items():
try:
cap = Cap(cap_name)
except ValueError:
logger.warning("unknown cap in tiers.yaml: %s", cap_name)
continue
caps[cap] = CapabilityLimits(
rpm=limits_dict.get("rpm"),
batch=limits_dict.get("batch"),
subscribe=limits_dict.get("subscribe"),
)
return CapabilitySet(caps)
def _is_transient(e: Exception) -> bool:
"""是否为"可重试的瞬时错误"——网络抖动 / 限流 / 服务端 5xx。
与权限/参数错误(403/401/400/404)区分:后者重试也无用,不重试。
用类名匹配而非 import SDK 异常,避免探测期对 SDK 内部耦合。
"""
cls = e.__class__.__name__
if cls in {
"RateLimitError", "InternalServerError", "APIError",
"ConnectionError", "TimeoutError", "ConnectError",
"ConnectTimeout", "ReadTimeout", "RemoteProtocolError",
"httpx.ConnectError", "httpx.TimeoutException",
}:
return True
# APIError 体系下,status_code 5xx/429 视为瞬时
status = getattr(e, "status_code", None)
if isinstance(status, int) and (status == 429 or status >= 500):
return True
return False
def _call_with_retry(fn, attempts: int = 3, backoff: float = 0.6) -> None:
"""调用 fn();对瞬时错误退避重试,权限/参数错误立即抛出。
attempts=总尝试次数(含首次)。返回 None,异常由调用方分类。
"""
last_exc: Exception | None = None
for i in range(attempts):
try:
fn()
return
except Exception as e: # noqa: BLE001
last_exc = e
# 权限/参数类错误:重试无意义,立即抛出交给 try_call 归类
if not _is_transient(e):
raise
# 瞬时错误:最后一轮不再 sleep
if i < attempts - 1:
time.sleep(backoff * (i + 1))
# 重试耗尽,抛出最后一次异常
assert last_exc is not None
raise last_exc
def _probe_real(tiers: dict) -> tuple[CapabilitySet, list[str]]:
"""逐 capability 试探。需要 API key。
**关键**:探测始终在付费端点上进行,用 key 鉴权验证有效性。
绝不能读旧 capabilities 缓存的档位来选服务器 —— 否则首次保存 key 时,
旧缓存是 none 档 → get_client() 返回 free 服务器 → free 服务器忽略 key →
乱填 key 也能拿到日K → 误判成 free 档(鸡生蛋蛋生鸡的循环依赖 bug)。
返回 (capset, probe_log)。
"""
from tickflow import TickFlow
from .client import _base_url, PAID_ENDPOINT
key = secrets_store.get_tickflow_key()
# 探测专用客户端:强制走付费端点验证 key。
# base_url 用用户自定义端点(若已配置测速切换),否则默认付费端点。
probe_base = _base_url() or PAID_ENDPOINT
tf = TickFlow(api_key=key, base_url=probe_base)
available: dict[Cap, CapabilityLimits] = {}
log: list[str] = []
def try_call(cap: Cap, fn, default_limits: dict[str, Any]) -> None:
try:
_call_with_retry(fn)
available[cap] = CapabilityLimits(
rpm=default_limits.get("rpm"),
batch=default_limits.get("batch"),
subscribe=default_limits.get("subscribe"),
)
log.append(f"{cap}")
except Exception as e: # noqa: BLE001
msg = str(e).lower()
cls = e.__class__.__name__
# PermissionError 类名 / HTTP 403 / 中英文权限关键词都算"明确无权限"
is_perm_denied = (
cls in {"PermissionError", "AuthorizationError"}
or "permission" in msg or "unauthorized" in msg
or "403" in msg or "forbidden" in msg
or "套餐" in msg or "权限" in msg or "需要" in msg
)
if is_perm_denied:
log.append(f"{cap}(无权限)")
else:
# 重试耗尽仍失败的瞬时错误 — 标记为疑似,而非直接判定"无此能力"
log.append(f"? {cap} ({cls}: {e})")
# 用各档默认上限作为占位(无 X-RateLimit-* 头时)
# 取所有档的并集,逐 cap 试探
all_caps_defaults: dict[str, dict[str, Any]] = {}
for tier in ("free", "starter", "pro", "expert"):
for cap_name, lim in tiers.get(tier, {}).items():
all_caps_defaults.setdefault(cap_name, lim)
def defaults(cap: Cap) -> dict[str, Any]:
return all_caps_defaults.get(str(cap), {})
# 全部用 keyword-only 形式调用,符合 SDK 真实签名
# quote.by_symbol
try_call(Cap.QUOTE_BY_SYMBOL,
lambda: tf.quotes.get(symbols=[_PROBE_SYMBOL], as_dataframe=False),
defaults(Cap.QUOTE_BY_SYMBOL))
# quote.batch — 批量行情(POST /v1/quotes)。用 get_by_symbols 试探。
try_call(Cap.QUOTE_BATCH,
lambda: tf.quotes.get_by_symbols([_PROBE_SYMBOL], as_dataframe=False),
defaults(Cap.QUOTE_BATCH))
# quote.pool — 用一个真实存在的 universe id 试探。
# universes.list() 在 Free 也开放,先拿任意一个 universe id 再用 get_by_universes 试。
def _probe_pool():
unis = tf.universes.list()
if not unis:
raise RuntimeError("no universes available")
first_id = unis[0]["id"] if isinstance(unis[0], dict) else getattr(unis[0], "id")
return tf.quotes.get_by_universes([first_id], as_dataframe=False)
try_call(Cap.QUOTE_POOL, _probe_pool, defaults(Cap.QUOTE_POOL))
# kline.daily.by_symbol — Free 也有
try_call(Cap.KLINE_DAILY_BY_SYMBOL,
lambda: tf.klines.get(_PROBE_SYMBOL, period="1d", count=1, as_dataframe=False),
defaults(Cap.KLINE_DAILY_BY_SYMBOL))
# kline.daily.batch
try_call(Cap.KLINE_DAILY_BATCH,
lambda: tf.klines.batch([_PROBE_SYMBOL], period="1d", count=1, as_dataframe=False),
defaults(Cap.KLINE_DAILY_BATCH))
# kline.minute.by_symbol
try_call(Cap.KLINE_MINUTE_BY_SYMBOL,
lambda: tf.klines.get(_PROBE_SYMBOL, period="1m", count=1, as_dataframe=False),
defaults(Cap.KLINE_MINUTE_BY_SYMBOL))
# kline.minute.batch
try_call(Cap.KLINE_MINUTE_BATCH,
lambda: tf.klines.batch([_PROBE_SYMBOL], period="1m", count=1, as_dataframe=False),
defaults(Cap.KLINE_MINUTE_BATCH))
# intraday
try_call(Cap.INTRADAY,
lambda: tf.klines.intraday(_PROBE_SYMBOL, count=1, as_dataframe=False),
defaults(Cap.INTRADAY))
# intraday.batch
try_call(Cap.INTRADAY_BATCH,
lambda: tf.klines.intraday_batch([_PROBE_SYMBOL], count=1, as_dataframe=False),
defaults(Cap.INTRADAY_BATCH))
# depth5 — 按标的查(单只)
try_call(Cap.DEPTH5,
lambda: tf.depth.get(_PROBE_SYMBOL),
defaults(Cap.DEPTH5))
# depth5.batch — 批量查(SDK 0.1.23+ 提供 depth.batch,对应官方 /v1/depth/batch 端点)
try_call(Cap.DEPTH5_BATCH,
lambda: tf.depth.batch([_PROBE_SYMBOL]),
defaults(Cap.DEPTH5_BATCH))
# financial — SDK 提供 income / balance_sheet / cash_flow / metrics / shares
# 用 metrics 探测(单据最小)
try_call(Cap.FINANCIAL,
lambda: tf.financials.metrics([_PROBE_SYMBOL], latest=True, as_dataframe=False),
defaults(Cap.FINANCIAL))
# adj_factor — 实际在 klines.ex_factors
try_call(Cap.ADJ_FACTOR,
lambda: tf.klines.ex_factors([_PROBE_SYMBOL], as_dataframe=False),
defaults(Cap.ADJ_FACTOR))
# websocket 不在探测期试连接(成本太高且阻塞),按档位默认推断
# 若 expert 的其他 cap 都通,则推断 websocket 也可用
if (Cap.FINANCIAL in available and Cap.INTRADAY_BATCH in available):
available[Cap.WEBSOCKET] = CapabilityLimits(
subscribe=defaults(Cap.WEBSOCKET).get("subscribe", 100),
)
log.append("✓ websocket (inferred from expert tier)")
return CapabilitySet(available), log
def detect_capabilities(force: bool = False) -> CapabilitySet:
"""探测当前 API Key 的能力集。"""
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if not force and cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
cached = json.load(f)
# schema 版本校验:旧缓存或缺版本号 → 过期,丢弃后重新探测
if cached.get("schema_version") == _CACHE_SCHEMA_VERSION:
return _capset_from_json(cached)
logger.info("capabilities 缓存 schema 版本过期(缓存=%s, 当前=%d), 重新探测",
cached.get("schema_version"), _CACHE_SCHEMA_VERSION)
tiers = _load_tiers_yaml()
if settings.use_free_mode:
# 无 key —— 归 none 档(走 free-api 服务器,仅历史日K)
capset = _tier_to_capset(tiers["none"])
_persist(capset, "None", log=["无 API Key(无档 · free-api 服务器)"], missing=[], extras=[])
return capset
# 有 API key — 真实探测
try:
capset, probe_log = _probe_real(tiers)
# 判定档位:无效 key → none,免费 key → free,付费 → starter/pro/expert
classified = _classify_tier(capset, tiers)
if classified.is_invalid:
# 无效 key(连单只日K都拿不到):归 none 档,标记要求清除 key
capset = _tier_to_capset(tiers["none"])
probe_log.append("⚠ Key 无效(单只日K也无法获取),判定为无档")
_persist(capset, "None", log=probe_log, missing=[], extras=[], invalid_key=True)
return capset
if classified.is_free:
# 免费有效 key:能力按 free 档(= none 档能力,走 free-api 服务器)
capset = _tier_to_capset(tiers["free"])
_persist(capset, "Free", log=probe_log + ["✓ 免费有效 key(运行时走 free-api 服务器)"], missing=[], extras=[])
return capset
# 付费档(starter+) — 探测出的能力即为真实可用
label, missing, extras = _compute_label_and_missing(capset, tiers)
capset = _override_limits_with_detected_tier(capset, label, tiers)
_persist(capset, label, log=probe_log, missing=missing, extras=extras)
return capset
except Exception as e:
logger.exception("detect_capabilities failed; using none baseline: %s", e)
capset = _tier_to_capset(tiers["none"])
_persist(capset, "None(探测失败)", log=[f"探测失败:{e}"], missing=[], extras=[])
return capset
# ===== Tier 代表性 capability(signature caps)=====
# 拥有**任意一个**即认作该档及以上。自上而下匹配。
# 这套设计的好处:单个 capability 探测的 transient 失败不会把整体档位"误降"。
TIER_SIGNATURES: dict[str, set[Cap]] = {
"expert": {Cap.FINANCIAL, Cap.INTRADAY_BATCH, Cap.WEBSOCKET},
"pro": {Cap.KLINE_MINUTE_BATCH, Cap.KLINE_MINUTE_BY_SYMBOL,
Cap.INTRADAY, Cap.DEPTH5, Cap.DEPTH5_BATCH},
"starter": {Cap.QUOTE_BATCH, Cap.KLINE_DAILY_BATCH,
Cap.ADJ_FACTOR, Cap.QUOTE_POOL},
# free / none 不需 signature — 由 _classify_tier 的分水岭逻辑判定
}
@dataclass(slots=True, frozen=True)
class TierClassification:
"""档位判定结果。
判定依据是"复权因子分水岭":
- 连单只日K都没有 → 无效 key(is_invalid),归 none 档
- 有单只日K、无复权因子 → 免费 key(is_free)
- 有复权因子 → 付费档(starter+),具体档位由 signature 决定
"""
tier: str # "none" / "free" / "starter" / "pro" / "expert"
is_invalid: bool # 无效 key(连单只日K都拿不到)
is_free: bool # 免费有效 key(有日K、无复权因子)
def _classify_tier(capset: CapabilitySet, tiers: dict) -> TierClassification:
"""根据探测出的能力集判定档位。
分水岭是 KLINE_DAILY_BY_SYMBOL(单只日K)与 ADJ_FACTOR(复权因子):
- 无单只日K → none(无效 key)
- 有日K无复权 → free(免费 key)
- 有复权因子 → 走 signature 判定 starter/pro/expert
"""
held = set(capset.all().keys())
# 1) 连单只日K都没有 → 无效 key
if Cap.KLINE_DAILY_BY_SYMBOL not in held:
return TierClassification(tier="none", is_invalid=True, is_free=False)
# 2) 有日K但无复权因子 → 免费 key
if Cap.ADJ_FACTOR not in held:
return TierClassification(tier="free", is_invalid=False, is_free=True)
# 3) 有复权因子 → 付费档,按 signature 自上而下判定
if held & TIER_SIGNATURES["expert"]:
base = "expert"
elif held & TIER_SIGNATURES["pro"]:
base = "pro"
elif held & TIER_SIGNATURES["starter"]:
base = "starter"
else:
# 有复权因子但无任何代表能力 — 兜底为 starter(复权本身是 starter 特征)
base = "starter"
return TierClassification(tier=base, is_invalid=False, is_free=False)
# 补丁友好命名(label 后缀用)
_CAP_ALIASES: dict[Cap, str] = {
Cap.KLINE_MINUTE_BATCH: "分钟K",
Cap.KLINE_MINUTE_BY_SYMBOL: "分钟K",
Cap.INTRADAY: "分时",
Cap.INTRADAY_BATCH: "批量分时",
Cap.DEPTH5: "五档",
Cap.DEPTH5_BATCH: "批量五档",
Cap.WEBSOCKET: "WS",
Cap.FINANCIAL: "财务",
Cap.ADJ_FACTOR: "复权",
Cap.QUOTE_BATCH: "批量行情",
Cap.QUOTE_POOL: "标的池",
Cap.KLINE_DAILY_BATCH: "日K批量",
}
def _override_limits_with_detected_tier(
capset: CapabilitySet, label: str, tiers: dict,
) -> CapabilitySet:
"""探测完成后,用判档对应的 limits 覆盖每个 cap 的速率/批量。
判档前每个 cap 用的是"所有档默认值的并集"(为了不漏数据),
判档后才知道用户真实档位,limits 用该档的实际值更准。
label 可能是 "Pro" / "Pro + 分钟K" / "Pro+" 等组合形式 — 取第一个词当作基线档名。
"""
base_name = label.split()[0].split("+")[0].strip().lower() # "Pro + 分钟K" → "pro"
tier_limits = tiers.get(base_name, {})
new_caps: dict[Cap, CapabilityLimits] = {}
for cap, _old_lim in capset.all().items():
spec = tier_limits.get(cap.value)
if spec:
new_caps[cap] = CapabilityLimits(
rpm=spec.get("rpm"),
batch=spec.get("batch"),
subscribe=spec.get("subscribe"),
)
else:
# 不在该档定义里(extras),用 expert 档兜底(最宽松)
expert_spec = tiers.get("expert", {}).get(cap.value, {})
new_caps[cap] = CapabilityLimits(
rpm=expert_spec.get("rpm"),
batch=expert_spec.get("batch"),
subscribe=expert_spec.get("subscribe"),
)
return CapabilitySet(new_caps)
def _tier_caps_set(tiers: dict, tier_name: str) -> set[Cap]:
"""读 tiers.yaml 的某档定义,转为 Cap 集合。"""
return {Cap(c) for c in tiers.get(tier_name, {}).keys() if c in {x.value for x in Cap}}
def _compute_label_and_missing(
capset: CapabilitySet, tiers: dict,
) -> tuple[str, list[str], list[str]]:
"""返回 (label, missing_caps, extra_caps)。
label:档位标签。
missing_caps:本档**应有但未探测到**的 capability(用于诊断:可能是探测 bug 或权限丢失)。
extra_caps:超出本档的额外 capability(自定义组合)。
"""
held = set(capset.all().keys())
# 1) 完全匹配 — 干净命中某档
for tier_name in ["free", "starter", "pro", "expert"]:
if held == _tier_caps_set(tiers, tier_name):
return tier_name.capitalize(), [], []
# 2) 按 signature 自上而下判档
if held & TIER_SIGNATURES["expert"]:
base = "expert"
elif held & TIER_SIGNATURES["pro"]:
base = "pro"
elif held & TIER_SIGNATURES["starter"]:
base = "starter"
else:
base = "free"
base_caps = _tier_caps_set(tiers, base)
missing = sorted(c.value for c in (base_caps - held))
extras = base_caps and (held - base_caps) or set() # extras 是超出该档的部分
# 实际超出 = held 中"既不属于本档、也不属于本档下方任何档"的 cap
# 简化:extras = held - base_caps
extras_set = held - base_caps
# 3) 拼 label
if not extras_set:
# 完全在本档内(可能缺一两项 — 由 missing 反映)
return base.capitalize(), missing, []
# 补丁过多 → 用 "≈" 形式
if len(extras_set) > 3:
return f"{base.capitalize()}+", missing, sorted(c.value for c in extras_set)
suffix = sorted({_CAP_ALIASES.get(e, str(e)) for e in extras_set})
return f"{base.capitalize()} + " + " + ".join(suffix), missing, sorted(c.value for c in extras_set)
def _compute_label(capset: CapabilitySet, tiers: dict) -> str:
"""对外简化签名 — 只要 label。"""
label, _missing, _extras = _compute_label_and_missing(capset, tiers)
return label
def _persist(
capset: CapabilitySet,
label: str,
log: list[str] | None = None,
missing: list[str] | None = None,
extras: list[str] | None = None,
invalid_key: bool = False,
) -> None:
settings.data_dir.mkdir(parents=True, exist_ok=True)
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
payload = {
"schema_version": _CACHE_SCHEMA_VERSION,
"label": label,
"capabilities": capset.to_dict(),
"probe_log": log or [],
"missing_caps": missing or [], # 本档应有但未探测到
"extras_caps": extras or [], # 超出本档的额外能力
"invalid_key": invalid_key, # 探测出的 key 无效(连单只日K都拿不到)
}
with cache_path.open("w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
def _capset_from_json(data: dict[str, Any]) -> CapabilitySet:
caps: dict[Cap, CapabilityLimits] = {}
for cap_name, lim in data.get("capabilities", {}).items():
try:
cap = Cap(cap_name)
except ValueError:
continue
caps[cap] = CapabilityLimits(
rpm=lim.get("rpm"),
batch=lim.get("batch"),
subscribe=lim.get("subscribe"),
)
return CapabilitySet(caps)
def tier_label() -> str:
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
return json.load(f).get("label", "Unknown")
return "Unknown"
def probe_log() -> list[str]:
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
return json.load(f).get("probe_log", [])
return []
def missing_caps() -> list[str]:
"""本档应有但未探测到的 capability — 通常意味着探测有 bug 或权限边界。"""
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
return json.load(f).get("missing_caps", [])
return []
def extras_caps() -> list[str]:
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
return json.load(f).get("extras_caps", [])
return []
def is_invalid_key() -> bool:
"""最近一次探测是否判定 key 无效(连单只日K都拿不到)。
settings 层据此清除已存的 key,避免乱填的 key 被持久化。
"""
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
if cache_path.exists():
with cache_path.open(encoding="utf-8") as f:
return bool(json.load(f).get("invalid_key", False))
return False
def base_tier_name() -> str:
"""当前档位的基础名(小写): none / free / starter / pro / expert。
供 client 层判断"是否走 free-api 服务器"(none/free → free 服务器)。
"""
label = tier_label()
return label.split()[0].split("+")[0].strip().lower()
+158
View File
@@ -0,0 +1,158 @@
"""标的池(Universe)定义(§6.3)。
Phase 1 实现:
- 常用指数成份(沪深 300 / 中证 500 / 上证 50)用数据源 `quote.pool` 端点拉取并缓存
- 全 A 通过 instruments.batch 获取
- 自选池 = 用户的 watchlist
"""
from __future__ import annotations
import logging
from datetime import date
from pathlib import Path
from typing import Literal
import polars as pl
from app.config import settings
from app.tickflow.client import get_client
logger = logging.getLogger(__name__)
PoolId = Literal["CSI300", "CSI500", "SSE50", "CN_Equity_A", "CN_Index", "watchlist"]
# 数据源 universe id 是内部命名(见 tf.universes.list())。
# 没有官方对照表,启动时按名称模糊匹配从 universes.list() 里找。
# 常见名:沪深300 / 中证500 / 上证50 / 全 A
_POOL_NAME_HINTS = {
"CSI300": ["沪深300", "HS300", "CSI300"],
"CSI500": ["中证500", "ZZ500", "CSI500"],
"SSE50": ["上证50", "SH50", "SSE50"],
}
def _find_universe_id(hints: list[str]) -> str | None:
"""从 universes.list() 里按 name/id 子串匹配找一个 universe id。"""
try:
tf = get_client()
unis = tf.universes.list()
except Exception as e: # noqa: BLE001
logger.warning("universes.list failed: %s", e)
return None
for u in unis or []:
item = u if isinstance(u, dict) else {"id": getattr(u, "id", ""), "name": getattr(u, "name", "")}
haystack = (item.get("id", "") + " " + item.get("name", "")).lower()
for h in hints:
if h.lower() in haystack:
return item["id"]
return None
def _pool_cache_path(pool_id: str) -> Path:
return settings.data_dir / "pools" / f"{pool_id}.parquet"
def get_pool(pool_id: PoolId, refresh: bool = False) -> list[str]:
"""返回标的池里的 symbol 列表。"""
if pool_id == "watchlist":
return _load_watchlist()
cache = _pool_cache_path(pool_id)
if cache.exists() and not refresh:
df = pl.read_parquet(cache)
return df["symbol"].to_list()
symbols = _fetch_pool(pool_id)
if symbols:
cache.parent.mkdir(parents=True, exist_ok=True)
pl.DataFrame({"symbol": symbols, "as_of": [date.today()] * len(symbols)}).write_parquet(cache)
return symbols
def _fetch_pool(pool_id: PoolId) -> list[str]:
"""从数据源拉取池成份。
实现:先用 universes.list 找到 universe id,再 quotes.get_by_universes 拉成份。
"""
tf = get_client()
if pool_id in _POOL_NAME_HINTS:
uid = _find_universe_id(_POOL_NAME_HINTS[pool_id])
if not uid:
logger.warning("无法在数据源 universes 列表里匹配到 %s", pool_id)
return []
try:
df = tf.quotes.get_by_universes([uid], as_dataframe=True)
if df is not None and len(df) > 0 and "symbol" in df.columns:
return df["symbol"].astype(str).tolist()
except Exception as e: # noqa: BLE001
logger.warning("fetch pool %s via universe %s failed: %s", pool_id, uid, e)
if pool_id == "CN_Equity_A":
# 全 A — 优先直接用 CN_Equity_A universe (包含沪深京三市)
uid = _find_universe_id(["CN_Equity_A", "沪深京A股", "全A"])
if uid:
try:
df = tf.quotes.get_by_universes([uid], as_dataframe=True)
if df is not None and len(df) > 0 and "symbol" in df.columns:
return sorted(set(df["symbol"].astype(str).tolist()))
except Exception as e: # noqa: BLE001
logger.warning("fetch CN_Equity_A via universe %s failed: %s", uid, e)
# fallback: 聚合申万一级行业 (覆盖度较低, 缺北交所/新股)
try:
unis = tf.universes.list()
except Exception as e: # noqa: BLE001
logger.warning("universes.list failed: %s", e)
unis = []
sw1_ids = []
for u in unis or []:
item = u if isinstance(u, dict) else {"id": getattr(u, "id", "")}
uid = item.get("id", "")
if "SW1_" in uid:
sw1_ids.append(uid)
if sw1_ids:
try:
df = tf.quotes.get_by_universes(sw1_ids, as_dataframe=True)
if df is not None and "symbol" in df.columns:
return sorted(set(df["symbol"].astype(str).tolist()))
except Exception as e: # noqa: BLE001
logger.warning("aggregate SW1 fetch failed: %s", e)
if pool_id == "CN_Index":
uid = _find_universe_id(["CN_Index", "沪深指数", "指数"])
ids = [uid] if uid else ["CN_Index"]
try:
df = tf.quotes.get_by_universes(ids, as_dataframe=True)
if df is not None and len(df) > 0 and "symbol" in df.columns:
return sorted(set(df["symbol"].astype(str).tolist()))
except Exception as e: # noqa: BLE001
logger.warning("fetch CN_Index via universe %s failed: %s", ids, e)
return []
def _load_watchlist() -> list[str]:
"""读取用户自选(由 watchlist service 维护)。"""
path = settings.data_dir / "user_data" / "watchlist.parquet"
if not path.exists():
return []
df = pl.read_parquet(path)
if df.is_empty() or "symbol" not in df.columns:
return []
return df["symbol"].to_list()
# 兜底:Free 用户/无 API 时给一个小型可用集合,让 UI 不至于空白
DEMO_SYMBOLS = [
"600000.SH", # 浦发银行
"600036.SH", # 招商银行
"600519.SH", # 贵州茅台
"601318.SH", # 中国平安
"601398.SH", # 工商银行
"000001.SZ", # 平安银行
"000333.SZ", # 美的集团
"000651.SZ", # 格力电器
"000858.SZ", # 五粮液
"002594.SZ", # 比亚迪
]
+988
View File
@@ -0,0 +1,988 @@
"""Repository 层(§7.4)。
数据分层:
- DuckDB 视图: 冷查询(统计、元数据、用户自定义SQL)
- Polars 缓存: 热路径(enriched 最新日 ~5500行 + instruments ~5500行)
- Polars scan_parquet: 分钟K/历史日K (predicate pushdown)
缓存生命周期:
- startup 时不加载(数据可能为空)
- pipeline 完成后调用 refresh_cache()
- 服务层通过 get_enriched_latest() / get_instruments() 获取缓存
"""
from __future__ import annotations
import logging
import threading
from datetime import date
from pathlib import Path
import duckdb
import polars as pl
from app.config import settings
logger = logging.getLogger(__name__)
class DataStore:
"""唯一的存储入口 — 进程启动时创建。"""
def __init__(self, data_dir: Path | None = None) -> None:
self.data_dir = Path(data_dir or settings.data_dir)
self.data_dir.mkdir(parents=True, exist_ok=True)
# 关键子目录(§7.2)
for sub in (
"kline_daily",
"kline_daily_enriched",
"kline_index_daily",
"kline_index_enriched",
"kline_minute",
"adj_factor",
"financials",
"instruments",
"instruments_index",
"instruments_ext",
"kline_ext",
"pools",
"backtest_results",
"screener_results",
"ai_cache",
"user_data",
"depth5",
):
(self.data_dir / sub).mkdir(parents=True, exist_ok=True)
# 财务数据子目录
for sub in ("metrics", "income", "balance_sheet", "cash_flow"):
(self.data_dir / "financials" / sub).mkdir(parents=True, exist_ok=True)
# DuckDB 内存模式 — 不建 .db 文件(§7.1)
self.db = duckdb.connect(database=":memory:")
self._register_views()
def _register_views(self) -> None:
"""把 Parquet 目录挂载为 DuckDB 视图(§7.3)。"""
d = self.data_dir.as_posix()
statements = [
f"""CREATE OR REPLACE VIEW kline_daily AS
SELECT * FROM read_parquet('{d}/kline_daily/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_enriched AS
SELECT * FROM read_parquet('{d}/kline_daily_enriched/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_index_daily AS
SELECT * FROM read_parquet('{d}/kline_index_daily/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_index_enriched AS
SELECT * FROM read_parquet('{d}/kline_index_enriched/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_minute AS
SELECT * FROM read_parquet('{d}/kline_minute/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW adj_factor AS
SELECT * FROM read_parquet('{d}/adj_factor/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW instruments AS
SELECT * FROM read_parquet('{d}/instruments/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW instruments_index AS
SELECT * FROM read_parquet('{d}/instruments_index/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW instruments_ext AS
SELECT * FROM read_parquet('{d}/instruments_ext/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_ext AS
SELECT * FROM read_parquet('{d}/kline_ext/**/*.parquet', union_by_name=true)""",
# 财务数据视图
f"""CREATE OR REPLACE VIEW financials_metrics AS
SELECT * FROM read_parquet('{d}/financials/metrics/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW financials_income AS
SELECT * FROM read_parquet('{d}/financials/income/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW financials_balance_sheet AS
SELECT * FROM read_parquet('{d}/financials/balance_sheet/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW financials_cash_flow AS
SELECT * FROM read_parquet('{d}/financials/cash_flow/*.parquet', union_by_name=true)""",
# 五档盘口 sealed 真假涨停(独立旁路存储,不进 enriched)
f"""CREATE OR REPLACE VIEW depth5 AS
SELECT * FROM read_parquet('{d}/depth5/**/*.parquet', union_by_name=true)""",
]
for sql in statements:
try:
self.db.execute(sql)
except duckdb.IOException:
logger.debug("view registration skipped (no parquet yet): %s", sql[:60])
class KlineRepository:
"""日 K / 分钟 K 的读写入口。"""
def __init__(self, store: DataStore) -> None:
self.store = store
self.db = store.db
self._lock = threading.Lock()
# ---- Polars 缓存 ----
self._enriched_cache: pl.DataFrame | None = None # 最新一天 (~5500行)
self._enriched_cache_date: date | None = None
self._live_agg_cache: pl.DataFrame | None = None # 预计算聚合表 (~5500行)
self._live_agg_cache_date: date | None = None
self._instruments_cache: pl.DataFrame | None = None
# 完整 enriched 历史 (含所有指标, 供 filter_history 策略使用)
self._enriched_history_cache: pl.DataFrame | None = None # ~100万行
self._enriched_history_start: date | None = None
self._index_instruments_cache: pl.DataFrame | None = None
# parquet glob 路径
self._enriched_glob = str(store.data_dir / "kline_daily_enriched" / "**" / "*.parquet")
self._index_enriched_glob = str(store.data_dir / "kline_index_enriched" / "**" / "*.parquet")
self._minute_glob = str(store.data_dir / "kline_minute" / "**" / "*.parquet")
self._inst_glob = str(store.data_dir / "instruments" / "**" / "*.parquet")
self._index_inst_glob = str(store.data_dir / "instruments_index" / "**" / "*.parquet")
def execute_all(self, sql: str, params: list | None = None) -> list[tuple]:
"""线程安全的 SELECT → fetchall。DuckDB 单 connection 非线程安全,所有读路径须走此方法。"""
with self._lock:
return self.db.execute(sql, params or []).fetchall()
def execute_one(self, sql: str, params: list | None = None) -> tuple | None:
"""线程安全的 SELECT → fetchone。"""
with self._lock:
return self.db.execute(sql, params or []).fetchone()
# ================================================================
# Polars 缓存管理
# ================================================================
def refresh_cache(self) -> None:
"""刷新 Polars 缓存。在 pipeline 完成后、服务启动时调用。"""
self._refresh_instruments()
self._refresh_index_instruments()
self._refresh_enriched()
def clear_cache(self) -> None:
"""清空所有 Polars 内存缓存。
与 refresh_cache 的区别: refresh_cache 在磁盘无数据时会提前 return,
导致内存里的旧缓存残留 (clear 数据后看板仍显示旧数据的根因)。
本方法无条件清空, 供清除数据/重置场景调用。
"""
self._enriched_cache = None
self._enriched_cache_date = None
self._enriched_history_cache = None
self._enriched_history_start = None
self._live_agg_cache = None
self._live_agg_cache_date = None
self._instruments_cache = None
self._index_instruments_cache = None
def _refresh_enriched(self) -> None:
"""从 parquet 加载 enriched 最新日到内存 + 构建聚合表。
enriched parquet 仅存 14 列基础数据。启动时读入历史数据并即时计算完整指标,
将结果缓存在内存中供各服务使用。
优化: 扩大历史读取范围, 同时缓存完整历史 (含指标), 供 filter_history 策略直接复用。
"""
try:
latest = self._latest_enriched_date_duckdb()
if not latest:
# 磁盘已无数据: 必须清空内存缓存, 否则旧数据会残留
# (清数据后看板仍显示旧数据的根因)
self.clear_cache()
return
# Step 1: 直接读最新日期的分区文件 (仅 14 列)
enriched_dir = self.store.data_dir / "kline_daily_enriched"
ds = latest.isoformat() if hasattr(latest, "isoformat") else str(latest)
target_parquet = enriched_dir / f"date={ds}" / "part.parquet"
if not target_parquet.exists():
return
df_latest = pl.read_parquet(target_parquet)
if df_latest.is_empty():
return
# Step 2: 读近 300 天 14 列数据 → compute → filter(latest) → 缓存
# 300 日历天 ≈ 210 交易日, 覆盖 filter_history 最大 lookback(90) + warmup(60)
try:
from datetime import timedelta
from app.indicators.pipeline import compute_indicators, compute_signals, compute_limit_signals
start_full = latest - timedelta(days=300)
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close",
"volume", "amount", "raw_close", "raw_high", "raw_low"]
if c in df_latest.columns]
lf = (
pl.scan_parquet(self._enriched_glob)
.filter(pl.col("date") >= start_full)
.sort(["symbol", "date"])
)
df_hist = lf.select(read_cols).collect()
if not df_hist.is_empty():
instruments = self._instruments_cache if self._instruments_cache is not None else pl.DataFrame()
df_full = compute_indicators(df_hist)
df_full = compute_signals(df_full)
if instruments is not None and not instruments.is_empty():
df_full = compute_limit_signals(df_full, instruments)
# JOIN instruments 到完整历史 (filter_history/basic_filter 需要 name/股本等列)
if instruments is not None and not instruments.is_empty():
inst_cols = [c for c in ["name", "total_shares", "float_shares"]
if c in instruments.columns and c not in df_full.columns]
if inst_cols:
df_full = df_full.join(
instruments.select(["symbol", *inst_cols]).unique(subset=["symbol"]),
on="symbol",
how="left",
)
# 缓存完整历史 (含指标+必要基础信息) 供 filter_history/backtest 直接复用
self._enriched_history_cache = df_full
self._enriched_history_start = df_full["date"].min()
logger.info("enriched 历史缓存: %d rows, %s ~ %s",
len(df_full), self._enriched_history_start, latest)
# 只取最新一天作为 enriched_cache
df_today = df_full.filter(pl.col("date") == latest)
if not df_today.is_empty():
self._enriched_cache = df_today
self._enriched_cache_date = latest
# 构建盘中递推基准: 若最新分区是今天的实时盘中数据,
# 递推状态必须停在上一交易日, 不能把今天作为“昨日”。
self._build_live_agg(self._live_agg_baseline_date(latest))
logger.info("enriched 缓存已计算: %d 只, 日期 %s (即时计算)", len(df_today), latest)
return
except Exception as e: # noqa: BLE001
logger.warning("enriched 即时计算失败, 使用原始 14 列缓存: %s", e)
# 降级: 直接使用 14 列数据 + 构建 live_agg
self._enriched_cache = df_latest
self._enriched_cache_date = latest
self._build_live_agg(self._live_agg_baseline_date(latest))
logger.info("enriched 缓存已加载: %d 只, 日期 %s", len(df_latest), latest)
except Exception as e: # noqa: BLE001
logger.warning("enriched 缓存刷新失败: %s", e)
def _build_live_agg(self, latest: date) -> None:
"""从 OHLCV 即时计算递推状态 + 窗口聚合, 构建盘中实时聚合表。
优化: 优先使用 _enriched_history_cache (启动时已计算), 避免重复 compute_indicators。
"""
from datetime import timedelta
from app.indicators.pipeline import _ema_alpha
start_60d = latest - timedelta(days=90) # 日历90天 ≈ 60个交易日
# 优先使用已有的历史缓存 (避免重复 scan_parquet + compute_indicators)
if self._enriched_history_cache is not None and not self._enriched_history_cache.is_empty():
hist_all = self._enriched_history_cache
if "date" in hist_all.columns and hist_all["date"].min() <= start_60d:
# 从历史缓存中提取所需列 (历史缓存已有指标列)
base_cols = ["symbol", "date", "open", "high", "low", "close", "volume",
"raw_close", "raw_high", "raw_low"]
needed = [c for c in base_cols if c in hist_all.columns]
df_hist = hist_all.filter(
(pl.col("date") >= start_60d) & (pl.col("date") <= latest)
).select(needed).sort(["symbol", "date"])
# 用历史缓存的指标列提取最新日状态 (无需再次 compute_indicators)
state_source = hist_all.filter(pl.col("date") == latest)
state_cols = [
"symbol",
"ema5", "ema10", "ema20", "ema30", "ema60",
"macd_dea",
"kdj_k", "kdj_d",
"atr_14",
"close", "high", "low",
"annual_vol_20d",
]
existing_state = [c for c in state_cols if c in state_source.columns]
agg_a = state_source.select(existing_state)
else:
df_hist = pl.DataFrame()
agg_a = pl.DataFrame()
else:
# 降级: 读 parquet + compute_indicators
df_hist, agg_a = self._build_live_agg_from_parquet(latest, start_60d)
if df_hist.is_empty():
self._live_agg_cache = pl.DataFrame()
self._live_agg_cache_date = None
return
if agg_a.is_empty():
self._live_agg_cache = pl.DataFrame()
self._live_agg_cache_date = None
return
# 单独计算 _ema12 / _ema26 (compute_indicators 内部会 drop 掉)
df_ema = df_hist.sort(["symbol", "date"]).with_columns([
pl.col("close").ewm_mean(alpha=_ema_alpha(12), adjust=False).over("symbol").alias("_ema12"),
pl.col("close").ewm_mean(alpha=_ema_alpha(26), adjust=False).over("symbol").alias("_ema26"),
]).filter(pl.col("date") == latest).select("symbol", "_ema12", "_ema26")
agg_a = agg_a.join(df_ema, on="symbol", how="inner")
# 单独计算 RSI 状态列 (compute_indicators 内部会 drop 掉)
df_rsi_base = df_hist.sort(["symbol", "date"]).with_columns(
pl.col("close").diff().over("symbol").alias("_daily_delta")
)
gain = pl.when(pl.col("_daily_delta") > 0).then(pl.col("_daily_delta")).otherwise(0.0)
loss = pl.when(pl.col("_daily_delta") < 0).then(-pl.col("_daily_delta")).otherwise(0.0)
rsi_exprs = []
for n in (6, 14, 24):
a = 1.0 / n
rsi_exprs.append(gain.ewm_mean(alpha=a, adjust=False).over("symbol").alias(f"_rsi_avg_gain_{n}"))
rsi_exprs.append(loss.ewm_mean(alpha=a, adjust=False).over("symbol").alias(f"_rsi_avg_loss_{n}"))
df_rsi = (
df_rsi_base
.with_columns(rsi_exprs)
.filter(pl.col("date") == latest)
.select("symbol", *[f"_rsi_avg_gain_{n}" for n in (6, 14, 24)],
*[f"_rsi_avg_loss_{n}" for n in (6, 14, 24)])
)
agg_a = agg_a.join(df_rsi, on="symbol", how="inner")
# 前复权因子: adj_factor = close(复权) / raw_close(原始)
if "raw_close" in df_hist.columns:
adj_factor_df = (
df_hist.filter(pl.col("date") == latest)
.select("symbol", (pl.col("close") / pl.col("raw_close")).alias("_adj_factor"))
)
agg_a = agg_a.join(adj_factor_df, on="symbol", how="left")
if "_adj_factor" in agg_a.columns:
agg_a = agg_a.with_columns(pl.col("_adj_factor").fill_null(1.0))
# annual_vol_20d 递推状态: 最近 19 天日收益率的部分和 / 平方和
df_daily_pct = (
df_hist.sort(["symbol", "date"])
.with_columns(
pl.col("close").pct_change().over("symbol").alias("_daily_pct")
)
)
df_vol = df_daily_pct.group_by("symbol").agg([
pl.col("_daily_pct").tail(19).sum().alias("_vol_19d_pct_sum"),
(pl.col("_daily_pct") ** 2).tail(19).sum().alias("_vol_19d_pct_sq_sum"),
])
agg_a = agg_a.join(df_vol, on="symbol", how="left")
# 昨日连板数: 从 enriched parquet 取 (用于增量计算同向 +1)
lf = pl.scan_parquet(self._enriched_glob).filter(pl.col("date") == latest)
consec_cols = [c for c in ["symbol", "consecutive_limit_ups", "consecutive_limit_downs"]
if c in lf.collect_schema().names()]
if len(consec_cols) == 3:
consec_df = lf.select(consec_cols).collect()
if not consec_df.is_empty():
consec = consec_df.select(
"symbol",
pl.col("consecutive_limit_ups").alias("_prev_consec_up"),
pl.col("consecutive_limit_downs").alias("_prev_consec_down"),
)
agg_a = agg_a.join(consec, on="symbol", how="left")
# B类: 按 symbol 分组聚合 — 窗口统计
agg_b = (
df_hist.sort(["symbol", "date"])
.group_by("symbol")
.agg([
pl.col("close").tail(4).sum().alias("_ma5_partial_sum"),
pl.col("close").tail(9).sum().alias("_ma10_partial_sum"),
pl.col("close").tail(19).sum().alias("_ma20_partial_sum"),
pl.col("close").tail(29).sum().alias("_ma30_partial_sum"),
pl.col("close").tail(59).sum().alias("_ma60_partial_sum"),
pl.col("close").tail(19).sum().alias("_boll_partial_sum"),
(pl.col("close").tail(19) ** 2).sum().alias("_boll_partial_sq_sum"),
pl.col("high").tail(59).max().alias("_high_59d"),
pl.col("low").tail(59).min().alias("_low_59d"),
pl.col("close").tail(5).first().alias("_close_5d_ago"),
pl.col("close").tail(10).first().alias("_close_10d_ago"),
pl.col("close").tail(20).first().alias("_close_20d_ago"),
pl.col("close").tail(30).first().alias("_close_30d_ago"),
pl.col("close").tail(60).first().alias("_close_60d_ago"),
pl.col("volume").tail(4).sum().alias("_vol_ma5_partial_sum"),
pl.col("volume").tail(9).sum().alias("_vol_ma10_partial_sum"),
pl.col("low").tail(8).min().alias("_kdj_8d_low"),
pl.col("high").tail(8).max().alias("_kdj_8d_high"),
pl.col("close").tail(59).len().alias("_window_len"),
])
)
self._live_agg_cache = agg_a.join(agg_b, on="symbol", how="inner")
self._live_agg_cache_date = latest
def _live_agg_baseline_date(self, latest: date) -> date:
"""盘中递推基准日期。当天实时分区存在时使用上一可用交易日。"""
if latest != date.today():
return latest
try:
row = self.execute_one(
"SELECT max(date) FROM kline_enriched WHERE date < ?",
[latest],
)
if row and row[0]:
d = row[0]
return d if isinstance(d, date) else date.fromisoformat(str(d))
except Exception: # noqa: BLE001
pass
return latest
def _build_live_agg_from_parquet(self, latest: date, start_60d: date) -> tuple[pl.DataFrame, pl.DataFrame]:
"""降级路径: 从 parquet 读取数据并计算指标 (当 _enriched_history_cache 不可用时)。"""
from app.indicators.pipeline import compute_indicators
lf = (
pl.scan_parquet(self._enriched_glob)
.filter(pl.col("date") >= start_60d)
.filter(pl.col("date") <= latest)
.sort(["symbol", "date"])
)
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close", "volume",
"raw_close", "raw_high", "raw_low"]
if c in lf.collect_schema().names()]
df_hist = lf.select(read_cols).collect()
if df_hist.is_empty():
return df_hist, pl.DataFrame()
df_with_indicators = compute_indicators(df_hist)
state_cols = [
"symbol",
"ema5", "ema10", "ema20", "ema30", "ema60",
"macd_dea",
"kdj_k", "kdj_d",
"atr_14",
"close", "high", "low",
"annual_vol_20d",
]
existing_state = [c for c in state_cols if c in df_with_indicators.columns]
agg_a = df_with_indicators.filter(pl.col("date") == latest).select(existing_state)
return df_hist, agg_a
def _refresh_instruments(self) -> None:
"""加载 instruments 到内存。"""
try:
df = pl.scan_parquet(self._inst_glob).collect()
if not df.is_empty():
self._instruments_cache = df
logger.info("instruments 缓存已加载: %d", len(df))
except Exception as e: # noqa: BLE001
logger.warning("instruments 缓存刷新失败: %s", e)
def _refresh_index_instruments(self) -> None:
"""加载指数 instruments 到内存。"""
try:
df = pl.scan_parquet(self._index_inst_glob).collect()
if not df.is_empty():
self._index_instruments_cache = df
logger.info("index instruments 缓存已加载: %d", len(df))
except Exception as e: # noqa: BLE001
logger.debug("index instruments 缓存刷新跳过: %s", e)
def get_enriched_latest(self) -> tuple[pl.DataFrame, date | None]:
"""返回缓存的 enriched 最新日 DataFrame + 日期。如无缓存则懒加载。"""
if self._enriched_cache is None:
self._refresh_enriched()
if self._enriched_cache is None:
return pl.DataFrame(), self._enriched_cache_date
return self._enriched_cache, self._enriched_cache_date
def get_enriched_history(self, target_date: date, lookback_days: int) -> pl.DataFrame | None:
"""返回预计算的 enriched 历史数据 (仅 lookback 范围, 不含 warmup)。
warmup 部分在 _refresh_enriched 计算指标时已使用, 策略只需要最终的 lookback 窗口。
返回 ~33万行 (90日历天) 而非 ~107万行, filter_history 策略的 group_by 快 20x+。
"""
cache = self._enriched_history_cache
if cache is None or cache.is_empty():
return None
if "date" not in cache.columns:
return None
cache_max = cache["date"].max()
cache_min = cache["date"].min()
from datetime import timedelta
# 验证缓存覆盖完整范围 (含 warmup)
warmup_start = target_date - timedelta(days=(lookback_days + 60) * 2)
if cache_min > warmup_start or cache_max < target_date:
return None
# 只返回 lookback 范围 (日历天数 ≈ 2/3 交易日, 足够覆盖)
lookback_start = target_date - timedelta(days=lookback_days)
return cache.filter((pl.col("date") >= lookback_start) & (pl.col("date") <= target_date))
def get_enriched_range(
self,
start: date,
end: date,
symbols: list[str] | None = None,
columns: list[str] | None = None,
) -> pl.DataFrame | None:
"""从预计算 enriched 历史缓存返回完整区间;缓存不覆盖时返回 None。"""
if self._enriched_history_cache is None:
self._refresh_enriched()
cache = self._enriched_history_cache
if cache is None or cache.is_empty() or "date" not in cache.columns:
return None
cache_min = cache["date"].min()
cache_max = cache["date"].max()
if cache_min > start or cache_max < end:
return None
df = cache.filter((pl.col("date") >= start) & (pl.col("date") <= end))
if symbols is not None:
df = df.filter(pl.col("symbol").is_in(symbols))
if columns and not df.is_empty():
existing = [c for c in columns if c in df.columns]
if "symbol" not in existing and "symbol" in df.columns:
existing.insert(0, "symbol")
if "date" not in existing and "date" in df.columns:
existing.insert(1, "date")
df = df.select(existing)
return df.sort(["symbol", "date"])
def get_live_agg(self) -> pl.DataFrame:
"""返回盘中实时指标预计算聚合表。如无缓存则懒加载。"""
if self._live_agg_cache is None:
self._refresh_enriched()
if self._live_agg_cache is None:
return pl.DataFrame()
return self._live_agg_cache
def get_instruments(self) -> pl.DataFrame:
"""返回缓存的 instruments DataFrame。如无缓存则懒加载。"""
if self._instruments_cache is None:
self._refresh_instruments()
if self._instruments_cache is None:
return pl.DataFrame()
return self._instruments_cache
def get_index_instruments(self) -> pl.DataFrame:
"""返回缓存的指数 instruments DataFrame。如无缓存则懒加载。"""
if self._index_instruments_cache is None:
self._refresh_index_instruments()
if self._index_instruments_cache is None:
return pl.DataFrame()
return self._index_instruments_cache
def get_index_symbol_set(self) -> set[str]:
"""返回已缓存指数 symbol 集合。"""
df = self.get_index_instruments()
if df.is_empty() or "symbol" not in df.columns:
return set()
return set(df["symbol"].cast(pl.Utf8).to_list())
def enriched_latest_date(self) -> date | None:
"""返回缓存中的 enriched 最新日期。"""
return self._enriched_cache_date
# ================================================================
# 热路径: Polars 查询 (Chart / Screener / Signals / Intraday)
# ================================================================
def get_daily(
self,
symbol: str,
start: date,
end: date,
columns: list[str] | None = None,
) -> pl.DataFrame:
"""单股日K查询 — 从14列parquet读取后即时计算指标。"""
from datetime import timedelta
# 扩展范围用于指标预热 (MA60 需要 ~60 交易日 ≈ 120 日历日)
warmup_start = start - timedelta(days=150)
# 扫描14列 parquet
df = self._scan_daily_symbol(symbol, warmup_start, end, None)
if not df.is_empty():
df = self._compute_enriched_range(df)
# 尝试用缓存数据覆盖最新日 (盘中更准确)
cached, cache_date = self.get_enriched_latest()
if not df.is_empty() and cached is not None and not cached.is_empty() and cache_date:
if start <= cache_date <= end:
cached_part = self._filter_cached(cached, symbol, None)
if not cached_part.is_empty():
df = df.filter(pl.col("date") != cache_date)
common_cols = [c for c in df.columns if c in cached_part.columns]
df = pl.concat([df.select(common_cols), cached_part.select(common_cols)])
# 裁剪到请求范围
if not df.is_empty():
df = df.filter((pl.col("date") >= start) & (pl.col("date") <= end))
if columns and not df.is_empty():
existing = [c for c in columns if c in df.columns]
df = df.select(existing)
return df
def get_daily_batch(
self,
symbols: list[str],
start: date,
end: date,
columns: list[str] | None = None,
) -> pl.DataFrame:
"""批量日K查询。"""
cached, cache_date = self.get_enriched_latest()
if cached is not None and not cached.is_empty() and cache_date:
if start >= cache_date:
return self._filter_cached_batch(cached, symbols, columns)
# 回退 scan_parquet
return self._scan_daily_batch(symbols, start, end, columns)
def get_index_daily(
self,
symbol: str,
start: date,
end: date,
columns: list[str] | None = None,
) -> pl.DataFrame:
"""指数日K查询 — 从独立指数 enriched parquet 读取后即时计算通用指标。"""
from datetime import timedelta
warmup_start = start - timedelta(days=150)
df = self._scan_index_daily_symbol(symbol, warmup_start, end, None)
if not df.is_empty():
df = self._compute_index_enriched_range(df)
df = df.filter((pl.col("date") >= start) & (pl.col("date") <= end))
if columns and not df.is_empty():
existing = [c for c in columns if c in df.columns]
df = df.select(existing)
return df
def get_minute(
self,
symbol: str,
trade_date: date,
) -> pl.DataFrame:
"""分钟K查询 — Polars scan_parquet + predicate pushdown。"""
try:
return pl.scan_parquet(self._minute_glob).filter(
(pl.col("symbol") == symbol)
& (pl.col("datetime").dt.date() == trade_date)
).sort("datetime").collect()
except Exception as e: # noqa: BLE001
logger.warning("分钟K查询失败: %s", e)
return pl.DataFrame()
# ================================================================
# Polars 查询内部方法
# ================================================================
def _compute_enriched_range(self, df: pl.DataFrame) -> pl.DataFrame:
"""对14列enriched数据即时计算完整指标+信号。输入应含足够预热行数。"""
from app.indicators.pipeline import compute_indicators, compute_signals, compute_limit_signals, filter_halt_days
if df.is_empty() or df.height < 2:
return df
# 兜底过滤历史脏数据中的停牌日 (close 可能被填充为前收盘价)
df = filter_halt_days(df)
if df.is_empty() or df.height < 2:
return df
try:
df = compute_indicators(df)
df = compute_signals(df)
instruments = self.get_instruments()
df = compute_limit_signals(df, instruments)
except Exception as e: # noqa: BLE001
logger.warning("on-demand compute failed: %s", e)
return df
def _compute_index_enriched_range(self, df: pl.DataFrame) -> pl.DataFrame:
"""指数只计算通用技术指标和通用信号,跳过涨跌停/股本/市值逻辑。"""
from app.indicators.pipeline import compute_indicators, compute_signals
if df.is_empty() or df.height < 2:
return df
try:
df = compute_indicators(df)
df = compute_signals(df)
except Exception as e: # noqa: BLE001
logger.warning("index on-demand compute failed: %s", e)
return df
def _filter_cached(self, cached: pl.DataFrame, symbol: str, columns: list[str] | None) -> pl.DataFrame:
df = cached.filter(pl.col("symbol") == symbol)
if columns and not df.is_empty():
existing = [c for c in columns if c in df.columns]
df = df.select(existing)
return df
def _filter_cached_batch(self, cached: pl.DataFrame, symbols: list[str], columns: list[str] | None) -> pl.DataFrame:
df = cached.filter(pl.col("symbol").is_in(symbols))
if columns and not df.is_empty():
existing = [c for c in columns if c in df.columns]
df = df.select(existing)
return df.sort(["symbol", "date"])
def _scan_daily_symbol(self, symbol: str, start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
try:
lf = pl.scan_parquet(self._enriched_glob,
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
(pl.col("symbol") == symbol)
& (pl.col("date") >= start)
& (pl.col("date") <= end)
).sort("date")
if columns:
schema_names = lf.collect_schema().names()
existing = [c for c in columns if c in schema_names]
lf = lf.select(existing)
return lf.collect()
except Exception as e: # noqa: BLE001
logger.warning("日K查询失败: %s", e)
return pl.DataFrame()
def _scan_daily_batch(self, symbols: list[str], start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
try:
lf = pl.scan_parquet(self._enriched_glob,
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
(pl.col("symbol").is_in(symbols))
& (pl.col("date") >= start)
& (pl.col("date") <= end)
).sort(["symbol", "date"])
if columns:
schema_names = lf.collect_schema().names()
existing = [c for c in columns if c in schema_names]
lf = lf.select(existing)
return lf.collect()
except Exception as e: # noqa: BLE001
logger.warning("日K批量查询失败: %s", e)
return pl.DataFrame()
def _scan_index_daily_symbol(self, symbol: str, start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
try:
lf = pl.scan_parquet(self._index_enriched_glob,
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
(pl.col("symbol") == symbol)
& (pl.col("date") >= start)
& (pl.col("date") <= end)
).sort("date")
if columns:
schema_names = lf.collect_schema().names()
existing = [c for c in columns if c in schema_names]
lf = lf.select(existing)
return lf.collect()
except Exception as e: # noqa: BLE001
logger.warning("指数日K查询失败: %s", e)
return pl.DataFrame()
def _merge_cached_and_scan(
self,
cached: pl.DataFrame,
cache_date: date,
symbol: str,
start: date,
end: date,
columns: list[str] | None,
) -> pl.DataFrame:
"""合并缓存部分 + scan 历史部分。
历史部分用 strict < cache_date, 避免与缓存重复。
两部分 schema 可能不一致 (增量 vs 全量), concat 前对齐列。
"""
hist = self._scan_daily_symbol(symbol, start, cache_date, columns)
cached_part = self._filter_cached(cached, symbol, columns)
if hist.is_empty():
return cached_part
if cached_part.is_empty():
return hist
# 去重: 历史部分可能包含 cache_date, 去掉后再合并
hist = hist.filter(pl.col("date") < cache_date)
# 对齐列: 取交集, 统一类型
common_cols = [c for c in hist.columns if c in cached_part.columns]
hist = hist.select(common_cols)
cached_part = cached_part.select(common_cols)
# 统一类型: 历史可能是 Float64, 缓存可能是 Int64, 统一为 cast
for c in common_cols:
if hist[c].dtype != cached_part[c].dtype:
# 统一到更宽的类型
target = hist[c].dtype if hist.height > cached_part.height else cached_part[c].dtype
hist = hist.with_columns(pl.col(c).cast(target))
cached_part = cached_part.with_columns(pl.col(c).cast(target))
return pl.concat([hist, cached_part])
# ================================================================
# DuckDB 查询 (冷路径: 统计/元数据/自定义SQL)
# ================================================================
def latest_minute_date(self, symbol: str) -> date | None:
try:
with self._lock:
row = self.db.execute(
"SELECT max(CAST(datetime AS DATE)) FROM kline_minute WHERE symbol = ?",
[symbol],
).fetchone()
if row and row[0]:
return row[0] if isinstance(row[0], date) else date.fromisoformat(str(row[0]))
except duckdb.CatalogException:
pass
return None
def earliest_daily_date(self) -> date | None:
"""本地日K数据的最早日期。"""
try:
with self._lock:
res = self.db.execute(
"SELECT min(date) FROM kline_daily",
).fetchone()
if res and res[0]:
d = res[0]
return d if isinstance(d, date) else date.fromisoformat(str(d))
except Exception:
return None
return None
def earliest_minute_date(self) -> date | None:
"""本地分钟K数据的最早日期。"""
try:
with self._lock:
res = self.db.execute(
"SELECT min(CAST(datetime AS DATE)) FROM kline_minute",
).fetchone()
if res and res[0]:
d = res[0]
return d if isinstance(d, date) else date.fromisoformat(str(d))
except Exception:
return None
return None
def latest_daily_date(self) -> date | None:
"""本地日K数据的最新日期。"""
try:
with self._lock:
res = self.db.execute(
"SELECT max(date) FROM kline_daily",
).fetchone()
if res and res[0]:
d = res[0]
return d if isinstance(d, date) else date.fromisoformat(str(d))
except Exception:
return None
return None
def _latest_enriched_date_duckdb(self) -> date | None:
try:
with self._lock:
res = self.db.execute(
"SELECT max(date) FROM kline_enriched",
).fetchone()
if res and res[0]:
d = res[0]
return d if isinstance(d, date) else date.fromisoformat(str(d))
except Exception: # noqa: BLE001
return None
return None
# ================================================================
# 写入 (Pipeline / Sync)
# ================================================================
def append_daily(self, df: pl.DataFrame) -> None:
"""按日分区写入日K数据 (merge-upsert)。"""
if df.is_empty():
return
self._write_daily_partition(df, "kline_daily")
def append_enriched(self, df: pl.DataFrame) -> None:
"""按日分区写入 enriched 数据 (merge-upsert)。磁盘仅写入 14 列存储列。"""
if df.is_empty():
return
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
df_storage = df.select(storage_cols)
self._write_daily_partition(df_storage, "kline_daily_enriched")
def append_index_daily(self, df: pl.DataFrame) -> None:
"""按日分区写入指数日K数据 (merge-upsert)。"""
if df.is_empty():
return
self._write_daily_partition(df, "kline_index_daily")
def append_index_enriched(self, df: pl.DataFrame) -> None:
"""按日分区写入指数 enriched 数据。磁盘仅写入通用基础行情窄表。"""
if df.is_empty():
return
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
df_storage = df.select(storage_cols)
self._write_daily_partition(df_storage, "kline_index_enriched")
def save_index_instruments(self, df: pl.DataFrame) -> None:
"""保存指数标的维表。"""
if df.is_empty() or "symbol" not in df.columns:
return
out = self.store.data_dir / "instruments_index" / "instruments_index.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
df.unique(subset=["symbol"], keep="last").sort("symbol").write_parquet(out)
self._index_instruments_cache = None
self._refresh_index_instruments()
def refresh_index_views(self) -> None:
"""刷新指数相关 DuckDB 视图。"""
d = self.store.data_dir.as_posix()
statements = [
f"""CREATE OR REPLACE VIEW kline_index_daily AS
SELECT * FROM read_parquet('{d}/kline_index_daily/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW kline_index_enriched AS
SELECT * FROM read_parquet('{d}/kline_index_enriched/**/*.parquet', union_by_name=true)""",
f"""CREATE OR REPLACE VIEW instruments_index AS
SELECT * FROM read_parquet('{d}/instruments_index/**/*.parquet', union_by_name=true)""",
]
for sql in statements:
try:
with self._lock:
self.db.execute(sql)
except Exception as e: # noqa: BLE001
logger.debug("index view refresh skipped: %s", e)
def _write_daily_partition(self, df: pl.DataFrame, table: str) -> None:
"""按 date 分区写入 parquet,每个日期一个文件,支持 merge-upsert。"""
base = self.store.data_dir / table
for date_df in df.partition_by("date"):
dt = date_df["date"][0]
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
out = base / f"date={ds}" / "part.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
if out.exists():
existing = pl.read_parquet(out)
date_df = pl.concat([existing, date_df], how="diagonal_relaxed").unique(
subset=["symbol", "date"], keep="last"
)
date_df = date_df.sort(["symbol", "date"])
date_df.write_parquet(out)
def flush_live_daily(self, df: pl.DataFrame) -> None:
"""覆写当天 kline_daily 分区 (实时行情落盘, 非merge)。"""
if df.is_empty() or "date" not in df.columns:
return
base = self.store.data_dir / "kline_daily"
dt = df["date"][0]
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
out = base / f"date={ds}" / "part.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
df.sort(["symbol", "date"]).write_parquet(out)
def flush_live_enriched(self, df: pl.DataFrame) -> None:
"""覆写当天 kline_daily_enriched 分区 (实时 enriched 落盘, 非merge)。
内存缓存保留完整指标列供各服务使用,磁盘仅写入 14 列存储列。
"""
if df.is_empty() or "date" not in df.columns:
return
# 内存缓存: 保留完整 66 列
self._enriched_cache = df.sort(["symbol"])
dt = df["date"][0]
self._enriched_cache_date = dt
# 磁盘写入: 仅 14 列存储列
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
df_storage = df.select(storage_cols).sort(["symbol"])
base = self.store.data_dir / "kline_daily_enriched"
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
out = base / f"date={ds}" / "part.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
df_storage.write_parquet(out)
+66
View File
@@ -0,0 +1,66 @@
"""请求调度器(§5.6)。
按 capability 分别维护令牌桶。Phase 0:基础实现;Phase 1 接入批量合并、优先级队列。
"""
from __future__ import annotations
import asyncio
import time
from dataclasses import dataclass
from .capabilities import Cap, CapabilitySet
@dataclass
class _Bucket:
capacity: int # 每分钟令牌数
tokens: float
last_refill: float # 单位:秒
def consume(self, n: int = 1) -> float:
"""尝试消费 n 个令牌,返回需要等待的秒数(0 表示无需等待)。"""
now = time.monotonic()
elapsed = now - self.last_refill
# 60s 内补满 capacity,匀速补
refill = elapsed * (self.capacity / 60.0)
self.tokens = min(self.capacity, self.tokens + refill)
self.last_refill = now
if self.tokens >= n:
self.tokens -= n
return 0.0
deficit = n - self.tokens
# 还需多少秒才能补齐
wait = deficit / (self.capacity / 60.0)
# 不预扣,留给下一次再竞争(避免饿死优先级高的请求)
return wait
class Scheduler:
"""每个 capability 一个桶。"""
def __init__(self, capset: CapabilitySet) -> None:
self._capset = capset
self._buckets: dict[Cap, _Bucket] = {}
self._locks: dict[Cap, asyncio.Lock] = {}
for cap, lim in capset.all().items():
if lim.rpm:
self._buckets[cap] = _Bucket(
capacity=lim.rpm,
tokens=lim.rpm,
last_refill=time.monotonic(),
)
self._locks[cap] = asyncio.Lock()
async def acquire(self, cap: Cap, n: int = 1) -> None:
"""阻塞直到拿到 n 个令牌。无桶 = 不限流(由调用方保证)。"""
bucket = self._buckets.get(cap)
if bucket is None:
return
lock = self._locks[cap]
async with lock:
while True:
wait = bucket.consume(n)
if wait == 0:
return
await asyncio.sleep(wait)