Update app.py

app.py

@@ -1,62 +1,340 @@
 import os
-import sys
+import uuid
+import httpx
+import torch
 import logging
-from fastapi import FastAPI
+import json
+import asyncio
+from typing import Dict, Optional
+from fastapi import FastAPI, Request, BackgroundTasks, HTTPException, Depends
 from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import uvicorn
+from contextlib import asynccontextmanager
 
-# Setup logging to see what's happening
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Configuration - NOW WORKING!
+MODEL_ID = "google/gemma-1.1-2b-it"
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+API_KEY = os.getenv("API_KEY", "default-key-123")
+MAX_TOKENS = int(os.getenv("MAX_TOKENS", "450"))
+DEVICE = os.getenv("DEVICE", "cpu")
+PORT = int(os.getenv("PORT", "7860"))
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+# Security
+security = HTTPBearer()
 
-# Debug: Print ALL environment variables
-logger.info("=" * 50)
-logger.info("STARTING APP - CHECKING ENVIRONMENT")
-logger.info("=" * 50)
+# Job storage
+jobs: Dict[str, dict] = {}
 
-for key in sorted(os.environ.keys()):
-    value = os.environ.get(key)
-    if any(term in key.lower() for term in ['token', 'key', 'secret', 'pass', 'auth']):
-        # Hide sensitive values but show they exist
-        logger.info(f"ENV {key}: {'SET' if value else 'NOT SET'} (hidden for security)")
-    else:
-        logger.info(f"ENV {key}: {value}")
+class AIGenerator:
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+        self.loaded = False
+        self.load_error = None
+        
+    def load_model(self):
+        """Load the AI model with authentication"""
+        if self.loaded:
+            return True
+            
+        logger.info(f"🚀 Loading model: {MODEL_ID}")
+        
+        if not HF_TOKEN:
+            logger.error("❌ HF_TOKEN is not set!")
+            self.load_error = "HF_TOKEN environment variable is not set"
+            return False
+            
+        try:
+            # Load tokenizer with authentication
+            logger.info("📥 Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN  # Key change: use 'token' parameter
+            )
+            
+            # Set padding token
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                
+            logger.info("✅ Tokenizer loaded")
+            
+            # Load model with authentication
+            logger.info("📥 Loading model...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                torch_dtype=torch.float32,
+                token=HF_TOKEN,  # Key change: use 'token' parameter
+                device_map=None
+            )
+            
+            # Move to device
+            self.model = self.model.to(DEVICE)
+            self.model.eval()
+            
+            self.loaded = True
+            logger.info("🎉 Model loaded successfully!")
+            return True
+            
+        except Exception as e:
+            self.load_error = str(e)
+            logger.error(f"❌ Model loading failed: {str(e)}")
+            return False
 
-# Explicitly check for HF_TOKEN
-HF_TOKEN = os.getenv("HF_TOKEN")
-logger.info(f"HF_TOKEN explicitly: {'SET' if HF_TOKEN else 'NOT SET'}")
+# Global generator instance
+generator = AIGenerator()
 
-# Check other common names
-for var_name in ["HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"]:
-    value = os.getenv(var_name)
-    if value:
-        logger.info(f"Found in {var_name}: {'SET' if value else 'NOT SET'}")
+async def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """Verify API key"""
+    if credentials.credentials != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    return True
 
