Update app.py

app.py

@@ -1,9 +1,17 @@
-# === ЗАГРУЗКА МОДЕЛИ (ИСПРАВЛЕННАЯ ЧАСТЬ) ===
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import json
+from pathlib import Path
+from datetime import datetime
+import requests
+
+# === ЗАГРУЗКА МОДЕЛИ ===
 print("🚀 Загрузка модели...")
 model_id = "OpenRussianAI/OpenAirAI-X"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-# ВАЖНО: Убеждаемся, что у токенизатора есть pad_token
+# Важно: устанавливаем pad_token, если его нет, чтобы избежать ошибок при генерации
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
@@ -12,7 +20,43 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(device).eval()
 print(f"✅ Модель загружена на {device}")
 
-# === ГЕНЕРАЦИЯ ОТВЕТА (ПОЛНОСТЬЮ ПЕРЕПИСАННАЯ) ===
+# === ХРАНИЛИЩЕ ИСТОРИИ ===
+HISTORY_DIR = Path("chat_history")
+HISTORY_DIR.mkdir(exist_ok=True)
+
+def get_history_path(username):
+    safe_name = "".join(c for c in username if c.isalnum() or c in ('-', '_')).lower()
+    return HISTORY_DIR / f"{safe_name}.json"
+
+def load_history(username):
+    path = get_history_path(username)
+    if path.exists():
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    return {}
+
+def save_history(username, history_data):
+    path = get_history_path(username)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(history_data, f, ensure_ascii=False, indent=2)
+
+def check_if_pro(token):
+    if not token or not token.startswith("hf_"):
+        return False
+    try:
+        response = requests.get(
+            "https://huggingface.co/api/whoami-v2",
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=5
+        )
+        if response.status_code == 200:
+            data = response.json()
+            return data.get("isPro", False) or data.get("plan", {}).get("name") == "PRO"
+    except Exception as e:
+        print(f"Ошибка проверки PRO: {e}")
+    return False
+
+# === ГЕНЕРАЦИЯ ОТВЕТА (ИСПРАВЛЕННАЯ) ===
 def generate_response(message, history, username, current_chat_id):
     if not username:
         gr.Warning("Сначала введите имя пользователя!")
@@ -38,35 +82,34 @@ def generate_response(message, history, username, current_chat_id):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=128, # Уменьшим, чтобы модель не уходила в бред
+            max_new_tokens=128,          # Ограничиваем длину, чтобы модель не "уходила в бред"
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id, # Важно!
-            repetition_penalty=1.2 # Штраф за повторения, чтобы убрать цикл "Я OpenAI"
+            eos_token_id=tokenizer.eos_token_id, # Токен конца строки
+            repetition_penalty=1.2       # Штраф за повторения (убирает цикл "Я OpenAI")
         )
     
-    # Декодируем только НОВУЮ часть текста
+    # Декодируем ТОЛЬКО новую часть текста (ответ модели)
     generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
     response_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
     
-    # Очистка от мусора и обрезка
-    # Обычно модель может начать ответ с пробела или новой строки
+    # Очистка ответа
     ai_response = response_text.strip()
     
-    # Если модель все же написала "Пользователь:" или "AI:" в конце, обрезаем
+    # Если модель начала писать следующий вопрос или метку, обрезаем
     stop_words = ["Пользователь:", "User:", "\n\n"]
     for stop_word in stop_words:
         if stop_word in ai_response:
             ai_response = ai_response.split(stop_word)[0].strip()
             
-    # Если ответ пустой после очистки
     if not ai_response:
         ai_response = "..."
 
     history = history + [{"role": "assistant", "content": ai_response}]
     
+    # Сохранение истории
     if username and current_chat_id:
         history_data = load_history(username)
         if current_chat_id not in history_data:
@@ -79,4 +122,235 @@ def generate_response(message, history, username, current_chat_id):
         save_history(username, history_data)
     
     chat_list_choices = get_chat_list(username)
