fix rust tests

arch/tools/cli/main.py

@@ -14,15 +14,18 @@ from cli.utils import (
 )
 from cli.core import (
     start_arch_modelserver,
+    start_brightstaff,
     stop_arch_modelserver,
     start_arch,
-    stop_arch,
+    stop_docker_container,
     download_models_from_hf,
 )
 from cli.consts import (
     ARCHGW_DOCKER_IMAGE,
+    BRIGHTSTAFF_DOCKER_IMAGE,
     KATANEMO_DOCKERHUB_REPO,
     SERVICE_NAME_ARCHGW,
+    SERVICE_NAME_BRIGHTSTAFF,
     SERVICE_NAME_MODEL_SERVER,
     SERVICE_ALL,
 )
@@ -40,6 +43,7 @@ logo = r"""
 
 # Command to build archgw and model_server Docker images
 ARCHGW_DOCKERFILE = "./arch/Dockerfile"
+BRIGHTSTAFF_DOCKERFILE = "./arch/Dockerfile.brightstaff"
 MODEL_SERVER_BUILD_FILE = "./model_server/pyproject.toml"
 
 
@@ -51,6 +55,19 @@ def get_version():
         return "version not found"
 
 
+def verify_service_name(service):
+    """Verify if the service name is valid."""
+    if service not in [
+        SERVICE_NAME_ARCHGW,
+        SERVICE_NAME_MODEL_SERVER,
+        SERVICE_NAME_BRIGHTSTAFF,
+        SERVICE_ALL,
+    ]:
+        print(f"Error: Invalid service {service}. Exiting")
+        sys.exit(1)
+    return True
+
+
 @click.group(invoke_without_command=True)
 @click.option("--version", is_flag=True, help="Show the archgw cli version and exit.")
 @click.pass_context
@@ -75,9 +92,8 @@ def main(ctx, version):
 )
 def build(service):
     """Build Arch from source. Must be in root of cloned repo."""
-    if service not in [SERVICE_NAME_ARCHGW, SERVICE_NAME_MODEL_SERVER, SERVICE_ALL]:
-        print(f"Error: Invalid service {service}. Exiting")
-        sys.exit(1)
+    verify_service_name(service)
+
     # Check if /arch/Dockerfile exists
     if service == SERVICE_NAME_ARCHGW or service == SERVICE_ALL:
         if os.path.exists(ARCHGW_DOCKERFILE):
@@ -108,6 +124,35 @@ def build(service):
 
     click.echo("archgw image built successfully.")
 
+    if service == SERVICE_NAME_BRIGHTSTAFF or service == SERVICE_ALL:
+        if os.path.exists(BRIGHTSTAFF_DOCKERFILE):
+            click.echo("Building brightstaff image...")
+            try:
+                subprocess.run(
+                    [
+                        "docker",
+                        "build",
+                        "-f",
+                        BRIGHTSTAFF_DOCKERFILE,
+                        "-t",
+                        f"{KATANEMO_DOCKERHUB_REPO}:brightstaff_latest",
+                        "-t",
+                        f"{BRIGHTSTAFF_DOCKER_IMAGE}",
+                        ".",
+                        "--add-host=host.docker.internal:host-gateway",
+                    ],
+                    check=True,
+                )
+                click.echo("brightstaff image built successfully.")
+            except subprocess.CalledProcessError as e:
+                click.echo(f"Error building brightstaff image: {e}")
+                sys.exit(1)
+        else:
+            click.echo("Error: Dockerfile not found in /arch")
+            sys.exit(1)
+
+    click.echo("brightstaff image built successfully.")
+
     """Install the model server dependencies using Poetry."""
     if service == SERVICE_NAME_MODEL_SERVER or service == SERVICE_ALL:
         # Check if pyproject.toml exists
@@ -146,9 +191,7 @@ def build(service):
 )
 def up(file, path, service, foreground):
     """Starts Arch."""
-    if service not in [SERVICE_NAME_ARCHGW, SERVICE_NAME_MODEL_SERVER, SERVICE_ALL]:
-        log.info(f"Error: Invalid service {service}. Exiting")
-        sys.exit(1)
+    verify_service_name(service)
 
     if service == SERVICE_ALL and foreground:
         # foreground can only be specified when starting individual services
@@ -233,10 +276,13 @@ def up(file, path, service, foreground):
 
     if service == SERVICE_NAME_ARCHGW:
         start_arch(arch_config_file, env, foreground=foreground)
+    if service == SERVICE_NAME_BRIGHTSTAFF:
+        start_brightstaff(arch_config_file, env, foreground=foreground)
     else:
         download_models_from_hf()
         start_arch_modelserver(foreground)
         start_arch(arch_config_file, env, foreground=foreground)
+        start_brightstaff(arch_config_file, env, foreground=foreground)
 
 
 @click.command()
@@ -255,10 +301,10 @@ def down(service):
     if service == SERVICE_NAME_MODEL_SERVER:
         stop_arch_modelserver()
     elif service == SERVICE_NAME_ARCHGW:
-        stop_arch()
+        stop_docker_container()
     else:
         stop_arch_modelserver()
-        stop_arch()
+        stop_docker_container()
 
 
 @click.command()