Add explicit error messages

app.py

@@ -47,8 +47,10 @@ with demo:
     gr.Markdown("## Client side")
 
     gr.Markdown("### Step 1: Generate the keys. ")
-    # TODO: Re-initialize the key message once generated and sent
     keygen_button = gr.Button("Generate the keys and send evaluation key to the server.")
+    evaluation_key = gr.Textbox(
+        label="Evaluation key representation:", max_lines=2, interactive=False
+    )
     client_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
     
     gr.Markdown("### Step 2: Infos. ")
@@ -133,7 +135,7 @@ with demo:
     get_output_button = gr.Button("Receive the encrypted output from the server.")
 
     encrypted_output_representation = gr.Textbox(
-        label="Encrypted output representation: ", max_lines=1, interactive=False
+        label="Encrypted output representation: ", max_lines=2, interactive=False
     )
 
     gr.Markdown("### Step 6: Decrypt the output.")
@@ -146,7 +148,7 @@ with demo:
     # Button generate the keys
     keygen_button.click(
         keygen_send,
-        outputs=[client_id, keygen_button],
+        outputs=[client_id, evaluation_key, keygen_button],
     )
 
     # Button to pre-process, generate the key, encrypt and send the user inputs from the client 

backend.py

@@ -53,7 +53,7 @@ def shorten_bytes_object(bytes_object, limit=500):
 
 
 def clean_temporary_files(n_keys=20):
-    """Clean keys and encrypted inputs.
+    """Clean older keys and encrypted files.
 
     A maximum of n_keys keys and associated temporary files are allowed to be stored. Once this 
     limit is reached, the oldest files are deleted.
@@ -73,16 +73,15 @@ def clean_temporary_files(n_keys=20):
             user_ids.append(key_dir.name)
             shutil.rmtree(key_dir)
 
-    # Get all the encrypted objects in the temporary folder
+    # Get all the encrypted files in the temporary folders
     client_files = CLIENT_FILES.iterdir()
     server_files = SERVER_FILES.iterdir()
 
-    # Delete all files related to the ids whose keys were deleted
-    for user_directory in chain(client_files, server_files):
+    # Delete all files related to the IDs whose keys were deleted
+    for user_dir in chain(client_files, server_files):
         for user_id in user_ids:
-            if user_id in user_directory.name:
-                for client_server_file in user_directory.iterdir():
-                    client_server_file.unlink()
+            if user_id in user_dir.name:
+                shutil.rmtree(user_dir)
 
 
 def _get_client(client_id):
@@ -122,11 +121,44 @@ def _get_client_file_path(name, client_id, client_type=None):
     return dir_path / f"{name}{client_type_suffix}"
 
 
+def _send_to_server(client_id, client_type, file_name):
+    """Send the encrypted inputs or the evaluation key to the server.
+
+    Args:
+        client_id (int): The client ID to consider.
+        client_type (Optional[str]): The type of client to consider (either 'user', 'bank', 'third_party' or 
+            None).
+        file_name (str): File name to send (either 'evaluation_key' or 'encrypted_inputs').
+    """
+    # Get the paths to the encrypted inputs
+    encrypted_file_path = _get_client_file_path(file_name, client_id, client_type)
+
+    # Define the data and files to post
+    data = {
+        "client_id": client_id,
+        "client_type": client_type,
+        "file_name": file_name,
+    }
+
+    files = [
+        ("files", open(encrypted_file_path, "rb")),
+    ]
+
+    # Send the encrypted inputs or evaluation key to the server
+    url = SERVER_URL + "send_file"
+    with requests.post(
+        url=url,
+        data=data,
+        files=files,
+    ) as response:
+        return response.ok
+
+
 def keygen_send():
     """Generate the private and evaluation key, and send the evaluation key to the server.
     
     Returns:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
     """
     # Clean temporary files
     clean_temporary_files()
@@ -154,48 +186,18 @@ def keygen_send():
 
     # Send the evaluation key to the server
     _send_to_server(client_id, None, file_name)
-    
-    return client_id, gr.update(value="Keys are generated and sent ✅")
-
-
-def _send_to_server(client_id, client_type, file_name):
-    """Send the encrypted inputs or the evaluation key to the server.
-
-    Args:
-        client_id (int): The client ID to consider.
-        client_type (Optional[str]): The type of client to consider (either 'user', 'bank', 'third_party' or 
-            None).
-        file_name (str): File name to send (either 'evaluation_key' or 'encrypted_inputs').
-    """
-    # Get the paths to the encrypted inputs
-    encrypted_file_path = _get_client_file_path(file_name, client_id, client_type)
 
