improve imports

agenta-backend/agenta_backend/routers/app_router.py

@@ -22,19 +22,28 @@
     AppVariantOutput,
     AddVariantFromImagePayload,
     EnvironmentOutput,
+    Image
 )
 from agenta_backend.models import converters
 
-if os.environ["FEATURE_FLAG"] not in ["ee"]:
-    from agenta_backend.services import docker_utils
-
 if os.environ["FEATURE_FLAG"] in ["cloud", "ee", "demo"]:
     from agenta_backend.ee.services.selectors import (
         get_user_and_org_id,
     )  # noqa pylint: disable-all
 else:
     from agenta_backend.services.selectors import get_user_and_org_id
 
+if os.environ["FEATURE_FLAG"] in ["cloud"]:
+    from agenta_backend.ee.services import (
+        lambda_deployment_manager as deployment_manager,
+    )  # noqa pylint: disable-all
+elif os.environ["FEATURE_FLAG"] in ["ee"]:
+    from agenta_backend.ee.services import (
+        deployment_manager
+    )  # noqa pylint: disable-all
+else:
+    from agenta_backend.services import deployment_manager
+
 router = APIRouter()
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -222,12 +231,16 @@ async def add_variant_from_image(
     """
 
     if os.environ["FEATURE_FLAG"] not in ["cloud", "ee"]:
+        image = Image(
+            docker_id=payload.docker_id,
+            tags=payload.tags,
+        )
         if not payload.tags.startswith(settings.registry):
             raise HTTPException(
                 status_code=500,
                 detail="Image should have a tag starting with the registry name (agenta-server)",
             )
-        elif docker_utils.find_image_by_docker_id(payload.docker_id) is None:
+        elif deployment_manager.validate_image(image) is False:
             raise HTTPException(status_code=404, detail="Image not found")
 
     try:

agenta-backend/agenta_backend/routers/container_router.py

@@ -24,10 +24,6 @@
 else:
     from agenta_backend.services import container_manager
 
-
-if os.environ["FEATURE_FLAG"] in ["oss", "cloud"]:
-    from agenta_backend.services.docker_utils import restart_container
-
 import logging
 
 logger = logging.getLogger(__name__)
@@ -97,7 +93,7 @@ async def restart_docker_container(
         container_id = deployment.container_id
 
         logger.debug(f"Restarting container with id: {container_id}")
-        restart_container(container_id)
+        container_manager.restart_container(container_id)
         return {"message": "Please wait a moment. The container is now restarting."}
     except Exception as ex:
         return JSONResponse({"message": str(ex)}, status_code=500)

agenta-backend/agenta_backend/services/container_manager.py

@@ -19,6 +19,7 @@
 from agenta_backend.models.db_models import (
     AppDB,
 )
+from agenta_backend.services import docker_utils
 
 client = docker.from_env()
 
@@ -188,3 +189,11 @@ async def get_image_details_from_docker_hub(
             f"{repo_owner}/{repo_name}:{image_name}"
         )
         return image_details["Id"]
+
+def restart_container(container_id: str):
+    """Restart docker container.
+
+    Args:
+        container_id (str): The id of the container to restart.
+    """
+    docker_utils.restart_container(container_id)