Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,85 +1,126 @@
 import streamlit as st
 import re
 import datetime
+import sqlite3
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch 
 
-# --- 1. Configuration et Design (Rose & Blanc/Manga) ---
-
-st.set_page_config(
-    page_title="Charlotte-APY 💖",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# Style personnalisé pour l'ambiance Manga/Rose/Blanc
-st.markdown(
-    """
-    <style>
-    /* Généralités */
-    .stApp {
-        background-color: #FFF0F5; /* Rose Pâle (Blush) */
-        color: #333333; /* Texte Gris Foncé */
-    }
-    
-    /* Titre (Tonalité Manga) */
-    h1 {
-        color: #FF69B4; /* Rose Vif */
-        text-shadow: 2px 2px 4px #F08080; /* Ombre douce */
-        font-family: 'Comic Sans MS', cursive, sans-serif; /* Police ludique */
-        text-align: center;
-    }
-    
-    /* Boutons */
-    .stButton>button {
-        background-color: #FFB6C1; /* Rose Clair */
-        color: #8B0000; /* Rouge Fonce pour le texte */
-        border: 2px solid #FF69B4;
-        border-radius: 10px;
-        transition: all 0.2s;
-        font-weight: bold;
-    }
-    .stButton>button:hover {
-        background-color: #FF69B4; /* Rose Vif au survol */
-        color: white;
-    }
+# --- 0. CONFIGURATION ET INITIALISATION DE LA BASE DE DONNÉES ---
+
+DB_NAME = "charlotte_apy.db"
+MODEL_NAME = "Clemylia/Tiny-charlotte"
+MAX_QUOTA = 600
+MAX_TOKENS_PER_RESPONSE = 100
+
+# ----------------------------------------------------
+# A. LOGIQUE DE BASE DE DONNÉES SQLite (Mise à jour)
+# ----------------------------------------------------
+
+def init_db():
+    """Initialise la base de données et la table api_keys."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS api_keys (
+            key_id TEXT PRIMARY KEY,
+            quota_remaining INTEGER,
+            max_quota INTEGER,
+            date_last_use TEXT
+        )
+    """)
+    conn.commit()
+    conn.close()
+
+def get_all_keys():
+    """Récupère toutes les clés d'API."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("SELECT key_id, quota_remaining, max_quota, date_last_use FROM api_keys")
+    rows = cursor.fetchall()
+    conn.close()
     
-    /* Conteneurs et zones de texte */
-    .stTextInput>div>div>input, .stTextArea>div>div>textarea {
-        background-color: white;
-        border: 1px solid #FFB6C1;
-        border-radius: 5px;
-    }
+    keys = {}
+    for row in rows:
+        key_id, quota_remaining, max_quota, date_last_use = row
+        keys[key_id] = {
+            'quota_remaining': quota_remaining,
+            'max_quota': max_quota,
+            'date_last_use': date_last_use
+        }
+    return keys
+
+def get_key_data(key_id):
+    """Récupère les données d'une seule clé."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("SELECT quota_remaining, max_quota, date_last_use FROM api_keys WHERE key_id = ?", (key_id,))
+    row = cursor.fetchone()
+    conn.close()
     
-    /* En-têtes secondaires */
-    h2, h3 {
-        color: #FF1493; /* Rose Profond */
-        border-bottom: 2px solid #FFC0CB; /* Ligne sous les titres */
-        padding-bottom: 5px;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-st.title("💖 Charlotte-APY : Votre Portail d'API Tiny-Charlotte 🎀")
-
-# --- 2. Initialisation de l'État de Session ---
+    if row:
+        return {'quota_remaining': row[0], 'max_quota': row[1], 'date_last_use': row[2]}
+    return None
+
+def add_key_to_db(key_id, max_quota=MAX_QUOTA):
+    """Ajoute une nouvelle clé à la base de données."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    today = datetime.date.today().isoformat()
+    try:
+        cursor.execute("""
+            INSERT INTO api_keys (key_id, quota_remaining, max_quota, date_last_use) 
+            VALUES (?, ?, ?, ?)
+        """, (key_id, max_quota, max_quota, today))
+        conn.commit()
+        conn.close()
+        return True
+    except sqlite3.IntegrityError:
+        conn.close()
+        return False
+
+def delete_key_from_db(key_id):
+    """Supprime une clé de la base de données."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM api_keys WHERE key_id = ?", (key_id,))
+    conn.commit()
+    conn.close()
+
+def update_key_quota_in_db(key_id, new_remaining_quota, new_date_last_use):
+    """Met à jour le quota et la date d'utilisation."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("""
+        UPDATE api_keys 
+        SET quota_remaining = ?, date_last_use = ? 
+        WHERE key_id = ?
+    """, (new_remaining_quota, new_date_last_use, key_id))
+    conn.commit()
+    conn.close()
+
+# Nouvelle fonction d'administration : réinitialiser le quota
+def reset_key_quota_in_db(key_id):
+    """Réinitialise le quota d'une clé au maximum et met à jour la date."""
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    today = datetime.date.today().isoformat()
+    cursor.execute("""
+        UPDATE api_keys 
+        SET quota_remaining = max_quota, date_last_use = ? 
+        WHERE key_id = ?
+    """, (today, key_id))
+    conn.commit()
+    conn.close()
 
-if 'api_keys' not in st.session_state:
-    # Format : {'clé_api': {'quota_remaining': 600, 'max_quota': 600, 'date_last_use': '2025-11-17'}}
-    st.session_state['api_keys'] = {} 
+# Initialisation au démarrage
+init_db()
 
-# --- 3. LOGIQUE DE VALIDATION DE CLÉ ---
+# ----------------------------------------------------
+# B. LOGIQUE DE VALIDATION ET CHARGEMENT DU MODÈLE
+# ----------------------------------------------------
 
 def validate_key(key_str):
-    """
-    Vérifie si la clé d'API respecte les contraintes :
-    1. Commence par 'Tn-charlotte'.
-    2. Contient au moins 5 chiffres.
-    3. Contient au moins 7 lettres (a-z ou A-Z).
-    4. Est unique.
-    """
+    """Vérifie si la clé respecte les contraintes de format."""
     
     if not key_str.startswith("Tn-charlotte"):
         return False, "La clé doit obligatoirement commencer par **Tn-charlotte**."
@@ -91,30 +132,20 @@ def validate_key(key_str):
     num_letters = len(re.findall(r'[a-zA-Z]', key_str))
     if num_letters < 7:
         return False, f"La clé doit contenir au moins **7 lettres** (actuel : {num_letters})."
-
-    if key_str in st.session_state.api_keys:
-        return False, "Cette clé d'API existe déjà."
         
     return True, "Clé valide !"
 
-# --- 4. CHARGEMENT DU MODÈLE ---
-
 @st.cache_resource
 def load_tiny_charlotte():
-    """
-    Charge le modèle tiny-charlotte et le tokenizer en utilisant le cache Streamlit.
-    """
-    model_name = "Clemylia/Tiny-charlotte"
-    
+    """Charge le modèle tiny-charlotte et le tokenizer en utilisant le cache Streamlit."""
     try:
-        st.sidebar.info(f"⏳ Chargement du modèle {model_name}...")
+        st.sidebar.info(f"⏳ Chargement du modèle {MODEL_NAME}...")
         
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         
-        # Utiliser un dispositif si possible (CUDA si disponible, sinon CPU)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
         
         st.sidebar.success(f"✅ Modèle Tiny-Charlotte chargé sur **{device}**.")
         return model, tokenizer, device
@@ -123,15 +154,16 @@ def load_tiny_charlotte():
         st.sidebar.error(f"❌ Erreur de chargement du modèle : {e}")
         return None, None, None
 
-# --- 5. INFÉRENCE ET GESTION DU QUOTA ---
+# ----------------------------------------------------
+# C. INFÉRENCE ET GESTION DU QUOTA
+# ----------------------------------------------------
 
 def run_inference(api_key, prompt, model, tokenizer, device):
     """
-    Exécute l'inférence, vérifie/réinitialise le quota journalier (600 tokens) 
-    et limite la génération à 128 tokens.
+    Exécute l'inférence avec le modèle, gère le quota via SQLite.
     """
     today = datetime.date.today().isoformat()
-    key_data = st.session_state.api_keys.get(api_key)
+    key_data = get_key_data(api_key)
     
     if not key_data:
         return "Erreur: Clé d'API non valide ou non trouvée.", 0
@@ -140,57 +172,71 @@ def run_inference(api_key, prompt, model, tokenizer, device):
     if key_data['date_last_use'] != today:
         key_data['quota_remaining'] = key_data['max_quota']
         key_data['date_last_use'] = today
-        st.session_state.api_keys[api_key] = key_data
-        st.success(f"🎉 Quota réinitialisé pour la clé **{api_key}** : 600 tokens disponibles.")
+        update_key_quota_in_db(api_key, key_data['quota_remaining'], today)
+        st.success(f"🎉 Quota réinitialisé automatiquement pour la clé **{api_key}** : {key_data['max_quota']} tokens disponibles aujourd'hui.")
 
     # 2. Vérification du Quota
-    MAX_TOKENS_PER_RESPONSE = 100
-    
     if key_data['quota_remaining'] < MAX_TOKENS_PER_RESPONSE:
-        return f"🚫 **Quota journalier atteint** ({key_data['quota_remaining']} / 600). Veuillez réessayer demain.", 0
+        return f"🚫 **Quota journalier atteint** ({key_data['quota_remaining']} / {MAX_QUOTA}). Veuillez réessayer demain.", 0
         
     if model is None or tokenizer is None:
         return "Erreur interne: Le modèle n'est pas prêt.", 0
 
     # 3. Exécution de la Génération
     try:
-        # Encodage du prompt et transfert vers le dispositif (CPU/CUDA)
         input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
         
-        # Génération de la réponse
         output = model.generate(
             input_ids,
             max_length=input_ids.shape[1] + MAX_TOKENS_PER_RESPONSE, 
-            do_sample=True,
-            top_k=50,
-            top_p=0.95,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id # Important pour la génération
+            do_sample=True, top_k=50, top_p=0.95, num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id 
         )
         
-        # Décodage
         response_text = tokenizer.decode(output[0], skip_special_tokens=True)
-        
-        # Calcul des tokens utilisés (Tokens générés - Tokens du prompt)
         tokens_generated = output.shape[1] - input_ids.shape[1]
         
     except Exception as e:
         st.error(f"Erreur lors de la génération: {e}")
         return "Erreur d'inférence. Problème avec le modèle.", 0
 
-    # 4. Mise à Jour du Quota
-    key_data['quota_remaining'] -= tokens_generated
-    st.session_state.api_keys[api_key] = key_data
+    # 4. Mise à Jour du Quota DANS LA BASE DE DONNÉES
+    new_remaining = key_data['quota_remaining'] - tokens_generated
+    update_key_quota_in_db(api_key, new_remaining, today)
 
     return response_text, tokens_generated
 
-# --- 6. INTERFACES UTILISATEUR ---
+# ----------------------------------------------------
+# D. INTERFACE UTILISATEUR STREAMLIT
+# ----------------------------------------------------
+
+# Configuration Streamlit et Style
+st.set_page_config(page_title="Charlotte-APY 💖", layout="wide", initial_sidebar_state="expanded")
+
+st.markdown(
+    """
+    <style>
+    /* ... (CSS pour le style Rose & Manga) ... */
+    .stApp { background-color: #FFF0F5; color: #333333; }
+    h1 { color: #FF69B4; text-shadow: 2px 2px 4px #F08080; font-family: 'Comic Sans MS', cursive, sans-serif; text-align: center; }
+    .stButton>button { background-color: #FFB6C1; color: #8B0000; border: 2px solid #FF69B4; border-radius: 10px; font-weight: bold; transition: all 0.2s; }
+    .stButton>button:hover { background-color: #FF69B4; color: white; }
+    .stTextInput>div>div>input, .stTextArea>div>div>textarea { background-color: white; border: 1px solid #FFB6C1; border-radius: 5px; }
+    h2, h3 { color: #FF1493; border-bottom: 2px solid #FFC0CB; padding-bottom: 5px; }
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+st.title("💖 Charlotte-APY : Votre Portail d'API Tiny-Charlotte 🎀")
+
 
 def generate_api_key_ui():
     """Interface de création de clé d'API."""
     st.subheader("🔑 Créer une Nouvelle Clé d'API")
     
-    if len(st.session_state.api_keys) >= 6:
+    current_keys = get_all_keys()
+    
+    if len(current_keys) >= 6:
         st.warning("❌ **Limite Atteinte !** Vous gérez déjà le maximum de 6 clés d'API.")
         return
         
@@ -205,14 +251,16 @@ def generate_api_key_ui():
         if submitted:
             is_valid, message = validate_key(new_key)
             
+            if get_key_data(new_key) is not None:
+                 is_valid = False
+                 message = "Cette clé d'API existe déjà (base de données)."
+            
             if is_valid:
-                st.session_state.api_keys[new_key] = {
-                    'quota_remaining': 600,
-                    'max_quota': 600,
-                    'date_last_use': datetime.date.today().isoformat()
-                }
-                st.success(f"✅ Clé d'API **{new_key}** créée ! Quota : 600 tokens/jour.")
-                st.rerun() 
+                if add_key_to_db(new_key):
+                    st.success(f"✅ Clé d'API **{new_key}** créée ! Quota : {MAX_QUOTA} tokens/jour.")
+                    st.rerun() 
+                else:
+                    st.error("Échec de l'ajout à la base de données.")
             else:
                 st.error(f"🚫 **Erreur de validation :** {message}")
 
@@ -220,47 +268,68 @@ def manage_api_keys_ui():
     """Interface de gestion et suppression des clés d'API."""
     st.subheader("🗂️ Gérer Vos Clés d'API")
     
-    keys_list = list(st.session_state.api_keys.keys())
+    keys_data_dict = get_all_keys()
+    keys_list = list(keys_data_dict.keys())
     
     if not keys_list:
         st.info("Vous n'avez pas encore de clés d'API.")
         return
 
-    # Affichage des clés dans un tableau
-    keys_data = []
-    for key, data in st.session_state.api_keys.items():
-        keys_data.append({
+    keys_data_list = []
+    for key, data in keys_data_dict.items():
+        keys_data_list.append({
             "Clé d'API": key,
             "Tokens Restants": f"{data['quota_remaining']} / {data['max_quota']}",
             "Dernière Utilisation (Réinitialisation)": data['date_last_use']
         })
         
-    st.dataframe(keys_data, use_container_width=True, hide_index=True)
+    st.dataframe(keys_data_list, use_container_width=True, hide_index=True)
     
     st.markdown("---") 
 
     # Section de Suppression
     st.write("### 🗑️ Supprimer une Clé")
-    key_to_delete = st.selectbox("Sélectionnez la clé à supprimer :", [""] + keys_list)
+    key_to_delete = st.selectbox("Sélectionnez la clé à supprimer :", [""] + keys_list, key="delete_select")
     
     if st.button("💔 Supprimer la Clé Sélectionnée", disabled=(key_to_delete == "")):
-        if key_to_delete in st.session_state.api_keys:
-            del st.session_state.api_keys[key_to_delete]
+        if key_to_delete:
+            delete_key_from_db(key_to_delete)
             st.success(f"🗑️ Clé **{key_to_delete}** supprimée avec succès.")
             st.rerun()
+
+def admin_quota_ui():
+    """Interface d'administration pour la réinitialisation manuelle du quota."""
+    st.subheader("⚙️ Administration : Réinitialisation du Quota")
+    
+    keys_list = list(get_all_keys().keys())
+    
+    if not keys_list:
+        st.info("Aucune clé d'API à administrer.")
+        return
         
+    key_to_reset = st.selectbox("Sélectionnez la clé à réinitialiser :", [""] + keys_list, key="reset_select")
+    
+    if st.button("🔄 Réinitialiser le Quota à 600"):
+        if key_to_reset:
+            reset_key_quota_in_db(key_to_reset)
+            st.success(f"✅ Quota pour la clé **{key_to_reset}** réinitialisé à {MAX_QUOTA} tokens.")
+            st.rerun()
+        else:
+            st.warning("Veuillez sélectionner une clé.")
+
+
 def test_api_ui(model, tokenizer, device):
     """Interface pour tester l'API."""
     
     st.subheader("🧪 Tester l'API Tiny-Charlotte")
     
-    keys_list = list(st.session_state.api_keys.keys())
+    keys_list = list(get_all_keys().keys())
     
     if not keys_list:
         st.warning("Créez une clé d'API avant de pouvoir tester l'inférence.")
         return
         
-    selected_key = st.selectbox("Sélectionnez votre clé d'API :", keys_list)
+    selected_key = st.selectbox("Sélectionnez votre clé d'API :", keys_list, key="inference_select")
     
     with st.form("inference_form"):
         prompt = st.text_area("Votre Requête pour Tiny-Charlotte", height=100)
@@ -272,10 +341,10 @@ def test_api_ui(model, tokenizer, device):
                 return
                 
             if model is None:
-                 st.error("Le modèle n'a pas pu être chargé. Vérifiez votre connexion ou la librairie transformers.")
+                 st.error("Le modèle n'a pas pu être chargé.")
                  return
                  
-            with st.spinner("Appel de l'API Charlotte-APY en cours... (Vérification du quota et inférence)"):
+            with st.spinner("Appel de l'API Charlotte-APY en cours..."):
                 
                 response, tokens_used = run_inference(selected_key, prompt, model, tokenizer, device)
                 
@@ -283,30 +352,75 @@ def test_api_ui(model, tokenizer, device):
                 st.info(response)
                 
                 # Affichage des métriques de quota
-                if selected_key in st.session_state.api_keys:
-                    remaining = st.session_state.api_keys[selected_key]['quota_remaining']
+                updated_data = get_key_data(selected_key)
+                
+                if updated_data:
+                    remaining = updated_data['quota_remaining']
                     
                     if tokens_used > 0:
-                        st.success(f"Tokens utilisés : **{tokens_used}**. Tokens restants pour cette clé aujourd'hui : **{remaining}** / 600.")
-                    elif remaining < 128:
-                        # Si le quota est bas, c'est l'erreur de quota qui est apparue
-                         st.error(f"Tokens restants : **{remaining}** / 600. Vous n'avez pas assez de tokens pour une réponse complète (128).")
+                        st.success(f"Tokens utilisés : **{tokens_used}**. Tokens restants pour cette clé aujourd'hui : **{remaining}** / {MAX_QUOTA}.")
+                    elif remaining < MAX_TOKENS_PER_RESPONSE:
+                         st.error(f"Tokens restants : **{remaining}** / {MAX_QUOTA}. Vous n'avez pas assez de tokens pour une réponse complète ({MAX_TOKENS_PER_RESPONSE}).")
                     else:
                         st.warning("Aucun token utilisé (erreur de modèle ou quota atteint).")
+                else:
+                    st.error("Impossible de récupérer les données de la clé après l'inférence.")
 
+def api_documentation_ui():
+    """Fournit la documentation pour l'utilisation externe de l'API (simulée)."""
+    st.subheader("📖 Documentation de l'API Tiny-Charlotte")
+    
+    st.markdown(f"""
+    L'authentification se fait via la clé d'API dans l'en-tête `Authorization: Bearer <votre_clé>`.
+    
+    * **Limite :** {MAX_QUOTA} tokens par clé par jour.
+    * **Réponse max :** {MAX_TOKENS_PER_RESPONSE} tokens.
+    
+    ### 🌐 Endpoint (Simulé) : `POST https://api.charlotte-apy.com/v1/inference`
+    """)
+    
+    code_example = """
+import requests
+import json
+
+API_URL = "https://api.charlotte-apy.com/v1/inference" 
+YOUR_API_KEY = "Tn-charlotte_Ma_Cle_12345ABCDEFG" 
+
+payload = {
+    "prompt": "Peux-tu me donner un conseil sur l'espoir ?",
+}
+
+headers = {
+    "Content-Type": "application/json",
+    "Authorization": f"Bearer {YOUR_API_KEY}" 
+}
+
+try:
+    response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
+    response.raise_for_status() 
+    data = response.json()
+    print("Réponse de Tiny-Charlotte :", data.get("generated_text"))
+    
+except requests.exceptions.RequestException as e:
+    print(f"Erreur lors de l'appel API: {e}")
+"""
+    st.code(code_example, language="python")
 
-# --- 7. DISPOSITION PRINCIPALE DE L'APPLICATION ---
 
+# --- DISPOSITION PRINCIPALE DE L'APPLICATION ---
 def main_app():
     
-    # Le modèle est chargé une seule fois au début de l'exécution
+    # 1. Chargement du Modèle (mis en cache)
     model, tokenizer, device = load_tiny_charlotte()
 
-    # Colonnes pour la création de clé
+    # 2. Gestion des Clés et Quota (Colonnes)
     col1, col2 = st.columns([1, 1])
 
     with col1:
         generate_api_key_ui()
+    
+    with col2:
+        admin_quota_ui() # Nouvelle section d'administration
         
     st.markdown("---") 
     
@@ -314,10 +428,14 @@ def main_app():
     
     st.markdown("---") 
     
-    # Le test de l'API dépend du modèle chargé
+    # 3. Test de l'API
     test_api_ui(model, tokenizer, device)
+    
+    st.markdown("---") 
+    
+    # 4. Documentation
+    api_documentation_ui()
 
 
-# Exécution de l'application
 if __name__ == "__main__":
     main_app()