feat: add login using API key

This commit is contained in:
Tadashi
2025-01-02 11:44:46 +07:00
parent fcc44b73ce
commit b052b38c2d
6 changed files with 95 additions and 10 deletions

View File

@@ -441,7 +441,11 @@ div.markmap {
}
#google-login {
width: 250px;
max-width: 450px;
}
#user-api-key-wrapper {
max-width: 450px;
}
#login-row {

View File

@@ -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")

View File

@@ -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,

View File

@@ -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),

View File

@@ -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:

View File

@@ -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("<h1 style='text-align:center;'>Welcome to Kotaemon</h1>")
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(