From 3bde47e56206ef2233814c19a50f0d6599927b11 Mon Sep 17 00:00:00 2001
From: Angel Nunez Mencias <git@angelnu.com>
Date: Sun, 1 Jun 2025 22:23:49 +0200
Subject: [PATCH] add keycloak auth to playground ui

---
 llama_stack/distribution/ui/README.md        |  3 ++
 llama_stack/distribution/ui/app.py           | 32 +++++++++++++++++++-
 llama_stack/distribution/ui/modules/api.py   |  2 ++
 llama_stack/distribution/ui/requirements.txt |  1 +
 pyproject.toml                               |  1 +
 5 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/llama_stack/distribution/ui/README.md b/llama_stack/distribution/ui/README.md
index 51c2d2bc2..0e96690ec 100644
--- a/llama_stack/distribution/ui/README.md
+++ b/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)            |
\ No newline at end of file
diff --git a/llama_stack/distribution/ui/app.py b/llama_stack/distribution/ui/app.py
index 441f65d20..a7a13bd99 100644
--- a/llama_stack/distribution/ui/app.py
+++ b/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()
diff --git a/llama_stack/distribution/ui/modules/api.py b/llama_stack/distribution/ui/modules/api.py
index 11455ed46..44969ea71 100644
--- a/llama_stack/distribution/ui/modules/api.py
+++ b/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", ""),
diff --git a/llama_stack/distribution/ui/requirements.txt b/llama_stack/distribution/ui/requirements.txt
index 53a1e7bf3..dac342fe1 100644
--- a/llama_stack/distribution/ui/requirements.txt
+++ b/llama_stack/distribution/ui/requirements.txt
@@ -3,3 +3,4 @@ llama-stack-client>=0.2.1
 pandas
 streamlit
 streamlit-option-menu
+streamlit-keycloak
diff --git a/pyproject.toml b/pyproject.toml
index f79857c3d..bec5cc487 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,6 +49,7 @@ ui = [
     "pandas",
     "llama-stack-client>=0.2.9",
     "streamlit-option-menu",
+    "streamlit-keycloak",
 ]
 
 [dependency-groups]