Update app.py

app.py

@@ -1,23 +1,20 @@
+# app.py - Stanley AI Backend for Hugging Face
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import time
-import re
 import logging
-from threading import Thread
-import queue
 import io
 import base64
 import random
+import sqlite3
+import secrets
+from datetime import datetime
 from PIL import Image, ImageDraw, ImageFont
-import os
 import json
-import hashlib
-import secrets
-from datetime import datetime, timedelta
-import sqlite3
 from functools import wraps
+import os
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -27,47 +24,102 @@ app = Flask(__name__)
 CORS(app)
 
 # ============================================================================
-# DATABASE SETUP (Lightweight for Hugging Face)
+# DATABASE SETUP
 # ============================================================================
 
 def init_db():
     """Initialize SQLite database for API keys"""
-    conn = sqlite3.connect('/tmp/api_keys.db')
-    c = conn.cursor()
-    
-    # Create API keys table
-    c.execute('''CREATE TABLE IF NOT EXISTS api_keys
-                 (id TEXT PRIMARY KEY,
-                  name TEXT,
-                  api_key TEXT UNIQUE,
-                  created TIMESTAMP,
-                  last_used TIMESTAMP,
-                  requests INTEGER DEFAULT 0,
-                  credits INTEGER DEFAULT 1000,
-                  is_active BOOLEAN DEFAULT 1,
-                  user_id TEXT,
-                  tier TEXT DEFAULT 'free')''')
+    try:
+        conn = sqlite3.connect('/tmp/api_keys.db')
+        c = conn.cursor()
+        
+        # Create API keys table
+        c.execute('''CREATE TABLE IF NOT EXISTS api_keys
+                     (id TEXT PRIMARY KEY,
+                      name TEXT,
+                      api_key TEXT UNIQUE,
+                      created TIMESTAMP,
+                      last_used TIMESTAMP,
+                      requests INTEGER DEFAULT 0,
+                      credits INTEGER DEFAULT 1000,
+                      is_active BOOLEAN DEFAULT 1,
+                      user_id TEXT,
+                      tier TEXT DEFAULT 'free')''')
+        
+        # Create invite codes table
+        c.execute('''CREATE TABLE IF NOT EXISTS invite_codes
+                     (code TEXT PRIMARY KEY,
+                      created TIMESTAMP,
+                      expires TIMESTAMP,
+                      used BOOLEAN DEFAULT 0,
+                      used_by TEXT)''')
+        
+        # Create usage stats table
+        c.execute('''CREATE TABLE IF NOT EXISTS usage_stats
+                     (date TEXT,
+                      api_key TEXT,
+                      endpoint TEXT,
+                      credits_used INTEGER)''')
+        
+        conn.commit()
+        conn.close()
+        logger.info("✅ Database initialized successfully")
+    except Exception as e:
+        logger.error(f"❌ Database initialization error: {e}")
+
+init_db()
+
+# ============================================================================
+# MODEL LOADING (Optimized for Hugging Face)
+# ============================================================================
+
+model = None
+tokenizer = None
+model_loaded = False
+
+# System Prompt for Stanley AI
+STANLEY_AI_SYSTEM = """You are STANLEY AI - an advanced artificial intelligence assistant created by Stanley Samwel Owino.
+You are knowledgeable, helpful, and occasionally use Kiswahili phrases when culturally appropriate.
+Provide comprehensive, thoughtful responses that demonstrate deep understanding."""
+
+def load_model():
+    """Load the AI model with optimizations for Hugging Face"""
+    global model, tokenizer, model_loaded
     
-    # Create invite codes table
-    c.execute('''CREATE TABLE IF NOT EXISTS invite_codes
-                 (code TEXT PRIMARY KEY,
-                  created TIMESTAMP,
-                  expires TIMESTAMP,
-                  used BOOLEAN DEFAULT 0,
-                  used_by TEXT)''')
+    if model_loaded:
+        return
     
-    # Create usage stats table
-    c.execute('''CREATE TABLE IF NOT EXISTS usage_stats
-                 (date TEXT,
-                  api_key TEXT,
-                  endpoint TEXT,
-                  credits_used INTEGER)''')
+    logger.info("🚀 Loading STANLEY AI Model...")
     
