nygaardcodecommentclassification/api/API.py

"""Main script: REST API initialization and endpoints.

This module defines the FastAPI application for:
- Retrieving available ML models and supported programming languages
- Classifying code comments using multi-label classification

The API is designed to be client-agnostic and supports concurrent requests
through asynchronous endpoint handlers and background thread execution
for CPU-bound ML inference tasks.
"""

import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from datetime import datetime
from functools import wraps
from http import HTTPStatus
import json
from typing import Any, Callable, Dict

from fastapi import FastAPI, Request, Response
from fastapi.responses import RedirectResponse
from fastapi.middleware.cors import CORSMiddleware

from nygaardcodecommentclassification import config
from nygaardcodecommentclassification.api.controllers import PredictionController
from nygaardcodecommentclassification.api.schemas import PredictionRequest

# ---------------------------------------------------------------------------
# Global Resources
# ---------------------------------------------------------------------------
# Initialize the prediction controller (models loaded from MLflow on startup)
controller = PredictionController()

# Thread pool for CPU-bound ML inference tasks
# This prevents blocking the async event loop during model predictions
_executor = ThreadPoolExecutor(max_workers=4)

# ---------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI) -> Any:
    """Async context manager for application lifecycle events.

    This handles:
    - Startup: Load all ML models into memory for fast inference
    - Shutdown: Release model resources and clear GPU memory if applicable

    Args:
        app: The FastAPI application instance

    Yields:
        None: Control back to the application after startup is complete
    """
    # Startup: load models into memory
    controller.startup()
    yield
    # Shutdown: release resources
    controller.shutdown()
    _executor.shutdown(wait=True)


# ---------------------------------------------------------------------------
# FastAPI Application Definition
# ---------------------------------------------------------------------------
app = FastAPI(
    title="Nygaard Code Comment Classification API",
    description="""
    Multi-label classification API for code comments.
    """,
    version="1.0",
    lifespan=lifespan,
)

# ---------------------------------------------------------------------------
# CORS Middleware Configuration
# ---------------------------------------------------------------------------
# Enable Cross-Origin Resource Sharing (CORS) for client-agnostic access.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Restrict to specific domains in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------------------------------------------------------------------------
# Response Decorator
# ---------------------------------------------------------------------------
def construct_response(f: Callable) -> Callable:
    """Decorator to enforce a consistent JSON response structure across all endpoints.

    This decorator wraps endpoint functions to provide:
    - Uniform response format with timestamp, method, URL, status, and data
    - Centralized error handling for ValueError (client errors) and Exception (server errors)
    - Automatic HTTP status code mapping

    Args:
        f: The endpoint function to wrap

    Returns:
        Wrapped function that returns a standardized response dict

    Response Structure:
        {
            "timestamp": "ISO 8601 timestamp",
            "method": "HTTP method (GET, POST, etc.)",
            "url": "Full request URL",
            "status-code": "HTTP status code",
            "message": "Status message or error description",
            "data": "Response payload (if successful)"
        }
    """

    @wraps(f)
    async def wrap(request: Request, *args, **kwargs) -> Dict[str, Any]:
        # Initialize response with request metadata
        response_struct: Dict[str, Any] = {
            "timestamp": datetime.now().isoformat(),
            "method": request.method,
            "url": str(request.url),
        }

        try:
            # Execute the wrapped endpoint function
            results = await f(request, *args, **kwargs)

            # If function returns a dict with status/message/data, use it directly
            if isinstance(results, dict) and "status-code" in results:
                response_struct.update(results)
            else:
                # Fallback for simple returns without explicit status
                response_struct["status-code"] = HTTPStatus.OK
                response_struct["message"] = HTTPStatus.OK.phrase
                response_struct["data"] = results

        except ValueError as e:
            # Client errors: invalid input, unsupported language/model, etc.
            response_struct["status-code"] = HTTPStatus.BAD_REQUEST
            response_struct["message"] = str(e)
        except Exception as e:
            # Server errors: model failures, configuration issues, etc.
            response_struct["status-code"] = HTTPStatus.INTERNAL_SERVER_ERROR
            response_struct["message"] = f"Internal Server Error: {str(e)}"

        return response_struct

    return wrap


# ---------------------------------------------------------------------------
# API Endpoints
# ---------------------------------------------------------------------------


@app.get("/models", tags=["Info"])
@construct_response
async def _get_models(request: Request) -> Dict[str, Any]:
    """Retrieve the list of available ML models grouped by language.

    Returns:
        Dict containing:
        - status-code: HTTP 200 on success
        - message: Status description
        - data: Dict mapping languages to available model types

    Example Response:
        {
            "java": ["catboost"],
            "python": ["catboost"],
            "pharo": ["catboost"]
        }
    """
    data = controller.get_models_info()
    return {"status-code": HTTPStatus.OK, "message": "Available models retrieved", "data": data}


@app.get("/languages", tags=["Info"])
@construct_response
async def _get_languages(request: Request) -> Dict[str, Any]:
    """Retrieve the list of supported programming languages.

    Returns the programming languages for which code comment classification
    is available. Each language has its own trained model.

    Returns:
        Dict containing:
        - status-code: HTTP 200 on success
        - message: Status description
        - data: Dict with "languages" key containing list of supported languages

    Example Response:
        {
            "languages": ["java", "python", "pharo"]
        }
    """
    data = {"languages": config.LANGUAGES}
    return {"status-code": HTTPStatus.OK, "message": "Supported languages retrieved", "data": data}


@app.post("/predict", tags=["Prediction"])
@construct_response
async def _predict(
    request: Request, response: Response, payload: PredictionRequest
) -> Dict[str, Any]:
    """Classify code comments using multi-label classification.

    This endpoint performs ML inference to classify code comments into
    multiple categories.

    Args:
        request: The FastAPI request object
        response: The FastAPI response object
        payload: PredictionRequest containing:
            - texts: List of code comments
            - class_names: List of class names corresponding to each comment
            - language: Programming language ("java", "python", "pharo")
            - model_type: Model to use (default: "catboost")

    Returns:
        Dict containing:
        - status-code: HTTP 200 on success, 400 on invalid input, 500 on error
        - message: Status description
        - data: Dict with model_used, language, and results list

    Example Request:
        POST /predict
        {
            "texts": ["This method calculates fibonacci", "this is a deprecated function"],
            "class_names": ["MathUtils", "Utils"],
            "language": "java",
            "model_type": "catboost"
        }

    Example Response:
        {
            "results": [
                {"text": "This method calculates fibonacci", "class_name": "MathUtils", "labels": ["summary"]},
                {"text": "this is a deprecated function", "class_name": "Utils", "labels": ["deprecation"]}
            ]
        }
    """
    loop = asyncio.get_event_loop()
    results = await loop.run_in_executor(
        _executor,
        controller.predict,
        payload.texts,
        payload.class_names,
        payload.language,
        payload.model_type,
    )

    response.headers["X-model"] = payload.model_type
    response.headers["X-language"] = payload.language

    # Collect all predicted labels
    all_labels = [label for result in results for label in result["labels"]]
    response.headers["X-predicted-labels"] = json.dumps(all_labels)

    return {
        "status-code": HTTPStatus.OK,
        "message": "Prediction successful",
        "data": {
            "model_used": payload.model_type,
            "language": payload.language,
            "results": results,
        },
    }


@app.get("/", tags=["Info"])
async def _root(request: Request) -> RedirectResponse:
    """Root endpoint redirecting to API documentation.
    Returns:
        Redirect response to the auto-generated API docs at /docs
    """
    return RedirectResponse(url="/docs")


# ---------------------------------------------------------------------------
# Entry Point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)