Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import sqlite3
 import secrets
 import string
 import torch
@@ -6,8 +7,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # --- CONFIGURATION DU MODÈLE ---
 MODEL_ID = "Finisha-F-scratch/Charlotte-amity"
-
-print("Chargement du tokenizer et du modèle...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID, 
@@ -15,22 +14,60 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-# --- LOGIQUE DE VALIDATION ---
-def validate_key(key):
-    # Format: tn- + 4 lettres + 4 chiffres (Total 11 caractères)
-    if key.startswith("tn-") and len(key) == 11:
+# --- GESTION DE LA BASE DE DONNÉES (SQLITE) ---
+DB_NAME = "charlotte_api.db"
+
+def init_db():
+    conn = sqlite3.connect(DB_NAME)
+    c = conn.cursor()
+    # Table pour stocker les clés et le nombre de requêtes
+    c.execute('''CREATE TABLE IF NOT EXISTS api_keys 
+                 (key TEXT PRIMARY KEY, requests_count INTEGER)''')
+    conn.commit()
+    conn.close()
+
+init_db()
+
+def create_key():
+    """Génère une clé tn-xxxx0000 et l'enregistre"""
+    letters = ''.join(secrets.choice(string.ascii_lowercase) for _ in range(4))
+    digits = ''.join(secrets.choice(string.digits) for _ in range(4))
+    new_key = f"tn-{letters}{digits}"
+    
+    conn = sqlite3.connect(DB_NAME)
+    c = conn.cursor()
+    try:
+        c.execute("INSERT INTO api_keys (key, requests_count) VALUES (?, ?)", (new_key, 0))
+        conn.commit()
+        return new_key
+    except:
+        return "Erreur lors de la génération. Réessayez."
+    finally:
+        conn.close()
+
+def validate_and_log(key):
+    """Vérifie la clé et incrémente le compteur"""
+    conn = sqlite3.connect(DB_NAME)
+    c = conn.cursor()
+    c.execute("SELECT requests_count FROM api_keys WHERE key=?", (key,))
+    result = c.fetchone()
+    if result is not None:
+        c.execute("UPDATE api_keys SET requests_count = requests_count + 1 WHERE key=?", (key,))
+        conn.commit()
+        conn.close()
         return True
+    conn.close()
     return False
 
-# --- FONCTION API PRINCIPALE ---
-def predict(api_key, message):
+# --- FONCTION D'INFÉRENCE (L'API) ---
+def charlotte_api(api_key, message):
     """
-    Cette fonction est l'entrée de ton API.
+    Point d'entrée principal pour les appels externes.
     """
-    if not validate_key(api_key):
-        return "Erreur: Clé API invalide. Format requis: tn-xxxx0000"
+    if not validate_and_log(api_key):
+        return "❌ Erreur : Clé API 'tn-' invalide ou non trouvée."
 
-    # Encodage avec respect de la fenêtre de contexte de 128
+    # Limitation stricte au contexte de 128 tokens
     inputs = tokenizer(message, return_tensors="pt", truncation=True, max_length=128).to(model.device)
     
     with torch.no_grad():
@@ -38,31 +75,40 @@ def predict(api_key, message):
             **inputs,
             max_new_tokens=100,
             do_sample=True,
-            temperature=0.7,
+            temperature=0.8,
             pad_token_id=tokenizer.eos_token_id
         )
     
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # On ne renvoie que la réponse de Charlotte (sans le prompt)
+    # On renvoie uniquement la réponse générée
     response = full_text[len(message):].strip()
     return response
 
-# --- INTERFACE ET POINTS D'ENTRÉE API ---
-with gr.Blocks() as demo:
-    gr.Markdown("# 🌸 Charlotte-Amity API Server")
-    
-    key_input = gr.Textbox(label="Clé API")
-    msg_input = gr.Textbox(label="Message")
-    output = gr.Textbox(label="Réponse")
-    
-    submit_btn = gr.Button("Envoyer")
+# --- INTERFACE GRADIO ---
+with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+    gr.Markdown("# 🌸 Charlotte-Amity : Système API Personnel")
     
-    # L'argument api_name est CRUCIAL pour l'appel externe
-    submit_btn.click(
-        fn=predict, 
-        inputs=[key_input, msg_input], 
-        outputs=output,
-        api_name="chat" 
+    with gr.Tab("🔑 Gérer mes clés"):
+        gr.Markdown("Cliquez ci-dessous pour générer une clé d'accès unique.")
+        key_output = gr.Textbox(label="Votre nouvelle clé API", interactive=False)
+        generate_btn = gr.Button("Générer ma clé tn-")
+        generate_btn.click(create_key, outputs=key_output)
+        gr.Info("Note: Gardez précieusement votre clé pour vos appels externes.")
+
+    with gr.Tab("💬 Test Rapide"):
+        api_input = gr.Textbox(label="Clé API (tn-...)")
+        msg_input = gr.Textbox(label="Votre message")
+        text_output = gr.Textbox(label="Réponse de Charlotte")
+        run_btn = gr.Button("Envoyer")
+        run_btn.click(charlotte_api, inputs=[api_input, msg_input], outputs=text_output)
+
+    # L'EXPOSITION DE L'API POUR L'EXTÉRIEUR
+    # Cette ligne permet d'appeler la fonction via l'URL du Space
+    api_interface = gr.Interface(
+        fn=charlotte_api,
+        inputs=[gr.Textbox(label="api_key"), gr.Textbox(label="message")],
+        outputs="text",
+        api_name="chat" # C'est le nom à utiliser dans le client Python
     )
 
 if __name__ == "__main__":