Update app.py

app.py

@@ -79,4 +79,163 @@ async def load_model():
     except Exception as e:
         logger.warning(f"Predictor weights not loaded: {e}")
 
-    predictor = Kr
+    predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512)
+    logger.info("Kronos NQ+ES v2 ready.")
+
+
+class CandleOut(BaseModel):
+    timestamp: str
+    open: float
+    high: float
+    low: float
+    close: float
+
+
+class ForecastResponse(BaseModel):
+    instrument: str
+    timeframe: str
+    generated_at: str
+    historical: list[CandleOut]
+    forecast_mean: list[CandleOut]
+    forecast_upper: list[CandleOut]
+    forecast_lower: list[CandleOut]
+    direction: str
+    confidence: float
+    volatility_ratio: float
+
+
+def fetch_candles(ticker: str, interval: str, period: str) -> pd.DataFrame:
+    raw = yf.download(ticker, period=period, interval=interval, progress=False)
+    if raw.empty:
+        raise HTTPException(status_code=502, detail=f"No data for {ticker}")
+    df = raw.reset_index()
+    if hasattr(df.columns, 'levels'):
+        df.columns = [c[0] if isinstance(c, tuple) else c for c in df.columns]
+    rename = {c: c.lower() for c in df.columns}
+    df.rename(columns=rename, inplace=True)
+    if "datetime" in df.columns:
+        df.rename(columns={"datetime": "timestamp"}, inplace=True)
+    elif "date" in df.columns:
+        df.rename(columns={"date": "timestamp"}, inplace=True)
+    df["timestamp"] = pd.to_datetime(df["timestamp"])
+    return df[["timestamp", "open", "high", "low", "close"]].dropna()
+
+
+def run_forecast(df: pd.DataFrame, forecast_bars: int, n_samples: int = 10):
+    lookback = min(len(df), 400)
+    x_df = df.tail(lookback).reset_index(drop=True)
+
+    freq = pd.infer_freq(x_df["timestamp"])
+    if freq is None:
+        delta = x_df["timestamp"].iloc[-1] - x_df["timestamp"].iloc[-2]
+        future_ts = pd.Series([x_df["timestamp"].iloc[-1] + delta * (i + 1) for i in range(forecast_bars)])
+    else:
+        future_ts = pd.Series(pd.date_range(
+            start=x_df["timestamp"].iloc[-1],
+            periods=forecast_bars + 1,
+            freq=freq,
+        )[1:])
+
+    samples = []
+    for _ in range(n_samples):
+        pred_df = predictor.predict(
+            df=x_df[["open", "high", "low", "close"]],
+            x_timestamp=x_df["timestamp"],
+            y_timestamp=future_ts,
+            pred_len=forecast_bars,
+            T=0.3,
+            top_p=0.5,
+            sample_count=1,
+        )
+        samples.append(pred_df[["open", "high", "low", "close"]].values)
+
+    samples = np.array(samples)
+    mean = samples.mean(axis=0)
+    upper = np.percentile(samples, 90, axis=0)
+    lower = np.percentile(samples, 10, axis=0)
+
+    return mean, upper, lower, future_ts
+
+
+def calc_direction(mean_candles: np.ndarray, last_close: float):
+    final_close = mean_candles[-1, 3]
+    pct_change = (final_close - last_close) / last_close * 100
+
+    if pct_change > 0.10:
+        return "BULLISH", min(abs(pct_change) * 30, 95)
+    elif pct_change < -0.10:
+        return "BEARISH", min(abs(pct_change) * 30, 95)
+    else:
+        return "NEUTRAL", max(50 - abs(pct_change) * 150, 10)
+
+
+def calc_vol_ratio(mean_candles: np.ndarray, hist_df: pd.DataFrame):
+    pred_ranges = mean_candles[:, 1] - mean_candles[:, 2]
+    hist_ranges = (hist_df["high"] - hist_df["low"]).tail(len(mean_candles)).values
+    if hist_ranges.mean() == 0:
+        return 1.0
+    return float(pred_ranges.mean() / hist_ranges.mean())
+
+
+def candles_to_list(arr: np.ndarray, timestamps) -> list[CandleOut]:
+    out = []
+    for i, ts in enumerate(timestamps):
+        out.append(CandleOut(
+            timestamp=str(ts),
+            open=round(float(arr[i, 0]), 2),
+            high=round(float(arr[i, 1]), 2),
+            low=round(float(arr[i, 2]), 2),
+            close=round(float(arr[i, 3]), 2),
+        ))
+    return out
+
+
+@app.get("/forecast", response_model=ForecastResponse)
+async def get_forecast(
+    instrument: str = Query("NQ", pattern="^(NQ|ES)$"),
+    timeframe: str = Query("1h", pattern="^(5m|1h)$"),
+):
+    if predictor is None:
+        raise HTTPException(status_code=503, detail="Model still loading")
+
+    ticker = TICKER_MAP[instrument]
+    tf_cfg = TIMEFRAME_MAP[timeframe]
+
+    df = fetch_candles(ticker, tf_cfg["interval"], tf_cfg["period"])
+    logger.info(f"Fetched {len(df)} candles for {instrument} @ {timeframe}")
+
+    mean, upper, lower, future_ts = run_forecast(df, tf_cfg["forecast_bars"])
+
+    last_close = float(df["close"].iloc[-1])
+    direction, confidence = calc_direction(mean, last_close)
+    vol_ratio = calc_vol_ratio(mean, df)
+
+    hist_tail = df.tail(50)
+    historical = [
+        CandleOut(
+            timestamp=str(row.timestamp),
+            open=round(float(row.open), 2),
+            high=round(float(row.high), 2),
+            low=round(float(row.low), 2),
+            close=round(float(row.close), 2),
+        )
+        for row in hist_tail.itertuples()
+    ]
+
+    return ForecastResponse(
+        instrument=instrument,
+        timeframe=timeframe,
+        generated_at=datetime.utcnow().isoformat() + "Z",
+        historical=historical,
+        forecast_mean=candles_to_list(mean, future_ts),
+        forecast_upper=candles_to_list(upper, future_ts),
+        forecast_lower=candles_to_list(lower, future_ts),
+        direction=direction,
+        confidence=round(confidence, 1),
+        volatility_ratio=round(vol_ratio, 2),
+    )
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model_loaded": predictor is not None}