Merge pull request #198 from rishab11250/feature/hf-token-endpoint

backend/app/routes/auth.py

@@ -14,6 +14,7 @@ from app.database import get_db
 from app.models import User, ApiKey
 from app.schemas import (
     GoogleLoginRequest,
+    HFTokenUpdate,
     RefreshRequest,
     TokenResponse,
     UpdatePassword,
@@ -279,6 +280,34 @@ def get_me(user: User = Depends(get_current_user)):
     """
     return UserResponse.model_validate(user)
 
+@router.put("/hf-token", response_model=UserResponse)
+def update_hf_token(
+    payload: HFTokenUpdate,
+    user: User = Depends(get_current_user),
+    db: Session = Depends(get_db),
+):
+    """Update the HuggingFace token for the authenticated user.
+
+    Stores the provided HF token in the user's profile so it can be used
+    for HuggingFace API calls (e.g. InferenceClient) in place of the
+    globally configured ``HF_TOKEN`` environment variable.
+
+    Args:
+        payload: HFTokenUpdate object containing the new ``hf_token`` value.
+        user: The currently authenticated user, obtained from the
+            ``get_current_user`` dependency.
+        db: SQLAlchemy database session, obtained from the dependency.
+
+    Returns:
+        UserResponse: The updated user profile including the new ``hf_token``
+        field.
+    """
+    user.hf_token = payload.hf_token
+    db.commit()
+    db.refresh(user)
+    return UserResponse.model_validate(user)
+
+
 @router.put("/update")
 def update_user_info(payload:UserUpdate,
                     user: User = Depends(get_current_user),

backend/tests/test_auth.py

@@ -79,3 +79,39 @@ def test_refresh_token_success(client, refresh_token):
     assert payload["access_token"]
     assert payload["refresh_token"]
     assert payload["token_type"] == "bearer"
+
+
+def test_update_hf_token_success(client, auth_headers):
+    response = client.put(
+        "/api/v1/auth/hf-token",
+        json={"hf_token": "hf_new_token_value"},
+        headers=auth_headers,
+    )
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert payload["hf_token"] == "hf_new_token_value"
+
+
+def test_update_hf_token_requires_auth(client):
+    response = client.put(
+        "/api/v1/auth/hf-token",
+        json={"hf_token": "hf_unauth"},
+    )
+
+    assert response.status_code in (401, 403)
+
+
+def test_hf_token_appears_in_user_response(client, auth_headers, user, db_session):
+    # First update the token
+    put_resp = client.put(
+        "/api/v1/auth/hf-token",
+        json={"hf_token": "hf_persist_token"},
+        headers=auth_headers,
+    )
+    assert put_resp.status_code == 200
+
+    # Then verify it shows up in GET /me
+    me_resp = client.get("/api/v1/auth/me", headers=auth_headers)
+    assert me_resp.status_code == 200
+    assert me_resp.json()["hf_token"] == "hf_persist_token"