ylff/server.py

import logging
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Import existing routers
from ylff.routers import (
    training,
    validation,
    inference,
    profiling,
    jobs,
    models,
    ingest,
)

logger = logging.getLogger(__name__)

def create_app() -> FastAPI:
    """Create and configure the YLFF API application."""
    app = FastAPI(
        title="YLFF API",
        description="API for You Learn From Failure training pipeline",
        version="0.1.0",
    )

    # Configure CORS - Allow all for local development convenience
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],  # For Next.js dev server
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # Include routers
    app.include_router(training.router, prefix="/api/v1")
    app.include_router(jobs.router, prefix="/api/v1")
    app.include_router(models.router, prefix="/api/v1")
    app.include_router(ingest.router, prefix="/api/v1")
    
    # Validation router check
    try:
        app.include_router(validation.router, prefix="/api/v1")
    except AttributeError:
        logger.warning("Could not include validation router")

    try:
        app.include_router(inference.router, prefix="/api/v1/inference")
    except AttributeError:
        pass

    try:
        app.include_router(profiling.router, prefix="/api/v1/profiling")
    except AttributeError:
        pass

    @app.get("/health")
    async def health_check():
        return {"status": "ok", "service": "ylff"}

    from fastapi.responses import JSONResponse
    @app.exception_handler(Exception)
    async def global_exception_handler(request: Request, exc: Exception):
        import traceback
        logger.error(f"Global Error: {str(exc)}\n{traceback.format_exc()}")
        return JSONResponse(
            status_code=500,
            content={"detail": f"Internal Server Error: {str(exc)}"},
        )

    return app

app = create_app()

def start_server(host: str = "0.0.0.0", port: int = 8000):
    """Run the API server."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    logger.info(f"Starting YLFF API server at http://{host}:{port}")
    uvicorn.run(app, host=host, port=port, log_level="info", access_log=True)