-    conn.commit()
-    conn.close()
-    logger.info("✅ Database initialized")
+    try:
+        # Use a smaller model for Hugging Face compatibility
+        model_name = "Qwen/Qwen2.5-0.5B-Instruct"
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True
+        )
+        
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        model_loaded = True
+        logger.info("✅ Model loaded successfully!")
+        
+    except Exception as e:
+        logger.error(f"❌ Error loading model: {e}")
+        model_loaded = False
 
-init_db()
+# Load model on startup
+load_model()
 
 # ============================================================================
 # API KEY MANAGEMENT FUNCTIONS
@@ -76,19 +128,23 @@ init_db()
 def authenticate_api_key():
     """Authenticate API key from request header"""
     auth_header = request.headers.get('Authorization')
-    if not auth_header:
+    if not auth_header or not auth_header.startswith('Bearer '):
         return None
     
-    # Extract token (Bearer <token>)
-    token = auth_header.replace('Bearer ', '').strip()
-    
-    conn = sqlite3.connect('/tmp/api_keys.db')
-    c = conn.cursor()
-    c.execute('''SELECT * FROM api_keys WHERE api_key = ? AND is_active = 1''', (token,))
-    key_data = c.fetchone()
-    conn.close()
+    api_key = auth_header.replace('Bearer ', '').strip()
     
-    return key_data
+    try:
+        conn = sqlite3.connect('/tmp/api_keys.db')
+        c = conn.cursor()
+        c.execute('''SELECT * FROM api_keys 
+                     WHERE api_key = ? AND is_active = 1''', (api_key,))
+        key_data = c.fetchone()
+        conn.close()
+        
+        return key_data
+    except Exception as e:
+        logger.error(f"Authentication error: {e}")
+        return None
 
 def check_credits(api_key_data):
     """Check if API key has sufficient credits"""
@@ -96,40 +152,45 @@ def check_credits(api_key_data):
 
 def update_usage(api_key, endpoint, credits_used=1):
     """Update usage statistics"""
-    conn = sqlite3.connect('/tmp/api_keys.db')
-    c = conn.cursor()
-    
-    # Update key usage
-    c.execute('''UPDATE api_keys 
-                 SET requests = requests + 1, 
-                     credits = credits - ?, 
-                     last_used = CURRENT_TIMESTAMP 
-                 WHERE api_key = ?''', 
-              (credits_used, api_key))
-    
-    # Record in stats
-    today = datetime.now().strftime('%Y-%m-%d')
-    c.execute('''INSERT INTO usage_stats (date, api_key, endpoint, credits_used)
-                 VALUES (?, ?, ?, ?)''',
-              (today, api_key, endpoint, credits_used))
-    
-    conn.commit()
-    conn.close()
+    try:
+        conn = sqlite3.connect('/tmp/api_keys.db')
+        c = conn.cursor()
+        
+        # Update key usage
+        c.execute('''UPDATE api_keys 
+                     SET requests = requests + 1, 
+                         credits = credits - ?, 
+                         last_used = CURRENT_TIMESTAMP 
+                     WHERE api_key = ?''', 
+                  (credits_used, api_key))
+        
+        # Record in stats
+        today = datetime.now().strftime('%Y-%m-%d')
+        c.execute('''INSERT INTO usage_stats (date, api_key, endpoint, credits_used)
+                     VALUES (?, ?, ?, ?)''',
+                  (today, api_key, endpoint, credits_used))
+        
+        conn.commit()
+        conn.close()
+        
+    except Exception as e:
+        logger.error(f"Usage update error: {e}")
 
 def generate_api_key(user_id="anonymous", name="My API Key", tier="free"):
     """Generate a new API key"""
-    # Create secure key
-    key_id = f"key_{secrets.token_hex(8)}"
-    api_key = f"stanley_{secrets.token_urlsafe(32)}"
-    
-    conn = sqlite3.connect('/tmp/api_keys.db')
-    c = conn.cursor()
-    
     try:
