Update to print disallowed endpoints in-place in the model list

Update configuration handling to put all clients in `clients` with backwards-compat. parsing

Troubleshoot radio button rendering

Refactor permissions configuration to support other oauth methods

app.py

@@ -1,5 +1,7 @@
 import os
 import json
+from time import sleep
+
 import gradio as gr
 
 import uvicorn
@@ -11,7 +13,7 @@ from starlette.responses import RedirectResponse
 from authlib.integrations.starlette_client import OAuth, OAuthError
 from fastapi import FastAPI, Request
 
-from shared import Client
+from shared import Client, User, OAuthProvider
 
 app = FastAPI()
 config = {}
@@ -51,41 +53,44 @@ def init_config():
     global clients
     global llm_host_names
     config = json.loads(os.environ['CONFIG'])
-    reserved_keys = ("huggingface_text", "allowed_domains_override")
-    for name in config:
-        if name in reserved_keys:
-            continue
-        model_personas = config[name].get("personas", {})
+    client_config = config.get("clients") or config
+    for name in client_config:
+        model_personas = client_config[name].get("personas", {})
         client = Client(
-            api_url=os.environ.get(config[name]['api_url'],
-                                   config[name]['api_url']),
-            api_key=os.environ.get(config[name]['api_key'],
-                                   config[name]['api_key']),
+            api_url=os.environ.get(client_config[name]['api_url'],
+                                   client_config[name]['api_url']),
+            api_key=os.environ.get(client_config[name]['api_key'],
+                                   client_config[name]['api_key']),
             personas=model_personas
         )
         clients[name] = client
-    llm_host_names = list(config.keys())
+    llm_host_names = list(client_config.keys())
 
 
-def get_allowed_models(user_domain: str) -> List[str]:
+def get_allowed_models(user: User) -> List[str]:
     """
     Get a list of allowed endpoints for a specified user domain. Allowed domains
     are configured in each model's configuration and may optionally be overridden
     in the Gradio demo configuration.
-    :param user_domain: User domain (i.e. neon.ai, google.com, guest)
-    :return: List of allowed endpoints from configuration
+    :param user: User to get permissions for
+    :return: List of allowed endpoints from configuration (including empty
+        strings for disallowed endpoints)
     """
-    overrides = config.get("allowed_domains_override", {})
+    overrides = config.get("permissions_override", {})
     allowed_endpoints = []
     for client in clients:
-        allowed_domains = overrides.get(client,
-                                        clients[client].config.inference.allowed_domains)
-        if allowed_domains is None:
-            # Allowed domains not specified; model is public
+        permission = overrides.get(client,
+                                   clients[client].config.inference.permissions)
+        if not permission:
+            # Permissions not specified (None or empty dict); model is public
             allowed_endpoints.append(client)
-        elif user_domain in allowed_domains:
-            # User domain is in the allowed domain list
+        elif user.oauth == OAuthProvider.GOOGLE and user.permissions_id in \
+                permission.get("google_domains", []):
+            # Google oauth  domain is in the allowed domain list
             allowed_endpoints.append(client)
+        else:
+            allowed_endpoints.append("")
+            print(f"No permission to access {client}")
     return allowed_endpoints
 
 
