Update token_cleaner.py

token_cleaner.py

@@ -25,86 +25,44 @@ logger = logging.getLogger(__name__)
 
 class TokenCleaner:
     def __init__(self, hf_token, dataset_name="HIRO12121212/videoinutoken"):
-        """
-        Initialize Token Cleaner
-        
-        Args:
-            hf_token: HuggingFace API token
-            dataset_name: Dataset name containing videoinu tokens
-        """
         self.hf_token = hf_token
         self.dataset_name = dataset_name
         self.hf_api = HfApi()
         logger.info(f"Token Cleaner initialized for dataset: {dataset_name}")
 
     def decode_jwt_token(self, token):
-        """
-        Decode JWT token to extract expiration time
-        
-        Args:
-            token: JWT token string
-            
-        Returns:
-            dict: Decoded token data or None if invalid
-        """
         try:
-            # JWT format: header.payload.signature
             parts = token.split('.')
             if len(parts) != 3:
                 logger.warning("Invalid JWT format (not 3 parts)")
                 return None
-                
-            # Decode payload (second part)
             payload = parts[1]
-            # Add padding if needed (JWT base64 doesn't use padding)
             padding = 4 - (len(payload) % 4)
             if padding != 4:
                 payload += '=' * padding
-                
-            # Decode base64
             decoded_bytes = base64.urlsafe_b64decode(payload)
             decoded_json = json.loads(decoded_bytes)
             return decoded_json
-            
         except Exception as e:
             logger.error(f"Error decoding JWT token: {str(e)}")
             return None
 
     def is_token_expired(self, token_data):
-        """
-        Check if token is expired based on 'exp' field
-        
-        Args:
-            token_data: Decoded JWT data
-            
-        Returns:
-            bool: True if expired, False otherwise
-        """
         if not token_data or 'exp' not in token_data:
-            return True  # Consider invalid tokens as expired
-            
+            return True
         exp_timestamp = token_data['exp']
         current_timestamp = int(time.time())
         is_expired = current_timestamp >= exp_timestamp
-        
         if is_expired:
             exp_date = datetime.fromtimestamp(exp_timestamp)
             logger.info(f"Token expired on {exp_date} (email: {token_data.get('email', 'unknown')})")
-            
         return is_expired
 
     def fetch_all_tokens(self):
-        """
-        Fetch all tokens from HuggingFace dataset
-        
-        Returns:
-            list: Array of token data
-        """
         try:
             download_url = f"https://huggingface.co/datasets/{self.dataset_name}/resolve/main/videoinu_tokens.json"
             headers = {"Authorization": f"Bearer {self.hf_token}"}
             response = requests.get(download_url, headers=headers, timeout=10)
-            
             if response.status_code == 200:
                 tokens_array = json.loads(response.text)
                 logger.info(f"✅ Fetched {len(tokens_array)} tokens from dataset")
@@ -112,27 +70,16 @@ class TokenCleaner:
             else:
                 logger.error(f"Failed to fetch tokens: {response.status_code}")
                 return []
-                
         except Exception as e:
             logger.error(f"Error fetching tokens: {str(e)}")
             return []
 
     def update_dataset(self, tokens_array):
-        """
-        Update dataset with cleaned tokens
-        
-        Args:
-            tokens_array: Array of valid tokens
-            
-        Returns:
-            bool: True if successful
-        """
         try:
             import tempfile
             with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                 json.dump(tokens_array, f, indent=2)
             temp_path = f.name
-            
             self.hf_api.upload_file(
                 path_or_fileobj=temp_path,
                 path_in_repo="videoinu_tokens.json",
@@ -140,36 +87,22 @@ class TokenCleaner:
                 repo_type="dataset",
                 token=self.hf_token
             )
-            
             os.unlink(temp_path)
             logger.info(f"✅ Updated dataset with {len(tokens_array)} tokens")
             return True
-            
         except Exception as e:
             logger.error(f"Error updating dataset: {str(e)}")
             return False
 
     def clean_expired_tokens(self):
-        """
-        Main cleaning function - removes expired tokens from dataset
-        
-        Returns:
-            dict: Statistics about cleaning operation
-        """
         logger.info("=" * 80)
         logger.info("Starting token cleaning cycle...")
         logger.info("=" * 80)
         
-        # Fetch all tokens
         all_tokens = self.fetch_all_tokens()
         if not all_tokens:
             logger.warning("No tokens found in dataset")
-            return {
-                "total": 0,
-                "valid": 0,
-                "expired": 0,
-                "invalid": 0
-            }
+            return {"total": 0, "valid": 0, "expired": 0, "invalid": 0}
             
         total_count = len(all_tokens)
         valid_tokens = []
@@ -178,20 +111,17 @@ class TokenCleaner:
         
         logger.info(f"Checking {total_count} tokens...")
         
-        # Check each token
         for i, token_entry in enumerate(all_tokens):
             token = token_entry.get('token', '')
             email = token_entry.get('email', 'unknown')
             credits = token_entry.get('credits_remaining', 0)
             
-            # Decode token
             decoded = self.decode_jwt_token(token)
             if decoded is None:
                 logger.warning(f"[{i+1}/{total_count}] Invalid token format: {email}")
                 invalid_count += 1
                 continue
                 
