Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
nicarq committed Jan 18, 2024
1 parent 79b4a4a commit 1440714
Show file tree
Hide file tree
Showing 6 changed files with 47 additions and 259 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,4 @@ db/.secret
./storage/*
vector_fs_db
storage
.vscode/*
5 changes: 0 additions & 5 deletions .vscode/settings.json

This file was deleted.

2 changes: 1 addition & 1 deletion shinkai-libs/shinkai-message-pyo3/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion shinkai-libs/shinkai-message-pyo3/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "shinkai_message_pyo3"
version = "0.1.5"
version = "0.1.6"
edition = "2018"
authors = ["Nico Arqueros <[email protected]>"]

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use shinkai_message_primitives::{
};

use crate::shinkai_pyo3_utils::{pyo3_job_scope::PyJobScope, pyo3_agent_llm_interface::PyAgentLLMInterface, pyo3_serialized_agent::PySerializedAgent, pyo3_shinkai_name::PyShinkaiName};
use super::{shinkai_builder_pyo3::PyShinkaiMessageBuilder, encryption_method_pyo3::PyEncryptionMethod};
use super::{shinkai_builder_pyo3::PyShinkaiMessageBuilder, encryption_method_pyo3::PyEncryptionMethod, message_schema_type_pyo3::PyMessageSchemaType};

#[pymodule]
fn shinkai_message_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
Expand All @@ -20,6 +20,7 @@ fn shinkai_message_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<PyAgentLLMInterface>()?;
m.add_class::<PySerializedAgent>()?;
m.add_class::<PyShinkaiName>()?;
m.add_class::<PyMessageSchemaType>()?;
// Add any other classes or functions you want to expose to Python here
Ok(())
}
Expand All @@ -38,254 +39,3 @@ pub struct PyExternalMetadata {
pub struct PyInternalMetadata {
pub inner: InternalMetadata,
}

// #[pymethods]
// impl PyShinkaiMessage {
// #[new]
// #[args(kwargs = "**")]
// fn new(kwargs: Option<&PyDict>) -> PyResult<Self> {
// // Initialize your PyShinkaiMessage here using kwargs if necessary
// // For now, let's just use the default
// Ok(Self {
// inner: Default::default(),
// })
// }

// #[getter]
// fn get_external_metadata(&self) -> PyResult<PyExternalMetadata> {
// Ok(PyExternalMetadata {
// inner: self.inner.external_metadata.clone(),
// })
// }

// #[setter]
// fn set_external_metadata(&mut self, py_external_metadata: PyExternalMetadata) {
// self.inner.external_metadata = py_external_metadata.inner;
// }

// // ... rest of your methods here, replacing [self](file:///Users/nijaar/shinkai/develop/shinkai-node/src/network/node_error.rs#11%2C13-11%2C13) with `self.inner`
// }

// // Previous Approach 2

// #[pymethods]
// impl ShinkaiMessage {
// #[getter]
// fn get_body(&self) -> PyResult<MessageBody> {
// Ok(self.body.clone())
// }

// #[setter]
// fn set_body(&mut self, body: MessageBody) {
// self.body = body;
// }

// #[getter]
// fn get_external_metadata(&self) -> PyResult<ExternalMetadata> {
// Ok(self.external_metadata.clone())
// }

// #[setter]
// fn set_external_metadata(&mut self, external_metadata: ExternalMetadata) {
// self.external_metadata = external_metadata;
// }

// #[getter]
// fn get_encryption(&self) -> PyResult<EncryptionMethod> {
// Ok(self.encryption.clone())
// }

// #[setter]
// fn set_encryption(&mut self, encryption: EncryptionMethod) {
// self.encryption = encryption;
// }

// #[getter]
// fn get_version(&self) -> PyResult<ShinkaiVersion> {
// Ok(self.version.clone())
// }

// #[setter]
// fn set_version(&mut self, version: ShinkaiVersion) {
// self.version = version;
// }
// }

// Previous Approach
//
// #[pyclass]
// #[derive(Clone)]
// struct PyEncryptionMethod {
// value: EncryptionMethod,
// }

// #[pymethods]
// impl PyEncryptionMethod {
// #[new]
// fn new(value: String) -> Self {
// let method = match value.as_str() {
// "DiffieHellmanChaChaPoly1305" | "default" => EncryptionMethod::DiffieHellmanChaChaPoly1305,
// _ => EncryptionMethod::None,
// };

// PyEncryptionMethod { value: method }
// }

// #[getter]
// fn value(&self) -> String {
// match self.value {
// EncryptionMethod::DiffieHellmanChaChaPoly1305 => String::from("DiffieHellmanChaChaPoly1305"),
// EncryptionMethod::None => String::from("None"),
// }
// }
// }

// #[pyclass]
// #[derive(Clone)]
// struct PyShinkaiVersion {
// value: ShinkaiVersion,
// }

// #[pyclass]
// #[derive(Clone)]
// pub struct PyMessageBody {
// value: MessageBody,
// }

// #[pyclass]
// #[derive(Clone)]
// struct PyInternalMetadata {
// #[pyo3(get, set)]
// sender_subidentity: String,
// #[pyo3(get, set)]
// recipient_subidentity: String,
// #[pyo3(get, set)]
// inbox: String,
// #[pyo3(get, set)]
// signature: String,
// #[pyo3(get, set)]
// encryption: PyEncryptionMethod, // Assuming you have this class defined
// }

// #[pyclass]
// #[derive(Clone)]
// struct PyExternalMetadata {
// #[pyo3(get, set)]
// sender: String,
// #[pyo3(get, set)]
// recipient: String,
// #[pyo3(get, set)]
// scheduled_time: String,
// #[pyo3(get, set)]
// signature: String,
// #[pyo3(get, set)]
// other: String,
// }

// #[pyclass]
// #[derive(Clone)]
// struct PyShinkaiMessage {
// #[pyo3(get, set)]
// body: PyMessageBody,
// #[pyo3(get, set)]
// external_metadata: PyExternalMetadata,
// #[pyo3(get, set)]
// encryption: PyEncryptionMethod,
// #[pyo3(get, set)]
// version: PyShinkaiVersion,
// }

// #[pymodule]
// fn shinkai_message_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
// #[pymethods]
// impl PyInternalMetadata {
// #[new]
// fn new(
// sender_subidentity: String,
// recipient_subidentity: String,
// inbox: String,
// signature: String,
// encryption: PyEncryptionMethod,
// ) -> Self {
// PyInternalMetadata {
// sender_subidentity,
// recipient_subidentity,
// inbox,
// signature,
// encryption,
// }
// }
// }

// #[pymethods]
// impl PyExternalMetadata {
// #[new]
// fn new(sender: String, recipient: String, scheduled_time: String, signature: String, other: String) -> Self {
// PyExternalMetadata {
// sender,
// recipient,
// scheduled_time,
// signature,
// other,
// }
// }
// }

// #[pymethods]
// impl PyShinkaiVersion {
// #[new]
// #[args(value = "String::from(\"V1_0\")")]
// fn new(value: String) -> Self {
// let version = match value.as_str() {
// "V1_0" => ShinkaiVersion::V1_0,
// _ => ShinkaiVersion::Unsupported,
// };

// PyShinkaiVersion { value: version }
// }

// #[getter]
// fn value(&self) -> String {
// match self.value {
// ShinkaiVersion::V1_0 => String::from("V1_0"),
// ShinkaiVersion::Unsupported => String::from("Unsupported"),
// }
// }
// }

// Ok(())
// }

// #[pymethods]
// impl PyMessageBody {
// #[new]
// fn new(value: String, body: PyObject) -> PyResult<Self> {
// let py = body.py();
// let value = match value.as_str() {
// "encrypted" => MessageBody::Encrypted(body.extract::<EncryptedShinkaiBody>(py)?),
// "unencrypted" => MessageBody::Unencrypted(body.extract::<ShinkaiBody>(py)?),
// _ => return Err(PyValueError::new_err("Invalid value")),
// };

// Ok(PyMessageBody { value })
// }

// #[getter]
// fn value(&self) -> PyResult<String> {
// match &self.value {
// MessageBody::Encrypted(_) => Ok(String::from("encrypted")),
// MessageBody::Unencrypted(_) => Ok(String::from("unencrypted")),
// }
// }

// #[getter]
// fn body(&self) -> PyResult<PyObject> {
// let gil_guard = Python::acquire_gil();
// let py = gil_guard.python();

// match &self.value {
// MessageBody::Encrypted(body) => Ok(body.into_py(py)),
// MessageBody::Unencrypted(body) => Ok(body.into_py(py)),
// }
// }
// }
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,48 @@ def test_get_all_inboxes_for_profile(self):
self.assertEqual(result_json["encryption"], "DiffieHellmanChaChaPoly1305")
self.assertEqual(result_json["version"], "V1_0")

def test_custom_job_message(self):
my_encryption_sk_string = '7008829b80ae4350cf049e48d8bce4714e216b674fff0bf34f97f7b98d928d3f'
my_identity_sk_string = 'b6baf0fa268f993c57223d5db96e5e1de776fcb0195ee6137f33de9d8d9dd749'
receiver_public_key = '798cbd64d78c4a0fba338b2a6349634940dc4e5b601db1029e02c41e0fe05679'
node = "@@node1.shinkai"
job_id = "job1"
content = "Job content"
files_inbox = ""
data = json.dumps({
"job_id": job_id,
"content": content,
"files_inbox": files_inbox
}) # Simulating JobMessage data
sender = node
sender_subidentity = ""
recipient = node
recipient_subidentity = "main/agent/agent_1"
other = "job_inbox::job1::false"
schema = shinkai_message_pyo3.PyMessageSchemaType("JobMessageSchema")

result = shinkai_message_pyo3.PyShinkaiMessageBuilder.create_custom_shinkai_message_to_node(
my_encryption_sk_string,
my_identity_sk_string,
receiver_public_key,
data,
sender,
sender_subidentity,
recipient,
recipient_subidentity,
other,
schema
)

# print("Result:", result)
result_json = json.loads(result)

self.assertTrue("encrypted" in result_json["body"])
self.assertEqual(result_json["external_metadata"]["sender"], sender)
self.assertEqual(result_json["external_metadata"]["recipient"], recipient)
self.assertEqual(result_json["external_metadata"]["other"], other)
self.assertEqual(result_json["encryption"], "DiffieHellmanChaChaPoly1305")
self.assertEqual(result_json["version"], "V1_0")

if __name__ == '__main__':
unittest.main()

0 comments on commit 1440714

Please sign in to comment.