-@app.get("/")
-async def root():
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan manager - preload model on startup"""
+    logger.info("🚀 Starting AI API Server...")
+    logger.info(f"📊 Config: Model={MODEL_ID}, Device={DEVICE}, MaxTokens={MAX_TOKENS}")
+    
+    # Try to preload model (non-blocking)
+    try:
+        generator.load_model()
+    except Exception as e:
+        logger.warning(f"Model preloading failed, will load on first request: {e}")
+    
+    yield
+
+app = FastAPI(lifespan=lifespan)
+
+def generate_text(prompt: str, max_tokens: int = None) -> str:
+    """Generate text based on prompt"""
+    try:
+        if not generator.loaded:
+            if not generator.load_model():
+                raise Exception(f"Model failed to load: {generator.load_error}")
+        
+        logger.info(f"📝 Generating text for prompt: '{prompt[:50]}...'")
+        
+        # Tokenize
+        inputs = generator.tokenizer(
+            prompt, 
+            return_tensors="pt", 
+            truncation=True,
+            max_length=512
+        )
+        
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+        
+        # Generate
+        with torch.no_grad():
+            outputs = generator.model.generate(
+                **inputs,
+                max_new_tokens=max_tokens or MAX_TOKENS,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.8,
+                pad_token_id=generator.tokenizer.pad_token_id,
+                repetition_penalty=1.1
+            )
+        
+        # Decode
+        generated_text = generator.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Remove prompt if included
+        if prompt in generated_text:
+            generated_text = generated_text.replace(prompt, "").strip()
+        
+        logger.info(f"✅ Generated {len(generated_text)} characters")
+        return generated_text
+        
+    except Exception as e:
+        logger.error(f"❌ Generation failed: {str(e)}")
+        raise
+
+@app.post("/api/generate-sync")
+async def generate_sync(
+    request: Request,
+    auth: bool = Depends(verify_api_key)
+):
+    """
+    Synchronous text generation
+    Body: {"prompt": "your text", "max_tokens": 100}
+    """
+    try:
+        data = await request.json()
+        
+        if not data.get("prompt"):
+            raise HTTPException(status_code=400, detail="Prompt is required")
+            
+        prompt = data["prompt"]
+        max_tokens = data.get("max_tokens")
+        
+        logger.info(f"📥 Sync request: '{prompt[:50]}...'")
+        
+        generated_text = generate_text(prompt, max_tokens)
+        
+        return JSONResponse({
+            "status": "success",
+            "result": generated_text,
+            "prompt": prompt,
+            "text_length": len(generated_text),
+            "model": MODEL_ID
+        })
+        
+    except Exception as e:
+        logger.error(f"❌ Sync generation error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/generate")
+async def generate_async(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    auth: bool = Depends(verify_api_key)
+):
+    """
+    Asynchronous text generation (for longer tasks)
+    Body: {"prompt": "your text", "max_tokens": 100, "callback_url": "optional"}
+    """
+    try:
+        data = await request.json()
+        job_id = str(uuid.uuid4())
+        
+        if not data.get("prompt"):
+            raise HTTPException(status_code=400, detail="Prompt is required")
+            
+        prompt = data["prompt"]
+        max_tokens = data.get("max_tokens")
+        callback_url = data.get("callback_url")
+        
+        logger.info(f"📥 Async request {job_id}")
+        
+        jobs[job_id] = {
+            "status": "processing",
+            "prompt": prompt
+        }
+        
+        # Process in background
+        background_tasks.add_task(
+            process_job_async,
+            job_id,
+            prompt,
+            max_tokens,
+            callback_url
+        )
+        
+        return JSONResponse({
+            "job_id": job_id,
+            "status": "queued",
+            "message": "Generation started",
+            "model": MODEL_ID
+        })
+        
+    except Exception as e:
+        logger.error(f"❌ Async request error: {str(e)}")
+        raise HTTPException(status_code=400, detail=str(e))
+
+async def process_job_async(job_id: str, prompt: str, max_tokens: int = None, callback_url: str = None):
+    """Background processing for async jobs"""
+    try:
+        logger.info(f"🔄 Processing async job {job_id}")
+        
+        generated_text = generate_text(prompt, max_tokens)
+        
+        jobs[job_id] = {
+            "status": "complete",
+            "result": generated_text,
+            "prompt": prompt,
+            "text_length": len(generated_text)
+        }
+        
+        logger.info(f"✅ Completed async job {job_id}")
+        
+        # Send callback if provided
+        if callback_url:
+            try:
+                async with httpx.AsyncClient(timeout=30.0) as client:
+                    await client.post(
+                        callback_url,
+                        json={
+                            "job_id": job_id,
+                            "status": "complete",
+                            "result": generated_text,
+                            "prompt": prompt
+                        }
+                    )
+            except Exception as e:
+                logger.error(f"❌ Callback failed: {e}")
+                
+    except Exception as e:
+        error_msg = str(e)
+        logger.error(f"❌ Async job {job_id} failed: {error_msg}")
+        jobs[job_id] = {
+            "status": "failed",
+            "error": error_msg,
+            "prompt": prompt
+        }
+
+@app.get("/api/status/{job_id}")
+async def get_status(job_id: str, auth: bool = Depends(verify_api_key)):
+    """Check job status"""
+    if job_id not in jobs:
+        raise HTTPException(status_code=404, detail="Job not found")
+    
+    return JSONResponse(jobs[job_id])
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
     return JSONResponse({
-        "message": "AI API Debug Version",
-        "status": "running",
-        "hf_token_set": bool(os.getenv("HF_TOKEN")),
-        "all_env_vars": {k: 'SET' if v else 'NOT SET' for k, v in os.environ.items() if 'TOKEN' in k or 'KEY' in k}
+        "status": "healthy",
+        "model_loaded": generator.loaded,
+        "model": MODEL_ID,
+        "device": DEVICE,
+        "max_tokens": MAX_TOKENS
     })
 
-@app.get("/test-env")
-async def test_env():
-    """Test if HF_TOKEN is available"""
-    hf_token = os.getenv("HF_TOKEN", "")
+@app.get("/model-info")
+async def model_info():
+    """Model information"""
     return JSONResponse({
-        "hf_token_exists": bool(hf_token),
-        "hf_token_length": len(hf_token) if hf_token else 0,
-        "hf_token_preview": f"{hf_token[:10]}..." if hf_token and len(hf_token) > 10 else hf_token,
-        "all_env_keys": list(os.environ.keys())
+        "model": MODEL_ID,
+        "loaded": generator.loaded,
+        "error": generator.load_error,
+        "device": DEVICE,
+        "requires_auth": True,
+        "token_available": bool(HF_TOKEN)
     })
 
-@app.get("/health")
-async def health():
-    return JSONResponse({"status": "healthy"})
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return JSONResponse({
+        "message": "🤖 AI Text Generation API",
+        "version": "1.0",
+        "model": MODEL_ID,
+        "status": "operational" if generator.loaded else "model_loading",
+        "endpoints": {
+            "generate_sync": "POST /api/generate-sync",
+            "generate_async": "POST /api/generate",
+            "check_status": "GET /api/status/{job_id}",
+            "health": "GET /health",
+            "model_info": "GET /model-info"
+        },
+        "usage": 'curl -X POST /api/generate-sync -H "Authorization: Bearer YOUR_KEY" -d \'{"prompt":"Hello"}\''
+    })
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=PORT,
+        log_level="info"
+    )