Upload models/msp/api/api_server_secure.py with huggingface_hub

models/msp/api/api_server_secure.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""
+Secure FastAPI server for serving Mistral 7B fine-tuned models
+Includes API key authentication like commercial services
+"""
+
+import os
+import sys
+import secrets
+from typing import Optional
+from fastapi import FastAPI, HTTPException, Header, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+import uvicorn
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from inference.inference_mistral7b import load_local_model, generate_with_local_model, get_device_info
+import torch
+
+# Configuration - Resolve model path relative to msp root
+_MODEL_BASE = Path(__file__).parent.parent / "mistral7b-finetuned-ahb2apb"
+DEFAULT_MODEL_PATH = str(_MODEL_BASE)
+
+# API Key authentication
+API_KEYS = set()
+API_KEY_FILE = "api_keys.txt"
+
+# Load or generate API keys
+def load_api_keys():
+    """Load API keys from file or create default"""
+    global API_KEYS
+    
+    if os.path.exists(API_KEY_FILE):
+        with open(API_KEY_FILE, 'r') as f:
+            API_KEYS = {line.strip() for line in f if line.strip()}
+    else:
+        # Generate a default API key
+        default_key = secrets.token_urlsafe(32)
+        with open(API_KEY_FILE, 'w') as f:
+            f.write(default_key + '\n')
+        API_KEYS = {default_key}
+        print(f"\n🔑 Generated default API key: {default_key}")
+        print(f"   Save this key! Store it in: {API_KEY_FILE}")
+    
+    print(f"✓ Loaded {len(API_KEYS)} API key(s)")
+
+def verify_api_key(api_key: str = Header(None)):
+    """Verify API key in request header"""
+    if api_key is None:
+        raise HTTPException(
+            status_code=401,
+            detail="API key required. Add header: 'X-API-Key: your-api-key'"
+        )
+    if api_key not in API_KEYS:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+    return api_key
+
+# Global model and tokenizer
+model = None
+tokenizer = None
+device_info = None
+
+app = FastAPI(
+    title="Mistral 7B AHB2APB API (Secure)",
+    description="Secure API for serving the fine-tuned Mistral 7B model for AHB2APB conversion",
+    version="1.0.0"
+)
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, restrict this!
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Security scheme
+security = HTTPBearer(auto_error=False)
+
+# Request/Response models
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_length: Optional[int] = 512
+    temperature: Optional[float] = 0.7
+
+class GenerateResponse(BaseModel):
+    response: str
+    model: str
+    max_length: int
+    temperature: float
+
+class HealthResponse(BaseModel):
+    status: str
+    model_loaded: bool
+    device: str
+    model_path: str
+    authentication: str
+
+@app.on_event("startup")
+async def startup():
+    """Load model and API keys on startup"""
+    global model, tokenizer, device_info
+    load_api_keys()
+    
+    model_path = os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH)
+    
+    print(f"\nLoading model from: {model_path}")
+    print("=" * 70)
+    
+    try:
+        device_info = get_device_info()
+        model, tokenizer = load_local_model(model_path)
+        print(f"\n✓ Model loaded successfully on {device_info['device']}!")
+        print(f"✓ API server ready (authentication enabled)")
+        print("=" * 70)
+    except Exception as e:
+        print(f"\n✗ Error loading model: {e}")
+        sys.exit(1)
+
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    """Health check endpoint (no auth required)"""
+    return HealthResponse(
+        status="healthy" if model is not None else "error",
+        model_loaded=model is not None,
+        device=device_info["device"] if device_info else "unknown",
+        model_path=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
+        authentication="enabled"
+    )
+
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "name": "Mistral 7B AHB2APB API (Secure)",
+        "version": "1.0.0",
+        "status": "running",
+        "authentication": "API key required",
+        "model": os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
+        "endpoints": {
+            "health": "/health",
+            "generate": "/api/generate (requires API key)",
+            "docs": "/docs"
+        }
+    }
+
+@app.post("/api/generate", response_model=GenerateResponse)
+async def generate(request: GenerateRequest, api_key: str = Depends(verify_api_key)):
+    """Generate text from a prompt (requires API key)"""
+    if model is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        response = generate_with_local_model(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=request.prompt,
+            max_length=request.max_length or 512,
+            temperature=request.temperature or 0.7
+        )
+        
+        return GenerateResponse(
+            response=response,
+            model=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
+            max_length=request.max_length or 512,
+            temperature=request.temperature or 0.7
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Start Secure Mistral 7B API server")
+    parser.add_argument("--model-path", type=str, default=DEFAULT_MODEL_PATH)
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--reload", action="store_true")
+    
+    args = parser.parse_args()
+    os.environ["MODEL_PATH"] = args.model_path
+    
+    print(f"\n🔒 Starting Secure Mistral 7B AHB2APB API Server")
+    print(f"   Model: {args.model_path}")
+    print(f"   Host: {args.host}")
+    print(f"   Port: {args.port}\n")
+    
+    # Change to api directory for proper module resolution
+    import os
+    os.chdir(os.path.dirname(os.path.abspath(__file__)))
+    uvicorn.run("api_server_secure:app", host=args.host, port=args.port, reload=args.reload)
+