Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,261 @@
+import os
+import secrets
+import hashlib
+from typing import Optional, Dict, Any
+from datetime import datetime, timedelta
+import logging
+
+from fastapi import FastAPI, HTTPException, Depends, Security, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import uvicorn
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="LLM AI Agent API",
+    description="Secure AI Agent API with Local LLM deployment",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
+
+# CORS middleware for cross-origin requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Configure this for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Security
+security = HTTPBearer()
+
+# Configuration
+class Config:
+    # API Keys - In production, use environment variables
+    API_KEYS = {
+        os.getenv("API_KEY_1", "your-secure-api-key-1"): "user1",
+        os.getenv("API_KEY_2", "your-secure-api-key-2"): "user2",
+        # Add more API keys as needed
+    }
+    
+    # Model configuration
+    MODEL_NAME = os.getenv("MODEL_NAME", "microsoft/DialoGPT-medium")  # Lightweight model for free tier
+    MAX_LENGTH = int(os.getenv("MAX_LENGTH", "512"))
+    TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
+    TOP_P = float(os.getenv("TOP_P", "0.9"))
+    
+    # Rate limiting (requests per minute per API key)
+    RATE_LIMIT = int(os.getenv("RATE_LIMIT", "10"))
+
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+text_generator = None
+
+# Request/Response models
+class ChatRequest(BaseModel):
+    message: str = Field(..., min_length=1, max_length=1000, description="Input message for the AI agent")
+    max_length: Optional[int] = Field(None, ge=10, le=2048, description="Maximum response length")
+    temperature: Optional[float] = Field(None, ge=0.1, le=2.0, description="Response creativity (0.1-2.0)")
+    system_prompt: Optional[str] = Field(None, max_length=500, description="Optional system prompt")
+
+class ChatResponse(BaseModel):
+    response: str
+    model_used: str
+    timestamp: str
+    tokens_used: int
+    processing_time: float
+
+class HealthResponse(BaseModel):
+    status: str
+    model_loaded: bool
+    timestamp: str
+    version: str
+
+# Rate limiting storage (in production, use Redis)
+request_counts: Dict[str, Dict[str, int]] = {}
+
+def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)) -> str:
+    """Verify API key authentication"""
+    api_key = credentials.credentials
+    
+    if api_key not in Config.API_KEYS:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    return Config.API_KEYS[api_key]
+
+def check_rate_limit(api_key: str) -> bool:
+    """Simple rate limiting implementation"""
+    current_minute = datetime.now().strftime("%Y-%m-%d-%H-%M")
+    
+    if api_key not in request_counts:
+        request_counts[api_key] = {}
+    
+    if current_minute not in request_counts[api_key]:
+        request_counts[api_key][current_minute] = 0
+    
+    if request_counts[api_key][current_minute] >= Config.RATE_LIMIT:
+        return False
+    
+    request_counts[api_key][current_minute] += 1
+    return True
+
+@app.on_event("startup")
+async def load_model():
+    """Load the LLM model on startup"""
+    global model, tokenizer, text_generator
+    
+    try:
+        logger.info(f"Loading model: {Config.MODEL_NAME}")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)
+        
+        # Add padding token if it doesn't exist
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Load model with optimizations for free tier
+        model = AutoModelForCausalLM.from_pretrained(
+            Config.MODEL_NAME,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            low_cpu_mem_usage=True
+        )
+        
+        # Create text generation pipeline
+        text_generator = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device=0 if torch.cuda.is_available() else -1
+        )
+        
+        logger.info("Model loaded successfully!")
+        
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise e
+
+@app.get("/", response_model=HealthResponse)
+async def root():
+    """Health check endpoint"""
+    return HealthResponse(
+        status="healthy",
+        model_loaded=model is not None,
+        timestamp=datetime.now().isoformat(),
+        version="1.0.0"
+    )
+
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    """Detailed health check"""
+    return HealthResponse(
+        status="healthy" if model is not None else "model_not_loaded",
+        model_loaded=model is not None,
+        timestamp=datetime.now().isoformat(),
+        version="1.0.0"
+    )
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat(
+    request: ChatRequest,
+    user: str = Depends(verify_api_key)
+):
+    """Main chat endpoint for AI agent interaction"""
+    start_time = datetime.now()
+    
+    # Check rate limiting
+    api_key = None  # In a real implementation, you'd extract this from the token
+    # if not check_rate_limit(api_key):
+    #     raise HTTPException(
+    #         status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+    #         detail="Rate limit exceeded. Please try again later."
+    #     )
+    
+    if model is None or tokenizer is None:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Model not loaded. Please try again later."
+        )
+    
+    try:
+        # Prepare input
+        input_text = request.message
+        if request.system_prompt:
+            input_text = f"System: {request.system_prompt}\nUser: {request.message}\nAssistant:"
+        
+        # Generate response
+        max_length = request.max_length or Config.MAX_LENGTH
+        temperature = request.temperature or Config.TEMPERATURE
+        
+        # Generate text
+        generated = text_generator(
+            input_text,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=Config.TOP_P,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1,
+            truncation=True
+        )
+        
+        # Extract response
+        response_text = generated[0]['generated_text']
+        if input_text in response_text:
+            response_text = response_text.replace(input_text, "").strip()
+        
+        # Calculate processing time
+        processing_time = (datetime.now() - start_time).total_seconds()
+        
+        # Count tokens (approximate)
+        tokens_used = len(tokenizer.encode(response_text))
+        
+        return ChatResponse(
+            response=response_text,
+            model_used=Config.MODEL_NAME,
+            timestamp=datetime.now().isoformat(),
+            tokens_used=tokens_used,
+            processing_time=processing_time
+        )
+        
+    except Exception as e:
+        logger.error(f"Error generating response: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Error generating response: {str(e)}"
+        )
+
+@app.get("/models")
+async def get_model_info(user: str = Depends(verify_api_key)):
+    """Get information about the loaded model"""
+    return {
+        "model_name": Config.MODEL_NAME,
+        "model_loaded": model is not None,
+        "max_length": Config.MAX_LENGTH,
+        "temperature": Config.TEMPERATURE,
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }
+
+if __name__ == "__main__":
+    # For local development
+    uvicorn.run(
+        "app:app",
+        host="0.0.0.0",
+        port=int(os.getenv("PORT", "7860")),  # Hugging Face Spaces uses port 7860
+        reload=False
+    )