-    # Define the data and files to post
-    data = {
-        "client_id": client_id,
-        "client_type": client_type,
-        "file_name": file_name,
-    }
-
-    files = [
-        ("files", open(encrypted_file_path, "rb")),
-    ]
-
-    # Send the encrypted inputs or evaluation key to the server
-    url = SERVER_URL + "send_file"
-    with requests.post(
-        url=url,
-        data=data,
-        files=files,
-    ) as response:
-        return response.ok
+    # Create a truncated version of the evaluation key for display
+    evaluation_key_short = shorten_bytes_object(evaluation_key)
+    
+    return client_id, evaluation_key_short, gr.update(value="Keys are generated and evaluation key is sent ✅")
 
 
 def _encrypt_send(client_id, inputs, client_type):
     """Encrypt the given inputs for a specific client and send it to the server.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
         inputs (numpy.ndarray): The inputs to encrypt.
         client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
     
@@ -203,6 +205,8 @@ def _encrypt_send(client_id, inputs, client_type):
         client_id, encrypted_inputs_short (int, bytes): Integer ID representing the current client 
             and a byte short representation of the encrypted input to send. 
     """
+    if client_id == "":
+        raise gr.Error("Please generate the keys first.")
 
     # Retrieve the client instance
     client = _get_client(client_id)
@@ -236,7 +240,7 @@ def pre_process_encrypt_send_user(client_id, *inputs):
     """Pre-process, encrypt and send the user inputs for a specific client to the server.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
     
     Returns:
@@ -284,7 +288,7 @@ def pre_process_encrypt_send_bank(client_id, *inputs):
     """Pre-process, encrypt and send the user inputs for a specific client to the server.
 
     Args:        
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
     
     Returns:
@@ -300,7 +304,7 @@ def pre_process_encrypt_send_third_party(client_id, *inputs):
     """Pre-process, encrypt and send the user inputs for a specific client to the server.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
 
     Returns:
@@ -326,10 +330,11 @@ def run_fhe(client_id):
     """Run the model on the encrypted inputs previously sent using FHE.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
     """
 
-    # TODO : add a warning for users to send all client types' inputs 
+    if client_id == "":
+        raise gr.Error("Please generate the keys first.")
 
     data = {
         "client_id": client_id,
@@ -344,18 +349,22 @@ def run_fhe(client_id):
         if response.ok:
             return response.json()
         else:
-            raise gr.Error("Please wait for the inputs to be sent to the server.")
+            raise gr.Error("Please send the inputs from all three parties to the server first.")
 
 
 def get_output(client_id):
     """Retrieve the encrypted output.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
     
     Returns:
         encrypted_output_short (bytes): A byte short representation of the encrypted output. 
     """
+
+    if client_id == "":
+        raise gr.Error("Please generate the keys first.")
+
     data = {
         "client_id": client_id,
     }
@@ -381,24 +390,27 @@ def get_output(client_id):
 
             return encrypted_output_short
         else:
-            raise gr.Error("Please wait for the FHE execution to be completed.")
+            raise gr.Error("Please run the FHE execution first and wait for it to be completed.")
 
 
 def decrypt_output(client_id):
     """Decrypt the result.
 
     Args:
-        client_id (int): The current client ID to consider.
+        client_id (str): The current client ID to consider.
 
     Returns:
         output(numpy.ndarray): The decrypted output
-
     """
+
+    if client_id == "":
+        raise gr.Error("Please generate the keys first.")
+
     # Get the encrypted output path
     encrypted_output_path = _get_client_file_path("encrypted_output", client_id)
 
     if not encrypted_output_path.is_file():
-        raise gr.Error("Please run the FHE execution first.")
+        raise gr.Error("Please receive the outputs from the server first.")
 
     # Load the encrypted output as bytes
     with encrypted_output_path.open("rb") as encrypted_output_file:
@@ -412,5 +424,5 @@ def decrypt_output(client_id):
 
     # Determine the predicted class
     output = numpy.argmax(output_proba, axis=1)
-
+    
     return output