New Endpoints for better Registration UX (#72)

* checkpointing

* checkpointing

* update

* working

* checkpointing

* first_device_needs_registration_code

* checkpointing

* sync

* add new builder method
update node ports

* - wasm updated
- registration endpoint updated
- added new boilerplate for easy tests

* update tests and fix pyo3

* creates venv if not present

.github/Dockerfile.wasm

@@ -3,6 +3,34 @@ FROM rust:bookworm
 WORKDIR /app/shinkai-libs
 COPY ./shinkai-libs .
 
+# Install pyenv dependencies
+RUN apt-get update && apt-get install -y \
+    make build-essential libssl-dev zlib1g-dev libbz2-dev \
+    libreadline-dev libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev \
+    xz-utils tk-dev libffi-dev liblzma-dev git
+
+# Install pyenv
+RUN curl https://pyenv.run | bash
+
+# Set environment variables for pyenv
+ENV PYENV_ROOT /root/.pyenv
+ENV PATH $PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
+
+# Install Python 3.10.6 with pyenv and set it as default
+RUN pyenv install 3.10.6
+RUN pyenv global 3.10.6
+
+# Check python version
+RUN python --version
+
+# Create a Python virtual environment and install maturin
+RUN python -m venv /venv && \
+    . /venv/bin/activate && \
+    pip install maturin
+
+# Activate the virtual environment in subsequent RUN commands
+ENV PATH="/venv/bin:$PATH"
+
 #RUN WASM tests
 RUN cd shinkai-message-wasm && cargo build
 
@@ -26,4 +54,5 @@ RUN node -v
 RUN curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
 #RUN wasm-pack test --node
 COPY .github/run-wasm*.sh /entrypoints/
+COPY .github/run-python*.sh /entrypoints/
 RUN chmod 755 /entrypoints/*.sh

.github/run-python-tests.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+cd /app/shinkai-libs/shinkai-message-pyo3
+
+# Initialize pyenv
+set -e
+export PYENV_ROOT="$HOME/.pyenv"
+export PATH="$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
+eval "$(pyenv init -)"
+
+# Get the path of the Python interpreter inside the Docker container
+python_path=$(which python)
+
+# Set the PYO3_PYTHON and PYTHON_SYS_EXECUTABLE environment variables
+export PYO3_PYTHON="$python_path"
+export PYTHON_SYS_EXECUTABLE="$python_path"
+
+# Print the Python version
+python --version
+
+# Create a virtual environment if it doesn't exist
+if [ ! -d "./venv" ]
+then
+    python -m venv venv
+fi
+
+# Activate your virtual environment
+source ./venv/bin/activate
+
+# Run maturin develop and capture its output
+output=$(maturin build -i python)
+
+# If maturin develop is successful, extract the path of the built wheel file
+if [ $? -eq 0 ]; then
+    echo "Maturin build successful"
+    wheel_file=$(ls target/wheels/*.whl)
+    echo "Wheel file: $wheel_file"
+
+    # Update the installed package using the built wheel file
+    echo "Running pip install --upgrade \"$wheel_file\"..."
+    pip_output=$(pip install --upgrade "$wheel_file")
+
+    # Run the tests and print their output
+    echo "Running tests..."
+    test_output=$(python -m unittest tests.test_shinkai_message_pyo3)
+    echo "$test_output"
+else
+    echo "maturin develop failed"
+fi

.github/workflows/main.yaml

@@ -43,6 +43,9 @@ jobs:
       - name: Run tests shinkai-message-wasm cargo tests
         run: docker run --rm --entrypoint /entrypoints/run-wasm-pack-tests.sh testing_image_wasm:${SHORT_SHA}
 
+      - name: Run tests shinkai-message-pyo3 python tests
+        run: docker run --rm --entrypoint /entrypoints/run-python-tests.sh testing_image_wasm:${SHORT_SHA}        
+
       - name: Remove image
         if: always()
         run: docker rmi testing_image_wasm:${SHORT_SHA} || true

.gitignore

@@ -25,3 +25,5 @@ qr_code_profile_*.png
 qr_code_device_*.png
 .venv/
 shinkai-libs/shinkai-message-pyo3/tests/__pycache__
+output.wav
+shinkai-libs/shinkai-message-pyo3/venv

Cargo.toml

@@ -47,7 +47,7 @@ mupdf = "0.4.2"
 qrcode = "0.12"
 image = "0.23"
 urlencoding = "1.1.1"
-hex = "=0.4.2"
+hex = "=0.4.3"
 
 
 [dependencies.rocksdb]

Dockerfile.shinkai-app

@@ -1,5 +1,11 @@
 FROM rust:bookworm
 
+# Install Python3 and pip
+RUN apt-get update && apt-get install -y python3 python3-pip
+
+# Install maturin
+RUN pip3 install maturin
+
 # Install nvm 
 RUN rm /bin/sh && ln -s /bin/bash /bin/sh
 ENV NVM_DIR /usr/local/nvm
@@ -24,3 +30,6 @@ COPY . .
 
 WORKDIR /app/shinkai-libs/shinkai-message-wasm
 RUN wasm-pack build
+
+WORKDIR /app/shinkai-libs/shinkai-message-pyo3
+RUN maturin build

README.md

@@ -84,4 +84,10 @@ You need to compile the wasm library from `shinkai-message-wasm` and copy the re
 npm run test.unit
 ```
 
+### Shinkai PYO3 Tests
+
+
+### Further CI Development
+
+Use `act -j test-wasm -P self-hosted=nektos/act-environments-ubuntu:18.04 --container-architecture linux/amd64` to run the tests locally in a docker container. This is useful for debugging CI issues.
 

scripts/run_node1.sh

@@ -1,31 +1,16 @@
 #!/bin/bash
 
 export NODE_IP="0.0.0.0"
-export NODE_PORT="8080"
+export NODE_PORT="9552"
 export NODE_API_IP="0.0.0.0"
-export NODE_API_PORT="13013"
+export NODE_API_PORT="9550"
 export IDENTITY_SECRET_KEY="df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119"
 export ENCRYPTION_SECRET_KEY="d83f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81159"
 export PING_INTERVAL_SECS="0"
 export GLOBAL_IDENTITY_NAME="@@node1.shinkai"
-export RUST_LOG=warn,error,info,debug
+export RUST_LOG=warn,error,info
 export STARTING_NUM_QR_PROFILES="1"
 export STARTING_NUM_QR_DEVICES="1"
-
-if [ -z "$1" ]
-then
-  echo "No argument supplied for CONNECT_PK, using empty string"
-  export CONNECT_PK=""
-else
-  export CONNECT_PK=$1
-fi
-
-if [ -z "$2" ]
-then
-  echo "No argument supplied for CONNECT_ADDR, using empty string"
-  export CONNECT_ADDR=""
-else
-  export CONNECT_ADDR=$2
-fi
+export FIRST_DEVICE_NEEDS_REGISTRATION_CODE="false"
 
 cargo run

shinkai-libs/shinkai-message-primitives/Cargo.toml

@@ -22,7 +22,7 @@ anyhow = "1.0.72"
 thiserror = "1.0.44"
 console_log = "0.2"
 log = "0.4.14"
-hex = "=0.4.2"
+hex = "=0.4.3"
 
 [dependencies.serde]
 version = "1.0.188"

shinkai-libs/shinkai-message-primitives/src/shinkai_utils/shinkai_message_builder.rs

@@ -599,6 +599,53 @@ impl ShinkaiMessageBuilder {
         )
     }
 
+    pub fn initial_registration_with_no_code_for_device(
+        my_device_encryption_sk: EncryptionStaticKey,
+        my_device_signature_sk: SignatureStaticKey,
+        profile_encryption_sk: EncryptionStaticKey,
+        profile_signature_sk: SignatureStaticKey,
+        registration_name: String,
+        sender_subidentity: String,
+        sender: ProfileName,
+        receiver: ProfileName,
+    ) -> Result<ShinkaiMessage, &'static str> {
+        let my_device_signature_pk = ed25519_dalek::PublicKey::from(&my_device_signature_sk);
+        let my_device_encryption_pk = x25519_dalek::PublicKey::from(&my_device_encryption_sk);
+        let profile_signature_pk = ed25519_dalek::PublicKey::from(&profile_signature_sk);
+        let profile_encryption_pk = x25519_dalek::PublicKey::from(&profile_encryption_sk);
+        let other = encryption_public_key_to_string(my_device_encryption_pk);
+
+        let identity_type = "device".to_string();
+        let permission_type = "admin".to_string();
+
+        let registration_code = RegistrationCode {
+            code: "".to_string(),
+            registration_name: registration_name.clone(),
+            device_identity_pk: signature_public_key_to_string(my_device_signature_pk),
+            device_encryption_pk: other.clone(),
+            profile_identity_pk: signature_public_key_to_string(profile_signature_pk),
+            profile_encryption_pk: encryption_public_key_to_string(profile_encryption_pk),
+            identity_type,
+            permission_type,
+        };
+
+        let body = serde_json::to_string(&registration_code).map_err(|_| "Failed to serialize data to JSON")?;
+        let other = encryption_public_key_to_string(my_device_encryption_pk.clone());
+
+        ShinkaiMessageBuilder::new(my_device_encryption_sk, my_device_signature_sk, my_device_encryption_pk)
+            .message_raw_content(body)
+            .body_encryption(EncryptionMethod::None)
+            .internal_metadata_with_schema(
+                sender_subidentity,
+                "".to_string(),
+                "".to_string(),
+                MessageSchemaType::UseRegistrationCode,
+                EncryptionMethod::None,
+            )
+            .external_metadata_with_other(receiver.clone(), sender, other)
+            .build()
+    }
+
     pub fn get_all_inboxes_for_profile(
         my_subidentity_encryption_sk: EncryptionStaticKey,
         my_subidentity_signature_sk: SignatureStaticKey,
@@ -1036,6 +1083,63 @@ mod tests {
         assert!(verify_signature(&my_identity_pk, &message_clone).unwrap())
     }
 
+    #[test]
+    fn test_initial_registration_with_no_code_for_device() {
+        let (my_device_identity_sk, my_device_identity_pk) = unsafe_deterministic_signature_keypair(0);
+        let (my_device_encryption_sk, my_device_encryption_pk) = unsafe_deterministic_encryption_keypair(0);
+
+        let (profile_identity_sk, profile_identity_pk) = unsafe_deterministic_signature_keypair(1);
+        let (profile_encryption_sk, profile_encryption_pk) = unsafe_deterministic_encryption_keypair(1);
+
+        let recipient = "@@other_node.shinkai".to_string();
+        let sender = recipient.clone();
+        let sender_subidentity = "main".to_string();
+
+        let registration_name = "registration_name".to_string();
+
+        let message_result = ShinkaiMessageBuilder::initial_registration_with_no_code_for_device(
+            my_device_encryption_sk.clone(),
+            my_device_identity_sk,
+            profile_encryption_sk,
+            profile_identity_sk,
+            registration_name.clone(),
+            sender_subidentity.clone(),
+            sender.clone(),
+            recipient.clone(),
+        );
+        println!("message_result: {:?}", message_result);
+        assert!(message_result.is_ok());
+        let message = message_result.unwrap();
+
+        assert_eq!(message.encryption, EncryptionMethod::None);
+
+        if let MessageBody::Unencrypted(shinkai_body) = message.body {
+            if let MessageData::Unencrypted(shinkai_data) = shinkai_body.message_data {
+                // Parse the decrypted content from a JSON string to a RegistrationCode struct
+                let parsed_content: Result<RegistrationCode, _> =
+                    serde_json::from_str(&shinkai_data.message_raw_content);
+                match &parsed_content {
+                    Ok(registration_code) => {
+                        println!("Parsed content: {:?}", registration_code);
+                    }
+                    Err(e) => {
+                        eprintln!("Failed to parse content: {:?}", e);
+                    }
+                }
+
+                let registration_code = parsed_content.unwrap();
+                assert_eq!(registration_code.code, "");
+                assert_eq!(registration_code.registration_name, registration_name);
+                assert_eq!(registration_code.permission_type, "admin");
+                assert_eq!(registration_code.identity_type, "device");
+            }
+            assert_eq!(shinkai_body.internal_metadata.sender_subidentity, sender_subidentity);
+        }
+
+        let external_metadata = message.external_metadata;
+        assert_eq!(external_metadata.sender, recipient);
+    }
+
     #[test]
     fn test_builder_missing_fields() {
         let (my_identity_sk, my_identity_pk) = unsafe_deterministic_signature_keypair(0);

shinkai-libs/shinkai-message-pyo3/Cargo.toml

@@ -19,7 +19,7 @@ x25519-dalek = "1.2.0"
 ed25519-dalek = "1.0.1"
 rand_os = "0.2.2"
 rand = "0.8"
-hex = "=0.4.2"
+hex = "=0.4.3"
 
 [dependencies.pyo3]
 version = "0.19.2"

shinkai-libs/shinkai-message-pyo3/run_tests.sh

@@ -1,10 +1,16 @@
 #!/bin/bash
 
+# Create a virtual environment if it doesn't exist
+if [ ! -d "./venv" ]
+then
+    python -m venv venv
+fi
+
 # Activate your virtual environment
 source ./venv/bin/activate
 
 # Run maturin develop and capture its output
-output=$(maturin develop)
+output=$(maturin build -i python)
 
 # If maturin develop is successful, extract the path of the built wheel file
 if [ $? -eq 0 ]; then

shinkai-libs/shinkai-message-pyo3/src/shinkai_pyo3_wrapper/shinkai_builder_pyo3.rs

@@ -7,7 +7,7 @@ use pyo3::{prelude::*, pyclass, types::PyDict, PyResult};
 use shinkai_message_primitives::{
     schemas::{agents::serialized_agent::SerializedAgent, inbox_name::InboxName, registration_code::RegistrationCode},
     shinkai_message::shinkai_message_schemas::{
-        APIAddAgentRequest, APIGetMessagesFromInboxRequest, APIReadUpToTimeRequest, IdentityPermissions, JobCreation,
+        APIAddAgentRequest, APIGetMessagesFromInboxRequest, APIReadUpToTimeRequest, IdentityPermissions, JobCreationInfo,
         JobScope, MessageSchemaType, RegistrationCodeRequest, RegistrationCodeType, JobMessage,
     },
     shinkai_utils::{
@@ -818,7 +818,7 @@ impl PyShinkaiMessageBuilder {
                 }
             };
 
-            let job_creation = JobCreation { scope: scope.inner };
+            let job_creation = JobCreationInfo { scope: scope.inner };
             let body = match serde_json::to_string(&job_creation) {
                 Ok(body) => body,
                 Err(e) => return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string().clone())),