Update app.py

app.py

@@ -1,6 +1,9 @@
 import os
+import json
 from openai import OpenAI
 import gradio as gr
+from werkzeug.security import generate_password_hash, check_password_hash
+from huggingface_hub import HfApi, HfFolder
 
 # ================== YOUR SETTINGS ==================
 API_KEY = os.environ.get("API_KEY")
@@ -9,6 +12,10 @@ client = OpenAI(
     api_key=API_KEY
 )
 
+# Hugging Face API
+hf_api = HfApi()
+hf_token = HfFolder.get_token()
+
 # Auto-fetch available models
 print("Fetching available models...")
 try:
@@ -30,7 +37,7 @@ def is_model_available(model_id):
     try:
         test_completion = client.chat.completions.create(
             model=model_id,
-            messages=[{"role": "user", "content": "Say 'test'"}],
+            messages=[{"role": "user", "content": "Say \'test\'"}],
             max_tokens=16384,
             stream=False
         )
@@ -45,9 +52,117 @@ if not available_models:
     available_models = [default_model]  # Fallback to default model if no models are available
 default_model = available_models[0] if available_models else None
 
+# User data storage
+def register_user(username, password):
+    try:
+        user_data = {
+            "password": generate_password_hash(password),
+            "history": []
+        }
+        hf_api.upload_file(
+            path_or_fileobj=json.dumps(user_data),
+            path_in_repo=f"{username}.json",
+            repo_id="your-repo-id",
+            repo_type="dataset",
+            token=hf_token
+        )
+        return "Registration successful!"
+    except Exception as e:
+        return f"Registration failed: {e}"
+
+def login_user(username, password):
+    try:
+        user_data = hf_api.download_file(
+            repo_id="your-repo-id",
+            repo_type="dataset",
+            path_in_repo=f"{username}.json",
+            token=hf_token
+        )
+        with open(user_data, 'r') as f:
+            data = json.load(f)
+        if check_password_hash(data["password"], password):
+            return data["history"], "Login successful!"
+        else:
+            return None, "Incorrect password."
+    except Exception as e:
+        return None, f"Login failed: {e}"
+
+def save_history(username, history):
+    try:
+        user_data = hf_api.download_file(
+            repo_id="your-repo-id",
+            repo_type="dataset",
+            path_in_repo=f"{username}.json",
+            token=hf_token
+        )
+        with open(user_data, 'r') as f:
+            data = json.load(f)
+        data["history"] = history
+        hf_api.upload_file(
+            path_or_fileobj=json.dumps(data),
+            path_in_repo=f"{username}.json",
+            repo_id="your-repo-id",
+            repo_type="dataset",
+            token=hf_token
+        )
+        return "History saved!"
+    except Exception as e:
+        return f"Failed to save history: {e}"
+
+def load_model(selected_model):
+    try:
+        test_completion = client.chat.completions.create(
+            model=selected_model,
+            messages=[{"role": "user", "content": "Say \'test\'"}],
+            max_tokens=16384,
+            stream=False
+        )
+        if test_completion.choices[0].message.content:
+            return selected_model, f"✅ Loaded & Ready: {selected_model}"
+    except Exception as e:
+        return None, f"❌ Unavailable: {str(e)}\\nTry another model or NVIDIA support."
+
+def chat_stream(message, history, current_model, username):
+    if not current_model:
+        history.append({"role": "assistant", "content": "⚠️ Load a model first!"})
+        return history
+
+    history.append({"role": "user", "content": message})
+    yield history
+
+    history.append({"role": "assistant", "content": ""})
+    yield history
+
+    messages = []
+    for msg in history[:-1]:
+        content = msg["content"] if isinstance(msg["content"], str) else str(msg["content"])
+        messages.append({"role": msg["role"], "content": content})
+
+    full_response = ""
+    try:
+        stream = client.chat.completions.create(
+            model=current_model,
+            messages=messages,
+            temperature=0.7,
+            top_p=0.8,
+            max_tokens=4096,
+            stream=True
+        )
+
+        for chunk in stream:
+            if chunk.choices and chunk.choices[0].delta.content is not None:
+                full_response += chunk.choices[0].delta.content
+                history[-1]["content"] = full_response
+                yield history
+    except Exception as e:
+        history[-1]["content"] = f"❌ Error: {str(e)}\\nSwitch models."
+        yield history
+
+    save_history(username, history)
+
 # ================== GUI ==================
 with gr.Blocks(title="NVIDIA AI Chat") as demo:
-    gr.Markdown("# 🟢 NVIDIA AI Chat\nFixed format errors + manual load with test.")
+    gr.Markdown("# 🟢 NVIDIA AI Chat\\nFixed format errors + manual load with test.")
 
     current_model = gr.State(None)
     status_label = gr.Markdown("⚠️ No model loaded.")
@@ -80,7 +195,17 @@ with gr.Blocks(title="NVIDIA AI Chat") as demo:
     load_btn.click(load_model, model_dropdown, [current_model, status_label])
 
     # Send (Enter or button)
-    msg.submit(chat_stream, [msg, chatbot, current_model], chatbot).then(lambda: "", None, msg)
-    send_btn.click(chat_stream, [msg, chatbot, current_model], chatbot).then(lambda: "", None, msg)
+    msg.submit(chat_stream, [msg, chatbot, current_model, "username"], chatbot).then(lambda: "", None, msg)
+    send_btn.click(chat_stream, [msg, chatbot, current_model, "username"], chatbot).then(lambda: "", None, msg)
+
+    # Authentication
+    with gr.Row():
+        username = gr.Textbox(label="Username")
+        password = gr.Password(label="Password")
+        register_btn = gr.Button("Register")
+        login_btn = gr.Button("Login")
+
+    register_btn.click(register_user, [username, password], status_label)
+    login_btn.click(login_user, [username, password], [chatbot, status_label])
 
 demo.launch(auth=("bsis3b", "classmate"))