141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
"""
|
|
FastAPI routes for Sanguo Quant API
|
|
"""
|
|
from fastapi import APIRouter, HTTPException, Depends, WebSocket, Query, Header
|
|
from pydantic import BaseModel
|
|
from .schemas import CtaBacktestRequest, OptimizeRequest, FactorAnalysisRequest
|
|
from .auth import verify_token as verify_token_impl, verify_password, create_token
|
|
from .ws import manager
|
|
|
|
|
|
router = APIRouter()
|
|
_orchestrator = None
|
|
_auth_config = {"username": "admin", "password_hash": "", "jwt_secret": "x", "expire_minutes": 60}
|
|
|
|
|
|
def set_orchestrator(orch):
|
|
"""Set the global orchestrator instance"""
|
|
global _orchestrator
|
|
_orchestrator = orch
|
|
|
|
|
|
def get_orchestrator():
|
|
"""Get the global orchestrator instance"""
|
|
return _orchestrator
|
|
|
|
|
|
def set_auth_config(cfg):
|
|
"""Set authentication configuration"""
|
|
_auth_config.update(cfg)
|
|
|
|
|
|
class LoginRequest(BaseModel):
|
|
"""Login request schema"""
|
|
username: str
|
|
password: str
|
|
|
|
|
|
async def verify_token(authorization: str | None = Header(None)):
|
|
"""Dependency to verify JWT token from Authorization header"""
|
|
if authorization is None:
|
|
raise HTTPException(status_code=401, detail="Missing authorization header")
|
|
|
|
if not authorization.startswith("Bearer "):
|
|
raise HTTPException(status_code=401, detail="Invalid authorization header format")
|
|
|
|
token = authorization.split(" ")[1]
|
|
return verify_token_impl(token)
|
|
|
|
|
|
@router.post("/auth/login")
|
|
def login(req: LoginRequest):
|
|
"""Authenticate user and return JWT token"""
|
|
if req.username != _auth_config["username"] or not verify_password(req.password, _auth_config["password_hash"]):
|
|
raise HTTPException(status_code=401, detail="用户名或密码错误")
|
|
return {"token": create_token(req.username)}
|
|
|
|
|
|
@router.post("/backtest/cta", dependencies=[Depends(verify_token)])
|
|
async def submit_cta(req: CtaBacktestRequest):
|
|
"""Submit CTA backtest task"""
|
|
tid = await get_orchestrator().submit_cta(
|
|
strategy_class=req.strategy,
|
|
symbol=req.symbol,
|
|
params=req.params,
|
|
start=req.start,
|
|
end=req.end,
|
|
cfg=None
|
|
)
|
|
return {"task_id": tid}
|
|
|
|
|
|
@router.post("/backtest/optimize", dependencies=[Depends(verify_token)])
|
|
async def submit_optimize(req: OptimizeRequest):
|
|
"""Submit optimization task"""
|
|
tid = await get_orchestrator().submit_optimize(
|
|
strategy_class=req.strategy,
|
|
symbol=req.symbol,
|
|
grid=req.grid,
|
|
start=req.start,
|
|
end=req.end,
|
|
cfg=None
|
|
)
|
|
return {"task_id": tid}
|
|
|
|
|
|
@router.post("/factor/analyze", dependencies=[Depends(verify_token)])
|
|
async def submit_factor(req: FactorAnalysisRequest):
|
|
"""Submit factor analysis task"""
|
|
tid = await get_orchestrator().submit_factor(
|
|
symbols=req.symbols,
|
|
factor_names=req.factor_names,
|
|
start=req.start,
|
|
end=req.end,
|
|
cfg=None,
|
|
output_dir="/tmp/factor"
|
|
)
|
|
return {"task_id": tid}
|
|
|
|
|
|
@router.get("/task/{task_id}", dependencies=[Depends(verify_token)])
|
|
def get_status(task_id: str):
|
|
"""Get task status"""
|
|
s = get_orchestrator().get_status(task_id)
|
|
if s is None:
|
|
raise HTTPException(status_code=404, detail="task not found")
|
|
stage = get_orchestrator().pool.get_stage(task_id)
|
|
return {
|
|
"task_id": task_id,
|
|
"status": s.value if hasattr(s, "value") else str(s),
|
|
"stage": stage or ""
|
|
}
|
|
|
|
|
|
@router.get("/task/{task_id}/result", dependencies=[Depends(verify_token)])
|
|
def get_result(task_id: str):
|
|
"""Get task result"""
|
|
r = get_orchestrator().get_result(task_id)
|
|
if r is None:
|
|
raise HTTPException(status_code=404, detail="result not ready")
|
|
return {"task_id": task_id, "statistics": r.statistics}
|
|
|
|
|
|
@router.websocket("/ws/task/{task_id}")
|
|
async def task_ws(websocket: WebSocket, task_id: str, token: str = Query(...)):
|
|
"""WebSocket endpoint for task status updates"""
|
|
try:
|
|
verify_token(token)
|
|
except Exception:
|
|
await websocket.close(code=4401)
|
|
return
|
|
|
|
await websocket.accept()
|
|
manager.connect(task_id, websocket)
|
|
|
|
try:
|
|
while True:
|
|
await websocket.receive_text() # Keep connection alive
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
manager.disconnect(task_id, websocket) |