@@ -107,7 +112,7 @@ def get_login_button(request: gr.Request) -> gr.Button:
     :param request: Gradio request to evaluate
     :return: Button for either login or logout action
     """
-    user = get_user(request)
+    user = get_user(request).username
     print(f"Getting login button for {user}")
 
     if user == "guest":
@@ -116,15 +121,39 @@ def get_login_button(request: gr.Request) -> gr.Button:
         return gr.Button(f"Logout {user}", link="/logout")
 
 
-def get_user(request: Request) -> str:
+def get_user(request: Request) -> User:
     """
     Get a unique user email address for the specified request
     :param request: FastAPI Request object with user session data
     :return: String user email address or "guest"
     """
+    # {'iss': 'https://accounts.google.com',
+    #  'azp': '***.apps.googleusercontent.com',
+    #  'aud': '***.apps.googleusercontent.com',
+    #  'sub': '###',
+    #  'hd': 'neon.ai',
+    #  'email': 'daniel@neon.ai',
+    #  'email_verified': True,
+    #  'at_hash': '***',
+    #  'nonce': '***',
+    #  'name': 'Daniel McKnight',
+    #  'picture': 'https://lh3.googleusercontent.com/a/***',
+    #  'given_name': '***',
+    #  'family_name': '***',
+    #  'iat': ###,
+    #  'exp': ###}
     if not request:
-        return "guest"
-    user = request.session.get('user', {}).get('email') or "guest"
+        return User(OAuthProvider.NONE, "guest", "")
+
+    user_dict = request.session.get("user", {})
+    if user_dict.get("iss") == "https://accounts.google.com":
+        user = User(OAuthProvider.GOOGLE, user_dict["email"], user_dict["hd"])
+    elif user_dict:
+        print(f"Unknown user session data: {user_dict}")
+        user = User(OAuthProvider.NONE, "guest", "")
+    else:
+        user = User(OAuthProvider.NONE, "guest", "")
+    print(user)
     return user
 
 
@@ -232,25 +261,25 @@ def get_model_options(request: gr.Request) -> List[gr.Radio]:
         # `user` is a valid Google email address or 'guest'
         user = get_user(request.request)
     else:
-        user = "guest"
-    print(f"Getting models for {user}")
+        user = User(OAuthProvider.NONE, "guest", "")
+    print(f"Getting models for {user.username}")
 
-    domain = "guest" if user == "guest" else user.split('@')[1]
-    allowed_llm_host_names = get_allowed_models(domain)
+    allowed_llm_host_names = get_allowed_models(user)
 
     radio_infos = [f"{name} ({clients[name].vllm_model_name})"
+                   if name in clients else "Not Authorized"
                    for name in allowed_llm_host_names]
     # Components
-    radios = [gr.Radio(choices=clients[name].personas.keys(),
+    radios = [gr.Radio(choices=clients[name].personas.keys() if name in clients else [],
                        value=None, label=info) for name, info
               in zip(allowed_llm_host_names, radio_infos)]
 
     # Select the first available option by default
     radios[0].value = list(clients[allowed_llm_host_names[0]].personas.keys())[0]
     print(f"Set default persona to {radios[0].value} for {allowed_llm_host_names[0]}")
-    # Ensure we always have the same number of rows
-    while len(radios) < len(llm_host_names):
-        radios.append(gr.Radio(choices=[], value=None, label="Not Authorized"))
+    # # Ensure we always have the same number of rows
+    # while len(radios) < len(llm_host_names):
+    #     radios.append(gr.Radio(choices=[], value=None, label="Not Authorized"))
     return radios
 
 
@@ -271,6 +300,17 @@ def init_gradio() -> gr.Blocks:
         @gr.on(triggers=[blocks.load, *[radio.input for radio in radios]],
                inputs=[radio_state, *radios], outputs=[radio_state, *radios])
         def radio_click(state, *new_state):
+            """
+            Handle any state changes that require re-rendering radio buttons
+            :param state: Previous radio state representation (before selection)
+            :param new_state: Current radio state (including selection)
+            :return: Desired new state (current option selected, previous option
+                deselected)
+            """
+            # Login and model options are triggered on load. This sleep is just
+            # a hack to make sure those events run before this logic to select
+            # the default model
+            sleep(0.1)
             try:
                 changed_index = next(i for i in range(len(state))
                                      if state[i] != new_state[i])
@@ -326,5 +366,5 @@ if __name__ == "__main__":
     init_config()
     init_oauth()
     blocks = init_gradio()
-    app = gr.mount_gradio_app(app, blocks, '/', auth_dependency=get_user)
+    app = gr.mount_gradio_app(app, blocks, '/')
     uvicorn.run(app, host='0.0.0.0', port=7860)

shared.py

@@ -1,3 +1,6 @@
+from dataclasses import dataclass
+from enum import IntEnum
+
 import yaml
 
 from typing import Dict, Optional, List
@@ -8,6 +11,18 @@ from huggingface_hub.utils import EntryNotFoundError
 from openai import OpenAI
 
 
+class OAuthProvider(IntEnum):
+    NONE = 0
+    GOOGLE = 1
+
+
+@dataclass
+class User:
+    oauth: OAuthProvider
+    username: str
+    permissions_id: str
+
+
 class PileConfig(BaseModel):
     file2persona: Dict[str, str]
     file2prefix: Dict[str, str]
@@ -17,7 +32,7 @@ class PileConfig(BaseModel):
 
 class InferenceConfig(BaseModel):
     chat_template: str
-    allowed_domains: Optional[List[str]] = None
+    permissions: Dict[str, list] = {}
 
 
 class RepoConfig(BaseModel):