Update app.py

app.py

@@ -1,148 +1,238 @@
+# app.py
 import os
-import requests
-import gradio as gr
-import time
 import json
+import time
 import hashlib
-from markdown import markdown
+import requests
+from urllib.parse import urlencode
+from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
+from fastapi.middleware.wsgi import WSGIMiddleware
+import gradio as gr
 
 # ======================
-# 🔐 USER AUTH SYSTEM
+# CONFIG
 # ======================
-
 DB_PATH = "users.json"
+LOGO = "logo.png"  # upload to your Space root (optional)
+# Set these environment variables in your Space secrets if you want Google OAuth:
+GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
+GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
+# Replace <YOUR_SPACE> below with your space name OR set REDIRECT_URI env var
+REDIRECT_URI = os.getenv("REDIRECT_URI") or "https://<YOUR_SPACE>.hf.space/login/callback"
+
+# Gemini API key (required)
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise ValueError("Please set GEMINI_API_KEY in environment variables (Settings → Secrets)")
 
-def hash_pw(pw):
+# ======================
+# UTIL: users DB
+# ======================
+def _hash_pw(pw: str) -> str:
     return hashlib.sha256(pw.encode()).hexdigest()
 
 def load_users():
     if not os.path.exists(DB_PATH):
         return {}
-    return json.load(open(DB_PATH, "r"))
+    try:
+        with open(DB_PATH, "r") as f:
+            return json.load(f)
+    except Exception:
+        return {}
 
-def save_users(data):
-    json.dump(data, open(DB_PATH, "w"), indent=4)
+def save_users(d):
+    with open(DB_PATH, "w") as f:
+        json.dump(d, f, indent=4)
 
-def signup(username, password):
+def signup_user(username, password):
     users = load_users()
     if username in users:
         return "❌ Хэрэглэгч байна!"
-    users[username] = hash_pw(password)
+    users[username] = _hash_pw(password)
     save_users(users)
     return "✅ Амжилттай бүртгэгдлээ!"
 
-def login(username, password):
+def check_login(username, password):
     users = load_users()
     if username not in users:
         return "❌ Хэрэглэгч олдсонгүй!"
-    if users[username] != hash_pw(password):
+    if users[username] != _hash_pw(password):
         return "❌ Нууц үг буруу!"
     return "SUCCESS"
 
 # ======================
-# 🔑 Gemini API
+# Gemini: simple API call
 # ======================
-
-gemini_api_key = os.getenv("GEMINI_API_KEY")
-if gemini_api_key is None:
-    raise ValueError("⚠️ GEMINI_API_KEY тохируулаагүй байна!")
-
-def gemini_response(prompt):
+def gemini_response(prompt: str) -> str:
     try:
-        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={gemini_api_key}"
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
         headers = {"Content-Type": "application/json"}
-        data = {"contents":[{"parts":[{"text":prompt}]}]}
-        res = requests.post(url, headers=headers, json=data)
-        res_json = res.json()
-        return res_json["candidates"][0]["content"]["parts"][0]["text"]
+        data = {"contents": [{"parts": [{"text": prompt}]}]}
+        r = requests.post(url, headers=headers, json=data, timeout=30)
+        r.raise_for_status()
+        j = r.json()
+        return j["candidates"][0]["content"]["parts"][0]["text"]
     except Exception as e:
         return f"⚠️ Gemini error: {e}"
 
 # ======================
-# 🌙 CHAT FUNCTION
+# Chat function (Gradio messages format)
 # ======================
-
 def chat_fn(message, chat_history):
     chat_history = chat_history or []
+    # append user
     chat_history.append({"role": "user", "content": message})
-    time.sleep(0.5)
+    # call model
+    time.sleep(0.2)
     reply = gemini_response(message)
+    # append assistant
     chat_history.append({"role": "assistant", "content": reply})
     return chat_history, chat_history
 
 # ======================
-# 🌟 UI CSS
+# Optional Google OAuth (FastAPI endpoints)
 # ======================
