login via fastAPI

Dockerfile

@@ -42,7 +42,7 @@ RUN curl -L https://huggingface.co/HuggingFaceTB/simplewiki-pruned-text-350k/res
 
 ENV WIKISPEEDIA_DB_PATH=/home/user/app/wikihop.db
 
-ENV VITE_API_BASE=""
+ENV VITE_ENV=production
 
 CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]
 

api.py

@@ -1,14 +1,16 @@
+from base64 import b64encode
 import sqlite3
 import json
 import os
 from typing import Tuple, List, Optional
 from functools import lru_cache
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 import uvicorn
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, RedirectResponse
+import requests
 
 app = FastAPI(title="WikiSpeedia API")
 
@@ -21,14 +23,20 @@ app.add_middleware(
     allow_headers=["*"],  # Allows all headers
 )
 
+CLIENT_SECRET = os.getenv("HUGGINGFACE_CLIENT_SECRET")
+IS_PROD = os.getenv("VITE_ENV") == "production"
+print("CLIENT_SECRET:", CLIENT_SECRET)
+print("IS_PROD:", IS_PROD)
 class ArticleResponse(BaseModel):
     title: str
     links: List[str]
 
+
 class HealthResponse(BaseModel):
     status: str
     article_count: int
 
+
 class SQLiteDB:
     def __init__(self, db_path: str):
         """Initialize the database with path to SQLite database"""
@@ -60,24 +68,25 @@ class SQLiteDB:
         self.cursor.execute("SELECT title FROM core_articles")
         return [row[0] for row in self.cursor.fetchall()]
 
+
 # Initialize database connection
 db = SQLiteDB(
     os.getenv("WIKISPEEDIA_DB_PATH", "/Users/jts/daily/wikihop/db/data/wikihop.db")
 )
 
+
 @app.get("/health", response_model=HealthResponse)
 async def health_check():
     """Health check endpoint that returns the article count"""
-    return HealthResponse(
-        status="healthy",
-        article_count=db._article_count
-    )
+    return HealthResponse(status="healthy", article_count=db._article_count)
+
 
 @app.get("/get_all_articles", response_model=List[str])
 async def get_all_articles():
     """Get all articles"""
     return db.get_all_articles()
 
+
 @app.get("/get_article_with_links/{article_title}", response_model=ArticleResponse)
 async def get_article(article_title: str):
     """Get article and its links by title"""
@@ -87,6 +96,49 @@ async def get_article(article_title: str):
     return ArticleResponse(title=title, links=links)
 
 
