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 )