-
-css = """
-.gradio-container {font-family: 'Segoe UI', sans-serif;}
-textarea, input {border-radius:15px; padding:10px;}
-button {border-radius:15px; padding:10px 20px; font-weight:bold;}
-"""
-
-logo_path = "logo.png"
-
-# ======================
-# 🚀 FULL APP UI
-# ======================
-
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
-
-    gr.Markdown("## 🔐 Welcome to ZeppFusion Login System")
-
-    # Tabs: Login / Signup
-    with gr.Tabs() as tabs:
-
-        with gr.Tab("Login"):
-            log_user = gr.Textbox(label="Username")
-            log_pw = gr.Textbox(label="Password", type="password")
-            log_status = gr.Textbox(label="Status", interactive=False)
-            log_btn = gr.Button("Login")
-
-        with gr.Tab("Signup"):
-            reg_user = gr.Textbox(label="New Username")
-            reg_pw = gr.Textbox(label="New Password", type="password")
-            reg_status = gr.Textbox(label="Status", interactive=False)
-            reg_btn = gr.Button("Signup")
-
-    # Chat UI (hidden initially)
-    with gr.Group(visible=False) as chat_ui:
-        with gr.Row():
-            gr.Image(value=logo_path, show_label=False)
-            gr.Markdown("## ZeppFusion AI Chat")
-        chatbot = gr.Chatbot(type="messages")
-        msg = gr.Textbox(placeholder="Таны асуулт...", lines=1)
-        clear = gr.Button("🗑️ Clear History")
-        msg.submit(chat_fn, [msg, chatbot], [chatbot, chatbot])
-        clear.click(lambda: [], None, chatbot)
-
-    # ======================
-    # BUTTON ACTIONS
-    # ======================
-
-    # Signup
-    reg_btn.click(
-        signup,
-        inputs=[reg_user, reg_pw],
-        outputs=reg_status
-    )
-
-    # Login
-    def login_action(username, password):
-        res = login(username, password)
-        if res == "SUCCESS":
-            return "✅ Амжилттай нэвтэрлээ!", gr.update(visible=False), gr.update(visible=True)
-        return res, gr.update(visible=True), gr.update(visible=False)
-
-    log_btn.click(
-        login_action,
-        inputs=[log_user, log_pw],
-        outputs=[log_status, tabs, chat_ui]
-    )
+fastapi_app = FastAPI()
+
+if GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET:
+    @fastapi_app.get("/login")
+    def start_login():
+        params = {
+            "client_id": GOOGLE_CLIENT_ID,
+            "redirect_uri": REDIRECT_URI,
+            "response_type": "code",
+            "scope": "openid email profile",
+            "access_type": "offline",
+            "prompt": "consent"
+        }
+        url = f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(params)}"
+        return RedirectResponse(url)
+
+    @fastapi_app.get("/login/callback")
+    def callback(code: str = None):
+        if not code:
+            return RedirectResponse("/")
+        data = {
+            "code": code,
+            "client_id": GOOGLE_CLIENT_ID,
+            "client_secret": GOOGLE_CLIENT_SECRET,
+            "redirect_uri": REDIRECT_URI,
+            "grant_type": "authorization_code"
+        }
+        token_res = requests.post("https://oauth2.googleapis.com/token", data=data).json()
+        id_token = token_res.get("id_token")
+        # NOTE: for simplicity we don't verify id_token here.
+        # You can decode/verify id_token to extract user's email and save session as needed.
+        return RedirectResponse("/")  # After login redirect to main app
 
 # ======================
-# RUN
+# Gradio UI (ChatGPT-like: sidebar + main)
 # ======================
+css = """
+/* Minimal ChatGPT-like style */
+:root { --bg:#0b0b0c; --panel:#0f1112; --muted:#9aa0a6; --accent:#10A37F; --user:#1e293b; }
+body { background: var(--bg); color: #e6eef3; }
+.gradio-container { max-width:1200px; margin:auto; padding:18px; }
+.sidebar { background: #0f1112; border-radius:8px; padding:14px; width:260px; min-height:600px; }
+.main { background: transparent; margin-left:20px; flex:1; }
+.card { background: var(--panel); border-radius:10px; padding:14px; }
+.msg-user { background: #0b1220; color:#e6eef3; border-radius:12px; padding:10px; display:inline-block; }
+.msg-assistant { background: rgba(16,163,127,0.12); color:#d2fff0; border-radius:12px; padding:10px; display:inline-block; }
+.input-bar { display:flex; gap:8px; align-items:center; }
+"""
 