+@app.get("/auth/callback")
+async def auth_callback(request: Request):
+
+    OAUTH_API_BASE = "https://huggingface.co/oauth/token"
+    CLIENT_ID = "a67ef241-fb7e-4300-a6bd-8430a7565c9a"
+
+    code = request.query_params.get("code")
+    if not code:
+        raise HTTPException(status_code=400, detail="No code provided")
+
+    response = requests.post(
+        OAUTH_API_BASE,
+        headers={
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Authorization": f"Basic {b64encode(f'{CLIENT_ID}:{CLIENT_SECRET}'.encode()).decode()}",
+        },
+        data={
+            "client_id": CLIENT_ID,
+            "code": code,
+            "grant_type": "authorization_code",
+            "redirect_uri": "http://localhost:8000/auth/callback" if not IS_PROD else "https://huggingfacetb-wikispeedia.hf.space/auth/callback",
+        },
+    )
+
+    # response.json() =
+    # {
+    #     "access_token": "hf_oauth_eyJhbGciOiJFZERTQSJ9.eyJzY29wZSI6WyJvcGVuaWQiLCJwcm9maWxlIiwiZW1haWwiLCJpbmZlcmVuY2UtYXBpIl0sImF1ZCI6Imh0dHBzOi8vaHVnZ2luZ2ZhY2UuY28iLCJvYXV0aEFwcCI6ImE2N2VmMjQxLWZiN2UtNDMwMC1hNmJkLTg0MzBhNzU2NWM5YSIsInNlc3Npb25JZCI6IjY3YTBkYjk3OWNmZDQ3ZGFkOGNmNDMwNyIsImlhdCI6MTc0NjIxOTEwOCwic3ViIjoiNjE3OGQ4NDIyNjczMjBhYmI5OWRmNzc2IiwiZXhwIjoxNzQ2MjQ3OTA4LCJpc3MiOiJodHRwczovL2h1Z2dpbmdmYWNlLmNvIn0.TNK7Nb2X22LHlFqleo6rzJjBngjTWpVIksE1Mw7m8vVxgr7CBbK_a1J4cW488n02391qqopcaNlZKFP8noZSAA",
+    #     "token_type": "bearer",
+    #     "expires_in": 28799,
+    #     "id_token": "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiI2MTc4ZDg0MjI2NzMyMGFiYjk5ZGY3NzYiLCJuYW1lIjoiSmFzb24gU3RpbGxlcm1hbiIsInByZWZlcnJlZF91c2VybmFtZSI6InN0aWxsZXJtYW4iLCJwcm9maWxlIjoiaHR0cHM6Ly9odWdnaW5nZmFjZS5jby9zdGlsbGVybWFuIiwicGljdHVyZSI6Imh0dHBzOi8vaHVnZ2luZ2ZhY2UuY28vYXZhdGFycy84NzM5NzA1ZWY3ZWFiYzk0NWExZWYzYzA3MTk2YWYxMy5zdmciLCJlbWFpbCI6Imphc29uLnQuc3RpbGxlcm1hbkBnbWFpbC5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiYXVkIjoiYTY3ZWYyNDEtZmI3ZS00MzAwLWE2YmQtODQzMGE3NTY1YzlhIiwiYXV0aF90aW1lIjoxNzQ2MjE5MTA4LCJpYXQiOjE3NDYyMTkxMDgsImV4cCI6MTc0NjIyMjcwOCwiaXNzIjoiaHR0cHM6Ly9odWdnaW5nZmFjZS5jbyJ9.pB7j-jkxxMG3GJNzipMNCsKQimk8_R0TcPrwi-Kln6qXcSccwGcWJvyMZvFRHjKB779UkMTzgCO-eY1CINX75KaRALLS_Eu0w448F_5LMixwpBXA6dntXBEdP69VLXakpXaPHjFY2HuvUN7fbE8e2_v4a-s7RRwHTDJIcxyH2Bd_OUpebFy1N6RNB_9MIL3jxXhsXyLNL2uDry0WIB52BJKBXB4EzE12HDGgNaWR6lrqr4nvjAExsGcTwarPhFSA5ndcbgh82vJxB3rVFhSU4iZ5AmMV1mDX6SgRVdPmWZPgTBwGeGlVN-OAHvLlNJ9FZ_i0qjrtA5IRU0o6ctKrfw",
+    #     "scope": "openid profile email inference-api",
+    #     "refresh_token": "hf_oauth__refresh_RiVshOppmioFVoxvYMXSPkMdyyzbyIqadj",
+    # }
+
+    print(response.json())
+
+    # redirect to the home page with access token and id token in the url
+    return RedirectResponse(url=f"/?access_token={response.json()['access_token']}&id_token={response.json()['id_token']}")
+
+    """Auth callback endpoint"""
+    return {"message": "Auth callback received"}
+
+
 # Mount the dist folder for static files
 app.mount("/", StaticFiles(directory="dist", html=True), name="static")
 

src/components/sign-in-with-hf-button.tsx

@@ -4,94 +4,52 @@ import { useEffect } from "react";
 import { jwtDecode } from "jwt-decode";
 
 const CLIENT_ID = "a67ef241-fb7e-4300-a6bd-8430a7565c9a";
-const REDIRECT_URI = "https://huggingfacetb-wikispeedia.hf.space";
+// const REDIRECT_URI = "https://huggingfacetb-wikispeedia.hf.space";
 // const REDIRECT_URI = "http://localhost:5173/auth/callback";
+const REDIRECT_URI = isProd ? "https://huggingfacetb-wikispeedia.hf.space/auth/callback" : "http://localhost:8000/auth/callback";
+
 const SCOPE = "openid%20profile%20email%20inference-api";
 const STATE = "1234567890";
 const SSO_URL = `https://huggingface.co/oauth/authorize?client_id=${CLIENT_ID}&redirect_uri=${REDIRECT_URI}&response_type=code&scope=${SCOPE}&prompt=consent&state=${STATE}`;
-const OAUTH_API_BASE = "https://huggingface.co/oauth/token";
-const CLIENT_SECRET = import.meta.env.VITE_HUGGINGFACE_CLIENT_SECRET; // THIS IS UNSAFE, must fix before real deploy
 
-import { oauthLoginUrl, oauthHandleRedirectIfPresent } from "@huggingface/hub";
+import { isProd } from "@/lib/constants";
 
 export const SignInWithHuggingFaceButton = () => {
   const [isSignedIn, setIsSignedIn] = useState(false);
   const [isLoading, setIsLoading] = useState(false);
   const [name, setName] = useState<string | null>(null);
 
-  // useEffect(() => {
-  //   const idToken = window.localStorage.getItem("huggingface_id_token");
-  //   const accessToken = window.localStorage.getItem("huggingface_access_token");
-
-  //   if (idToken && accessToken) {
-  //     const idTokenObject = JSON.parse(idToken);
-  //     if (idTokenObject.exp > Date.now() / 1000) {
-  //       setIsSignedIn(true);
-  //       setName(idTokenObject.name);
+  useEffect(() => {
+    // check if access_token and id_token are in the url
+    const accessTokenURL = new URLSearchParams(window.location.search).get("access_token");
+    const idTokenURL = new URLSearchParams(window.location.search).get("id_token");
 
-  //       return;
-  //     }
-  //   }
+    console.log(accessTokenURL, idTokenURL);
 
-  //   async function fetchToken() {
-  //     const code = new URLSearchParams(window.location.search).get("code");
-  //     if (code) {
-  //       // remove the code from the url
-  //       window.history.replaceState({}, "", window.location.pathname);
-  //       setIsLoading(true);
-  //       const response = await fetch(`${OAUTH_API_BASE}`, {
-  //         method: "POST",
-  //         headers: {
-  //           "Content-Type": "application/x-www-form-urlencoded",
-  //           Authorization: `Basic ${btoa(`${CLIENT_ID}:${CLIENT_SECRET}`)}`,
-  //         },
-  //         body: new URLSearchParams({
-  //           client_id: CLIENT_ID,
-  //           code,
-  //           grant_type: "authorization_code",
-  //           redirect_uri: REDIRECT_URI
-  //         }).toString(),
-  //       });
-  //       const data = await response.json();
-  //       window.localStorage.setItem("huggingface_access_token", data.access_token);
+    if (accessTokenURL && idTokenURL) {
+      // remove the access_token and id_token from the url
+      window.history.replaceState({}, "", window.location.pathname);
 
-  //       // parse the id_token
-  //       const idToken = jwtDecode(data.id_token);
-  //       console.log(idToken);
-  //       window.localStorage.setItem("huggingface_id_token", JSON.stringify(idToken));
-  //       setName(idToken.name);
-  //       setIsSignedIn(true);
-  //       setIsLoading(false);
-  //     }
-  //   }
+      window.localStorage.setItem("huggingface_access_token", accessTokenURL);
+      window.localStorage.setItem("huggingface_id_token", JSON.stringify(jwtDecode(idTokenURL)));
+    }
 
-  //   fetchToken();
-  // }, []);
+    // check if the user is already logged in
+    const idToken = window.localStorage.getItem("huggingface_id_token");
+    const accessToken = window.localStorage.getItem("huggingface_access_token");
 
-  useEffect(() => {
-    async function fetchToken() {
-      console.log("fetching token", window.location.href);
-      const oauthResult = await oauthHandleRedirectIfPresent();
+    if (idToken && accessToken) {
+      const idTokenObject = JSON.parse(idToken);
+      if (idTokenObject.exp > Date.now() / 1000) {
+        setIsSignedIn(true);
+        setName(idTokenObject.name);
 
-      if (!oauthResult) {
-        // If the user is not logged in, redirect to the login page
-        window.location.href = await oauthLoginUrl();
+        return;
       }
-
-      // You can use oauthResult.accessToken, oauthResult.accessTokenExpiresAt and oauthResult.userInfo
-      console.log(oauthResult);
     }
-
-    fetchToken();
+  
   }, []);
 
-  const handleLogin = async () => {
-    window.location.href =
-      (await oauthLoginUrl({
-        scopes: "inference-api,email",
-      })) + "&prompt=consent";
-  };
-
   if (isLoading) {
     return <div>Loading...</div>;
   }
@@ -101,7 +59,7 @@ export const SignInWithHuggingFaceButton = () => {
   }
 
   return (
-    <a onClick={handleLogin} href="#" rel="nofollow">
+    <a href={SSO_URL} rel="nofollow">
       <img
         src="https://huggingface.co/datasets/huggingface/badges/resolve/main/sign-in-with-huggingface-xl.svg"
         alt="Sign in with Hugging Face"

src/lib/constants.ts

@@ -1,3 +1,4 @@
-export const API_BASE = import.meta.env.VITE_API_BASE || "http://localhost:8000";
+export const isProd = import.meta.env.VITE_ENV === "production";
+export const API_BASE = isProd ? "" : "http://localhost:8000"; // we want this blank in production
 
 console.log("API_BASE", API_BASE);