"""智能选股域路由：自然语言选股（NDJSON 流式）+ 提问历史 + 全市场同步 + 个股预览。""" from __future__ import annotations import asyncio import json from datetime import datetime import pandas as pd from fastapi import APIRouter, Depends, HTTPException, Response from fastapi.responses import StreamingResponse from sqlalchemy import func, select, text from sqlalchemy.ext.asyncio import AsyncSession from .. import cache from .. import indicators as ind from ..auth import require_user from ..config import settings from ..data import fetcher, repository from ..data.aggregation import resample_bars from ..data.symbols import plain_code from ..db import async_session, get_session from ..models import AdjFactor, Candle, ScreenerQuery from ..schemas import ( CandleOut, PreviewInfoOut, PreviewResponse, ScreenerQueryListResponse, ScreenerQueryOut, ScreenerRunRequest, ScreenerRunResponse, ScreenerSyncRequest, ScreenerSyncStatus, ) from ..screener import engine, market_sync from ..screener.engine import DataNotReadyError from ..screener.llm import ScreenerError, parse_conditions from ._deps import ( ADJUST_MODES, FACTOR_STEP_SQL, FULL_MA_SET, INFO_SQL, adjust_bars, cached_json_response, raw_json, rows_to_bars, series_to_jsonable, ) router = APIRouter() # ---------- 智能选股 ---------- @router.post("/screener/run") async def screener_run( req: ScreenerRunRequest, session: AsyncSession = Depends(get_session), user=Depends(require_user), ) -> StreamingResponse: """自然语言 -> LLM 解析条件 -> 全市场筛选。也可直传 conditions 跳过 LLM（微调再跑）。 NDJSON 流式响应（每行一个 JSON 事件，前端逐行渲染进度）： {"type":"stage","key":"llm|date|prefilter|bars|filter_done|done","msg":"…","ms":123} {"type":"parsed","conditions":{…},"ms":456} LLM 解析出的结构化条件 {"type":"candidates","count":5400,"msg":"…","ms":…} SQL 预筛后的候选数 {"type":"progress","done":500,"total":5400} 逐股指标过滤进度 {"type":"result","result":{…ScreenerRunResponse…},"ms":…} {"type":"error","message":"…","code":400} 流中途失败（HTTP 已 200） 成功的提问（含解析出的条件与命中数）记录到 screener_queries，供历史一键重跑。 """ limit = settings.screener_default_limit async def gen(): try: if req.conditions: conds = req.conditions else: yield _ndjson({"type": "stage", "key": "llm", "msg": f"AI 解析条件中（{settings.llm_model}）…"}) conds = await parse_conditions(req.text) if not conds.indicator and not conds.snapshot: yield _ndjson({"type": "error", "code": 400, "message": "AI 未从描述中解析出任何筛选条件，请换种说法"}) return yield _ndjson({"type": "parsed", "conditions": conds.model_dump()}) result = None async for ev in engine.run_screen_events(session, conds, limit): if ev["type"] == "result": result = ev["result"] yield _ndjson({"type": "stage", "key": "done", "ms": ev.get("ms"), "msg": f"筛选完成：{result['total']} 只命中（数据基准 {result['trade_date']:%Y-%m-%d}）"}) else: yield _ndjson(ev) if result is None: yield _ndjson({"type": "error", "code": 500, "message": "选股流程未产出结果"}) return yield _ndjson({"type": "result", "result": ScreenerRunResponse(**result).model_dump(mode="json")}) # 相同文本 + 相同条件的上一条不重复记录（一键重跑场景） exists = ( await session.execute( select(ScreenerQuery.id).where( ScreenerQuery.user_id == user.id, ScreenerQuery.text == req.text.strip(), ScreenerQuery.conditions_json == json.dumps(conds.model_dump(), ensure_ascii=False), ) ) ).scalar_one_or_none() if exists is None: session.add(ScreenerQuery( user_id=user.id, text=req.text.strip(), conditions_json=json.dumps(conds.model_dump(), ensure_ascii=False), hit_count=result.get("total", 0), )) await session.commit() except DataNotReadyError as e: yield _ndjson({"type": "error", "code": 409, "message": str(e)}) except ValueError as e: # 未知指标/字段、条件为空 yield _ndjson({"type": "error", "code": 400, "message": str(e)}) except ScreenerError as e: code = 503 if "未配置 LLM_API_KEY" in str(e) else 502 yield _ndjson({"type": "error", "code": code, "message": str(e)}) except Exception as e: # noqa: BLE001 yield _ndjson({"type": "error", "code": 500, "message": f"选股失败: {e}"}) return StreamingResponse(gen(), media_type="application/x-ndjson", headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"}) def _ndjson(obj: dict) -> str: """dict -> NDJSON 行（json.dumps 保证 default=str 兜底 datetime 等）。""" return json.dumps(obj, ensure_ascii=False, default=str) + "\n" @router.get("/screener/queries", response_model=ScreenerQueryListResponse) async def screener_queries( limit: int = 20, session: AsyncSession = Depends(get_session), user=Depends(require_user), ) -> ScreenerQueryListResponse: """当前用户的提问历史（最新在前，含解析出的条件与命中数，可一键重跑）。""" limit = max(1, min(limit, 100)) rows = ( await session.execute( select(ScreenerQuery) .where(ScreenerQuery.user_id == user.id) .order_by(ScreenerQuery.created_at.desc()) .limit(limit) ) ).scalars().all() items = [] for r in rows: conds = None if r.conditions_json: try: from ..schemas import ScreenConditions conds = ScreenConditions.model_validate_json(r.conditions_json) except Exception: # noqa: BLE001 —— 旧格式/解析失败则只展示文本 conds = None items.append(ScreenerQueryOut( id=r.id, text=r.text, conditions=conds, hit_count=r.hit_count, created_at=r.created_at )) return ScreenerQueryListResponse(items=items) @router.delete("/screener/queries/{query_id}", status_code=204) async def screener_query_delete( query_id: int, session: AsyncSession = Depends(get_session), user=Depends(require_user), ) -> None: await session.execute( text("DELETE FROM screener_queries WHERE id = :i AND user_id = :u"), {"i": query_id, "u": user.id}, ) await session.commit() # ---------- 全市场数据同步 ---------- @router.post("/screener/sync", response_model=ScreenerSyncStatus) async def screener_sync_start( req: ScreenerSyncRequest, session: AsyncSession = Depends(get_session) ) -> ScreenerSyncStatus: """启动全市场数据同步（后台任务，立即返回状态）。""" try: await market_sync.start_sync(session, req.days, req.force) except ScreenerError as e: raise HTTPException(status_code=503, detail=str(e)) status = await market_sync.get_sync_status(session) return ScreenerSyncStatus(**{k: status.get(k) for k in ScreenerSyncStatus.model_fields}) @router.get("/screener/sync/status", response_model=ScreenerSyncStatus) async def screener_sync_status(session: AsyncSession = Depends(get_session)) -> ScreenerSyncStatus: """同步任务状态 + 数据实况（最新交易日/行数/ready）。""" status = await market_sync.get_sync_status(session) return ScreenerSyncStatus(**{k: status.get(k) for k in ScreenerSyncStatus.model_fields}) # ---------- 个股详情预览 ---------- @router.get("/screener/preview/{ts_code}", response_model=PreviewResponse) async def screener_preview( ts_code: str, limit: int = 500, adjust: str = "qfq", timeframe: str = "1d", mas: str = "5,10,20,60", zx: str = "10,20,30,60", end: str | None = None, session: AsyncSession = Depends(get_session), ) -> Response: """个股详情预览：日线（candles 不复权底座 + adj_factor 本地换算 bfq/qfq/hfq， 未缓存自动拉取，落后全市场最新交易日则强制刷新）+ 全套指标 + 最新截面信息卡。 timeframe 聚合到周/月/年（先复权再聚合）；mas 指定主图 MA 周期（逗号分隔）。 end=YYYY-MM-DD 时为「向前翻页」：返回该日之前最近 limit 根（含预热计算指标）， has_more 标记窗口前是否还有更早历史，前端据此继续向左滚动加载。""" if adjust not in ADJUST_MODES: raise HTTPException(status_code=400, detail=f"adjust 仅支持 {'/'.join(ADJUST_MODES)}") if timeframe not in ("1d", "1w", "1M", "1y"): raise HTTPException(status_code=400, detail="timeframe 仅支持 1d/1w/1M/1y") try: ma_periods = sorted({int(p) for p in mas.split(",") if p.strip().isdigit() and 1 <= int(p) <= 500}) except ValueError: raise HTTPException(status_code=400, detail="mas 格式应为逗号分隔的数字，如 5,10,20,60") if not ma_periods: ma_periods = [5, 10, 20, 60] try: zx_periods = sorted({int(p) for p in zx.split(",") if p.strip().isdigit() and 1 <= int(p) <= 500}) except ValueError: raise HTTPException(status_code=400, detail="zx 格式应为逗号分隔的数字，如 10,20,30,60") if not zx_periods: zx_periods = [10, 20, 30, 60] limit = max(30, min(limit, 5000)) end_dt: datetime | None = None if end: try: end_dt = datetime.strptime(end.strip()[:10], "%Y-%m-%d") except ValueError: raise HTTPException(status_code=400, detail="end 格式应为 YYYY-MM-DD") symbol = plain_code(ts_code) # --- 两级读缓存：历史窗口（end 翻页）只增不改，最新窗口每日由全市场同步推进； # 键含 ver:candles 版本号（同步完成后自增，旧缓存全部失效），TTL 兜底（cache.py）。 # 存序列化好的 JSON 直返（j: 前缀），跳过 json.loads + pydantic 校验/序列化（热路径数百 ms → 个位数）。 # 注：ma_periods 不参与缓存键 —— 前端已改为本地计算 MA，后端始终返回全量 MA 集合 cache_key = cache.digest( "preview", ts_code, timeframe, limit, adjust, end_dt.strftime("%Y-%m-%d") if end_dt else None, await cache.get_version("candles"), ) cached = await cached_json_response(f"pvj:{cache_key}") if cached is not None: return cached # --- 日线：candles(全量不复权底座)；未缓存拉取，落后于全市场最新交易日则强制刷新 --- # fetcher 只做「不复权」增量 upsert，底座口径恒为 bfq（TDX 全量 + Tushare 增量）， # 复权（qfq/hfq）读取时按 adj_factor 表本地换算。 # 每次只取「窗口 + 400 根预热」行（MA250/MACD EMA 在 400 根内充分收敛），不拉全量： # 首屏 ~500 根秒开，向左滚动时按 end 参数逐页向前翻。 frame_mult = {"1d": 1, "1w": 6, "1M": 24, "1y": 280}[timeframe] fetch_n = min(100000, limit * frame_mult + 400) source = "bfq" mode = "bfq" # 并发约定：注入 session 与 s2 各占一条连接，每次 gather 里每个 session 恰好跑一条查询 # （AsyncSession 单连接非并发安全），把 ~6 次串行 DB RTT 折叠成 2 个波次。 async with async_session() as s2: if end_dt is not None: # 向前翻页：取 end 之前的历史窗口，不触发同步（历史浏览）；max(ts) 用不到 rows = await repository.get_candles_before(session, symbol, "1d", before=end_dt, limit=fetch_n) global_latest = None else: # Wave 1：candles 窗口（注入 session）+ 全市场最新交易日（s2）并行 rows, global_latest = await asyncio.gather( repository.get_recent_candles(session, symbol, "1d", limit=fetch_n), s2.scalar(select(func.max(Candle.ts)).where(Candle.timeframe == "1d")), ) try: if not rows: await fetcher.sync_symbol(session, symbol, source="auto") rows = await repository.get_recent_candles(session, symbol, "1d", limit=fetch_n) elif global_latest is not None and rows[-1].ts.date() < global_latest.date(): await fetcher.sync_symbol(session, symbol, source="auto", force=True) rows = await repository.get_recent_candles(session, symbol, "1d", limit=fetch_n) except Exception: # noqa: BLE001 —— tushare/写库失败时回滚会话（否则毒化后兜底查询 500） await session.rollback() if not rows: rows = [] bars = rows_to_bars(rows) if not bars and end_dt is None: raise HTTPException(status_code=404, detail=f"无数据: {ts_code}（可先点「同步市场数据」）") # 信息卡取未聚合的日线最新 bar（聚合后 ts 是周期起点，不适用于「最新交易日」） last_daily = bars[-1] if bars else None prev_daily = bars[-2] if len(bars) > 1 else None # 翻页到底（end 之前无数据）：返回空页 + has_more=False，前端停止向前翻页 # --- Wave 2：复权因子（s2，覆盖索引 Index Only Scan）+ 信息卡（注入 session，LATERAL 一条）并行 --- async def _fetch_factors() -> list | None: if adjust == mode or not bars: return None # 只取因子「变化点」行（覆盖索引 Index Only Scan，免堆访问——adj_factor 堆碎片化 # 严重）；bisect 在阶梯函数上取值与日级序列逐字节一致 if end_dt is not None: # 分页：窗口 ≤ end 的变化点 + 全局最新因子（qfq 以最新因子归一） win = list((await s2.execute(FACTOR_STEP_SQL, {"code": ts_code, "upto": end_dt})).all()) if win: latest_f = (await s2.execute( select(AdjFactor.trade_date, AdjFactor.adj_factor) .where(AdjFactor.ts_code == ts_code) .order_by(AdjFactor.trade_date.desc()).limit(1) )).first() if latest_f is not None: win.append(latest_f) return win or None # 非分页：上界 global_latest（≥ 最新 bar），末项变化点即全局最新因子，比 # 「窗口 ≤ bars[-1].ts + 单独 latest」少一次查询 return list((await s2.execute( FACTOR_STEP_SQL, {"code": ts_code, "upto": global_latest} )).all()) or None factors, info_row = await asyncio.gather( _fetch_factors(), session.execute(INFO_SQL, {"code": ts_code, "target": last_daily.ts if last_daily else None}), ) # --- 复权换算：请求模式与底座模式不同时按 adj_factor 本地换算（无因子则维持原样） --- if factors: bars = adjust_bars(bars, factors, mode, adjust) mode = adjust source = adjust # --- 周期聚合：复权之后按日历聚合到周/月/年，指标在聚合后的序列上计算 --- bars = resample_bars(bars, timeframe) # --- 指标（在预热窗口上计算后截尾，保证预热正确；翻页到底的空页跳过） --- has_more = len(bars) > limit # 返回窗口之前还有更早历史（含预热行） indicators: dict[str, dict[str, list[float | None]]] = {} if bars: df = pd.DataFrame({"close": [b.close for b in bars], "high": [b.high for b in bars], "low": [b.low for b in bars]}) closes, highs, lows = df["close"], df["high"], df["low"] macd = ind.macd(closes) kdj = ind.kdj(highs, lows, closes) boll = ind.bollinger(closes) indicators = { # MA 始终返回全量集合（前端本地计算 MA，此处仅保留兼容；缓存键不依赖 ma_periods） "ma": {f"ma{p}": series_to_jsonable(ind.ma(closes, p)) for p in FULL_MA_SET}, "macd": { "dif": series_to_jsonable(macd["macd"]), "dea": series_to_jsonable(macd["signal"]), "hist": series_to_jsonable(macd["hist"]), }, "kdj": {k: series_to_jsonable(kdj[k]) for k in ("k", "d", "j")}, "rsi": { "rsi6": series_to_jsonable(ind.rsi(closes, 6)), "rsi12": series_to_jsonable(ind.rsi(closes, 12)), "rsi24": series_to_jsonable(ind.rsi(closes, 24)), }, "boll": {k: series_to_jsonable(boll[k]) for k in ("upper", "mid", "lower")}, "zx": { "short": series_to_jsonable(ind.ema2(closes)), "duokong": series_to_jsonable(ind.avg_ma(closes, tuple(zx_periods))), }, } limit = max(30, min(limit, len(bars))) for group in indicators.values(): for key in group: group[key] = group[key][-limit:] # --- 信息卡：Wave 2 已并行取回（stock_basic + 与行情同日对齐的快照、缺则最新日，见 INFO_SQL） --- row = info_row.first() if row is not None: m = row._mapping sb_name, sb_industry, sb_area, sb_market, sb_list_date = ( m["name"], m["industry"], m["area"], m["market"], m["list_date"] ) ds_turnover, ds_pe, ds_pb, ds_tmv, ds_cmv = ( m["turnover_rate"], m["pe_ttm"], m["pb"], m["total_mv"], m["circ_mv"] ) else: sb_name = sb_industry = sb_area = sb_market = sb_list_date = None ds_turnover = ds_pe = ds_pb = ds_tmv = ds_cmv = None def _yi(v) -> float | None: if v is None: return None v = float(v) return None if v != v else round(v / 1e4, 2) # 万元 -> 亿元 info = PreviewInfoOut( ts_code=ts_code, symbol=symbol, name=sb_name or ts_code, industry=sb_industry, area=sb_area, market=sb_market, list_date=sb_list_date, trade_date=last_daily.ts if last_daily else None, open=last_daily.open if last_daily else None, high=last_daily.high if last_daily else None, low=last_daily.low if last_daily else None, close=last_daily.close if last_daily else None, pre_close=prev_daily.close if prev_daily else None, pct_chg=((last_daily.close / prev_daily.close - 1) * 100) if last_daily and prev_daily and prev_daily.close else None, volume_hand=round(last_daily.volume / 100, 0) if last_daily else None, # 股 -> 手 amount_yi=round(last_daily.amount / 1e8, 2) if last_daily and last_daily.amount else None, # 元 -> 亿元 turnover_rate=ds_turnover, pe_ttm=ds_pe, pb=ds_pb, total_mv=_yi(ds_tmv), circ_mv=_yi(ds_cmv), ) candles = [ CandleOut(ts=b.ts, open=b.open, high=b.high, low=b.low, close=b.close, volume=b.volume, amount=b.amount, turnover=b.turnover) for b in bars[-limit:] ] resp = PreviewResponse(ts_code=ts_code, symbol=symbol, source=source, info=info, candles=candles, indicators=indicators, has_more=has_more) # 只序列化一次：本地（同步，120s）+ Redis（后台写，600s TTL 兜底跨进程/重启） raw = raw_json(resp) cache.local_set(f"pvj:{cache_key}", raw, ttl=120) cache.set_bg(f"pvj:{cache_key}", raw, ttl=600) return Response(content=raw, media_type="application/json")