Update app.py

app.py

@@ -238,6 +238,27 @@ class DatabaseManager:
             logger.error(f"Error getting all users stats: {e}")
             return []
 
+# ----------------- Hugging Face Auth -----------------
+class HuggingFaceAuth:
+    @staticmethod
+    async def validate_token(token: str) -> Tuple[bool, Optional[Dict]]:
+        """Validate Hugging Face token and return user info"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                headers = {"Authorization": f"Bearer {token}"}
+                async with session.get(
+                    "https://huggingface.co/api/whoami",
+                    headers=headers
+                ) as response:
+                    if response.status == 200:
+                        user_info = await response.json()
+                        return True, user_info
+                    else:
+                        return False, None
+        except Exception as e:
+            logger.error(f"Error validating HF token: {e}")
+            return False, None
+
 # ----------------- Model Manager -----------------
 class ModelManager:
     def __init__(self, config: Config):
@@ -293,7 +314,7 @@ class ModelManager:
             logger.error(f"❌ Error loading model: {e}")
             self.model_loaded = False
     
-    async def generate(self, request: GenerationRequest) -> Tuple[bool, str, int]:
+    async def generate(self, request: GenerationRequest, hf_token: str = None) -> Tuple[bool, str, int]:
         """Generate text with the model"""
         if not self.model_loaded:
             return False, "❌ Model not loaded", 0
@@ -380,6 +401,7 @@ class GemmaSaaSService:
         self.db = DatabaseManager(self.config)
         self.model_manager = ModelManager(self.config)
         self.analytics_manager = AnalyticsManager(self.db)
+        self.hf_auth = HuggingFaceAuth()
         self._validate_config()
     
     def _validate_config(self):
@@ -417,17 +439,33 @@ class GemmaSaaSService:
                 error="Internal service error"
             )
     
-    async def generate_text(self, prompt: str, api_key: str, **kwargs) -> APIResponse:
-        """Generate text with authentication and rate limiting"""
+    async def generate_text(self, prompt: str, hf_token: str, **kwargs) -> APIResponse:
+        """Generate text with Hugging Face authentication"""
         try:
-            # Validate API key
-            user = await self.db.validate_api_key(api_key)
-            if not user:
+            # Validate HF token
+            is_valid, user_info = await self.hf_auth.validate_token(hf_token)
+            if not is_valid or not user_info:
                 return APIResponse(
                     success=False, 
-                    error="⚠️ Invalid API key"
+                    error="⚠️ Invalid Hugging Face token"
                 )
             
+            # Get or create user
+            user_email = user_info.get('email', '')
+            user_name = user_info.get('name', user_info.get('preferred_username', 'User'))
+            
+            # Check if user exists
+            user = await self.db.validate_api_key(hf_token)  # We're using HF token as API key now
+            
+            if not user:
+                # Create new user
+                user_data = UserCreate(name=user_name, email=user_email, plan="free")
+                success, message, api_key = await self.db.create_user(user_data, hf_user_id=user_info.get('sub', ''))
+                if not success:
+                    return APIResponse(success=False, error=message)
+                
+                user = await self.db.validate_api_key(hf_token)
+            
             # Check rate limit
             can_make_request, requests_used = await self.db.check_rate_limit(user['id'])
             if not can_make_request:
@@ -442,7 +480,7 @@ class GemmaSaaSService:
             
             # Generate text
             request = GenerationRequest(prompt=prompt, **kwargs)
-            success, text, tokens_used = await self.model_manager.generate(request)
+            success, text, tokens_used = await self.model_manager.generate(request, hf_token)
             
             if success:
                 # Update usage statistics
@@ -470,14 +508,22 @@ class GemmaSaaSService:
                 error="Internal service error"
             )
     
-    async def get_user_stats(self, api_key: str) -> APIResponse:
-        """Get user statistics"""
+    async def get_user_stats(self, hf_token: str) -> APIResponse:
+        """Get user statistics using Hugging Face token"""
         try:
-            user = await self.db.validate_api_key(api_key)
+            # Validate HF token
+            is_valid, user_info = await self.hf_auth.validate_token(hf_token)
+            if not is_valid or not user_info:
+                return APIResponse(
+                    success=False, 
+                    error="Invalid Hugging Face token"
+                )
+            
+            user = await self.db.validate_api_key(hf_token)
             if not user:
                 return APIResponse(
                     success=False, 
-                    error="Invalid API key"
+                    error="User not found"
                 )
             
             # Generate analytics plot
@@ -505,11 +551,19 @@ class GemmaSaaSService:
                 error="Error retrieving stats"
             )
     
-    async def get_admin_stats(self, api_key: str) -> APIResponse:
+    async def get_admin_stats(self, hf_token: str) -> APIResponse:
         """Get admin statistics (only for admin users)"""
         try:
-            user = await self.db.validate_api_key(api_key)
-            if not user or user.get('email') != self.config.ADMIN_EMAIL:
+            # Validate HF token
+            is_valid, user_info = await self.hf_auth.validate_token(hf_token)
+            if not is_valid or not user_info:
+                return APIResponse(
+                    success=False, 
+                    error="Invalid Hugging Face token"
+                )
+            
+            user_email = user_info.get('email', '')
+            if user_email != self.config.ADMIN_EMAIL:
                 return APIResponse(
                     success=False, 
                     error="Unauthorized: Admin access required"
@@ -810,6 +864,25 @@ class GradioInterface:
         .dark-mode-plot {
             background: transparent !important;
         }
+        
+        .hf-login-container {
+            text-align: center;
+            padding: 2rem;
+        }
+        
+        .hf-login-btn {
+            background: linear-gradient(135deg, #ffd21e 0%, #ff9d0a 100%) !important;
+            color: #000 !important;
+            font-weight: 600 !important;
+            padding: 1rem 2rem !important;
+            border-radius: 12px !important;
+            margin: 1rem 0;
+        }
+        
+        .hf-login-btn:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 6px 20px rgba(255, 157, 10, 0.4) !important;
+        }
         """
     
     def create_header(self):
