Create app.py

app.py

@@ -0,0 +1,366 @@
+import os
+import io
+import csv
+import time
+import threading
+from datetime import datetime, timezone
+from typing import Dict, Optional, Tuple
+
+import gradio as gr
+from huggingface_hub import HfApi
+
+# -------------------------
+# Config
+# -------------------------
+DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID") # Space secret
+HF_TOKEN = os.environ["HF_TOKEN"]          # Space secret
+
+KEYS_PATH = "keys.csv"
+B_KEYS_PATH = "b_keys.csv"
+
+
+api = HfApi(token=HF_TOKEN)
+_lock = threading.Lock()
+
+# Simple cache to reduce Hub reads
+_cache = {
+    "ts": 0.0,
+    "keys": None,  # type: ignore
+    "b_keys": None,  # type: ignore
+}
+CACHE_TTL_SEC = 10.0
+
+# -------------------------
+# Helpers
+# -------------------------
+def _utc_now_iso() -> str:
+    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+
+def _download_csv(repo_id: str, path_in_repo: str) -> str:
+    # Download file bytes from the dataset repo, return text
+    file_bytes = api.hf_hub_download(
+        repo_id=repo_id,
+        filename=path_in_repo,
+        repo_type="dataset",
+        token=HF_TOKEN,
+    )
+    with open(file_bytes, "r", encoding="utf-8") as f:
+        return f.read()
+
+def _upload_csv(repo_id: str, path_in_repo: str, content: str, commit_message: str) -> None:
+    api.upload_file(
+        path_or_fileobj=io.BytesIO(content.encode("utf-8")),
+        path_in_repo=path_in_repo,
+        repo_id=repo_id,
+        repo_type="dataset",
+        commit_message=commit_message,
+        token=HF_TOKEN,
+    )
+
+def _parse_keys(csv_text: str) -> list:
+    keys = []
+    reader = csv.DictReader(io.StringIO(csv_text))
+    for r in reader:
+        if r["key"].strip():
+            keys.append({
+                "key": r["key"].strip(),
+                "sent": (r.get("sent") or "").strip(),
+            })
+    return keys
+
+def _parse_b_keys(csv_text: str) -> list:
+    keys = []
+    reader = csv.DictReader(io.StringIO(csv_text))
+    for r in reader:
+        if r["key"].strip():
+            keys.append({
+                "key": r["key"].strip(),
+                "count": int(r.get("count") or "0"),
+            })
+    return keys
+
+def _serialize_keys(keys: list) -> str:
+    out = io.StringIO()
+    fieldnames = ["key", "sent"]
+    w = csv.DictWriter(out, fieldnames=fieldnames)
+    w.writeheader()
+    for k in keys:
+        w.writerow({
+            "key": k["key"],
+            "sent": k.get("sent", ""),
+        })
+    return out.getvalue()
+
+def _serialize_b_keys(keys: list) -> str:
+    out = io.StringIO()
+    fieldnames = ["key", "count"]
+    w = csv.DictWriter(out, fieldnames=fieldnames)
+    w.writeheader()
+    for k in keys:
+        w.writerow({
+            "key": k["key"],
+            "count": k.get("count", 0),
+        })
+    return out.getvalue()
+
+def load_state(force: bool = False):
+    now = time.time()
+    if (not force) and _cache["keys"] is not None and (now - _cache["ts"] < CACHE_TTL_SEC):
+        return _cache["keys"], _cache["b_keys"]
+
+    keys_csv = _download_csv(DATASET_REPO_ID, KEYS_PATH)
+    b_keys_csv = _download_csv(DATASET_REPO_ID, B_KEYS_PATH)
+
+    keys = _parse_keys(keys_csv)
+    b_keys = _parse_b_keys(b_keys_csv)
+
+    _cache["ts"] = now
+    _cache["keys"] = keys
+    _cache["b_keys"] = b_keys
+    return keys, b_keys
+
+# -------------------------
+# 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]:
+    """
+    Returns the next available key.
+    Returns (key, status_message)
+    """
+    _ = request
+
+    with _lock:
+        keys, _ = load_state(force=True)
+
+        # Find available keys
+        available_keys = [k for k in keys if not k["sent"]]
+
+        if len(available_keys) < 1:
+            return "", "No available keys remaining."
+
+        # Send the next available key
+        key_to_send = available_keys[0]
+
+        # Mark key as sent
+        key_to_send["sent"] = _utc_now_iso()
+
+        # Save changes
+        updated_keys_csv = _serialize_keys(keys)
+
+        _upload_csv(
+            DATASET_REPO_ID,
+            KEYS_PATH,
+            updated_keys_csv,
+            commit_message=f"Key sent at {_utc_now_iso()}",
+        )
+
+        # refresh cache immediately
+        _cache["ts"] = 0.0
+
+        return key_to_send["key"], "Key sent successfully."
+
+def claim_b_key(
+    request: Optional[gr.Request] = None,
+) -> Tuple[str, str]:
+    """
+    Returns the  key and increments the usage count.
+    Returns (b_key, status_message)
+    """
+    _ = request
+
+    with _lock:
+        keys, b_keys = load_state(force=True)
+
+        if not b_keys or len(b_keys) == 0:
+            return "", "No keys available."
+
+        # Get the first key (you can add logic to rotate if needed)
+        b_key = b_keys[0]
+
+        # Increment the count
+        b_key["count"] += 1
+
+        # Save changes
+        updated_b_keys_csv = _serialize_b_keys(b_keys)
+
+        _upload_csv(
+            DATASET_REPO_ID,
+            B_KEYS_PATH,
+            updated_b_keys_csv,
+            commit_message=f"Key requested at {_utc_now_iso()}",
+        )
+
+        # refresh cache immediately
+        _cache["ts"] = 0.0
+
+        return b_key["key"], f"Key provided."
+
+def claim_key_authenticated(
+    request: Optional[gr.Request] = None,
+) -> Tuple[str, str]:
+    """
+    OAuth-protected endpoint that returns the next available key.
+    Requires user to be authenticated via Hugging Face OAuth.
+    Returns (key, status_message)
+    """
+    # Check if user is authenticated
+    user_info = _get_user_info(request)
+
+    if user_info is None:
+        return "", "Authentication required. Please sign in with your Hugging Face account."
+
+    username = user_info["username"]
+
+    with _lock:
+        keys, _ = load_state(force=True)
+
+        # Find available keys
+        available_keys = [k for k in keys if not k["sent"]]
+
+        if len(available_keys) < 1:
+            return "", "No available keys remaining."
+
+        # Send the next available key
+        key_to_send = available_keys[0]
+
+        # Mark key as sent with user info
+        key_to_send["sent"] = f"{_utc_now_iso()} (user: {username})"
+
+        # Save changes
+        updated_keys_csv = _serialize_keys(keys)
+
+        _upload_csv(
+            DATASET_REPO_ID,
+            KEYS_PATH,
+            updated_keys_csv,
+            commit_message=f"Key sent to {username} at {_utc_now_iso()}",
+        )
+
+        # refresh cache immediately
+        _cache["ts"] = 0.0
+
+        return key_to_send["key"], f"Key sent successfully to {username}."
+
+def claim_b_key_authenticated(
+    request: Optional[gr.Request] = None,
+) -> Tuple[str, str]:
+    """
+    OAuth-protected endpoint that returns the B key and increments the usage count.
+    Requires user to be authenticated via Hugging Face OAuth.
+    Returns (b_key, status_message)
+    """
+    # Check if user is authenticated
+    user_info = _get_user_info(request)
+
+    if user_info is None:
+        return "", "Authentication required. Please sign in with your Hugging Face account."
+
+    username = user_info["username"]
+
+    with _lock:
+        keys, b_keys = load_state(force=True)
+
+        if not b_keys or len(b_keys) == 0:
+            return "", "No keys available."
+
+        # Get the first key
+        b_key = b_keys[0]
+
+        # Increment the count
+        b_key["count"] += 1
+
+        # Save changes
+        updated_b_keys_csv = _serialize_b_keys(b_keys)
+
+        _upload_csv(
+            DATASET_REPO_ID,
+            B_KEYS_PATH,
+            updated_b_keys_csv,
+            commit_message=f"Key requested by {username} at {_utc_now_iso()}",
+        )
+
+        # refresh cache immediately
+        _cache["ts"] = 0.0
+
+        return b_key["key"], f"Key provided to {username}."
+
+# -------------------------
+# UI
+# -------------------------
+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.Tab("Service A"):
+        btn_a = gr.Button("Request", variant="secondary", size="sm")
+        out_a = gr.Textbox(label="Response", interactive=False, show_label=False)
+        status_a = gr.Textbox(label="", interactive=False, show_label=False)
+
+        btn_a.click(
+            fn=claim_key,
+            inputs=[],
+            outputs=[out_a, status_a],
+            api_name="claim_key",
+        )
+
+    with gr.Tab("Service B"):
+        btn_b = gr.Button("Request", variant="secondary", size="sm")
+        out_b = gr.Textbox(label="Response", interactive=False, show_label=False)
+        status_b = gr.Textbox(label="", interactive=False, show_label=False)
+
+        btn_b.click(
+            fn=claim_b_key,
+            inputs=[],
+            outputs=[out_b, status_b],
+            api_name="claim_b_key",
+        )
+
+    with gr.Tab("Service A (Authenticated)"):
+        gr.Markdown("### 🔐 Authentication Required")
+        gr.Markdown("This endpoint requires you to sign in with your Hugging Face account.")
+        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)
+
+        btn_a_auth.click(
+            fn=claim_key_authenticated,
+            inputs=[],
+            outputs=[out_a_auth, status_a_auth],
+            api_name="claim_key_authenticated",
+        )
+
+    with gr.Tab("Service B (Authenticated)"):
+        gr.Markdown("### 🔐 Authentication Required")
+        gr.Markdown("This endpoint requires you to sign in with your Hugging Face account.")
+        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)
+
+        btn_b_auth.click(
+            fn=claim_b_key_authenticated,
+            inputs=[],
+            outputs=[out_b_auth, status_b_auth],
+            api_name="claim_b_key_authenticated",
+        )
+
+demo.queue()
+demo.launch()