model_server.py

#!/usr/bin/env python3
"""
AI Coding Model Server
FastAPI server that hosts the 5B parameter coding model
"""

import torch
import spaces
import uvicorn
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
import logging
import os
import asyncio
import threading
from contextlib import asynccontextmanager

# Import model components
from models import CodeModel
from utils import format_code_response, validate_code_syntax

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global model instance
code_model = None
model_loading = False

class ChatMessage(BaseModel):
    """Chat message model."""
    message: str = Field(..., description="User's message")
    history: List[Dict[str, str]] = Field(default_factory=list, description="Chat history")
    language: str = Field(default="python", description="Target programming language")
    temperature: float = Field(default=0.7, ge=0.1, le=1.0, description="Generation temperature")

class ChatResponse(BaseModel):
    """Chat response model."""
    choices: List[Dict[str, Dict[str, str]]] = Field(..., description="Generated responses")
    history: List[Dict[str, str]] = Field(..., description="Updated chat history")
    usage: Optional[Dict[str, int]] = Field(None, description="Token usage information")

class HealthResponse(BaseModel):
    """Health check response."""
    status: str
    model_loaded: bool
    model_name: str
    device: str
    memory_usage: Optional[Dict[str, Any]] = None

class ModelInfoResponse(BaseModel):
    """Model information response."""
    model_name: str
    parameter_count: str
    max_length: int
    device: str
    is_loaded: bool
    vocab_size: int

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan management."""
    # Startup
    logger.info("Starting up AI Coding Model Server...")
    await load_model()
    
    yield
    
    # Shutdown
    logger.info("Shutting down server...")

async def load_model():
    """Load the model in background."""
    global code_model, model_loading
    
    if code_model is not None or model_loading:
        return
    
    model_loading = True
    logger.info("Loading coding model...")
    
    try:
        # Load model in thread to avoid blocking
        loop = asyncio.get_event_loop()
        code_model = await loop.run_in_executor(None, CodeModel)
        
        if code_model.is_loaded:
            logger.info(f"✅ Model loaded successfully: {code_model.model_name}")
        else:
            logger.error("❌ Failed to load model")
            
    except Exception as e:
        logger.error(f"❌ Error loading model: {e}")
        code_model = None
    finally:
        model_loading = False

def create_app() -> FastAPI:
    """Create and configure the FastAPI application."""
    
    # Create FastAPI app with lifespan management
    app = FastAPI(
        title="AI Coding Model Server",
        description="FastAPI server hosting a 5B parameter coding model",
        version="1.0.0",
        lifespan=lifespan
    )
    
    # Add CORS middleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],  # Configure appropriately for production
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    @app.get("/", response_model=Dict[str, str])
    async def root():
        """Root endpoint."""
        return {
            "message": "AI Coding Model Server",
            "version": "1.0.0",
            "status": "running" if code_model and code_model.is_loaded else "loading"
        }
    
    @app.get("/health", response_model=HealthResponse)
    async def health_check():
        """Health check endpoint."""
        if model_loading:
            return HealthResponse(
                status="loading",
                model_loaded=False,
                model_name="Loading...",
                device="unknown"
            )
        
        if not code_model or not code_model.is_loaded:
            raise HTTPException(status_code=503, detail="Model not loaded")
        
        # Get memory usage if available
        memory_info = None
        if torch.cuda.is_available():
            memory_info = {
                "allocated": torch.cuda.memory_allocated() / 1024**3,  # GB
                "cached": torch.cuda.memory_reserved() / 1024**3,     # GB
                "total": torch.cuda.get_device_properties(0).total_memory / 1024**3
            }
        
        return HealthResponse(
            status="healthy",
            model_loaded=True,
            model_name=code_model.model_name,
            device=code_model.device,
            memory_usage=memory_info
        )
    
    @app.get("/model/info", response_model=ModelInfoResponse)
    async def model_info():
        """Get detailed model information."""
        if not code_model:
            raise HTTPException(status_code=503, detail="Model not loaded")
        
        info = code_model.get_model_info()
        return ModelInfoResponse(**info)
    
    @app.post("/api/chat", response_model=ChatResponse)
    async def chat(request: ChatMessage):
        """Main chat endpoint."""
        if model_loading:
            raise HTTPException(status_code=503, detail="Model is still loading")
        
        if not code_model or not code_model.is_loaded:
            raise HTTPException(status_code=503, detail="Model not loaded")
        
        try:
            # Generate response using the model
            messages = request.history.copy()
            messages.append({"role": "user", "content": request.message})
            
            response_text = code_model.generate(
                messages=messages,
                temperature=request.temperature,
                max_new_tokens=2048,
                language=request.language
            )
            
            # Format the response
            formatted_response = format_code_response(response_text)
            
            # Update chat history
            new_history = request.history.copy()
            new_history.append({"role": "user", "content": request.message})
            new_history.append({"role": "assistant", "content": formatted_response})
            
            return ChatResponse(
                choices=[{"message": {"content": formatted_response}}],
                history=new_history
            )
            
        except Exception as e:
            logger.error(f"Chat error: {e}")
            raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")
    
    @app.post("/api/validate-code")
    async def validate_code(request: Dict[str, Any]):
        """Validate code syntax."""
        code = request.get("code", "")
        language = request.get("language", "python")
        
        if not code:
            raise HTTPException(status_code=400, detail="No code provided")
        
        validation_result = validate_code_syntax(code, language)
        return validation_result
    
    @app.get("/api/languages")
    async def get_supported_languages():
        """Get list of supported programming languages."""
        return {
            "languages": [
                "python", "javascript", "java", "cpp", "c", "go", "rust", 
                "typescript", "php", "ruby", "swift", "kotlin", "sql", 
                "html", "css", "bash", "powershell"
            ]
        }
    
    return app

def run_server(host: str = "0.0.0.0", port: int = 8000, reload: bool = False):
    """Run the FastAPI server."""
    app = create_app()
    
    console_info = f"""
🚀 AI Coding Model Server Starting...

📊 Server Info:
   • Host: {host}
   • Port: {port}
   • Model: Loading...
   • Device: {'CUDA' if torch.cuda.is_available() else 'CPU'}

🔗 Endpoints:
   • Health: http://{host}:{port}/health
   • Model Info: http://{host}:{port}/model/info
   • Chat: http://{host}:{port}/api/chat
   • API Docs: http://{host}:{port}/docs

💡 Usage:
   • Terminal client: python terminal_chatbot.py
   • API calls: POST to /api/chat with chat messages
   • Check status: GET /health

⚡ Server is ready! Press Ctrl+C to stop.
    """
    
    print(console_info)
    
    # Run server
    uvicorn.run(
        "model_server:create_app",
        host=host,
        port=port,
        reload=reload,
        log_level="info",
        access_log=True
    )

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="AI Coding Model Server")
    parser.add_argument("--host", default="0.0.0.0", help="Server host")
    parser.add_argument("--port", type=int, default=8000, help="Server port")
    parser.add_argument("--reload", action="store_true", help="Auto-reload on changes")
    
    args = parser.parse_args()
    
    run_server(
        host=args.host,
        port=args.port,
        reload=args.reload
    )