-    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)
+    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)
+
+# === УПРАВЛЕНИЕ ЧАТАМИ ===
+def new_chat(username):
+    if not username:
+        return [], None, gr.update(choices=[])
+    chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+    history_data = load_history(username)
+    history_data[chat_id] = {
+        "title": "Новый чат",
+        "created": datetime.now().isoformat(),
+        "messages": []
+    }
+    save_history(username, history_data)
+    chat_list_choices = get_chat_list(username)
+    return [], chat_id, gr.update(choices=chat_list_choices, value=None)
+
+def load_chat(chat_title, username, current_chat_id):
+    if not username or not chat_title:
+        return [], current_chat_id
+    history_data = load_history(username)
+    for cid, data in history_data.items():
+        if data["title"] == chat_title:
+            return data["messages"], cid
+    return [], current_chat_id
+
+def get_chat_list(username):
+    if not username:
+        return []
+    history_data = load_history(username)
+    sorted_chats = sorted(
+        history_data.items(),
+        key=lambda x: x[1].get("created", ""),
+        reverse=True
+    )
+    return [data["title"] for _, data in sorted_chats]
+
+def delete_chat(chat_title, username):
+    if not username or not chat_title:
+        return [], {}, None, gr.update(choices=[])
+    history_data = load_history(username)
+    for cid, data in list(history_data.items()):
+        if data["title"] == chat_title:
+            del history_data[cid]
+            break
+    save_history(username, history_data)
+    chat_list_choices = get_chat_list(username)
+    return [], history_data, None, gr.update(choices=chat_list_choices, value=None)
+
+# === CSS ===
+CUSTOM_CSS = """
+.main-header {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    padding: 20px;
+    border-radius: 12px;
+    margin-bottom: 20px;
+    text-align: center;
+}
+.main-header h1 { margin: 0; font-size: 2em; }
+.pro-badge {
+    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%);
+    color: #333;
+    padding: 4px 12px;
+    border-radius: 20px;
+    font-weight: bold;
+    display: inline-block;
+    margin-left: 10px;
+    box-shadow: 0 2px 8px rgba(255, 215, 0, 0.4);
+}
+.user-info {
+    padding: 10px;
+    background: white;
+    border-radius: 8px;
+    margin-bottom: 10px;
+    text-align: center;
+    font-weight: 600;
+}
+.privacy-banner {
+    background: linear-gradient(135deg, #10b981 0%, #059669 100%);
+    color: white;
+    padding: 15px 20px;
+    border-radius: 10px;
+    margin: 15px 0;
+    display: flex;
+    align-items: center;
+    gap: 12px;
+    font-size: 0.95em;
+    box-shadow: 0 2px 8px rgba(16, 185, 129, 0.3);
+}
+.privacy-banner .icon { font-size: 1.8em; }
+.privacy-banner strong { display: block; font-size: 1.1em; margin-bottom: 4px; }
+"""
+
+PRIVACY_BANNER = """
+<div class="privacy-banner">
+    <div class="icon">🔒</div>
+    <div>
+        <strong>Ваши данные в безопасности</strong>
+        Все сообщения обрабатываются локально. 
+        <b>Данные НЕ передаются в OpenRussianAI</b>. 
+        История хранится только внутри контейнера.
+    </div>
+</div>
+"""
+
+# === ИНТЕРФЕЙС ===
+with gr.Blocks(title="OpenAirAI-X Chat") as demo:
+    username_state = gr.State("")
+    history_state = gr.State({})
+    current_chat_id_state = gr.State(None)
+    
+    with gr.Column(visible=False) as main_interface:
+        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Русскоязычный AI-ассистент</p></div>")
+        gr.HTML(PRIVACY_BANNER)
+        
+        with gr.Row():
+            with gr.Column(scale=1, min_width=250):
+                user_info = gr.Textbox(label="Пользователь", value="Не авторизован", interactive=False, elem_classes="user-info")
+                pro_badge = gr.HTML('<div class="pro-badge">👑 PRO</div>', visible=False)
+                new_chat_btn = gr.Button("➕ Новый чат", variant="primary")
+                gr.Markdown("### 📚 История чатов")
+                chat_list = gr.Dropdown(choices=[], label="Ваши чаты", interactive=True, allow_custom_value=False, value=None)
+                delete_chat_btn = gr.Button("🗑️ Удалить выбранный чат", variant="stop", size="sm")
+                logout_btn = gr.Button("🚪 Выйти", size="sm")
+            
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(label="Диалог", height=500)
+                with gr.Row():
+                    msg_input = gr.Textbox(placeholder="Напишите сообщение...", lines=2, scale=5, show_label=False)
+                    send_btn = gr.Button("📤 О��править", variant="primary", scale=1)
+    
+    with gr.Column(visible=True) as login_screen:
+        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Введите данные для входа</p></div>")
+        gr.HTML(PRIVACY_BANNER)
+        gr.Markdown("### 🔐 Вход в систему\nВведите ваше имя пользователя. История чатов будет привязана к этому имени.")
+        username_input = gr.Textbox(label="Имя пользователя (HF Username)", placeholder="Например: RootLinux21")
+        token_input = gr.Textbox(label="HF Token (необязательно, для PRO)", placeholder="hf_...")
+        login_btn = gr.Button("🔑 Войти", variant="primary")
+    
+    # === ОБРАБОТЧИКИ ===
+    
+    def handle_login(username, token):
+        if not username.strip():
+            gr.Warning("Введите имя пользователя!")
+            return gr.update(), gr.update(), "Не авторизован", "", {}, None, gr.update(choices=[], value=None), gr.update(visible=False)
+        
+        is_pro = check_if_pro(token)
+        history_data = load_history(username)
+        chat_list_choices = get_chat_list(username)
+        
+        current_chat_id = None
+        if not history_data:
+            chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+            history_data[chat_id] = {
+                "title": "Новый чат",
+                "created": datetime.now().isoformat(),
+                "messages": []
+            }
+            save_history(username, history_data)
+            current_chat_id = chat_id
+            chat_list_choices = get_chat_list(username)
+        
+        pro_badge_html = '<div class="pro-badge">👑 PRO</div>'
+        display_name = f"👤 {username}"
+        
+        return (
+            gr.update(visible=True),
+            gr.update(visible=False),
+            display_name,
+            username,
+            history_data,
+            current_chat_id,
+            gr.update(choices=chat_list_choices, value=None),
+            gr.update(visible=is_pro, value=pro_badge_html)
+        )
+    
+    def handle_logout():
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True),
+            "Не авторизован",
+            "",
+            {},
+            None,
+            gr.update(choices=[], value=None),
+            gr.update(visible=False)
+        )
+    
+    login_btn.click(
+        fn=handle_login,
+        inputs=[username_input, token_input],
+        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
+    )
+    
+    send_btn.click(
+        fn=generate_response,
+        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
+        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
+    )
+    
+    msg_input.submit(
+        fn=generate_response,
+        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
+        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
+    )
+    
+    new_chat_btn.click(
+        fn=new_chat,
+        inputs=[username_state],
+        outputs=[chatbot, current_chat_id_state, chat_list]
+    )
+    
+    chat_list.change(
+        fn=load_chat,
+        inputs=[chat_list, username_state, current_chat_id_state],
+        outputs=[chatbot, current_chat_id_state]
+    )
+    
+    delete_chat_btn.click(
+        fn=delete_chat,
+        inputs=[chat_list, username_state],
+        outputs=[chatbot, history_state, current_chat_id_state, chat_list]
+    )
+    
+    logout_btn.click(
+        fn=handle_logout,
+        inputs=None,
+        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860, css=CUSTOM_CSS)