Generate a single key per user

app.py

@@ -19,9 +19,10 @@ from settings import (
     FAMILY_STATUS,
 )
 from backend import (
-    pre_process_keygen_encrypt_send_user,
-    pre_process_keygen_encrypt_send_bank,
-    pre_process_keygen_encrypt_send_third_party,
+    keygen_send,
+    pre_process_encrypt_send_user,
+    pre_process_encrypt_send_bank,
+    pre_process_encrypt_send_third_party,
     run_fhe,
     get_output,
     decrypt_output,
@@ -45,21 +46,26 @@ with demo:
 
     gr.Markdown("## Client side")
 
-    gr.Markdown("### Step 1: Infos. ")
+    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.")
+    client_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
+    
+    gr.Markdown("### Step 2: Infos. ")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### User")
-            gender = gr.Radio(["Female", "Male"], label="Gender")
+            gender = gr.Radio(["Female", "Male"], label="Gender", value="Female")
             bool_inputs = gr.CheckboxGroup(["Car", "Property", "Work phone", "Phone", "Email"], label="What do you own ?")
             num_children = gr.Slider(**CHILDREN_MIN_MAX, step=1, label="Number of children", info="How many children do you have (0 to 19) ?")
             num_family = gr.Slider(**FAMILY_MIN_MAX, step=1, label="Family", info="How many members does your family have? (1 to 20) ?")
             total_income = gr.Slider(**INCOME_MIN_MAX, label="Income", info="What's you total yearly income (in euros, 3780 to 220500) ?")
             age = gr.Slider(**AGE_MIN_MAX, step=1, label="Age", info="How old are you (20 to 68) ?")
-            income_type = gr.Dropdown(choices=INCOME_TYPES, label="Income type", info="What is your main type of income ?")
-            education_type = gr.Dropdown(choices=EDUCATION_TYPES, label="Education", info="What is your education background ?")
-            family_status = gr.Dropdown(choices=FAMILY_STATUS, label="Family", info="What is your family status ?")
-            occupation_type = gr.Dropdown(choices=OCCUPATION_TYPES, label="Occupation", info="What is your main occupation ?")
-            housing_type = gr.Dropdown(choices=HOUSING_TYPES, label="Housing", info="In what type of housing do you live ?")
+            income_type = gr.Dropdown(choices=INCOME_TYPES, value=INCOME_TYPES[0], label="Income type", info="What is your main type of income ?")
+            education_type = gr.Dropdown(choices=EDUCATION_TYPES, value=EDUCATION_TYPES[0], label="Education", info="What is your education background ?")
+            family_status = gr.Dropdown(choices=FAMILY_STATUS, value=FAMILY_STATUS[0], label="Family", info="What is your family status ?")
+            occupation_type = gr.Dropdown(choices=OCCUPATION_TYPES, value=OCCUPATION_TYPES[0], label="Occupation", info="What is your main occupation ?")
+            housing_type = gr.Dropdown(choices=HOUSING_TYPES, value=HOUSING_TYPES[0], label="Housing", info="In what type of housing do you live ?")
 
         with gr.Column():
             gr.Markdown("### Bank ")
@@ -67,50 +73,37 @@ with demo:
 
         with gr.Column():
             gr.Markdown("### Third party ")
-            employed = gr.Radio(["Yes", "No"], label="Is the person employed ?")
+            employed = gr.Radio(["Yes", "No"], label="Is the person employed ?", value="Yes")
             years_employed = gr.Slider(**EMPLOYED_MIN_MAX, step=1, label="Years of employment", info="How long have this person been employed (in years, 0 to 43) ?")
 
 
-    gr.Markdown("### Step 2: Keygen, encrypt  using FHE and send the inputs to the server.")
+    gr.Markdown("### Step 3: Encrypt  using FHE and send the inputs to the server.")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### User")
             encrypt_button_user = gr.Button("Encrypt the inputs and send to server.")
             
-            user_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
             encrypted_input_user = gr.Textbox(
                 label="Encrypted input representation:", max_lines=2, interactive=False
             )
-            # keys_user = gr.Textbox(
-            #     label="Keys representation:", max_lines=2, interactive=False
-            # )
-
 
 
         with gr.Column():
             gr.Markdown("### Bank ")
             encrypt_button_bank = gr.Button("Encrypt the inputs and send to server.")
 
-            bank_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
             encrypted_input_bank = gr.Textbox(
                 label="Encrypted input representation:", max_lines=2, interactive=False
             )
-            # keys_bank = gr.Textbox(
-            #     label="Keys representation:", max_lines=2, interactive=False
-            # )
 
 
         with gr.Column():
             gr.Markdown("### Third Party ")
             encrypt_button_third_party = gr.Button("Encrypt the inputs and send to server.")
 
-            third_party_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
             encrypted_input_third_party = gr.Textbox(
                 label="Encrypted input representation:", max_lines=2, interactive=False
             )
-            # keys_3 = gr.Textbox(
-            #     label="Keys representation:", max_lines=2, interactive=False
-            # )
 
     gr.Markdown("## Server side")
     gr.Markdown(
@@ -119,7 +112,7 @@ with demo:
         "the encrypted result to the client."
     )
 
-    gr.Markdown("### Step 6: Run FHE execution.")
+    gr.Markdown("### Step 4: Run FHE execution.")
     execute_fhe_button = gr.Button("Run FHE execution.")
     fhe_execution_time = gr.Textbox(
         label="Total FHE execution time (in seconds):", max_lines=1, interactive=False
@@ -131,7 +124,7 @@ with demo:
         "private key."
     )
 
-    gr.Markdown("### Step 7: Receive the encrypted output from the server.")
+    gr.Markdown("### Step 5: Receive the encrypted output from the server.")
     gr.Markdown(
         "The output displayed here is the encrypted result sent by the server, which has been "
         "decrypted using a different private key. This is only used to visually represent an "
@@ -143,55 +136,58 @@ with demo:
         label="Encrypted output representation: ", max_lines=1, interactive=False
     )
 
-    gr.Markdown("### Step 8: Decrypt the output.")
+    gr.Markdown("### Step 6: Decrypt the output.")
     decrypt_button = gr.Button("Decrypt the output")
 
     prediction_output = gr.Textbox(
         label="Credit card approval high risk: ", max_lines=1, interactive=False
     )
 
+    # Button generate the keys
+    keygen_button.click(
+        keygen_send,
+        outputs=[client_id, keygen_button],
+    )
+
     # Button to pre-process, generate the key, encrypt and send the user inputs from the client 
     # side to the server
     encrypt_button_user.click(
-        pre_process_keygen_encrypt_send_user,
-        inputs=[gender, bool_inputs, num_children, num_family, total_income, age, income_type, \
-                education_type, family_status, occupation_type, housing_type],
-        outputs=[user_id, encrypted_input_user],
+        pre_process_encrypt_send_user,
+        inputs=[client_id, gender, bool_inputs, num_children, num_family, total_income, age, \
+                income_type, education_type, family_status, occupation_type, housing_type],
+        outputs=[encrypted_input_user],
     )
 
     # Button to pre-process, generate the key, encrypt and send the bank inputs from the client 
     # side to the server
     encrypt_button_bank.click(
-        pre_process_keygen_encrypt_send_bank,
-        inputs=[account_length],
-        outputs=[bank_id, encrypted_input_bank],
+        pre_process_encrypt_send_bank,
+        inputs=[client_id, account_length],
+        outputs=[encrypted_input_bank],
     )
 
     # Button to pre-process, generate the key, encrypt and send the third party inputs from the 
     # client side to the server    
     encrypt_button_third_party.click(
-        pre_process_keygen_encrypt_send_third_party,
-        inputs=[employed, years_employed],
-        outputs=[third_party_id, encrypted_input_third_party],
+        pre_process_encrypt_send_third_party,
+        inputs=[client_id, employed, years_employed],
+        outputs=[encrypted_input_third_party],
     )
 
-    # TODO : ID should be unique
     # Button to send the encodings to the server using post method
-    execute_fhe_button.click(run_fhe, inputs=[user_id, bank_id, third_party_id], outputs=[fhe_execution_time])
+    execute_fhe_button.click(run_fhe, inputs=[client_id], outputs=[fhe_execution_time])
 
-    # TODO : ID should be unique
     # Button to send the encodings to the server using post method
     get_output_button.click(
         get_output, 
-        inputs=[user_id, bank_id, third_party_id], 
+        inputs=[client_id], 
         outputs=[encrypted_output_representation],
     )
 
-    # TODO : ID should be unique
     # Button to decrypt the output as the user
     decrypt_button.click(
         decrypt_output,
-        inputs=[user_id, bank_id, third_party_id],
+        inputs=[client_id],
         outputs=[prediction_output],
     )
 

backend.py

@@ -78,55 +78,110 @@ def clean_temporary_files(n_keys=20):
     server_files = SERVER_FILES.iterdir()
 
     # Delete all files related to the ids whose keys were deleted
-    for file in chain(client_files, server_files):
+    for user_directory in chain(client_files, server_files):
         for user_id in user_ids:
-            if user_id in file.name:
-                file.unlink()
+            if user_id in user_directory.name:
+                for client_server_file in user_directory.iterdir():
+                    client_server_file.unlink()
 
 
-def _get_client(client_id, client_type):
-    """Get the client API.
+def _get_client(client_id):
+    """Get the client instance.
 
     Args:
         client_id (int): The client ID to consider.
-        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').
 
     Returns:
-        FHEModelClient: The client API.
+        FHEModelClient: The client instance.
     """
-    key_dir = FHE_KEYS / f"{client_type}_{client_id}"
+    key_dir = FHE_KEYS / f"{client_id}"
 
     return MultiInputsFHEModelClient(DEPLOYMENT_PATH, key_dir=key_dir, nb_inputs=len(CLIENT_TYPES))
 
 
-def _keygen(client_id, client_type):
-    """Generate the private key associated to a client.
+def _get_client_file_path(name, client_id, client_type=None):
+    """Get the file path for the client.
 
     Args:
+        name (str): The desired file name (either 'evaluation_key', 'encrypted_inputs') of 
+            'encrypted_outputs').
         client_id (int): The client ID to consider.
-        client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
+        client_type (Optional[str]): The type of user to consider (either 'user', 'bank', 
+            'third_party' or None). Default to None, which is used for evaluation key and output.
+
+    Returns:
+        pathlib.Path: The file path.
+    """
+    client_type_suffix = "" 
+    if client_type is not None:
+        client_type_suffix = f"_{client_type}"
+
+    dir_path = CLIENT_FILES / f"{client_id}"
+    dir_path.mkdir(exist_ok=True)
+
+    return dir_path / f"{name}{client_type_suffix}"
+
+
+def _send_eval_key(client_id, evaluation_key_path):
+    """Send the evaluation key to the server.
+
+    Args:
+        client_id (int): The client ID to consider.
+        evaluation_key_path (Path): Path to the evaluation key to send.
+    """
+
+    # Define the data and files to post
+    data = {
+        "client_id": client_id,
+    }
+
+    files = [
+        ("files", open(evaluation_key_path, "rb")),
+    ]
+
+    # Send the evaluation key to the server
+    url = SERVER_URL + "send_eval_key"
+    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.
     """
     # Clean temporary files
     clean_temporary_files()
 
+    # Create an ID for the current client to consider
+    client_id = numpy.random.randint(0, 2**32)
+
     # Retrieve the client instance
-    client = _get_client(client_id, client_type)
+    client = _get_client(client_id)
 
-    # Generate a private key
+    # Generate the private and evaluation keys
     client.generate_private_and_evaluation_keys(force=True)
 
-    # Retrieve the serialized evaluation key. In this case, as circuits are fully leveled, this
-    # evaluation key is empty. However, for software reasons, it is still needed for proper FHE
-    # execution
+    # Retrieve the serialized evaluation key
     evaluation_key = client.get_serialized_evaluation_keys()
 
-    # Save evaluation_key as bytes in a file as it is too large to pass through regular Gradio
+    # Save evaluation key as bytes in a file as it is too large to pass through regular Gradio
     # buttons (see https://github.com/gradio-app/gradio/issues/1877)
-    evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type)
+    evaluation_key_path = _get_client_file_path("evaluation_key", client_id)
 
     with evaluation_key_path.open("wb") as evaluation_key_file:
         evaluation_key_file.write(evaluation_key)
 
+    # Send the evaluation key to the server
+    _send_eval_key(client_id, evaluation_key_path)
+    
+    return client_id, gr.update(value="Keys are generated and sent ✅")
+
 
 def _send_input(client_id, client_type):
     """Send the encrypted inputs as well as the evaluation key to the server.
@@ -135,8 +190,7 @@ def _send_input(client_id, client_type):
         client_id (int): The client ID to consider.
         client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
     """
-    # Get the paths to the evaluation key and encrypted inputs
-    evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type)
+    # Get the paths to the encrypted inputs
     encrypted_input_path = _get_client_file_path("encrypted_inputs", client_id, client_type)
 
     # Define the data and files to post
@@ -147,7 +201,6 @@ def _send_input(client_id, client_type):
 
     files = [
         ("files", open(encrypted_input_path, "rb")),
-        ("files", open(evaluation_key_path, "rb")),
     ]
 
     # Send the encrypted inputs and evaluation key to the server
@@ -160,24 +213,11 @@ def _send_input(client_id, client_type):
         return response.ok
 
 
-def _get_client_file_path(name, client_id, client_type):
-    """Get the correct temporary file path for the client.
-
-    Args:
-        name (str): The desired file name (either 'evaluation_key' or 'encrypted_inputs').
-        client_id (int): The client ID to consider.
-        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').
-
-    Returns:
-        pathlib.Path: The file path.
-    """
-    return CLIENT_FILES / f"{name}_{client_type}_{client_id}"
-
-
-def _keygen_encrypt_send(inputs, client_type):
-    """Encrypt the given inputs for a specific client.
+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.
         inputs (numpy.ndarray): The inputs to encrypt.
         client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
     
@@ -185,15 +225,9 @@ def _keygen_encrypt_send(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. 
     """
-    # Create an ID for the current client to consider
-    client_id = numpy.random.randint(0, 2**32)
-
-    _keygen(client_id, client_type)
 
     # Retrieve the client instance
-    client = _get_client(client_id, client_type)
-
-    # TODO : pre-process the data first 
+    client = _get_client(client_id)
 
     # Quantize, encrypt and serialize the inputs
     encrypted_inputs = client.quantize_encrypt_serialize_multi_inputs(
@@ -215,14 +249,14 @@ def _keygen_encrypt_send(inputs, client_type):
 
     _send_input(client_id, client_type)
 
-    # TODO: also return private key representation if possible
-    return client_id, encrypted_inputs_short
+    return encrypted_inputs_short
 
 
-def pre_process_keygen_encrypt_send_user(*inputs):
-    """Pre-process the given inputs for a specific client.
+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.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
     
     Returns:
@@ -263,13 +297,14 @@ def pre_process_keygen_encrypt_send_user(*inputs):
 
     preprocessed_user_inputs = PRE_PROCESSOR_USER.transform(user_inputs)
 
-    return _keygen_encrypt_send(preprocessed_user_inputs, "user")
+    return _encrypt_send(client_id, preprocessed_user_inputs, "user")
 
 
-def pre_process_keygen_encrypt_send_bank(*inputs):
-    """Pre-process the given inputs for a specific client.
+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.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
     
     Returns:
@@ -278,13 +313,14 @@ def pre_process_keygen_encrypt_send_bank(*inputs):
     """
     account_length = inputs[0]
     
-    return _keygen_encrypt_send(account_length, "bank")
+    return _encrypt_send(client_id, account_length, "bank")
     
 
-def pre_process_keygen_encrypt_send_third_party(*inputs):
-    """Pre-process the given inputs for a specific client.
+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.
         *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
 
     Returns:
@@ -303,24 +339,20 @@ def pre_process_keygen_encrypt_send_third_party(*inputs):
 
     preprocessed_third_party_inputs = PRE_PROCESSOR_THIRD_PARTY.transform(third_party_inputs)
 
-    return _keygen_encrypt_send(preprocessed_third_party_inputs, "third_party")
+    return _encrypt_send(client_id, preprocessed_third_party_inputs, "third_party")
 
 
-def run_fhe(user_id, bank_id, third_party_id):
+def run_fhe(client_id):
     """Run the model on the encrypted inputs previously sent using FHE.
 
     Args:
-        user_id (int): The user ID to consider.
-        bank_id (int): The bank ID to consider.
-        third_party_id (int): The third party ID to consider.
+        client_id (int): The current client ID to consider.
     """
 
     # TODO : add a warning for users to send all client types' inputs 
 
     data = {
-        "user_id": user_id,
-        "bank_id": bank_id,
-        "third_party_id": third_party_id,
+        "client_id": client_id,
     }
 
     # Trigger the FHE execution on the encrypted inputs previously sent
@@ -335,21 +367,17 @@ def run_fhe(user_id, bank_id, third_party_id):
             raise gr.Error("Please wait for the inputs to be sent to the server.")
 
 
-def get_output(user_id, bank_id, third_party_id):
+def get_output(client_id):
     """Retrieve the encrypted output.
 
     Args:
-        user_id (int): The user ID to consider.
-        bank_id (int): The bank ID to consider.
-        third_party_id (int): The third party ID to consider.
+        client_id (int): The current client ID to consider.
     
     Returns:
         encrypted_output_short (bytes): A byte short representation of the encrypted output. 
     """
     data = {
-        "user_id": user_id,
-        "bank_id": bank_id,
-        "third_party_id": third_party_id,
+        "client_id": client_id,
     }
 
     # Retrieve the encrypted output
@@ -363,7 +391,7 @@ def get_output(user_id, bank_id, third_party_id):
 
             # Save the encrypted output to bytes in a file as it is too large to pass through regular
             # Gradio buttons (see https://github.com/gradio-app/gradio/issues/1877)
-            encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")
+            encrypted_output_path = _get_client_file_path("encrypted_output", client_id)
 
             with encrypted_output_path.open("wb") as encrypted_output_file:
                 encrypted_output_file.write(encrypted_output)
@@ -376,20 +404,18 @@ def get_output(user_id, bank_id, third_party_id):
             raise gr.Error("Please wait for the FHE execution to be completed.")
 
 
-def decrypt_output(user_id, bank_id, third_party_id):
+def decrypt_output(client_id):
     """Decrypt the result.
 
     Args:
-        user_id (int): The user ID to consider.
-        bank_id (int): The bank ID to consider.
-        third_party_id (int): The third party ID to consider.
+        client_id (int): The current client ID to consider.
 
     Returns:
         output(numpy.ndarray): The decrypted output
 
     """
     # Get the encrypted output path
-    encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")
+    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.")
@@ -399,7 +425,7 @@ def decrypt_output(user_id, bank_id, third_party_id):
         encrypted_output_proba = encrypted_output_file.read()
 
     # Retrieve the client API
-    client = _get_client(user_id, "user")
+    client = _get_client(client_id)
 
     # Deserialize, decrypt and post-process the encrypted output
     output_proba = client.deserialize_decrypt_dequantize(encrypted_output_proba)

server.py

@@ -5,25 +5,34 @@ from typing import List
 from fastapi import FastAPI, File, Form, UploadFile
 from fastapi.responses import JSONResponse, Response
 
-from settings import DEPLOYMENT_PATH, SERVER_FILES
+from settings import DEPLOYMENT_PATH, SERVER_FILES, CLIENT_TYPES
 from utils.client_server_interface import MultiInputsFHEModelServer
 
 # Load the server
 FHE_SERVER = MultiInputsFHEModelServer(DEPLOYMENT_PATH)
 
 
-def _get_server_file_path(name, client_id, client_type):
-    """Get the correct temporary file path for the server.
+def _get_server_file_path(name, client_id, client_type=None):
+    """Get the file path for the server.
 
     Args:
-        name (str): The desired file name (either 'evaluation_key' or 'encrypted_inputs').
+        name (str): The desired file name (either 'evaluation_key', 'encrypted_inputs' or 
+            'encrypted_outputs').
         client_id (int): The client ID to consider.
-        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').
+        client_type (Optional[str]): The type of user to consider (either 'user', 'bank', 
+            'third_party' or None). Default to None, which is used for evaluation key and output.
 
     Returns:
         pathlib.Path: The file path.
     """
-    return SERVER_FILES / f"{name}_{client_type}_{client_id}"
+    client_type_suffix = "" 
+    if client_type is not None:
+        client_type_suffix = f"_{client_type}"
+
+    dir_path = SERVER_FILES / f"{client_id}"
+    dir_path.mkdir(exist_ok=True)
+
+    return dir_path / f"{name}{client_type_suffix}"
 
 
 # Initialize an instance of FastAPI
@@ -35,6 +44,20 @@ def root():
     return {"message": "Welcome to Credit Card Approval Prediction server!"}
 
 
+@app.post("/send_eval_key")
+def send_eval_key(
+    client_id: str = Form(),
+    files: List[UploadFile] = File(),
+):
+    """Send the evaluation key to the server."""
+    # Retrieve the evaluation key path
+    evaluation_key_path = _get_server_file_path("evaluation_key", client_id)
+    
+    # Write the file using the above path
+    with evaluation_key_path.open("wb") as evaluation_key:
+        evaluation_key.write(files[0].file.read())
+
+
 @app.post("/send_input")
 def send_input(
     client_id: str = Form(),
@@ -42,44 +65,33 @@ def send_input(
     files: List[UploadFile] = File(),
 ):
     """Send the inputs to the server."""
-    # Retrieve the encrypted inputs and the evaluation key paths
+    # Retrieve the encrypted inputs
     encrypted_inputs_path = _get_server_file_path("encrypted_inputs", client_id, client_type)
-    evaluation_key_path = _get_server_file_path("evaluation_key", client_id, client_type)
     
-    # Write the files using the above paths
-    with encrypted_inputs_path.open("wb") as encrypted_inputs, evaluation_key_path.open(
-        "wb"
-    ) as evaluation_key:
+    # Write the file using the above path
+    with encrypted_inputs_path.open("wb") as encrypted_inputs:
         encrypted_inputs.write(files[0].file.read())
-        evaluation_key.write(files[1].file.read())
 
 
 @app.post("/run_fhe")
 def run_fhe(
-    user_id: str = Form(),
-    bank_id: str = Form(),
-    third_party_id: str = Form(),
+    client_id: str = Form(),
 ):
     """Execute the model on the encrypted inputs using FHE."""
-    # Retrieve the evaluation key (from the user, as all evaluation keys should be the same)
-    evaluation_key_path = _get_server_file_path("evaluation_key", user_id, "user")
-
-    # Get the encrypted inputs
-    encrypted_user_inputs_path = _get_server_file_path("encrypted_inputs", user_id, "user")
-    encrypted_bank_inputs_path = _get_server_file_path("encrypted_inputs", bank_id, "bank")
-    encrypted_third_party_inputs_path = _get_server_file_path("encrypted_inputs", third_party_id, "third_party")
-    with (
-        evaluation_key_path.open("rb") as evaluation_key_file,
-        encrypted_user_inputs_path.open("rb") as encrypted_user_inputs_file,
-        encrypted_bank_inputs_path.open("rb") as encrypted_bank_inputs_file,
-        encrypted_third_party_inputs_path.open("rb") as encrypted_third_party_inputs_file,
-    ):
+    # Get the evaluation key
+    evaluation_key_path = _get_server_file_path("evaluation_key", client_id)
+
+    with evaluation_key_path.open("rb") as evaluation_key_file:
         evaluation_key = evaluation_key_file.read()
-        encrypted_user_input = encrypted_user_inputs_file.read()
-        encrypted_bank_input = encrypted_bank_inputs_file.read()
-        encrypted_third_party_input = encrypted_third_party_inputs_file.read()
 
-    encrypted_inputs = (encrypted_user_input, encrypted_bank_input, encrypted_third_party_input)
+    # Get the encrypted inputs from all parties
+    encrypted_inputs = []
+    for client_type in CLIENT_TYPES:
+        encrypted_inputs_path = _get_server_file_path("encrypted_inputs", client_id, client_type)
+
+        with encrypted_inputs_path.open("rb") as encrypted_inputs_file:
+            encrypted_input = encrypted_inputs_file.read()
+            encrypted_inputs.append(encrypted_input)
 
     # Run the FHE execution
     start = time.time()
@@ -87,7 +99,7 @@ def run_fhe(
     fhe_execution_time = round(time.time() - start, 2)
 
     # Retrieve the encrypted output path
-    encrypted_output_path = _get_server_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")
+    encrypted_output_path = _get_server_file_path("encrypted_output", client_id)
 
     # Write the file using the above path
     with encrypted_output_path.open("wb") as output_file:
@@ -98,13 +110,11 @@ def run_fhe(
 
 @app.post("/get_output")
 def get_output(
-    user_id: str = Form(),
-    bank_id: str = Form(),
-    third_party_id: str = Form(),
+    client_id: str = Form(),
 ):
     """Retrieve the encrypted output."""
     # Retrieve the encrypted output path
-    encrypted_output_path = _get_server_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")
+    encrypted_output_path = _get_server_file_path("encrypted_output", client_id)
 
     # Read the file using the above path
     with encrypted_output_path.open("rb") as encrypted_output_file: