main.py

import logging
from contextlib import asynccontextmanager
from datetime import datetime
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from config import config
from services.text_service import text_service
from services.vision_service import vision_service
from routers import text_router, vision_router

# Logging Setup
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger("main")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifecycle manager"""
    logger.info("=" * 60)
    logger.info("🤖 STARTING SMOLLM2 MULTIMODAL API")
    logger.info("=" * 60)
    
    try:
        # Initialize text service
        logger.info("Initializing Text Service...")
        await text_service.initialize()
        
        # Initialize vision service
        logger.info("Initializing Vision Service...")
        await vision_service.initialize()
        
        logger.info("=" * 60)
        logger.info("✓ All services initialized successfully")
        logger.info("=" * 60)
        
    except Exception as e:
        logger.critical(f"Startup failed: {e}")
        raise
    
    yield
    
    # Cleanup
    logger.info("Shutting down services...")
    await text_service.cleanup()
    await vision_service.cleanup()
    logger.info("Shutdown complete")

# Create FastAPI application
app = FastAPI(
    title="SmolLM2 Multimodal API",
    version="3.0",
    description="Production-ready API for SmolLM2 text and vision models",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(text_router.router)
app.include_router(vision_router.router)

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "name": "SmolLM2 Multimodal API",
        "version": "3.0",
        "endpoints": {
            "text": "/v1/text/chat/completions",
            "vision": "/v1/vision/analyze",
            "health": "/health"
        },
        "docs": "/docs"
    }

@app.get("/health")
async def health_check():
    """Comprehensive health check"""
    return {
        "status": "healthy",
        "services": {
            "text": text_service.is_ready(),
            "vision": vision_service.is_ready()
        },
        "timestamp": datetime.utcnow().isoformat()
    }

@app.get("/ping")
async def ping():
    """Simple ping endpoint"""
    all_ready = text_service.is_ready() and vision_service.is_ready()
    
    if not all_ready:
        return JSONResponse(
            status_code=503,
            content={"status": "initializing", "ready": False}
        )
    
    return {"status": "pong", "ready": True}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "main:app",
        host=config.HOST,
        port=config.PORT,
        log_level="info"
    )