@@ -847,6 +920,24 @@ class GradioInterface:
             
             return gr.HTML(examples_html)
     
+    def create_hf_login_component(self):
+        with gr.Column(elem_classes=["card"]):
+            gr.Markdown("### 🔐 Hugging Face Login")
+            gr.HTML("""
+            <div class="hf-login-container">
+                <p>To use this platform, please login with your Hugging Face account:</p>
+                <a href="https://huggingface.co/login" target="_blank" class="btn hf-login-btn">Login with Hugging Face</a>
+                <p>After logging in, go to your <a href="https://huggingface.co/settings/tokens" target="_blank" style="color: #ffd21e;">Hugging Face settings</a> to get your access token.</p>
+            </div>
+            """)
+            hf_token_input = gr.Textbox(
+                label="Hugging Face Access Token",
+                type="password",
+                placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+                elem_classes=["input"]
+            )
+            return hf_token_input
+
     async def create_interface(self):
         """Create the enhanced Gradio interface"""
         with gr.Blocks(
@@ -866,12 +957,7 @@ class GradioInterface:
                     with gr.Row():
                         with gr.Column(scale=1):
                             gr.Markdown("### ⚙️ Configuration")
-                            api_key_playground = gr.Textbox(
-                                label="🔑 API Key",
-                                type="password",
-                                placeholder="Enter your API key...",
-                                elem_classes=["input"]
-                            )
+                            hf_token_input = self.create_hf_login_component()
                             
                             with gr.Accordion("Advanced Settings", open=False):
                                 max_tokens_input = gr.Slider(
@@ -932,59 +1018,31 @@ class GradioInterface:
                             )
                             
                             # Stats display
-                            generation_stats = gr.JSON(
-                                label="📊 Generation Statistics",
-                                visible=False
-                            )
+                            generation_stats = gr.HTML()
                 
                 # Profile Tab
                 with gr.Tab("👤 Profile", elem_classes=["card"]):
                     with gr.Row():
                         with gr.Column(scale=1):
-                            gr.Markdown("### 🆕 Create Account")
-                            name_input = gr.Textbox(
-                                label="👤 Full Name",
-                                elem_classes=["input"]
-                            )
-                            email_input = gr.Textbox(
-                                label="📧 Email Address",
-                                elem_classes=["input"]
-                            )
-                            plan_input = gr.Dropdown(
-                                choices=["free", "pro", "enterprise"],
-                                value="free",
-                                label="📋 Plan",
+                            gr.Markdown("### 🔐 Authentication")
+                            profile_hf_token = gr.Textbox(
+                                label="Hugging Face Access Token",
+                                type="password",
+                                placeholder="Enter your Hugging Face token...",
                                 elem_classes=["input"]
                             )
-                            
-                            create_btn = gr.Button(
-                                "✨ Create API Key",
-                                elem_classes=["btn", "btn-primary"],
-                                variant="primary"
-                            )
-                            
-                            creation_status = gr.HTML()
-                            api_key_display = gr.Textbox(
-                                label="🔑 Your API Key",
-                                interactive=False,
-                                visible=False,
-                                elem_classes=["input"]
+                            profile_login_btn = gr.Button(
+                                "🔓 Login",
+                                elem_classes=["btn", "btn-primary"]
                             )
+                            profile_status = gr.HTML()
                         
                         with gr.Column(scale=1):
                             gr.Markdown("### 📊 Account Statistics")
-                            stats_api_key = gr.Textbox(
-                                label="🔑 API Key",
-                                type="password",
-                                placeholder="Enter API key to view stats",
-                                elem_classes=["input"]
-                            )
-                            
                             refresh_stats_btn = gr.Button(
                                 "🔄 Refresh Stats",
                                 elem_classes=["btn", "btn-secondary"]
                             )
-                            
                             user_stats_display = gr.HTML()
                 
                 # Analytics Tab
@@ -993,10 +1051,10 @@ class GradioInterface:
                     
                     with gr.Row():
                         with gr.Column(scale=1):
-                            analytics_api_key = gr.Textbox(
-                                label="🔑 API Key",
+                            analytics_hf_token = gr.Textbox(
+                                label="Hugging Face Access Token",
                                 type="password",
-                                placeholder="Enter API key to view analytics",
+                                placeholder="Enter your Hugging Face token...",
                                 elem_classes=["input"]
                             )
                             refresh_analytics_btn = gr.Button(
@@ -1012,10 +1070,10 @@ class GradioInterface:
                     
                     with gr.Row():
                         with gr.Column(scale=1):
-                            admin_api_key = gr.Textbox(
-                                label="🔑 Admin API Key",
+                            admin_hf_token = gr.Textbox(
+                                label="Hugging Face Access Token",
                                 type="password",
-                                placeholder="Enter admin API key",
+                                placeholder="Enter admin Hugging Face token...",
                                 elem_classes=["input"]
                             )
                             refresh_admin_btn = gr.Button(
@@ -1030,13 +1088,13 @@ class GradioInterface:
             self.create_footer()
             
             # Event Handlers
-            async def handle_generation(prompt, api_key, max_tokens, temperature, top_k, top_p, repetition_penalty):
-                if not api_key.strip():
-                    return "⚠️ Please enter your API key", {}, False
+            async def handle_generation(prompt, hf_token, max_tokens, temperature, top_k, top_p, repetition_penalty):
+                if not hf_token.strip():
+                    return "⚠️ Please enter your Hugging Face token", ""
                 
                 response = await self.service.generate_text(
                     prompt=prompt,
-                    api_key=api_key,
+                    hf_token=hf_token,
                     max_tokens=max_tokens,
                     temperature=temperature,
                     top_k=top_k,
@@ -1063,48 +1121,27 @@ class GradioInterface:
                     """
                     return (
                         response.data["generated_text"],
-                        stats_html,
-                        True
+                        stats_html
                     )
                 else:
                     return (
                         response.error,
-                        f'<div class="alert alert-error">{response.error}</div>',
-                        False
+                        f'<div class="alert alert-error">{response.error}</div>'
                     )
             
-            async def handle_user_creation(name, email, plan):
-                if not name or not name.strip():
-                    return (
-                        f'<div class="alert alert-error">❌ Name is required</div>',
-                        "",
-                        False
-                    )
+            async def handle_profile_login(hf_token):
+                if not hf_token.strip():
+                    return f'<div class="alert alert-error">❌ Please enter your Hugging Face token</div>'
                 
-                if not email or not email.strip():
-                    return (
-                        f'<div class="alert alert-error">❌ Email is required</div>',
-                        "",
-                        False
-                    )
-                
-                response = await self.service.create_user(name, email, plan)
+                response = await self.service.get_user_stats(hf_token)
                 
                 if response.success:
-                    return (
-                        f'<div class="alert alert-success">✅ Account created successfully! Your API key is below.</div>',
-                        response.data["api_key"],
-                        True
-                    )
+                    return f'<div class="alert alert-success">✅ Successfully logged in!</div>'
                 else:
-                    return (
-                        f'<div class="alert alert-error">❌ {response.error}</div>',
-                        "",
-                        False
-                    )
+                    return f'<div class="alert alert-error">❌ {response.error}</div>'
             
-            async def handle_stats_refresh(api_key):
-                response = await self.service.get_user_stats(api_key)
+            async def handle_stats_refresh(hf_token):
+                response = await self.service.get_user_stats(hf_token)
                 
                 if response.success:
                     stats = response.data
@@ -1140,8 +1177,8 @@ class GradioInterface:
                 else:
                     return f'<div class="alert alert-error">❌ {response.error}</div>'
             
-            async def handle_analytics_refresh(api_key):
-                response = await self.service.get_user_stats(api_key)
+            async def handle_analytics_refresh(hf_token):
+                response = await self.service.get_user_stats(hf_token)
                 
                 if response.success:
                     return response.data["plot"]
@@ -1153,8 +1190,8 @@ class GradioInterface:
                         font=dict(color='white')
                     )
             
-            async def handle_admin_refresh(api_key):
-                response = await self.service.get_admin_stats(api_key)
+            async def handle_admin_refresh(hf_token):
+                response = await self.service.get_admin_stats(hf_token)
                 
                 if response.success:
                     stats = response.data
@@ -1227,39 +1264,39 @@ class GradioInterface:
             generate_btn.click(
                 fn=handle_generation,
                 inputs=[
-                    prompt_input, api_key_playground, max_tokens_input,
+                    prompt_input, hf_token_input, max_tokens_input,
                     temperature_input, top_k_input, top_p_input, repetition_penalty_input
                 ],
-                outputs=[output_text, generation_stats, generation_stats]
+                outputs=[output_text, generation_stats]
             )
             
             clear_btn.click(
-                fn=lambda: ("", "", False),
+                fn=lambda: ("", ""),
                 inputs=[],
-                outputs=[prompt_input, output_text, generation_stats]
+                outputs=[prompt_input, output_text]
             )
             
-            create_btn.click(
-                fn=handle_user_creation,
-                inputs=[name_input, email_input, plan_input],
-                outputs=[creation_status, api_key_display, api_key_display]
+            profile_login_btn.click(
+                fn=handle_profile_login,
+                inputs=[profile_hf_token],
+                outputs=[profile_status]
             )
             
             refresh_stats_btn.click(
                 fn=handle_stats_refresh,
-                inputs=[stats_api_key],
+                inputs=[profile_hf_token],
                 outputs=[user_stats_display]
             )
             
             refresh_analytics_btn.click(
                 fn=handle_analytics_refresh,
-                inputs=[analytics_api_key],
+                inputs=[analytics_hf_token],
                 outputs=[analytics_plot]
             )
             
             refresh_admin_btn.click(
                 fn=handle_admin_refresh,
-                inputs=[admin_api_key],
+                inputs=[admin_hf_token],
                 outputs=[admin_plot, admin_stats_display]
             )