Update loader.py

loader.py

@@ -1,74 +1,131 @@
-# loader.py — Public proxy loader for private static Hugging Face Space
+# loader.py — public proxy for PRIVATE Static Space via huggingface.co/resolve & raw
 import os
+import mimetypes
+from urllib.parse import urljoin
 from fastapi import FastAPI, Request
-from fastapi.responses import Response, HTMLResponse, PlainTextResponse
+from fastapi.responses import Response, HTMLResponse, PlainTextResponse, JSONResponse
 import httpx
 
 app = FastAPI()
 
-# ========= CONFIG =========
-PRIVATE_STATIC_URL = os.getenv(
-    "PRIVATE_STATIC_URL",
-    "https://cheekeong2025-iip-grading.static.hf.space"  # ✅ fallback URL
-)
-PRIVATE_SPACE_ID = os.getenv("PRIVATE_SPACE_ID")
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-if not PRIVATE_SPACE_ID:
-    raise RuntimeError(
-        "❌ PRIVATE_SPACE_ID is not set.\n"
-        "Go to your Hugging Face Space → Settings → Repository secrets\n"
-        "and add PRIVATE_SPACE_ID=cheekeong2025/iip-grading"
-    )
-
-if not PRIVATE_STATIC_URL:
-    raise RuntimeError(
-        "❌ PRIVATE_STATIC_URL is not set.\n"
-        "Go to your Hugging Face Space → Settings → Repository secrets\n"
-        "and add PRIVATE_STATIC_URL=https://<your-private-static-space>.hf.space"
-    )
+# --- Config from Secrets ---
+SPACE_ID = os.getenv("PRIVATE_SPACE_ID")          # e.g. "cheekeong2025/IIP_Grading"
+HF_TOKEN = os.getenv("HF_TOKEN")                  # read-access token for that Space
+REVISION = os.getenv("REVISION", "main")          # branch/tag/commit
+# ---------------------------
 
+if not SPACE_ID:
+    raise RuntimeError("Set PRIVATE_SPACE_ID (e.g. 'cheekeong2025/IIP_Grading') in Settings → Repository secrets.")
 if not HF_TOKEN:
-    print("[loader] ⚠️ Warning: HF_TOKEN not set — attempting anonymous access.")
+    raise RuntimeError("Set HF_TOKEN (read access to the private Space) in Settings → Repository secrets.")
 
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
-# ===========================
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
+BASE_RESOLVE = f"https://huggingface.co/spaces/{SPACE_ID}/resolve/{REVISION}/"
+BASE_RAW     = f"https://huggingface.co/spaces/{SPACE_ID}/raw/{REVISION}/"
 
