Add Google OAuth Login flow.

.env

@@ -30,3 +30,11 @@ DUCKDB_USE_VIEWS=0
 # HF_USERNAME=
 # The default repo to deploy to for a staging demo. Can be overridden by a command line flag.
 # HF_STAGING_DEMO_REPO='HF_ORG/HF_REPO_NAME'
+
+# For Google-login. This is generated from the Google Cloud Console for a web client.
+# See: https://developers.google.com/identity/protocols/oauth2
+GOOGLE_CLIENT_ID='279475920249-i8llm8vbos1vj5m1qocir8narb3r0enu.apps.googleusercontent.com'
+# The client secret of the above client.
+# GOOGLE_CLIENT_SECRET=
+# A random string for oauth sessions.
+# LILAC_OAUTH_SECRET_KEY=

README.md

@@ -81,6 +81,25 @@ To run the docker image locally:
 docker run -p 5432:5432 lilac_blueprint
 ```
 
+#### Authentication
+
+Authentication is done via Google login. A Google Client token should be created
+from the Google API Console. Details can be found [here](https://developers.google.com/identity/protocols/oauth2).
+
+By default, the Lilac google client is used. The secret can be found in Google
+Cloud console, and should be defined under `GOOGLE_CLIENT_SECRET` in .env.local.
+
+For the session middleware, a random string should be created and defined as `LILAC_OAUTH_SECRET_KEY` in .env.local.
+
+You can generate a random secret key with:
+
+```py
+import string
+import random
+key = ''.join(random.choices(string.ascii_uppercase + string.digits, k=64))
+print(f"LILAC_OAUTH_SECRET_KEY='{key}'")
+```
+
 ### Configuration
 
 To use various API's, API keys need to be provided. Create a file named `.env.local` in the root, and add variables that are listed in `.env` with your own values.

mypy.ini

@@ -116,3 +116,7 @@ follow_imports = skip
 [mypy-langdetect.*]
 ignore_missing_imports = True
 follow_imports = skip
+
+[mypy-authlib.*]
+ignore_missing_imports = True
+follow_imports = skip

pyproject.toml

@@ -9,30 +9,33 @@ version = "0.0.1"
 [tool.poetry.dependencies]
 
 ### Required dependencies. ###
-dask = "^2023.3.2" 
-datasets = "^2.12.0" 
-distributed = "^2023.3.2.1" 
-duckdb = "^0.8.1" 
-fastapi = "^0.98.0" 
-gcsfs = "^2023.4.0" 
-google-cloud-storage = "^2.5.0" 
-gunicorn = "^20.1.0" 
-joblib = "^1.3.1" 
+authlib = "^1.2.1"
+dask = "^2023.3.2"
+datasets = "^2.12.0"
+distributed = "^2023.3.2.1"
+duckdb = "^0.8.1"
+fastapi = "^0.98.0"
+gcsfs = "^2023.4.0"
+google-cloud-storage = "^2.5.0"
+gunicorn = "^20.1.0"
+httpx = "^0.24.1"
+itsdangerous = "^2.1.2"
+joblib = "^1.3.1"
 openai-function-call = "^0.0.5" # Wraps OpenAI functions with Pydantic models.
 orjson = "^3.8.10" # Fast JSON serialization: https://fastapi.tiangolo.com/advanced/custom-response/#use-orjsonresponse
 pillow = "^9.3.0" # Image processing.
-psutil = "^5.9.5" 
-pyarrow = "^9.0.0" 
-pydantic = "^1.10.11" 
-python = "~3.9" 
-python-dotenv = "^1.0.0" 
-requests = "^2.28.1" 
-scikit-learn = "^1.3.0" 
-tenacity = "^8.2.2" 
-tqdm = "^4.65.0" 
-types-psutil = "^5.9.5.12" 
-typing-extensions = "^4.7.1" 
-uvicorn = {extras = ["standard"], version = "^0.22.0"} 
+psutil = "^5.9.5"
+pyarrow = "^9.0.0"
+pydantic = "^1.10.11"
+python = "~3.9"
+python-dotenv = "^1.0.0"
+requests = "^2.28.1"
+scikit-learn = "^1.3.0"
+tenacity = "^8.2.2"
+tqdm = "^4.65.0"
+types-psutil = "^5.9.5.12"
+typing-extensions = "^4.7.1"
+uvicorn = {extras = ["standard"], version = "^0.22.0"}
 
 ### Optional dependencies. ###
 
@@ -61,6 +64,7 @@ regex = "^2023.6.3"
 # For language detection.
 langdetect = {version = "^1.0.9", optional = true}
 
+
 [tool.poetry.extras]
 all = [
   "cohere",
@@ -97,29 +101,29 @@ optional = true
 
 [tool.poetry.group.dev.dependencies]
 bokeh = ">=2.4.2,<3" # Required for Dask monitoring.
-click = "^8.1.3" 
-google-api-python-client-stubs = "^1.13.0" 
-httpx = "^0.24.0" 
-huggingface-hub = "^0.15.1" 
-isort = "^5.12.0" 
-matplotlib = "^3.7.1" 
-mypy = "^1.0.0" 
-notebook = "^6.5.4" 
-pytest = "^7.1.3" 
-pytest-asyncio = "^0.20.2" 
-pytest-cov = "^4.0.0" 
-pytest-mock = "^3.10.0" 
-ruff = "^0.0.219" 
-setuptools = "^65.5.0" 
-toml = "^0.10.2" 
-types-Pillow = "^9.3.0.4" 
-types-cachetools = "^5.3.0.5" 
-types-regex = "^2023.6.3.0" 
-types-requests = "^2.28.11.5" 
-types-tqdm = "^4.65.0.0" 
-watchdog = {extras = ["watchmedo"], version = "^3.0.0"} 
-wheel = "^0.37.1" 
-yapf = "^0.32.0" 
+click = "^8.1.3"
+google-api-python-client-stubs = "^1.13.0"
+httpx = "^0.24.0"
+huggingface-hub = "^0.15.1"
+isort = "^5.12.0"
+matplotlib = "^3.7.1"
+mypy = "^1.0.0"
+notebook = "^6.5.4"
+pytest = "^7.1.3"
+pytest-asyncio = "^0.20.2"
+pytest-cov = "^4.0.0"
+pytest-mock = "^3.10.0"
+ruff = "^0.0.219"
+setuptools = "^65.5.0"
+toml = "^0.10.2"
+types-Pillow = "^9.3.0.4"
+types-cachetools = "^5.3.0.5"
+types-regex = "^2023.6.3.0"
+types-requests = "^2.28.11.5"
+types-tqdm = "^4.65.0.0"
+watchdog = {extras = ["watchmedo"], version = "^3.0.0"}
+wheel = "^0.37.1"
+yapf = "^0.32.0"
 
 [tool.poetry.scripts]
 deploy-hf = "scripts.deploy_hf:main"

src/auth.py

@@ -1,5 +1,7 @@
 """Authentication and ACL configuration."""
 
+from typing import Optional
+
 from pydantic import BaseModel
 
 from .config import CONFIG
@@ -32,6 +34,21 @@ class UserAccess(BaseModel):
   concept: ConceptUserAccess
 
 
+class UserInfo(BaseModel):
+  """User information."""
+  email: str
+  name: str
+  given_name: str
+  family_name: str
+
+
+class AuthenticationInfo(BaseModel):
+  """Authentication information for the user."""
+  user: Optional[UserInfo]
+  access: UserAccess
+  auth_enabled: bool
+
+
 def get_user_access() -> UserAccess:
   """Get the user access."""
   auth_enabled = CONFIG.get('LILAC_AUTH_ENABLED', False)

src/router_google_login.py

@@ -0,0 +1,62 @@
+"""Router for Google OAuth2 login."""
+
+from urllib.parse import urlparse, urlunparse
+
+from authlib.integrations.starlette_client import OAuth, OAuthError
+from fastapi import APIRouter, Request, Response
+from fastapi.responses import HTMLResponse
+from starlette.config import Config
+from starlette.responses import RedirectResponse
+
+from .config import CONFIG
+from .router_utils import RouteErrorHandler
+
+router = APIRouter(route_class=RouteErrorHandler)
+
+GOOGLE_CLIENT_ID = CONFIG.get('GOOGLE_CLIENT_ID', None)
+GOOGLE_CLIENT_SECRET = CONFIG.get('GOOGLE_CLIENT_SECRET', None)
+LILAC_AUTH_ENABLED = CONFIG.get('LILAC_AUTH_ENABLED', False)
+if LILAC_AUTH_ENABLED:
+  if GOOGLE_CLIENT_ID is None or GOOGLE_CLIENT_SECRET is None:
+    raise ValueError(
+      'Missing `GOOGLE_CLIENT_ID` or `GOOGLE_CLIENT_SECRET` when `LILAC_AUTH_ENABLED=true`')
+  SECRET_KEY = CONFIG.get('LILAC_OAUTH_SECRET_KEY', None)
+  if not SECRET_KEY:
+    raise ValueError('Missing `LILAC_OAUTH_SECRET_KEY` when `LILAC_AUTH_ENABLED=true`')
+
+  # Set up oauth
+  oauth = OAuth(
+    Config(environ={
+      'GOOGLE_CLIENT_ID': GOOGLE_CLIENT_ID,
+      'GOOGLE_CLIENT_SECRET': GOOGLE_CLIENT_SECRET
+    }))
+  oauth.register(
+    name='google',
+    server_metadata_url='https://accounts.google.com/.well-known/openid-configuration',
+    client_kwargs={'scope': 'openid email profile'},
+  )
+
+
+@router.get('/login')
+async def login(request: Request, origin_url: str) -> RedirectResponse:
+  """Redirects to Google OAuth login page."""
+  auth_path = urlunparse(urlparse(origin_url)._replace(path='/google/auth'))
+  return await oauth.google.authorize_redirect(request, auth_path)
+
+
+@router.get('/auth')
+async def auth(request: Request) -> Response:
+  """Handles the Google OAuth callback."""
+  try:
+    token = await oauth.google.authorize_access_token(request)
+  except OAuthError as error:
+    return HTMLResponse(f'<h1>{error}</h1>')
+  request.session['user'] = token['userinfo']
+  return RedirectResponse(url='/')
+
+
+@router.get('/logout')
+def logout(request: Request) -> RedirectResponse:
+  """Logs the user out."""
+  request.session.pop('user', None)
+  return RedirectResponse(url='/')