diff --git a/libs/ktem/ktem/assets/css/main.css b/libs/ktem/ktem/assets/css/main.css index fcef5341..87d173ce 100644 --- a/libs/ktem/ktem/assets/css/main.css +++ b/libs/ktem/ktem/assets/css/main.css @@ -441,7 +441,11 @@ div.markmap { } #google-login { - width: 250px; + max-width: 450px; +} + +#user-api-key-wrapper { + max-width: 450px; } #login-row { diff --git a/libs/ktem/ktem/index/file/ui.py b/libs/ktem/ktem/index/file/ui.py index 7ab4cf29..842747e3 100644 --- a/libs/ktem/ktem/index/file/ui.py +++ b/libs/ktem/ktem/index/file/ui.py @@ -1194,13 +1194,16 @@ class FileIndexPage(BasePage): request: gr.Request, ): if KH_DEMO_MODE: + user = None + if request is None: + raise ValueError("This feature is not available") + try: import gradiologin as grlogin user = grlogin.get_user(request) except (ImportError, AssertionError): - user = None - + pass if not user: raise ValueError("Please sign-in to use this feature") diff --git a/libs/ktem/ktem/pages/chat/__init__.py b/libs/ktem/ktem/pages/chat/__init__.py index 0b7dcdcf..424078d0 100644 --- a/libs/ktem/ktem/pages/chat/__init__.py +++ b/libs/ktem/ktem/pages/chat/__init__.py @@ -247,6 +247,14 @@ function() { MINDMAP_HTML_EXPORT_TEMPLATE.replace("\n", "").replace('"', '\\"'), ) +fetch_api_key_js = """ +function(_, __) { + api_key = getStorage('google_api_key', ''); + console.log('session API key:', api_key); + return [api_key, _]; +} +""" + class ChatPage(BasePage): def __init__(self, app): @@ -263,6 +271,7 @@ class ChatPage(BasePage): ) self._info_panel_expanded = gr.State(value=True) self._command_state = gr.State(value=None) + self._user_api_key = gr.Text(value="", visible=False) def on_building_ui(self): with gr.Row(): @@ -452,6 +461,7 @@ class ChatPage(BasePage): self.chat_panel.text_input, self.chat_panel.chatbot, self._app.user_id, + self._user_api_key, self._app.settings_state, self.chat_control.conversation_id, self.chat_control.conversation_rn, @@ -590,6 +600,10 @@ class ChatPage(BasePage): ) if KH_DEMO_MODE: + self.chat_control.btn_demo_logout.click( + fn=None, + js=self.chat_control.logout_js, + ) self.chat_control.btn_new.click( fn=lambda: self.chat_control.select_conv("", None), outputs=[ @@ -858,6 +872,7 @@ class ChatPage(BasePage): chat_input, chat_history, user_id, + user_api_key, settings, conv_id, conv_name, @@ -866,6 +881,7 @@ class ChatPage(BasePage): ): """Submit a message to the chatbot""" if KH_DEMO_MODE: + print("API key", user_api_key) try: import gradiologin as grlogin @@ -904,6 +920,7 @@ class ChatPage(BasePage): True, settings, user_id, + request=None, ) elif file_names: for file_name in file_names: @@ -1031,7 +1048,13 @@ class ChatPage(BasePage): def _on_app_created(self): if KH_DEMO_MODE: self._app.app.load( + fn=lambda x: x, + inputs=[self._user_api_key], + outputs=[self._user_api_key], + js=fetch_api_key_js, + ).then( fn=self.chat_control.toggle_demo_login_visibility, + inputs=[self._user_api_key], outputs=[ self.chat_control.cb_suggest_chat, self.chat_control.btn_new, diff --git a/libs/ktem/ktem/pages/chat/control.py b/libs/ktem/ktem/pages/chat/control.py index c02cec5d..eb46dbde 100644 --- a/libs/ktem/ktem/pages/chat/control.py +++ b/libs/ktem/ktem/pages/chat/control.py @@ -21,6 +21,14 @@ if not os.path.isdir(ASSETS_DIR): ASSETS_DIR = "libs/ktem/ktem/assets/icons" +logout_js = """ +function () { + removeFromStorage('google_api_key'); + window.location.href = "/logout"; +} +""" + + def is_conv_name_valid(name): """Check if the conversation name is valid""" errors = [] @@ -37,6 +45,7 @@ class ConversationControl(BasePage): def __init__(self, app): self._app = app + self.logout_js = logout_js self.on_building_ui() def on_building_ui(self): @@ -158,7 +167,6 @@ class ConversationControl(BasePage): ) self.btn_demo_logout = gr.Button( "Sign-out", - link="/logout", min_width=120, size="sm", scale=1, @@ -429,7 +437,7 @@ class ConversationControl(BasePage): gr.Info("Chat suggestions updated.") - def toggle_demo_login_visibility(self, request: gr.Request): + def toggle_demo_login_visibility(self, user_api_key, request: gr.Request): try: import gradiologin as grlogin @@ -437,7 +445,7 @@ class ConversationControl(BasePage): except (ImportError, AssertionError): user = None - if user: + if user or user_api_key: return [ gr.update(visible=True), gr.update(visible=True), diff --git a/libs/ktem/ktem/pages/help.py b/libs/ktem/ktem/pages/help.py index 6a7cf86a..2ecdf7eb 100644 --- a/libs/ktem/ktem/pages/help.py +++ b/libs/ktem/ktem/pages/help.py @@ -3,9 +3,11 @@ from pathlib import Path import gradio as gr import requests +from decouple import config from theflow.settings import settings KH_DEMO_MODE = getattr(settings, "KH_DEMO_MODE", False) +HF_SPACE_URL = config("HF_SPACE_URL", default="") def get_remote_doc(url: str) -> str: @@ -65,12 +67,14 @@ class HelpPage: with gr.Accordion("Create Your Own Space"): gr.Markdown( "This is a demo with limited functionality. " - "Use Duplicate space button to install Kotaemon " + "Use **Create space** button to install Kotaemon " "in your own space with all features " "(including upload and manage your private " "documents securely)." ) - gr.DuplicateButton( + gr.Button( + value="Create Your Own Space", + link=HF_SPACE_URL, variant="primary", size="lg", ) @@ -84,7 +88,7 @@ class HelpPage: f"{self.remote_content_url}/v{self.app_version}/docs/usage.md" ) if user_guide_md: - with gr.Accordion("User Guide"): + with gr.Accordion("User Guide", open=not KH_DEMO_MODE): gr.Markdown(user_guide_md) if self.app_version: diff --git a/sso_app_demo.py b/sso_app_demo.py index 69871564..ff50b50c 100644 --- a/sso_app_demo.py +++ b/sso_app_demo.py @@ -25,6 +25,30 @@ GOOGLE_CLIENT_SECRET = config("GOOGLE_CLIENT_SECRET", default="") SECRET_KEY = config("SECRET_KEY", default="default-secret-key") +save_api_key_js = """ +function(api_key) { + setStorage('google_api_key', api_key); + window.location.href = "/app"; +} +""" + +global_js = """ +function () { + // store info in local storage + globalThis.setStorage = (key, value) => { + localStorage.setItem(key, value) + } + globalThis.getStorage = (key, value) => { + item = localStorage.getItem(key); + return item ? item : value; + } + globalThis.removeFromStorage = (key) => { + localStorage.removeItem(key) + } +} +""" + + def add_session_middleware(app): config_data = { "GOOGLE_CLIENT_ID": GOOGLE_CLIENT_ID, @@ -93,8 +117,9 @@ async def auth(request: Request): with gr.Blocks( theme=KotaemonTheme(), css=gradio_app._css, + js=global_js, ) as login_demo: - with gr.Row(elem_id="login-row"): + with gr.Column(elem_id="login-row"): gr.Markdown("

Welcome to Kotaemon

") gr.Button( "Login with Google", @@ -102,6 +127,24 @@ with gr.Blocks( variant="primary", elem_id="google-login", ) + with gr.Accordion( + "Or use your own Gemini API key", + elem_id="user-api-key-wrapper", + open=False, + ): + api_key_input = gr.Textbox( + placeholder="API Key", + label="Enter your Gemini API key", + ) + api_key_save_btn = gr.Button( + "Save", + ) + + api_key_save_btn.click( + fn=lambda _: True, + inputs=[api_key_input], + js=save_api_key_js, + ) app = gr.mount_gradio_app(app, login_demo, path="/login-app") app = gr.mount_gradio_app(