main.py

"""FutureQuery forecast-service — FastAPI ensemble API (TimesFM + Chronos-2).

Run it::

    cd forecast-service
    pip install -r requirements.txt
    python main.py        # serves http://localhost:8008

Endpoints
---------
POST /forecast        -> ensemble forecast (TimesFM + Chronos-2) for a price series
GET  /calibration     -> Brier score, reliability curve, directional accuracy
POST /resolve         -> record a market's realised outcome for calibration
GET  /health          -> model availability / liveness
POST /pipeline/start  -> run the full scan->triage->ensemble->synthesise loop
POST /pipeline/resume -> approve / override the human checkpoint (browser-driven)

Models are loaded once on startup (not per request). The first run downloads
weights from HuggingFace (~2GB) — expected. If a model can't be loaded the
service still starts and falls back to a naive forecast, so the API never hard
-fails (see models.py).
"""

from __future__ import annotations

import logging
import os
from contextlib import asynccontextmanager
from typing import List, Literal, Optional

import numpy as np
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

import calibration
import ensemble
from models import ChronosModel, TimesFMModel
from pipeline_api import router as pipeline_router

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
log = logging.getLogger("forecast.main")

PORT = int(os.environ.get("FORECAST_PORT", "8008"))
SERVICE_VERSION = "1.1.0"
MARKET_CONTEXT_VERSION = "predict_trust_layer.v1"

# Holds the loaded models for the lifetime of the process.
STATE: dict = {"timesfm": None, "chronos": None}


@asynccontextmanager
async def lifespan(app: FastAPI):
    log.info("Loading forecasting models (first run downloads ~2GB)...")
    STATE["timesfm"] = TimesFMModel()
    STATE["chronos"] = ChronosModel()
    calibration.init_db()
    log.info(
        "Models ready — TimesFM available=%s, Chronos available=%s",
        STATE["timesfm"].available,
        STATE["chronos"].available,
    )
    yield
    STATE["timesfm"] = None
    STATE["chronos"] = None


app = FastAPI(title="FutureQuery forecast-service", version=SERVICE_VERSION, lifespan=lifespan)

# Allow the Next.js dev server / browser tools to call the API directly.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(pipeline_router)


# --- request / response models -------------------------------------------

class ForecastRequest(BaseModel):
    prices: List[float] = Field(default_factory=list, description="historical yes_prices")
    covariates: Optional[List[List[float]]] = Field(
        default=None, description="optional related series (past covariates)"
    )
    question_type: Literal["numeric", "event"] = "event"
    horizon: int = Field(default=5, ge=1, le=64)
    question: Optional[str] = None  # optional label, stored for calibration


class ResolveRequest(BaseModel):
    id: str
    outcome: float = Field(ge=0.0, le=1.0, description="1.0 = YES happened, 0.0 = NO")


# --- endpoints ------------------------------------------------------------

@app.get("/health")
def health() -> dict:
    tf = STATE.get("timesfm")
    ch = STATE.get("chronos")
    return {
        "status": "ok",
        "service": "futurequery-forecast-service",
        "version": SERVICE_VERSION,
        "revision": _service_revision(),
        "supports_market_context": True,
        "market_context_version": MARKET_CONTEXT_VERSION,
        "pipeline": {
            "supports_market_context": True,
            "context_source": "predict_trust_layer",
            "market_context_version": MARKET_CONTEXT_VERSION,
        },
        "timesfm": {
            "available": bool(tf and tf.available),
            "label": tf.label if tf else None,
        },
        "chronos2": {
            "available": bool(ch and ch.available),
            "label": ch.label if ch else None,
        },
    }


@app.post("/forecast")
def forecast(req: ForecastRequest) -> dict:
    prices = [float(p) for p in (req.prices or []) if _finite(p)]
    horizon = int(req.horizon)
    qtype = req.question_type

    # --- guards: refuse to forecast when we shouldn't --------------------
    if len(prices) == 0:
        warning = (
            "No price history for event market"
            if qtype == "event"
            else "No data provided"
        )
        payload = ensemble.short_series_response(prices, horizon, qtype, warning)
        return _finalize(payload, req, prices)

    if len(prices) < ensemble.MIN_POINTS:
        warning = (
            f"Series too short for reliable forecast "
            f"(have {len(prices)}, need >= {ensemble.MIN_POINTS})"
        )
        payload = ensemble.short_series_response(prices, horizon, qtype, warning)
        return _finalize(payload, req, prices)

    # --- run both models on startup-loaded handles -----------------------
    tf_model = STATE.get("timesfm") or TimesFMModel()
    ch_model = STATE.get("chronos") or ChronosModel()

    timesfm_fc = tf_model.forecast(prices, horizon, req.covariates)
    chronos_fc = ch_model.forecast(prices, horizon, req.covariates)

    payload = ensemble.full_response(prices, timesfm_fc, chronos_fc, qtype, horizon)
    return _finalize(payload, req, prices)


@app.get("/calibration")
def get_calibration() -> dict:
    return calibration.compute_calibration()


@app.post("/resolve")
def resolve(req: ResolveRequest) -> dict:
    found = calibration.record_resolution(req.id, req.outcome)
    return {"ok": found, "id": req.id}


# --- helpers --------------------------------------------------------------

def _finalize(payload: dict, req: ForecastRequest, prices: List[float]) -> dict:
    """Attach a calibration id and persist the forecast for later scoring."""
    headline = ensemble.headline_probability(payload)
    p_market = prices[-1] if prices else None
    trend = payload.get("ensemble", {}).get("trend")
    brier = payload.get("ensemble", {}).get("brier_estimate")
    try:
        fid = calibration.record_forecast(
            question=req.question or "",
            question_type=req.question_type,
            horizon=int(req.horizon),
            p_forecast=headline,
            p_market=p_market,
            trend=trend,
            brier_estimate=brier,
            use_forecast=bool(payload.get("use_forecast", False)),
        )
        payload["id"] = fid
    except Exception as exc:  # never let logging break a forecast
        log.warning("could not persist forecast for calibration: %s", exc)
        payload["id"] = None
    return payload


def _service_revision() -> str | None:
    for name in (
        "FORECAST_SERVICE_COMMIT",
        "RENDER_GIT_COMMIT",
        "RAILWAY_GIT_COMMIT_SHA",
        "FLY_MACHINE_VERSION",
        "SOURCE_VERSION",
        "GIT_COMMIT",
    ):
        value = os.environ.get(name)
        if value:
            return value[:64]
    return None


def _finite(value) -> bool:
    try:
        return np.isfinite(float(value))
    except (TypeError, ValueError):
        return False


if __name__ == "__main__":
    import uvicorn

    print("\n" + "=" * 56, flush=True)
    print("  forecast-service is READY when the next line says", flush=True)
    print("  'Uvicorn running on http://127.0.0.1:8008'", flush=True)
    print("=" * 56 + "\n", flush=True)
    # bind loopback (127.0.0.1): avoids the Windows firewall prompt and
    # localhost/IPv6 mismatches that can stop the server from coming up.
    # 127.0.0.1 locally; set FORECAST_HOST=0.0.0.0 in a container so the host can reach it.
    uvicorn.run(
        "main:app",
        host=os.environ.get("FORECAST_HOST", "127.0.0.1"),
        port=PORT,
        reload=False,
    )