add keycloak auth to playground ui

llama_stack/distribution/ui/README.md

@@ -48,3 +48,6 @@ uv run --with ".[ui]" streamlit run llama_stack/distribution/ui/app.py
 | TOGETHER_API_KEY           | API key for Together provider      | (empty string)            |
 | SAMBANOVA_API_KEY          | API key for SambaNova provider     | (empty string)            |
 | OPENAI_API_KEY             | API key for OpenAI provider        | (empty string)            |
+| KEYCLOAK_URL               | URL for keycloak authentication    | (empty string)            |
+| KEYCLOAK_REALM             | Keycloak realm                     | default                   |
+| KEYCLOAK_CLIENT_ID         | Client ID for keycloak auth        | (empty string)            |

llama_stack/distribution/ui/app.py

@@ -50,6 +50,36 @@ def main():
     )
     pg.run()
 
+def main2():
+    from dataclasses import asdict
+    st.subheader(f"Welcome {keycloak.user_info['preferred_username']}!")
+    st.write(f"Here is your user information:")
+    st.write(asdict(keycloak))
+
+def get_access_token() -> str|None:
+    return st.session_state.get('access_token')
 
 if __name__ == "__main__":
-    main()
+    
+    from streamlit_keycloak import login
+    import os
+    
+    keycloak_url = os.environ.get("KEYCLOAK_URL")
+    keycloak_realm = os.environ.get("KEYCLOAK_REALM", "default")
+    keycloak_client_id = os.environ.get("KEYCLOAK_CLIENT_ID")
+    
+    if keycloak_url and keycloak_client_id:
+        keycloak = login(
+            url=keycloak_url,
+            realm=keycloak_realm,
+            client_id=keycloak_client_id,
+            auto_refresh=True,
+            init_options={"checkLoginIframe": False},
+        )
+
+        if keycloak.authenticated:
+            st.session_state['access_token'] = keycloak.access_token
+            main()
+    # TBD - add other authentications
+    else:
+        main()

llama_stack/distribution/ui/modules/api.py

@@ -7,11 +7,13 @@
 import os
 
 from llama_stack_client import LlamaStackClient
+from llama_stack.distribution.ui.app import get_access_token
 
 
 class LlamaStackApi:
     def __init__(self):
         self.client = LlamaStackClient(
+            api_key=get_access_token(),
             base_url=os.environ.get("LLAMA_STACK_ENDPOINT", "http://localhost:8321"),
             provider_data={
                 "fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),

llama_stack/distribution/ui/requirements.txt

@@ -3,3 +3,4 @@ llama-stack-client>=0.2.1
 pandas
 streamlit
 streamlit-option-menu
+streamlit-keycloak

pyproject.toml

@@ -49,6 +49,7 @@ ui = [
     "pandas",
     "llama-stack-client>=0.2.9",
     "streamlit-option-menu",
+    "streamlit-keycloak",
 ]
 
 [dependency-groups]