Update nygaardcodecommentclassification/api/API.py

nygaardcodecommentclassification/api/API.py

@@ -27,8 +27,8 @@ from nygaardcodecommentclassification.api.schemas import PredictionRequest
 # ---------------------------------------------------------------------------
 # Global Resources
 # ---------------------------------------------------------------------------
-# Initialize the prediction controller with the models directory
-controller = PredictionController(config.MODELS_DIR)
+# 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
@@ -56,6 +56,7 @@ async def lifespan(app: FastAPI):
     controller.shutdown()
     _executor.shutdown(wait=True)
 
+
 # ---------------------------------------------------------------------------
 # FastAPI Application Definition
 # ---------------------------------------------------------------------------
@@ -108,6 +109,7 @@ def construct_response(f: Callable) -> Callable:
             "data": "Response payload (if successful)"
         }
     """
+
     @wraps(f)
     async def wrap(request: Request, *args, **kwargs) -> Dict[str, Any]:
         # Initialize response with request metadata
@@ -120,7 +122,7 @@ def construct_response(f: Callable) -> Callable:
         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)
@@ -144,12 +146,11 @@ def construct_response(f: Callable) -> Callable:
     return wrap
 
 
-
-
 # ---------------------------------------------------------------------------
 # API Endpoints
 # ---------------------------------------------------------------------------
 
+
 @app.get("/models", tags=["Info"])
 @construct_response
 async def _get_models(request: Request) -> Dict[str, Any]:
@@ -169,11 +170,7 @@ async def _get_models(request: Request) -> Dict[str, Any]:
         }
     """
     data = controller.get_models_info()
-    return {
-        "status-code": HTTPStatus.OK,
-        "message": "Available models retrieved",
-        "data": data
-    }
+    return {"status-code": HTTPStatus.OK, "message": "Available models retrieved", "data": data}
 
 
 @app.get("/languages", tags=["Info"])
@@ -196,11 +193,7 @@ async def _get_languages(request: Request) -> Dict[str, Any]:
         }
     """
     data = {"languages": config.LANGUAGES}
-    return {
-        "status-code": HTTPStatus.OK,
-        "message": "Supported languages retrieved",
-        "data": data
-    }
+    return {"status-code": HTTPStatus.OK, "message": "Supported languages retrieved", "data": data}
 
 
 @app.post("/predict", tags=["Prediction"])
@@ -214,7 +207,8 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
     Args:
         request: The FastAPI request object
         payload: PredictionRequest containing:
-            - texts: List of code comments to classify
+            - 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")
 
@@ -228,6 +222,7 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
         POST /predict
         {
             "texts": ["This method calculates fibonacci", "this is a deprecated function"],
+            "class_names": ["MathUtils", "Utils"],
             "language": "java",
             "model_type": "catboost"
         }
@@ -235,8 +230,8 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
     Example Response:
         {
             "results": [
-                {"text": "This method calculates fibonacci", "labels": ["summary"]},
-                {"text": "this is a deprecated function", "labels": ["deprecation"]}
+                {"text": "This method calculates fibonacci", "class_name": "MathUtils", "labels": ["summary"]},
+                {"text": "this is a deprecated function", "class_name": "Utils", "labels": ["deprecation"]}
             ]
         }
     """
@@ -245,8 +240,9 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
         _executor,
         controller.predict,
         payload.texts,
+        payload.class_names,
         payload.language,
-        payload.model_type
+        payload.model_type,
     )
 
     return {
@@ -255,8 +251,8 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
         "data": {
             "model_used": payload.model_type,
             "language": payload.language,
-            "results": results
-        }
+            "results": results,
+        },
     }
 
 
@@ -265,4 +261,5 @@ async def _predict(request: Request, payload: PredictionRequest) -> Dict[str, An
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7680)
+
+    uvicorn.run(app, host="0.0.0.0", port=7860)