Demo routes for webpage added + call_limits + file_limits

main.py

@@ -1,11 +1,17 @@
-import os
-import requests
+import os, httpx
 from fastapi import FastAPI, Request, HTTPException, UploadFile, File, Form
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from slowapi.middleware import SlowAPIMiddleware
 
 # --- Secrets set in the PUBLIC Space ---
-UPSTREAM_URL = os.environ.get("hf_url", "").rstrip("/")   # e.g., https://datfid-org-datfid-api.hf.space
-HF_TOKEN     = os.environ.get("hf_token")                 # HF access token to access your PRIVATE Space
+UPSTREAM_URL = os.environ.get("hf_url", "").rstrip("/")             # url to access Private Space
+HF_TOKEN     = os.environ.get("hf_token")                           # HF access token to access Private Space
+DEMO_FORWARD_URL   = os.getenv("DEMO_FORWARD_URL", "").rstrip("/")  # url to acces demo space 
+DATFID_DEMO_TOKEN  = os.getenv("DATFID_DEMO_TOKEN", "")             # token to access demo space
 
 if not HF_TOKEN:
     raise RuntimeError("Missing secret 'hf_token' in public Space.")
@@ -14,6 +20,59 @@ if not UPSTREAM_URL.startswith("https://"):
 
 app = FastAPI(title="DATFID Public Proxy", docs_url="/docs", redoc_url=None)
 
+# Prefer X-Forwarded-For (HF sits behind a proxy)
+def client_ip(request: Request):
+    xff = request.headers.get("x-forwarded-for")
+    return xff.split(",")[0].strip() if xff else (request.client.host or "0.0.0.0")
+
+limiter = Limiter(key_func=client_ip)  # or get_remote_address
+app.state.limiter = limiter
+app.add_middleware(SlowAPIMiddleware)
+
+@app.exception_handler(RateLimitExceeded)
+async def ratelimit_handler(request: Request, exc: RateLimitExceeded):
+    return JSONResponse(status_code=429, content={"detail": "Too many requests, slow down."})
+
+SDK_MAX_BODY_BYTES = int(os.getenv("SDK_MAX_BODY_BYTES", "25000000"))  # 25MB default for demo
+SDK_MAX_BODY_BYTES_extended = int(os.getenv("SDK_MAX_BODY_BYTES_extended", "125000000")) # 125MB default for prod 
+
+# Early global body-size guard (runs before routes)
+@app.middleware("http")
+async def limit_body_size(request: Request, call_next):
+    cl = request.headers.get("content-length")
+
+    # demo routes use the smaller cap, others use extended
+    path = request.url.path or ""
+    cap = SDK_MAX_BODY_BYTES if "-demo" in path else SDK_MAX_BODY_BYTES_extended
+
+    if cl and int(cl) > cap:
+        return JSONResponse(
+            {"detail": f"Payload too large (> {cap} bytes)"},
+            status_code=413,
+        )
+    return await call_next(request)
+
+# CORS for browser
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "https://datfid.com",
+        "https://www.datfid.com"
+    ],
+    # Optional: allow Vercel preview domains
+    # allow_origin_regex=r"^https:\/\/.*\.vercel\.app$",
+    allow_methods=["POST", "OPTIONS"],
+    allow_headers=["Content-Type"],   # no Authorization header needed from browser
+    allow_credentials=False,
+    max_age=86400,
+)
+
+# to ensure we don’t leak hop-by-hop headers
+def _filter_resp_headers(h):
+    # pass through useful headers but strip hop-by-hop
+    allowed = {"content-type", "content-disposition", "content-length"}
+    return {k: v for k, v in h.items() if k.lower() in allowed}
+
 def _extract_user_token(req: Request) -> str | None:
     """
     Read user's DATFID token from Authorization: Bearer <dt+...>.
@@ -24,7 +83,7 @@ def _extract_user_token(req: Request) -> str | None:
         return None
     return auth.split(" ", 1)[1].strip()
 
-def _forward(path: str, method: str = "GET", json_body=None, user_token: str | None = None):
+async def _forward(path: str, method: str = "GET", json_body=None, user_token: str | None = None):
     """
     Forward request to the PRIVATE Space:
       - 'Authorization: Bearer <HF_TOKEN>' to pass HF private gate