+        # Create secure key
+        key_id = f"key_{secrets.token_hex(8)}"
+        api_key = f"stanley_{secrets.token_urlsafe(32)}"
+        
+        conn = sqlite3.connect('/tmp/api_keys.db')
+        c = conn.cursor()
+        
         c.execute('''INSERT INTO api_keys 
                      (id, name, api_key, created, last_used, requests, credits, is_active, user_id, tier)
                      VALUES (?, ?, ?, CURRENT_TIMESTAMP, NULL, 0, 1000, 1, ?, ?)''',
                   (key_id, name, api_key, user_id, tier))
+        
         conn.commit()
         
         # Get the created key
@@ -148,8 +209,8 @@ def generate_api_key(user_id="anonymous", name="My API Key", tier="free"):
             'tier': key_data[9]
         }
         
-    except sqlite3.IntegrityError:
-        conn.close()
+    except Exception as e:
+        logger.error(f"Key generation error: {e}")
         return None
 
 # ============================================================================
@@ -161,11 +222,11 @@ def require_api_key(credits_cost=1):
     def decorator(f):
         @wraps(f)
         def decorated_function(*args, **kwargs):
-            # Skip authentication for OPTIONS requests (CORS preflight)
+            # Skip for OPTIONS (CORS preflight)
             if request.method == 'OPTIONS':
                 return f(*args, **kwargs)
             
-            # Check for API key
+            # Authenticate
             key_data = authenticate_api_key()
             
             if not key_data:
@@ -201,152 +262,198 @@ def require_api_key(credits_cost=1):
     return decorator
 
 # ============================================================================
-# MODEL LOADING
-# ============================================================================
-
-model = None
-tokenizer = None
-model_loaded = False
-
-# Simplified System Prompt for Hugging Face
-STANLEY_AI_SYSTEM = """You are STANLEY AI - an advanced large language model created by Stanley Samwel Owino. You provide comprehensive, helpful responses with occasional Kiswahili phrases when culturally relevant."""
-
-def load_model():
-    global model, tokenizer, model_loaded
-    if model_loaded:
-        return
-    
-    logger.info("🚀 Loading STANLEY AI Model...")
-    model_name = "Qwen/Qwen2.5-0.5B-Instruct"  # Smaller model for Hugging Face
-    
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
-        )
-        
-        model_loaded = True
-        logger.info("✅ STANLEY AI Model loaded successfully!")
-        
-    except Exception as e:
-        logger.error(f"❌ Error loading model: {e}")
-        model_loaded = False
-
-load_model()
-
-# ============================================================================
-# RESPONSE GENERATION
+# AI RESPONSE GENERATION
 # ============================================================================
 
-response_cache = {}
-CACHE_SIZE = 100
-
-def generate_response(user_message):
-    """Generate response using the model"""
+def generate_response(user_message, max_tokens=512):
+    """Generate AI response using the model"""
     if not model_loaded:
         return "Model is loading. Please try again in a moment."
     
     try:
+        # Prepare messages for chat template
         messages = [
             {"role": "system", "content": STANLEY_AI_SYSTEM},
             {"role": "user", "content": user_message}
         ]
         
-        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tokenizer(text, return_tensors="pt").to(model.device)
+        # Apply chat template
+        text = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+        
+        inputs = tokenizer(text, return_tensors="pt")
         
+        if torch.cuda.is_available():
+            inputs = inputs.to('cuda')
+        
+        # Generate response
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=512,
+                max_new_tokens=max_tokens,
                 temperature=0.7,
                 do_sample=True,
                 top_p=0.9,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
         
-        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        # Decode response
+        response = tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:], 
+            skip_special_tokens=True
+        )
+        
         return response.strip()
         
     except Exception as e:
         logger.error(f"Generation error: {e}")
-        return f"Error generating response: {str(e)}"
+        return f"I encountered an error: {str(e)}"
 
 # ============================================================================
-# IMAGE GENERATION (SIMPLIFIED FOR HUGGING FACE)
+# IMAGE GENERATION
 # ============================================================================
 
 def generate_fallback_image(prompt, width=512, height=512):
     """Generate simple images using PIL for Hugging Face"""
     try:
-        # Create base image
-        img = Image.new('RGB', (width, height), color=(random.randint(50, 200), 
-                                                       random.randint(50, 200), 
-                                                       random.randint(50, 200)))
+        # Create base image with gradient background
+        img = Image.new('RGB', (width, height), color=(0, 0, 0))
         draw = ImageDraw.Draw(img)
         