+# Gradio app built as WSGI app to mount into FastAPI when running as main.
+with gr.Blocks(css=css, theme=gr.themes.Base()) as gr_app:
+    with gr.Row():
+        # Sidebar
+        with gr.Column(elem_id="left", scale=1):
+            with gr.Box():
+                gr.Markdown("## ZeppFusion")
+                gr.Markdown("**Your AI assistant, ChatGPT-style UI**")
+                # Login / Signup controls in sidebar
+                with gr.TabbedInterface([]):  # invisible placeholder to keep layout stable
+                    pass
+            with gr.Box(elem_classes="sidebar"):
+                gr.Markdown("### Chats")
+                chat_list = gr.Textbox(value="• New chat", interactive=False)
+                gr.Markdown("### Account")
+                # show google login button only if env provided
+                if GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET:
+                    google_btn = gr.Button("Login with Google")
+                else:
+                    google_btn = None
+                # Local login/signup widgets (compact)
+                username = gr.Textbox(label="Username", placeholder="username")
+                password = gr.Textbox(label="Password", placeholder="password", type="password")
+                login_status = gr.Textbox(label="Status", interactive=False)
+                btn_login = gr.Button("Login")
+                signup_user = gr.Textbox(label="New username")
+                signup_pw = gr.Textbox(label="New password", type="password")
+                signup_status = gr.Textbox(label="Signup status", interactive=False)
+                btn_signup = gr.Button("Signup")
+                # Dark mode toggle (just cosmetic)
+                dark_toggle = gr.Checkbox(label="Dark mode", value=True)
+
+        # Main content
+        with gr.Column(elem_classes="main", scale=3):
+            with gr.Box(elem_classes="card"):
+                with gr.Row():
+                    gr.Image(value=LOGO if os.path.exists(LOGO) else None, show_label=False, elem_id="logo")
+                    gr.Markdown("### ZeppFusion AI — Chat")
+                # Chat window
+                chatbot = gr.Chatbot(type="messages")
+                msg = gr.Textbox(placeholder="Таны асуулт...", lines=1)
+                clear_btn = gr.Button("Clear")
+                # helper text area
+                info = gr.Markdown("Login to start chatting (or use Google if configured).")
+
+    # Actions
+    def on_signup(u, p):
+        return signup_user, signup_user, signup_user, signup_user  # placeholder to satisfy outputs (not used)
+
+    # Actually hookup signup/login
+    btn_signup.click(lambda u, p: signup_user_func(u,p) if False else None, inputs=[], outputs=[])
+    # Instead, we'll define proper functions below and rebind.
+
+    # Define proper functions
+    def do_signup(u, p):
+        if not u or not p:
+            return "❌ Оролт дутуу!"
+        res = signup_user(u, p)
+        return res
+
+    def do_login(u, p):
+        if not u or not p:
+            return "❌ Оролт дутуу!"
+        res = check_login(u, p)
+        return ("✅ Амжилттай нэвтэрлээ!" , gr.update(visible=False), gr.update(visible=True)) if res == "SUCCESS" else (res, gr.update(visible=True), gr.update(visible=False))
+
+    # Rewire buttons to actual functions (need to find widgets by variable)
+    btn_signup.click(do_signup, inputs=[signup_user, signup_pw], outputs=[signup_status])
+    btn_login.click(do_login, inputs=[username, password], outputs=[login_status, gr_app, chatbot])
+
+    # If Google OAuth configured -> open /login path
+    if GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET:
+        def open_google():
+            # return a direct link that user can click to open the FastAPI /login endpoint;
+            # in HuggingFace space this will route to /login on same host.
+            return RedirectResponse("/login")
+        google_btn.click(lambda: "/login", None, None)  # clicking the button will navigate the browser if configured by the runtime
+
+    # Chat actions
+    msg.submit(chat_fn, inputs=[msg, chatbot], outputs=[chatbot, chatbot])
+    clear_btn.click(lambda: [], None, chatbot)
+
+# Mount Gradio into FastAPI if running as main
 if __name__ == "__main__":
-    app.launch()
+    import uvicorn
+    from fastapi.middleware.cors import CORSMiddleware
+    fastapi_app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+    fastapi_app.mount("/", WSGIMiddleware(gr_app))
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)