+def _join(b: str, p: str) -> str:
+    return urljoin(b, p.lstrip("/"))
+
+def _with_q(url: str, req: Request) -> str:
+    q = str(req.url.query or "")
+    return f"{url}?{q}" if q else url
+
+def _looks_file(p: str) -> bool:
+    return "." in p.split("/")[-1]
+
+def _mime(p: str, default="application/octet-stream"):
+    if p.endswith("/") or p == "":
+        return "text/html; charset=utf-8"
+    m, _ = mimetypes.guess_type(p)
+    return m or default
+
+async def _get(client: httpx.AsyncClient, url: str) -> httpx.Response:
+    r = await client.get(url, headers=HEADERS)
+    print(f"[proxy] GET {url} -> {r.status_code}")
+    return r
+
+async def _fetch(client: httpx.AsyncClient, path: str, req: Request) -> tuple[httpx.Response, str]:
+    """
+    Try in order (first 200 wins):
+      resolve:path
+      resolve:static/path
+      (dir) resolve:path/index.html
+      (dir) resolve:static/path/index.html
+      raw:path
+      raw:static/path
+      (dir) raw:path/index.html
+      (dir) raw:static/path/index.html
+    """
+    as_dir = not _looks_file(path)
+    tried = []
+    candidates = [
+        _join(BASE_RESOLVE, path),
+        _join(BASE_RESOLVE, f"static/{path.lstrip('/')}"),
+    ]
+    if as_dir:
+        candidates += [
+            _join(BASE_RESOLVE, path.rstrip("/") + "/index.html"),
+            _join(BASE_RESOLVE, f"static/{path.rstrip('/')}/index.html"),
+        ]
+    candidates += [
+        _join(BASE_RAW, path),
+        _join(BASE_RAW, f"static/{path.lstrip('/')}"),
+    ]
+    if as_dir:
+        candidates += [
+            _join(BASE_RAW, path.rstrip("/") + "/index.html"),
+            _join(BASE_RAW, f"static/{path.rstrip('/')}/index.html"),
+        ]
+
+    last = None
+    for u in candidates:
+        url = _with_q(u, req)
+        r = await _get(client, url)
+        tried.append(u)
+        if r.status_code == 200:
+            return r, " -> ".join(tried)
+        last = r
+    return last, " -> ".join(tried)
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "space": SPACE_ID, "revision": REVISION}
+
+@app.get("/_debug/fetch/{path:path}")
+async def debug_fetch(path: str, request: Request):
+    try:
+        async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
+            r, tried = await _fetch(client, path, request)
+        return JSONResponse({"status": r.status_code, "content_type": r.headers.get("content-type"), "tried": tried})
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
 
 @app.get("/")
-async def root():
-    """Serve the index.html from the private static Space."""
-    print(f"[loader] Fetching root from {PRIVATE_STATIC_URL}/index.html")
-    async with httpx.AsyncClient(timeout=None) as client:
-        try:
-            r = await client.get(f"{PRIVATE_STATIC_URL}/index.html", headers=HEADERS)
-            r.raise_for_status()
-        except Exception as e:
-            return PlainTextResponse(f"Error fetching root: {e}", status_code=500)
-    return HTMLResponse(r.text, status_code=r.status_code)
+async def root(request: Request):
+    try:
+        async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
+            r, tried = await _fetch(client, "index.html", request)
+            if r.status_code == 404:
+                r, tried2 = await _fetch(client, "", request)
+                print(f"[root] tried: {tried} || {tried2}")
+            else:
+                print(f"[root] tried: {tried}")
+    except Exception as e:
+        return PlainTextResponse(f"Error fetching root: {e}", status_code=500)
 
+    ctype = r.headers.get("content-type", "text/html; charset=utf-8")
+    if "text/html" in ctype:
+        return HTMLResponse(r.text, status_code=r.status_code, media_type=ctype)
+    return Response(content=r.content, status_code=r.status_code, media_type=ctype)
 
 @app.get("/{path:path}")
 async def proxy(path: str, request: Request):
-    """Proxy all static assets (CSS, JS, images, etc.) from the private Space."""
-    url = f"{PRIVATE_STATIC_URL.rstrip('/')}/{path}"
-    print(f"[loader] Fetching asset: {url}")
     try:
-        async with httpx.AsyncClient(timeout=None) as client:
-            r = await client.get(url, headers=HEADERS)
-            r.raise_for_status()
+        async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
+            r, tried = await _fetch(client, path, request)
+        print(f"[proxy] path={path} status={r.status_code}")
     except Exception as e:
-        return PlainTextResponse(f"Error fetching {url}: {e}", status_code=500)
+        return PlainTextResponse(f"Error fetching {path}: {e}", status_code=500)
 
-    content_type = r.headers.get("content-type", "text/plain")
-    return Response(content=r.content, media_type=content_type, status_code=r.status_code)
-
-
-@app.get("/health")
-async def health():
-    """Simple health check."""
-    return {
-        "status": "ok",
-        "source": PRIVATE_STATIC_URL,
-        "space_id": PRIVATE_SPACE_ID
-    }
+    ctype = r.headers.get("content-type") or _mime(path)
+    return Response(content=r.content, media_type=ctype, status_code=r.status_code)