-            # Check expiration
             if self.is_token_expired(decoded):
                 logger.info(f"[{i+1}/{total_count}] 🗑️ Removing expired token: {email} ({credits} credits)")
                 expired_count += 1
@@ -201,12 +131,10 @@ class TokenCleaner:
                 logger.info(f"[{i+1}/{total_count}] ✅ Valid token: {email} (expires: {exp_date}, {credits} credits)")
                 valid_tokens.append(token_entry)
         
-        # Update dataset if any tokens were removed
         if expired_count > 0 or invalid_count > 0:
             logger.info(f"\n🧹 Cleaning dataset...")
             logger.info(f" Removing {expired_count} expired tokens")
             logger.info(f" Removing {invalid_count} invalid tokens")
-            
             if self.update_dataset(valid_tokens):
                 logger.info(f"✅ Dataset cleaned successfully!")
             else:
@@ -218,7 +146,8 @@ class TokenCleaner:
             "total": total_count,
             "valid": len(valid_tokens),
             "expired": expired_count,
-            "invalid": invalid_count
+            "invalid": invalid_count,
+            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         }
         
         logger.info("=" * 80)
@@ -232,12 +161,6 @@ class TokenCleaner:
         return stats
 
     def run_continuous(self, interval_minutes=60):
-        """
-        Run token cleaner continuously
-        
-        Args:
-            interval_minutes: Minutes between cleaning cycles (default: 60 = 1 hour)
-        """
         logger.info(f"🚀 Token Cleaner started - checking every {interval_minutes} minutes")
         while True:
             try:
@@ -254,61 +177,68 @@ class TokenCleaner:
                 logger.info(f"Retrying in 1 minute...")
                 time.sleep(60)
 
-# Gradio interface functions
-def clean_now(cleaner):
+# Global cleaner instance
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logger.error("HF_TOKEN environment variable not set!")
+    exit(1)
+
+cleaner = TokenCleaner(
+    hf_token=HF_TOKEN,
+    dataset_name="HIRO12121212/videoinutoken"
+)
+
+# Start background thread
+threading.Thread(target=cleaner.run_continuous, args=(60,), daemon=True).start()
+
+# Gradio functions (no arguments needed)
+def clean_now():
     stats = cleaner.clean_expired_tokens()
     summary = (
-        f"Cleaning Summary:\n"
-        f"Total tokens checked: {stats['total']}\n"
-        f"Valid tokens: {stats['valid']}\n"
-        f"Expired tokens removed: {stats['expired']}\n"
-        f"Invalid tokens removed: {stats['invalid']}"
+        f"🧹 Cleaning completed at {stats['timestamp']}\n\n"
+        f"Total tokens checked: **{stats['total']}**\n"
+        f"Valid tokens remaining: **{stats['valid']}**\n"
+        f"Expired tokens removed: **{stats['expired']}**\n"
+        f"Invalid tokens removed: **{stats['invalid']}**"
     )
     return summary
 
-if __name__ == "__main__":
-    # Get HF token from environment
-    HF_TOKEN = os.getenv("HF_TOKEN")
-    if not HF_TOKEN:
-        logger.error("HF_TOKEN environment variable not set!")
-        exit(1)
+def get_status():
+    return "Background cleaner is running (checks every 60 minutes). Use 'Clean Now' for manual cleaning."
+
+# Gradio interface
+with gr.Blocks(title="Videoinu Token Cleaner") as demo:
+    gr.Markdown("""
+    # 🧹 Videoinu Token Cleaner Dashboard
     
-    # Create cleaner instance
-    cleaner = TokenCleaner(
-        hf_token=HF_TOKEN,
-        dataset_name="HIRO12121212/videoinutoken"
-    )
+    This Space automatically cleans expired tokens from the dataset every hour.
     
-    # Start background cleaner thread
-    threading.Thread(target=cleaner.run_continuous, args=(60,), daemon=True).start()
+    - Background cleaning is always running.
+    - Click **Clean Now** to trigger an immediate clean.
+    - Results appear here + full logs in Space logs.
+    """)
     
-    # Gradio interface
-    with gr.Blocks(title="Videoinu Token Cleaner") as demo:
-        gr.Markdown("""
-        # Videoinu Token Cleaner
-        
-        This app cleans expired tokens from the Hugging Face dataset.
-        
-        - Background cleaning runs every 60 minutes.
-        - Click the button below to trigger a manual clean.
-        
-        Logs are available in the console/space logs.
-        """)
-        
-        output = gr.Textbox(label="Cleaning Results")
-        clean_button = gr.Button("Clean Now")
-        
-        clean_button.click(
-            fn=clean_now,
-            inputs=[],
-            outputs=[output],
-            api_name="clean_now"
-        )
+    status = gr.Textbox(label="Status", value="Background cleaner active", interactive=False)
+    output = gr.Markdown(label="Last Cleaning Result")
+    
+    with gr.Row():
+        clean_button = gr.Button("Clean Now", variant="primary", scale=1)
+        refresh_button = gr.Button("Refresh Status", scale=1)
     
-    # Launch Gradio app
-    # Use port 7860 for Hugging Face Spaces
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    clean_button.click(
+        fn=clean_now,
+        outputs=output,
+        api_name="clean_now"
+    )
+    
+    refresh_button.click(
+        fn=get_status,
+        outputs=status
+    )
+
+# Launch the app
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    show_error=True
+)