Add Gradio frontend with HuggingFace login

frontend.py

@@ -0,0 +1,200 @@
+"""Gradio UI for llm-backend."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Iterator, List, Optional
+import typing
+
+import gradio as gr
+from huggingface_hub import HfApi, login
+
+from src.config import UPLOAD_DIR
+from src.db import list_sessions_info, reset_history
+from src.team import TeamChatSession
+
+
+# ---------------------------------------------------------------------------
+# Authentication helpers
+# ---------------------------------------------------------------------------
+
+def hf_login(token: str) -> str:
+    """Login to HuggingFace and return the username."""
+    login(token=token, new_session=True)
+    info = HfApi().whoami(token=token)
+    return info.get("name") or info.get("email")
+
+
+# ---------------------------------------------------------------------------
+# File helpers
+# ---------------------------------------------------------------------------
+
+def _vm_host_path(user: str, vm_path: str) -> Path:
+    rel = Path(vm_path).relative_to("/data")
+    base = Path(UPLOAD_DIR) / user
+    target = (base / rel).resolve()
+    if not str(target).startswith(str(base)):
+        raise ValueError("Invalid path")
+    return target
+
+
+def list_directory(user: str, path: str) -> list[dict[str, str | bool]]:
+    target = _vm_host_path(user, path)
+    if not target.exists() or not target.is_dir():
+        raise FileNotFoundError(path)
+    entries: list[dict[str, str | bool]] = []
+    for entry in sorted(target.iterdir()):
+        entries.append({"name": entry.name, "is_dir": entry.is_dir()})
+    return entries
+
+
+def read_file(user: str, path: str) -> str:
+    target = _vm_host_path(user, path)
+    if not target.exists() or target.is_dir():
+        raise FileNotFoundError(path)
+    return target.read_text()
+
+
+def write_file(user: str, path: str, content: str) -> None:
+    target = _vm_host_path(user, path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(content)
+
+
+def delete_path(user: str, path: str) -> None:
+    target = _vm_host_path(user, path)
+    if target.is_dir():
+        for child in target.iterdir():
+            if child.is_file():
+                child.unlink()
+            else:
+                delete_path(user, str(Path("/data") / child.relative_to(target)))
+        target.rmdir()
+    elif target.exists():
+        target.unlink()
+
+
+# ---------------------------------------------------------------------------
+# Chat helpers
+# ---------------------------------------------------------------------------
+
+async def chat_generator(user: str, session: str, prompt: str) -> Iterator[List[List[str]]]:
+    history: List[List[str]] = []
+    async with TeamChatSession(user=user, session=session) as chat:
+        history.append([prompt, ""])
+        resp = ""
+        async for part in chat.chat_stream(prompt):
+            resp += part
+            history[-1][1] = resp
+            yield history
+
+
+# ---------------------------------------------------------------------------
+# Gradio callbacks
+# ---------------------------------------------------------------------------
+
+async def send_message(user: str, session: str, message: str, history: list[list[str]] | None) -> typing.Iterator[tuple[list[list[str]], str]]:
+    history = history or []
+    async with TeamChatSession(user=user, session=session) as chat:
+        history.append([message, ""])
+        resp = ""
+        async for part in chat.chat_stream(message):
+            resp += part
+            history[-1][1] = resp
+            yield history, ""
+
+def load_sessions(user: str) -> list[str]:
+    info = list_sessions_info(user)
+    return [s["name"] for s in info]
+
+
+def remove_session(user: str, session: str) -> None:
+    reset_history(user, session)
+
+
+# ---------------------------------------------------------------------------
+# UI construction
+# ---------------------------------------------------------------------------
+
+def build_ui() -> gr.Blocks:
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        user_state = gr.State(str)
+        session_state = gr.State("default")
+        history_state = gr.State([])
+
+        with gr.Column() as login_col:
+            gr.Markdown("## HuggingFace Login")
+            token_box = gr.Textbox(type="password", label="HuggingFace token")
+            login_btn = gr.Button("Login")
+            login_status = gr.Markdown()
+
+        with gr.Row(visible=False) as main_row:
+            with gr.Column(scale=3):
+                session_drop = gr.Dropdown(label="Session", interactive=True)
+                new_session = gr.Textbox(label="New Session Name")
+                create_btn = gr.Button("Create Session")
+                chatbox = gr.Chatbot(type="messages")
+                msg = gr.Textbox(label="Message")
+                send_btn = gr.Button("Send")
+
+            with gr.Column(scale=2):
+                gr.Markdown("### Files")
+                dir_path = gr.Textbox(value="/data", label="Path")
+                refresh_btn = gr.Button("List")
+                file_list = gr.Dataframe(headers=["Name", "Is Dir"], datatype=["str", "bool"], interactive=False)
+                open_path = gr.Textbox(label="File Path")
+                open_btn = gr.Button("Open")
+                file_editor = gr.Textbox(label="Content", lines=10)
+                save_btn = gr.Button("Save")
+                delete_btn = gr.Button("Delete")
+
+        def do_login(token: str):
+            user = hf_login(token)
+            sessions = load_sessions(user)
+            return {user_state: user, session_drop: gr.Dropdown.update(choices=sessions), login_col: gr.Column.update(visible=False), main_row: gr.Row.update(visible=True)}
+
+        login_btn.click(do_login, inputs=token_box, outputs=[user_state, session_drop, login_col, main_row])
+
+        def create_session(user: str, name: str):
+            if not name:
+                return gr.Dropdown.update()
+            reset_history(user, name)
+            sessions = load_sessions(user)
+            return gr.Dropdown.update(value=name, choices=sessions)
+
+        create_btn.click(create_session, inputs=[user_state, new_session], outputs=session_drop)
+
+        def refresh_files(user: str, path: str):
+            entries = list_directory(user, path)
+            data = [[e["name"], e["is_dir"]] for e in entries]
+            return {file_list: gr.Dataframe.update(value=data), dir_path: path}
+
+        refresh_btn.click(refresh_files, inputs=[user_state, dir_path], outputs=[file_list, dir_path])
+
+        def open_file(user: str, path: str):
+            content = read_file(user, path)
+            return {file_editor: content, open_path: path}
+
+        open_btn.click(open_file, inputs=[user_state, open_path], outputs=[file_editor, open_path])
+
+        def save_file(user: str, path: str, content: str):
+            write_file(user, path, content)
+            return gr.Textbox.update()
+
+        save_btn.click(save_file, inputs=[user_state, open_path, file_editor], outputs=file_editor)
+
+        def delete_file(user: str, path: str):
+            delete_path(user, path)
+            return gr.Textbox.update(value="")
+
+        delete_btn.click(delete_file, inputs=[user_state, open_path], outputs=open_path)
+
+        send_btn.click(send_message, inputs=[user_state, session_drop, msg, history_state], outputs=[chatbox, history_state, msg])
+        msg.submit(send_message, inputs=[user_state, session_drop, msg, history_state], outputs=[chatbox, history_state, msg])
+
+    return demo
+
+
+if __name__ == "__main__":
+    ui = build_ui()
+    ui.launch()

requirements.txt

@@ -8,3 +8,6 @@ fastapi
 uvicorn
 python-multipart
 httpx
+
+gradio
+huggingface_hub