-        # Add decorative elements based on prompt
+        # Add gradient effect
+        for i in range(height):
+            r = int(50 + (i / height) * 100)
+            g = int(50 + (i / height) * 150)
+            b = int(100 + (i / height) * 100)
+            draw.line([(0, i), (width, i)], fill=(r, g, b))
+        
+        # Add decorative elements based on prompt keywords
         if any(word in prompt.lower() for word in ['sun', 'light', 'bright']):
-            draw.ellipse([width//4, height//4, 3*width//4, 3*height//4], 
-                        fill=(255, 255, 0))
-        elif any(word in prompt.lower() for word in ['moon', 'night', 'dark']):
-            draw.ellipse([width//3, height//3, 2*width//3, 2*height//3], 
-                        fill=(200, 200, 200))
+            # Draw a sun
+            draw.ellipse(
+                [width//2 - 50, height//2 - 50, width//2 + 50, height//2 + 50],
+                fill=(255, 255, 0),
+                outline=(255, 200, 0),
+                width=3
+            )
+        
+        elif any(word in prompt.lower() for word in ['moon', 'night', 'star']):
+            # Draw a moon and stars
+            draw.ellipse(
+                [width//2 - 40, height//2 - 40, width//2 + 40, height//2 + 40],
+                fill=(200, 200, 200)
+            )
+            
+            # Draw some stars
+            for _ in range(10):
+                x = random.randint(50, width - 50)
+                y = random.randint(50, height - 200)
+                size = random.randint(2, 6)
+                draw.rectangle(
+                    [x, y, x + size, y + size],
+                    fill=(255, 255, 255)
+                )
+        
+        elif any(word in prompt.lower() for word in ['tree', 'nature', 'forest']):
+            # Draw a simple tree
+            # Trunk
+            draw.rectangle(
+                [width//2 - 15, height//2, width//2 + 15, height - 50],
+                fill=(139, 69, 19)
+            )
+            # Leaves
+            draw.ellipse(
+                [width//2 - 60, height//2 - 100, width//2 + 60, height//2 + 20],
+                fill=(34, 139, 34)
+            )
         
-        # Add text label
+        # Add Stanley AI branding
         try:
-            # Try to use a system font
-            font = ImageFont.load_default()
-            text = f"AI: {prompt[:30]}..." if len(prompt) > 30 else f"AI: {prompt}"
+            # Try to load a font
+            font_size = min(width // 25, 20)
+            try:
+                font = ImageFont.truetype("arial.ttf", font_size)
+            except:
+                # Fallback to default font
+                font = ImageFont.load_default()
             
-            # Draw text background
-            text_bbox = draw.textbbox((0, 0), text, font=font)
+            # Add prompt text
+            display_text = f"AI Generated: {prompt[:40]}..." if len(prompt) > 40 else f"AI: {prompt}"
+            text_bbox = draw.textbbox((0, 0), display_text, font=font)
             text_width = text_bbox[2] - text_bbox[0]
-            text_height = text_bbox[3] - text_bbox[1]
             
             x = (width - text_width) // 2
-            y = height - text_height - 20
+            y = height - 40
             
-            draw.rectangle([x-5, y-5, x+text_width+5, y+text_height+5], 
-                          fill=(0, 0, 0, 128))
-            draw.text((x, y), text, fill=(255, 255, 255), font=font)
+            # Text background
+            draw.rectangle(
+                [x - 10, y - 10, x + text_width + 10, y + 30],
+                fill=(0, 0, 0, 150)
+            )
+            
+            # Text
+            draw.text(
+                (x, y),
+                display_text,
+                fill=(255, 255, 255),
+                font=font
+            )
+            
+        except Exception as font_error:
+            logger.warning(f"Font error: {font_error}")
+        
+        # Add Stanley AI watermark
+        watermark = "Stanley AI"
+        wm_font_size = min(width // 40, 16)
+        try:
+            wm_font = ImageFont.truetype("arial.ttf", wm_font_size)
         except:
-            pass
+            wm_font = ImageFont.load_default()
+        
+        wm_bbox = draw.textbbox((0, 0), watermark, font=wm_font)
+        wm_width = wm_bbox[2] - wm_bbox[0]
+        
+        draw.text(
+            (width - wm_width - 10, 10),
+            watermark,
+            fill=(255, 255, 255, 180),
+            font=wm_font
+        )
         
         # Convert to base64
         buffered = io.BytesIO()
-        img.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode()
+        img.save(buffered, format="PNG", optimize=True)
+        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        
         return f"data:image/png;base64,{img_str}"
         
     except Exception as e:
         logger.error(f"Image generation error: {e}")
-        return None
+        # Ultimate fallback - solid color with text
+        try:
+            img = Image.new('RGB', (width, height), color=(random.randint(50, 200), random.randint(50, 200), random.randint(50, 200)))
+            draw = ImageDraw.Draw(img)
+            draw.text((50, height//2), f"Prompt: {prompt[:30]}", fill=(255, 255, 255))
+            
+            buffered = io.BytesIO()
+            img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            return f"data:image/png;base64,{img_str}"
+        except:
+            return None
 
 # ============================================================================
-# FLASK ROUTES
+# FLASK ROUTES - PUBLIC ENDPOINTS
 # ============================================================================
 
 @app.route('/')
 def home():
-    return jsonify({
-        "message": "🚀 STANLEY AI API is running!", 
-        "version": "3.0",
-        "status": "active",
-        "features": ["Chat API", "Image Generation", "API Key System"],
-        "documentation": "/api/docs"
-    })
-
-# ============================================================================
-# PUBLIC ENDPOINTS (No API key required)
-# ============================================================================
+    """Home page - serves the frontend"""
+    return send_file('index.html')
 
 @app.route('/api/status')
 def api_status():
@@ -354,7 +461,9 @@ def api_status():
     return jsonify({
         "status": "online",
         "model_loaded": model_loaded,
-        "timestamp": datetime.now().isoformat()
+        "timestamp": datetime.now().isoformat(),
+        "version": "3.0",
+        "message": "🚀 Stanley AI API is running!"
     })
 
 @app.route('/api/docs')
@@ -362,13 +471,19 @@ def api_docs():
     """API documentation"""
     return jsonify({
         "endpoints": {
-            "POST /api/chat": "Chat with Stanley AI (requires API key)",
-            "POST /api/generate-image": "Generate images (requires API key, 5 credits)",
-            "POST /api/v1/generate-key": "Generate API key",
-            "GET /api/v1/keys": "List your API keys (requires API key)",
+            "GET /api/status": "Check API status (public)",
+            "POST /api/v1/generate-key": "Generate API key (public)",
+            "POST /api/chat": "Chat with AI (requires API key, 1 credit)",
+            "POST /api/generate-image": "Generate image (requires API key, 5 credits)",
+            "GET /api/v1/keys": "List API keys (requires API key)",
             "DELETE /api/v1/keys/<key_id>": "Revoke API key (requires API key)",
-            "GET /api/v1/usage": "Get usage stats (requires API key)",
-            "GET /api/status": "Check API status (public)"
+            "GET /api/v1/usage": "Get usage stats (requires API key)"
+        },
+        "authentication": "Use Bearer token: Authorization: Bearer YOUR_API_KEY",
+        "credits": {
+            "free_tier": "1000 credits on signup",
+            "chat": "1 credit per request",
+            "image": "5 credits per image"
         }
     })
 
@@ -382,7 +497,7 @@ def create_api_key():
     try:
         data = request.get_json() or {}
         name = data.get('name', 'My API Key')
-        user_id = data.get('user_id', 'anonymous')
+        user_id = data.get('user_id', f"user_{secrets.token_hex(8)}")
         
         key_data = generate_api_key(user_id, name)
         
@@ -400,6 +515,7 @@ def create_api_key():
             }), 500
             
     except Exception as e:
+        logger.error(f"Create key error: {e}")
         return jsonify({
             "error": f"Error generating key: {str(e)}",
             "status": "error"
@@ -408,7 +524,7 @@ def create_api_key():
 @app.route('/api/v1/keys', methods=['GET'])
 @require_api_key(credits_cost=0)
 def list_api_keys():
-    """List API keys for the authenticated user"""
+    """List API keys for authenticated user"""
     try:
         api_key = request.api_key_data['api_key']
         
@@ -416,7 +532,7 @@ def list_api_keys():
         c = conn.cursor()
         
         # Get user ID from current key
-        c.execute('''SELECT user_id FROM api_keys WHERE api_key = ?''', (api_key,))
+        c.execute('SELECT user_id FROM api_keys WHERE api_key = ?', (api_key,))
         user_data = c.fetchone()
         
         if not user_data:
@@ -451,6 +567,7 @@ def list_api_keys():
         })
         
     except Exception as e:
+        logger.error(f"List keys error: {e}")
         return jsonify({
             "error": f"Error listing keys: {str(e)}",
             "status": "error"
@@ -461,12 +578,13 @@ def list_api_keys():
 def revoke_api_key(key_id):
     """Revoke an API key"""
     try:
+        api_key = request.api_key_data['api_key']
+        
         conn = sqlite3.connect('/tmp/api_keys.db')
         c = conn.cursor()
         
         # Check if key belongs to user
-        c.execute('''SELECT user_id FROM api_keys WHERE api_key = ?''', 
-                 (request.api_key_data['api_key'],))
+        c.execute('SELECT user_id FROM api_keys WHERE api_key = ?', (api_key,))
         user_data = c.fetchone()
         
         if not user_data:
@@ -492,6 +610,7 @@ def revoke_api_key(key_id):
         })
         
     except Exception as e:
+        logger.error(f"Revoke key error: {e}")
         return jsonify({
             "error": f"Error revoking key: {str(e)}",
             "status": "error"
@@ -528,7 +647,7 @@ def get_usage():
         for row in c.fetchall():
             daily_usage.append({
                 'date': row[0],
-                'credits_used': row[1]
+                'credits_used': row[1] or 0
             })
         
         conn.close()
@@ -544,13 +663,14 @@ def get_usage():
         })
         
     except Exception as e:
+        logger.error(f"Usage stats error: {e}")
         return jsonify({
             "error": f"Error getting usage: {str(e)}",
             "status": "error"
         }), 500
 
 # ============================================================================
-# CHAT ENDPOINT
+# AI CHAT ENDPOINT
 # ============================================================================
 
 @app.route('/api/chat', methods=['POST'])
@@ -558,20 +678,27 @@ def get_usage():
 def chat():
     """Chat with Stanley AI"""
     try:
+        start_time = time.time()
         data = request.get_json()
-        user_message = data.get('message', '')
+        user_message = data.get('message', '').strip()
         
         if not user_message:
             return jsonify({"error": "Message is required"}), 400
         
         logger.info(f"Chat request: {user_message[:50]}...")
         
+        # Generate response
         response = generate_response(user_message)
         
+        # Calculate response time
+        response_time = round(time.time() - start_time, 2)
+        
         return jsonify({
             "response": response,
             "status": "success",
-            "credits_remaining": request.api_key_data['credits']
+            "credits_remaining": request.api_key_data['credits'],
+            "response_time": response_time,
+            "model": "Stanley AI"
         })
         
     except Exception as e:
@@ -586,26 +713,36 @@ def chat():
 # ============================================================================
 
 @app.route('/api/generate-image', methods=['POST'])
-@require_api_key(credits_cost=5)  # Image generation costs 5 credits
+@require_api_key(credits_cost=5)
 def generate_image():
     """Generate image from text prompt"""
     try:
+        start_time = time.time()
         data = request.get_json()
-        prompt = data.get('prompt', '')
+        prompt = data.get('prompt', '').strip()
         
         if not prompt:
             return jsonify({"error": "Prompt is required"}), 400
         
-        logger.info(f"Image generation request: {prompt[:50]}...")
+        logger.info(f"Image generation: {prompt[:50]}...")
+        
+        # Get optional parameters
+        width = min(data.get('width', 512), 1024)
+        height = min(data.get('height', 512), 1024)
         
-        image_data = generate_fallback_image(prompt)
+        # Generate image
+        image_data = generate_fallback_image(prompt, width, height)
         
         if image_data:
+            generation_time = round(time.time() - start_time, 2)
+            
             return jsonify({
                 "image": image_data,
                 "prompt": prompt,
                 "status": "success",
                 "credits_remaining": request.api_key_data['credits'],
+                "generation_time": generation_time,
+                "dimensions": f"{width}x{height}",
                 "quality": "basic"
             })
         else:
@@ -622,7 +759,7 @@ def generate_image():
         }), 500
 
 # ============================================================================
-# TEST ENDPOINT (No authentication required for testing)
+# TEST ENDPOINTS (No authentication)
 # ============================================================================
 
 @app.route('/api/test', methods=['GET'])
@@ -631,32 +768,51 @@ def test_endpoint():
     return jsonify({
         "message": "Stanley AI API is working!",
         "timestamp": datetime.now().isoformat(),
-        "version": "3.0"
+        "version": "3.0",
+        "status": "success"
     })
 
 @app.route('/api/test/chat', methods=['POST'])
 def test_chat():
     """Test chat without API key (limited)"""
     try:
-        data = request.get_json()
+        data = request.get_json() or {}
         message = data.get('message', 'Hello')
         
-        # Simple response without model
+        # Simple predefined responses for testing
         responses = [
             "Hello! I'm Stanley AI, your intelligent assistant.",
-            "I'm here to help you with your questions.",
-            "Welcome to the Stanley AI platform!",
-            "How can I assist you today?"
+            "Welcome to the Stanley AI platform! How can I help you today?",
+            "I'm here to assist you with any questions you might have.",
+            "As Stanley AI, I'm designed to provide helpful and comprehensive responses.",
+            "Feel free to ask me anything. I'm here to help!"
         ]
         
+        # Check for specific keywords
+        message_lower = message.lower()
+        
+        if 'stanley' in message_lower:
+            response = "Yes, I'm Stanley AI! Created by Stanley Samwel Owino to provide advanced AI assistance."
+        elif 'api' in message_lower or 'key' in message_lower:
+            response = "You can generate an API key in the API portal. Visit /api/v1/generate-key to get started."
+        elif 'image' in message_lower:
+            response = "I can generate images! Use the /api/generate-image endpoint with your API key."
+        elif 'swahili' in message_lower or 'kiswahili' in message_lower:
+            response = "Jambo! Yes, I have knowledge of Kiswahili language and culture. Karibu sana!"
+        else:
+            response = random.choice(responses)
+        
         return jsonify({
-            "response": random.choice(responses),
+            "response": response,
             "status": "success",
-            "note": "This is a test response. For full features, use API key."
+            "note": "This is a test response. For full features, generate an API key."
         })
         
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        return jsonify({
+            "error": str(e),
+            "status": "error"
+        }), 500
 
 # ============================================================================
 # ERROR HANDLERS
@@ -670,8 +826,17 @@ def not_found(error):
         "code": 404
     }), 404
 
+@app.errorhandler(405)
+def method_not_allowed(error):
+    return jsonify({
+        "error": "Method not allowed",
+        "status": "error",
+        "code": 405
+    }), 405
+
 @app.errorhandler(500)
 def internal_error(error):
+    logger.error(f"Internal server error: {error}")
     return jsonify({
         "error": "Internal server error",
         "status": "error",
@@ -679,7 +844,7 @@ def internal_error(error):
     }), 500
 
 # ============================================================================
-# MAIN
+# MAIN APPLICATION
 # ============================================================================
 
 if __name__ == '__main__':
@@ -687,9 +852,15 @@ if __name__ == '__main__':
     print("🚀 STANLEY AI API v3.0")
     print("🌍 Hugging Face Space Ready")
     print("🔑 API Key System: Enabled")
-    print("💬 Chat: Requires API key (1 credit per request)")
-    print("🎨 Image Generation: Requires API key (5 credits per image)")
-    print("📊 Database: SQLite (persistent)")
+    print("💬 Chat Endpoint: /api/chat (1 credit)")
+    print("🎨 Image Generation: /api/generate-image (5 credits)")
+    print("📊 Database: SQLite (persistent storage)")
     print("=" * 60)
     
-    app.run(debug=True, host='0.0.0.0', port=7860)
+    # Run the Flask app
+    app.run(
+        debug=True,
+        host='0.0.0.0',
+        port=7860,
+        threaded=True
+    )