Update app/main.py

app/main.py

@@ -8,7 +8,11 @@ from . import configs as config
 from .models import UploadResponse, ClassifyRequest, ClassifyResponse, SegmentRequest, SegmentResponse
 from .services.predictor import classify_image
 from .services.segmenter import segment_image
-from .services.routing import router as model_router
+from .services.routing import (
+    router as model_router,
+    resolve_classifier_version,
+    get_ab_status,
+)
 from .storage import (
     save_image,
     get_image_path,
@@ -27,6 +31,9 @@ app = FastAPI(
 
 app.add_middleware(MetricsMiddleware)
 
+# ✅ Include the router for model management endpoints
+app.include_router(model_router)
+
 ALLOWED_CONTENT_TYPES = {"image/jpeg", "image/png", "image/jpg"}
 MAX_FILE_SIZE = 10 * 1024 * 1024
 
@@ -44,7 +51,7 @@ async def models_status():
     return {
         "classification_loaded": classification_model is not None,
         "segmentation_loaded": segmentation_model is not None,
-        "ab_testing": model_router.get_ab_status(),
+        "ab_testing": get_ab_status(),  # ✅ استخدام الدالة المستقلة
         "loaded_versions": config.get_loaded_versions(),
         "storage_paths": {
             "images": str(config.IMAGES_DIR),
@@ -55,7 +62,7 @@ async def models_status():
 
 @app.get("/api/ab-status")
 async def ab_testing_status():
-    return model_router.get_ab_status()
+    return get_ab_status()  # ✅ استخدام الدالة المستقلة
 
 
 @app.post("/api/upload", status_code=201, response_model=UploadResponse)
@@ -115,7 +122,8 @@ async def classify(body: ClassifyRequest):
             status=existing.get("status", "completed"),
         )
 
-    resolved_uri = model_router.resolve_classifier_version(body.model_version)
+    # ✅ استخدام الدالة المستقلة
+    resolved_uri = resolve_classifier_version(body.model_version)
 
     try:
         prediction, confidence, actual_version = classify_image(
@@ -230,4 +238,4 @@ async def health():
     return {
         "status": "healthy",
         "predict_ready": config.get_classification_model() is not None,
-    }
+    }