@@ -38,7 +97,10 @@ def _forward(path: str, method: str = "GET", json_body=None, user_token: str | N
     if user_token:
         headers["X-API-Key"] = user_token
 
-    r = requests.request(method, url, headers=headers, json=json_body, timeout=600)
+    timeout = httpx.Timeout(600.0)
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        r = await client.request(method, url, headers=headers, json=json_body)
+
     ct = r.headers.get("content-type", "")
 
     if "application/json" in ct:
@@ -49,7 +111,7 @@ def _forward(path: str, method: str = "GET", json_body=None, user_token: str | N
     # Fallback: return short text envelope if non-JSON
     return JSONResponse(status_code=r.status_code, content={"text": r.text[:1000]})
 
-def _forward_stream(path: str, files=None, data=None, user_token: str | None = None, method: str = "POST"):
+async def _forward_stream(path: str, files=None, data=None, user_token: str | None = None, method: str = "POST"):
     url = f"{UPSTREAM_URL}{path}"
     headers = {
         "Authorization": f"Bearer {HF_TOKEN}",
@@ -58,20 +120,41 @@ def _forward_stream(path: str, files=None, data=None, user_token: str | None = N
     if user_token:
         headers["X-API-Key"] = user_token
 
-    r = requests.request(method, url, headers=headers, files=files, data=data, timeout=600)
-    # Pass through status, content-type, and content-disposition if present
-    media_type = r.headers.get("content-type", "application/octet-stream")
-    cd = r.headers.get("content-disposition")
-    resp = Response(content=r.content, status_code=r.status_code, media_type=media_type)
-    if cd:
-        resp.headers["Content-Disposition"] = cd
-    return resp
+    timeout = httpx.Timeout(600.0)
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
+            if resp.status_code >= 400:
+                text = await resp.aread()
+                return Response(content=text, status_code=resp.status_code, media_type=resp.headers.get("content-type","text/plain"))
+            return StreamingResponse(resp.aiter_raw(), status_code=resp.status_code, headers=_filter_resp_headers(resp.headers))
+
+# for demo
+async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | None, method: str = "POST"):
+    if not DEMO_FORWARD_URL or not HF_TOKEN or not DATFID_DEMO_TOKEN:
+        raise HTTPException(status_code=500, detail="Demo not configured.")
+
+    url = DEMO_FORWARD_URL + path
+    headers = {
+        # HF private-space gate:
+        "Authorization": f"Bearer {HF_TOKEN}",
+        # App-level demo token (checked by the private API):
+        "X-DATFID-Token": DATFID_DEMO_TOKEN,
+    }
+
+    timeout = httpx.Timeout(120.0)
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        # stream response back to the caller (so large CSVs don’t load fully into memory)
+        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
+            if resp.status_code >= 400:
+                text = await resp.aread()
+                raise HTTPException(status_code=resp.status_code, detail=text.decode(errors="ignore"))
+            return StreamingResponse(resp.aiter_raw(), status_code=resp.status_code, headers=_filter_resp_headers(resp.headers))
 
 @app.get("/")
 async def root(req: Request):
     # Forward to private root (private gate still needs HF token)
     user_token = _extract_user_token(req)  # optional here
-    return _forward("/", "GET", user_token=user_token)
+    return await _forward("/", "GET", user_token=user_token)
 
 @app.get("/secure-ping/")
 async def secure_ping(req: Request):
@@ -79,7 +162,7 @@ async def secure_ping(req: Request):
     user_token = _extract_user_token(req)
     if not user_token:
         raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
-    return _forward("/secure-ping/", "GET", user_token=user_token)
+    return await _forward("/secure-ping/", "GET", user_token=user_token)
 
 @app.post("/modelfit/")
 async def modelfit(req: Request):
@@ -87,7 +170,7 @@ async def modelfit(req: Request):
     if not user_token:
         raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
     body = await req.json()
-    return _forward("/modelfit/", "POST", json_body=body, user_token=user_token)
+    return await _forward("/modelfit/", "POST", json_body=body, user_token=user_token)
 
 @app.post("/modelforecast/")
 async def modelforecast(req: Request):
@@ -95,7 +178,7 @@ async def modelforecast(req: Request):
     if not user_token:
         raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
     body = await req.json()
-    return _forward("/modelforecast/", "POST", json_body=body, user_token=user_token)
+    return await _forward("/modelforecast/", "POST", json_body=body, user_token=user_token)
 
 @app.post("/modelfit-file/")
 async def modelfit_file(
@@ -115,8 +198,12 @@ async def modelfit_file(
     if not user_token:
         raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
 
+    raw = await file.read()
+    if len(raw) > SDK_MAX_BODY_BYTES_extended:
+        raise HTTPException(status_code=413, detail="Payload too large.")
+
     files = {
-        "file": (file.filename, await file.read(), file.content_type or "application/octet-stream"),
+        "file": (file.filename, raw, file.content_type or "application/octet-stream"),
     }
     data = {
         "id_col": id_col,
@@ -128,7 +215,7 @@ async def modelfit_file(
         "filter_by_significance": filter_by_significance,
         "meanvar_test": meanvar_test,
     }
-    return _forward_stream("/modelfit-file/", files=files, data=data, user_token=user_token, method="POST")
+    return await _forward_stream("/modelfit-file/", files=files, data=data, user_token=user_token, method="POST")
 
 @app.post("/modelforecast-file/")
 async def modelforecast_file(
@@ -138,8 +225,74 @@ async def modelforecast_file(
     user_token = _extract_user_token(req)
     if not user_token:
         raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
+    
+    raw = await df_forecast.read()
+    if len(raw) > SDK_MAX_BODY_BYTES_extended:
+        raise HTTPException(status_code=413, detail="Payload too large.")
 
     files = {
-        "df_forecast": (df_forecast.filename, await df_forecast.read(), df_forecast.content_type or "application/octet-stream"),
+        "df_forecast": (df_forecast.filename, raw, df_forecast.content_type or "application/octet-stream"),
     }
-    return _forward_stream("/modelforecast-file/", files=files, data=None, user_token=user_token, method="POST")
+    return await _forward_stream("/modelforecast-file/", files=files, data=None, user_token=user_token, method="POST")
+
+@limiter.limit("10/10minute")  # 10 calls per 10 minutes per IP
+@app.post("/modelfit-file-demo/")
+async def modelfit_file_demo(
+    request: Request,
+    file: UploadFile = File(...),
+    id_col: str = Form(...),
+    time_col: str = Form(...),
+    y: str = Form(...),
+    # optional knobs
+    lag_y: str = Form(""),
+    lagged_features: str = Form(""),
+    current_features: str = Form(""),
+    filter_by_significance: str = Form("false"),
+    meanvar_test: str = Form("false"),
+):
+    
+    # read once, size-guard it, then forward
+    raw = await file.read()
+    if len(raw) > SDK_MAX_BODY_BYTES:
+        raise HTTPException(status_code=413, detail="Payload too large.")
+    
+    files = {
+        "file": (file.filename, raw, file.content_type or "application/octet-stream"),
+    }
+    data = {
+        "id_col": id_col,
+        "time_col": time_col,
+        "y": y,
+        "lag_y": lag_y,
+        "lagged_features": lagged_features,
+        "current_features": current_features,
+        "filter_by_significance": filter_by_significance,
+        "meanvar_test": meanvar_test,
+    }
+    # the path goes to the private demo route
+    return await _forward_demo_stream("/modelfit-file-demo/", files=files, data=data, method="POST")
+
+@limiter.limit("10/10minute")
+@app.post("/modelforecast-file-demo/")
+async def modelforecast_file_demo(
+    request: Request,
+    df_forecast: UploadFile = File(...),
+):
+    
+    raw = await df_forecast.read()
+    if len(raw) > SDK_MAX_BODY_BYTES:
+        raise HTTPException(status_code=413, detail="Payload too large.")
+    
+    files = {
+        "df_forecast": (df_forecast.filename, raw, df_forecast.content_type or "application/octet-stream"),
+    }
+    return await _forward_demo_stream("/modelforecast-file-demo/", files=files, data=None, method="POST")
+
+@app.get("/health-demo-proxy")
+async def health_demo_proxy():
+    # convenience endpoint to test private-space + demo token hop
+    try:
+        return await _forward_demo_stream("/health-demo", files=None, data=None, method="GET")
+    except HTTPException as e:
+        # bubble up errors so you can diagnose missing tokens, wrong URL, etc.
+        raise e

requirements.txt

@@ -1,5 +1,6 @@
-fastapi
-uvicorn
-requests
-python-multipart
-gunicorn
+fastapi
+uvicorn
+httpx
+python-multipart
+slowapi
+limits