Quadratic-Labs/PrivateNearestNeighbors-FHE

1import network2from shapely.geometry import Point3import streamlit as st4import config5from utils import (6    set_up_server,7    set_up_client,8    display_encrypted,9    transform_point,10    process_result,11    display_map,12    init_session,13    add_to_server_side,14    add_to_client_side,15    display_client_side,16    display_server_side,17    restart_session,18)19from concrete import fhe20 21server = set_up_server()22client = set_up_client(server.client_specs.serialize())23restaurants, point_coordinates = network.get()24 25c1, c2, c3 = init_session()26 27with c1:28    display_client_side()29with c2:30    st_map = display_map(restaurants)31with c3:32    display_server_side()33 34 35# keys generation view36if "evaluation_key" not in st.session_state:37    with c1:38        if st.button("Generate keys"):39            with st.spinner("Generating keys"):40                client.keys.load_if_exists_generate_and_save_otherwise(41                    config.keys_filepath42                )43                st.session_state["evaluation_key"] = client.evaluation_keys.serialize()44                add_to_client_side(45                    "Encryption/decryption keys and evaluation keys are generated."46                )47                add_to_client_side("The evaluation key is sent to the server.")48                add_to_server_side(49                    f"Evaluation key: {display_encrypted(st.session_state['evaluation_key'])}"50                )51            st.rerun()52 53else:54    if "position" not in st.session_state:55        with c1:56            st.write("Select your starting position on the st_map")57 58        if st_map.get("last_clicked"):59            position = Point(60                st_map.get("last_clicked")["lng"], st_map.get("last_clicked")["lat"]61            )62            add_to_client_side(63                f"Selected starting point is: ({position.x}, {position.y})"64            )65            st.session_state["position"] = position66            st.rerun()67 68    if "position" in st.session_state and "encrypted_position" not in st.session_state:69        position = st.session_state["position"]70        with c1:71            if st.button("Encrypt and send inputs"):72                with st.spinner("Encrypting inputs"):73                    client.keys.load(config.keys_filepath)74 75                    x, y = transform_point(position.x, position.y)76                    encrypted_x, encrypted_y = client.encrypt(x, y)77                    encrypted_position = (78                        encrypted_x.serialize(),79                        encrypted_y.serialize(),80                    )81                    add_to_server_side(82                        f"Received starting point: ({display_encrypted(encrypted_position[0])},{display_encrypted(encrypted_position[1])})"83                    )84                    st.session_state["encrypted_position"] = encrypted_position85                st.rerun()86 87    if (88        "encrypted_position" in st.session_state89        and "encrypted_results" not in st.session_state90    ):91        with c3:92            encrypted_x = st.session_state["encrypted_position"][0]93            encrypted_y = st.session_state["encrypted_position"][1]94 95            if st.button(f"Find {config.number_of_neighbors} nearest restaurants"):96                with st.spinner("Computing"):97                    st.write("Please wait, it may take few minutes.")98                    deserialized_x = fhe.Value.deserialize(encrypted_x)99                    deserialized_y = fhe.Value.deserialize(encrypted_y)100                    deserialized_evaluation_keys = fhe.EvaluationKeys.deserialize(101                        st.session_state["evaluation_key"]102                    )103                    client.keys.load_if_exists_generate_and_save_otherwise(104                        config.keys_filepath105                    )106                    res = server.run(107                        deserialized_x,108                        deserialized_y,109                        evaluation_keys=deserialized_evaluation_keys,110                    )111                    ser_res = fhe.Value.serialize(res)112                    add_to_server_side(113                        f"The encrypted result {display_encrypted(ser_res)} is sent to the client."114                    )115                    add_to_client_side(116                        f"The received encrypted result is: {display_encrypted(ser_res)}."117                    )118                    st.session_state["encrypted_results"] = ser_res119                st.rerun()120 121    if (122        "encrypted_results" in st.session_state123        and "decrypted_result" not in st.session_state124    ):125        res = st.session_state["encrypted_results"]126        with c1:127            if st.button("Decrypt result"):128                with st.spinner("Computing"):129                    client.keys.load_if_exists_generate_and_save_otherwise(130                        config.keys_filepath131                    )132                    deser_res = fhe.Value.deserialize(res)133                    decrypted_result = client.decrypt(deser_res)134                    process_result(restaurants, decrypted_result)135                    st.session_state["decrypted_result"] = decrypted_result136                    st.rerun()137 138    if "decrypted_result" in st.session_state:139        with c1:140            restart_session()141