Update app.py

app.py

@@ -4,7 +4,7 @@ import csv
 import time
 import threading
 from datetime import datetime, timezone
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple
 
 import gradio as gr
 from huggingface_hub import HfApi
@@ -122,23 +122,6 @@ def load_state(force: bool = False):
 # -------------------------
 # Core API
 # -------------------------
-def _get_user_info(request: Optional[gr.Request]) -> Optional[Dict]:
-    """
-    Extract user info from OAuth token in the request.
-    Returns None if user is not authenticated.
-    """
-    if request is None:
-        return None
-
-    # In Gradio with OAuth enabled, user info is available via request.username
-    # and other OAuth data via request.kwargs
-    if hasattr(request, 'username') and request.username:
-        return {
-            "username": request.username,
-            "headers": dict(request.headers) if hasattr(request, 'headers') else {}
-        }
-    return None
-
 def claim_key(
     request: Optional[gr.Request] = None,
 ) -> Tuple[str, str]:
@@ -215,7 +198,7 @@ def claim_b_key(
         return b_key["key"], f"Key provided."
 
 def claim_key_authenticated(
-    request: Optional[gr.Request] = None,
+    profile: Optional[gr.OAuthProfile] = None,
 ) -> Tuple[str, str]:
     """
     OAuth-protected endpoint that returns the next available key.
@@ -223,12 +206,10 @@ def claim_key_authenticated(
     Returns (key, status_message)
     """
     # Check if user is authenticated
-    user_info = _get_user_info(request)
-
-    if user_info is None:
+    if profile is None:
         return "", "Authentication required. Please sign in with your Hugging Face account."
 
-    username = user_info["username"]
+    username = profile.username
 
     with _lock:
         keys, _ = load_state(force=True)
@@ -261,7 +242,7 @@ def claim_key_authenticated(
         return key_to_send["key"], f"Key sent successfully to {username}."
 
 def claim_b_key_authenticated(
-    request: Optional[gr.Request] = None,
+    profile: Optional[gr.OAuthProfile] = None,
 ) -> Tuple[str, str]:
     """
     OAuth-protected endpoint that returns the B key and increments the usage count.
@@ -269,12 +250,10 @@ def claim_b_key_authenticated(
     Returns (b_key, status_message)
     """
     # Check if user is authenticated
-    user_info = _get_user_info(request)
-
-    if user_info is None:
+    if profile is None:
         return "", "Authentication required. Please sign in with your Hugging Face account."
 
-    username = user_info["username"]
+    username = profile.username
 
     with _lock:
         keys, b_keys = load_state(force=True)
@@ -308,7 +287,10 @@ def claim_b_key_authenticated(
 # -------------------------
 with gr.Blocks(title="API") as demo:
     gr.Markdown("## API Service")
-    gr.Markdown("*Note: Authenticated endpoints require signing in with your Hugging Face account.*")
+
+    with gr.Row():
+        gr.Markdown("*Note: Authenticated endpoints require signing in with your Hugging Face account.*")
+        gr.LoginButton()
 
     with gr.Tab("Service A"):
         btn_a = gr.Button("Request", variant="secondary", size="sm")
@@ -336,7 +318,7 @@ with gr.Blocks(title="API") as demo:
 
     with gr.Tab("Service A (Authenticated)"):
         gr.Markdown("### 🔐 Authentication Required")
-        gr.Markdown("This endpoint requires you to sign in with your Hugging Face account.")
+        gr.Markdown("Click the 'Sign in with Hugging Face' button at the top to authenticate.")
         btn_a_auth = gr.Button("Request Key (OAuth)", variant="primary", size="sm")
         out_a_auth = gr.Textbox(label="Response", interactive=False, show_label=False)
         status_a_auth = gr.Textbox(label="", interactive=False, show_label=False)
@@ -350,7 +332,7 @@ with gr.Blocks(title="API") as demo:
 
     with gr.Tab("Service B (Authenticated)"):
         gr.Markdown("### 🔐 Authentication Required")
-        gr.Markdown("This endpoint requires you to sign in with your Hugging Face account.")
+        gr.Markdown("Click the 'Sign in with Hugging Face' button at the top to authenticate.")
         btn_b_auth = gr.Button("Request Key (OAuth)", variant="primary", size="sm")
         out_b_auth = gr.Textbox(label="Response", interactive=False, show_label=False)
         status_b_auth = gr.Textbox(label="", interactive=False, show_label=False)
@@ -364,3 +346,4 @@ with gr.Blocks(title="API") as demo:
 
 demo.queue()
 demo.launch()
+