kinfer 0.4.3__cp311-cp311-macosx_11_0_arm64.whl → 0.5.2__cp311-cp311-macosx_11_0_arm64.whl

kinfer/rust/src/types.rs

@@ -0,0 +1,86 @@
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ModelMetadata {
+    pub joint_names: Vec<String>,
+    pub num_commands: Option<usize>,
+    pub carry_size: Vec<usize>,
+}
+
+impl ModelMetadata {
+    pub fn model_validate_json(json: String) -> Result<Self, Box<dyn std::error::Error>> {
+        Ok(serde_json::from_str(&json)?)
+    }
+
+    pub fn to_json(&self) -> Result<String, Box<dyn std::error::Error>> {
+        Ok(serde_json::to_string(self)?)
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
+pub enum InputType {
+    JointAngles,
+    JointAngularVelocities,
+    ProjectedGravity,
+    Accelerometer,
+    Gyroscope,
+    Command,
+    Time,
+    Carry,
+}
+
+impl InputType {
+    pub fn get_name(&self) -> &str {
+        match self {
+            InputType::JointAngles => "joint_angles",
+            InputType::JointAngularVelocities => "joint_angular_velocities",
+            InputType::ProjectedGravity => "projected_gravity",
+            InputType::Accelerometer => "accelerometer",
+            InputType::Gyroscope => "gyroscope",
+            InputType::Command => "command",
+            InputType::Time => "time",
+            InputType::Carry => "carry",
+        }
+    }
+
+    pub fn get_shape(&self, metadata: &ModelMetadata) -> Vec<usize> {
+        match self {
+            InputType::JointAngles => vec![metadata.joint_names.len()],
+            InputType::JointAngularVelocities => vec![metadata.joint_names.len()],
+            InputType::ProjectedGravity => vec![3],
+            InputType::Accelerometer => vec![3],
+            InputType::Gyroscope => vec![3],
+            InputType::Command => vec![metadata.num_commands.unwrap_or(0)],
+            InputType::Time => vec![1],
+            InputType::Carry => metadata.carry_size.clone(),
+        }
+    }
+
+    pub fn from_name(name: &str) -> Result<Self, Box<dyn std::error::Error>> {
+        match name {
+            "joint_angles" => Ok(InputType::JointAngles),
+            "joint_angular_velocities" => Ok(InputType::JointAngularVelocities),
+            "projected_gravity" => Ok(InputType::ProjectedGravity),
+            "accelerometer" => Ok(InputType::Accelerometer),
+            "gyroscope" => Ok(InputType::Gyroscope),
+            "command" => Ok(InputType::Command),
+            "time" => Ok(InputType::Time),
+            "carry" => Ok(InputType::Carry),
+            _ => Err(format!("Unknown input type: {}", name).into()),
+        }
+    }
+
+    pub fn get_names() -> Vec<&'static str> {
+        vec![
+            "joint_angles",
+            "joint_angular_velocities",
+            "projected_gravity",
+            "accelerometer",
+            "gyroscope",
+            "command",
+            "time",
+            "carry",
+        ]
+    }
+}

kinfer/rust_bindings/rust_bindings.pyi

@@ -8,14 +8,21 @@ import typing
 
 class ModelProviderABC:
     def __new__(cls) -> ModelProviderABC: ...
-    def get_joint_angles(self, joint_names:typing.Sequence[builtins.str]) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_joint_angular_velocities(self, joint_names:typing.Sequence[builtins.str]) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_projected_gravity(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_time(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
+    def get_inputs(self, input_types:typing.Sequence[builtins.str], metadata:PyModelMetadata) -> builtins.dict[builtins.str, numpy.typing.NDArray[numpy.float32]]: ...
+    def take_action(self, action:numpy.typing.NDArray[numpy.float32], metadata:PyModelMetadata) -> None: ...
+
+class PyInputType:
+    def __new__(cls, input_type:builtins.str) -> PyInputType: ...
+    def get_name(self) -> builtins.str: ...
+    def get_shape(self, metadata:PyModelMetadata) -> builtins.list[builtins.int]: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
+
+class PyModelMetadata:
+    def __new__(self, joint_names:typing.Sequence[builtins.str], num_commands:typing.Optional[builtins.int], carry_size:typing.Sequence[builtins.int]) -> PyModelMetadata: ...
+    def to_json(self) -> builtins.str: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
 
 class PyModelProvider:
     ...
@@ -35,3 +42,5 @@ class PyModelRuntime:
 
 def get_version() -> builtins.str: ...
 
+def metadata_from_json(json:builtins.str) -> PyModelMetadata: ...
+

kinfer/rust_bindings/src/lib.rs

@@ -1,89 +1,210 @@
 use async_trait::async_trait;
 use kinfer::model::{ModelError, ModelProvider, ModelRunner};
 use kinfer::runtime::ModelRuntime;
+use kinfer::types::{InputType, ModelMetadata};
 use ndarray::{Array, Ix1, IxDyn};
 use numpy::{PyArray1, PyArrayDyn, PyArrayMethods};
 use pyo3::exceptions::PyNotImplementedError;
 use pyo3::prelude::*;
+use pyo3::types::{PyAny, PyAnyMethods};
 use pyo3::{pymodule, types::PyModule, Bound, PyResult, Python};
 use pyo3_stub_gen::define_stub_info_gatherer;
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods};
+use std::collections::HashMap;
+use std::hash::Hash;
 use std::sync::Arc;
 use std::sync::Mutex;
-use std::time::Instant;
+
+type StepResult = (Py<PyArrayDyn<f32>>, Py<PyArrayDyn<f32>>);
+
 #[pyfunction]
 #[gen_stub_pyfunction]
 fn get_version() -> String {
     env!("CARGO_PKG_VERSION").to_string()
 }
 
-#[pyclass(subclass)]
+#[pyclass]
 #[gen_stub_pyclass]
-pub struct ModelProviderABC;
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct PyInputType {
+    pub input_type: InputType,
+}
+
+impl From<InputType> for PyInputType {
+    fn from(input_type: InputType) -> Self {
+        Self { input_type }
+    }
+}
+
+impl From<PyInputType> for InputType {
+    fn from(input_type: PyInputType) -> Self {
+        input_type.input_type
+    }
+}
 
 #[gen_stub_pymethods]
 #[pymethods]
-impl ModelProviderABC {
+impl PyInputType {
     #[new]
-    fn new() -> Self {
-        ModelProviderABC
+    fn __new__(input_type: &str) -> PyResult<Self> {
+        let input_type = InputType::from_name(input_type).map_or_else(
+            |_| {
+                Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                    "Invalid input type: {} (must be one of {})",
+                    input_type,
+                    InputType::get_names().join(", "),
+                )))
+            },
+            Ok,
+        )?;
+        Ok(Self { input_type })
     }
 
-    fn get_joint_angles<'py>(
-        &self,
-        joint_names: Vec<String>,
-    ) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        let n = joint_names.len();
-        Err(PyNotImplementedError::new_err(format!(
-            "Must override get_joint_angles with {} joint names",
-            n
-        )))
+    fn get_name(&self) -> String {
+        self.input_type.get_name().to_string()
     }
 
-    fn get_joint_angular_velocities<'py>(
-        &self,
+    fn get_shape(&self, metadata: PyModelMetadata) -> Vec<usize> {
+        self.input_type.get_shape(&metadata.into())
+    }
+
+    fn __repr__(&self) -> String {
+        format!("InputType({})", self.get_name())
+    }
+
+    fn __eq__(&self, other: Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyInputType>() {
+            Ok(self == &other)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+#[pyclass]
+#[gen_stub_pyclass]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct PyModelMetadata {
+    #[pyo3(get, set)]
+    pub joint_names: Vec<String>,
+    #[pyo3(get, set)]
+    pub num_commands: Option<usize>,
+    #[pyo3(get, set)]
+    pub carry_size: Vec<usize>,
+}
+
+#[pymethods]
+#[gen_stub_pymethods]
+impl PyModelMetadata {
+    #[new]
+    fn __new__(
         joint_names: Vec<String>,
-    ) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        let n = joint_names.len();
-        Err(PyNotImplementedError::new_err(format!(
-            "Must override get_joint_angular_velocities with {} joint names",
-            n
-        )))
+        num_commands: Option<usize>,
+        carry_size: Vec<usize>,
+    ) -> Self {
+        Self {
+            joint_names,
+            num_commands,
+            carry_size,
+        }
     }
 
-    fn get_projected_gravity<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        Err(PyNotImplementedError::new_err(
-            "Must override get_projected_gravity",
-        ))
+    fn to_json(&self) -> PyResult<String> {
+        let metadata = ModelMetadata {
+            joint_names: self.joint_names.clone(),
+            num_commands: self.num_commands,
+            carry_size: self.carry_size.clone(),
+        }
+        .to_json()
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        Ok(metadata)
     }
 
-    fn get_accelerometer<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        Err(PyNotImplementedError::new_err(
-            "Must override get_accelerometer",
-        ))
+    fn __repr__(&self) -> PyResult<String> {
+        let json = self.to_json()?;
+        Ok(format!("ModelMetadata({:?})", json))
+    }
+
+    fn __eq__(&self, other: Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyModelMetadata>() {
+            Ok(self == &other)
+        } else {
+            Ok(false)
+        }
     }
+}
 
-    fn get_gyroscope<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        Err(PyNotImplementedError::new_err(
-            "Must override get_gyroscope",
-        ))
+#[pyfunction]
+#[gen_stub_pyfunction]
+fn metadata_from_json(json: &str) -> PyResult<PyModelMetadata> {
+    let metadata = ModelMetadata::model_validate_json(json.to_string()).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid model metadata: {}", e))
+    })?;
+    Ok(PyModelMetadata::from(metadata))
+}
+
+impl From<ModelMetadata> for PyModelMetadata {
+    fn from(metadata: ModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names,
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size,
+        }
     }
+}
 
-    fn get_command<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        Err(PyNotImplementedError::new_err("Must override get_command"))
+impl From<&ModelMetadata> for PyModelMetadata {
+    fn from(metadata: &ModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names.clone(),
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size.clone(),
+        }
     }
+}
 
-    fn get_time<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
-        Err(PyNotImplementedError::new_err("Must override get_time"))
+impl From<PyModelMetadata> for ModelMetadata {
+    fn from(metadata: PyModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names,
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size,
+        }
     }
+}
 
-    fn take_action<'py>(
+#[pyclass(subclass)]
+#[gen_stub_pyclass]
+struct ModelProviderABC;
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl ModelProviderABC {
+    #[new]
+    fn __new__() -> Self {
+        ModelProviderABC
+    }
+
+    fn get_inputs<'py>(
         &self,
-        joint_names: Vec<String>,
-        action: Bound<'py, PyArray1<f32>>,
+        input_types: Vec<String>,
+        metadata: PyModelMetadata,
+    ) -> PyResult<HashMap<String, Bound<'py, PyArrayDyn<f32>>>> {
+        Err(PyNotImplementedError::new_err(format!(
+            "Must override get_inputs with {} input types {:?} and metadata {:?}",
+            input_types.len(),
+            input_types,
+            metadata
+        )))
+    }
+
+    fn take_action(
+        &self,
+        action: Bound<'_, PyArray1<f32>>,
+        metadata: PyModelMetadata,
     ) -> PyResult<()> {
         let n = action.len()?;
-        assert_eq!(joint_names.len(), n);
+        assert_eq!(metadata.joint_names.len(), n); // TODO: this is wrong
         Err(PyNotImplementedError::new_err(format!(
             "Must override take_action with {} action",
             n
@@ -96,127 +217,62 @@ impl ModelProviderABC {
 #[derive(Clone)]
 struct PyModelProvider {
     obj: Arc<Py<ModelProviderABC>>,
-    start_time: Instant,
 }
 
 #[pymethods]
 impl PyModelProvider {
     #[new]
-    fn new(obj: Py<ModelProviderABC>) -> Self {
-        Self {
-            obj: Arc::new(obj),
-            start_time: Instant::now(),
-        }
+    fn __new__(obj: Py<ModelProviderABC>) -> Self {
+        Self { obj: Arc::new(obj) }
     }
 }
 
 #[async_trait]
 impl ModelProvider for PyModelProvider {
-    async fn get_joint_angles(
-        &self,
-        joint_names: &[String],
-    ) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = (joint_names,);
-            let result = obj.call_method(py, "get_joint_angles", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_joint_angular_velocities(
+    async fn get_inputs(
         &self,
-        joint_names: &[String],
-    ) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = (joint_names,);
-            let result = obj.call_method(py, "get_joint_angular_velocities", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_projected_gravity(&self) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = ();
-            let result = obj.call_method(py, "get_projected_gravity", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_accelerometer(&self) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = ();
-            let result = obj.call_method(py, "get_accelerometer", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_gyroscope(&self) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = ();
-            let result = obj.call_method(py, "get_gyroscope", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_command(&self) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
-            let obj = self.obj.clone();
-            let args = ();
-            let result = obj.call_method(py, "get_command", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
-        })
-        .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_time(&self) -> Result<Array<f32, IxDyn>, ModelError> {
-        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
+        input_types: &[InputType],
+        metadata: &ModelMetadata,
+    ) -> Result<HashMap<InputType, Array<f32, IxDyn>>, ModelError> {
+        let input_names: Vec<String> = input_types
+            .iter()
+            .map(|t| t.get_name().to_string())
+            .collect();
+        let result = Python::with_gil(|py| -> PyResult<HashMap<InputType, Array<f32, IxDyn>>> {
             let obj = self.obj.clone();
-            let args = ();
-            let result = obj.call_method(py, "get_time", args, None)?;
-            let array = result.extract::<Vec<f32>>(py)?;
-            Ok(Array::from_vec(array).into_dyn())
+            let args = (input_names.clone(), PyModelMetadata::from(metadata.clone()));
+            let result = obj.call_method(py, "get_inputs", args, None)?;
+            let dict: HashMap<String, Vec<f32>> = result.extract(py)?;
+            let mut arrays = HashMap::new();
+            for (i, name) in input_names.iter().enumerate() {
+                let array = dict.get(name).ok_or_else(|| {
+                    PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                        "Missing input: {}",
+                        name
+                    ))
+                })?;
+                arrays.insert(input_types[i], Array::from_vec(array.clone()).into_dyn());
+            }
+            Ok(arrays)
         })
         .map_err(|e| ModelError::Provider(e.to_string()))?;
-        Ok(args)
-    }
-
-    async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError> {
-        Ok(carry)
+        Ok(result)
     }
 
     async fn take_action(
         &self,
-        joint_names: Vec<String>,
         action: Array<f32, IxDyn>,
+        metadata: &ModelMetadata,
     ) -> Result<(), ModelError> {
         Python::with_gil(|py| -> PyResult<()> {
             let obj = self.obj.clone();
             let action_1d = action
                 .into_dimensionality::<Ix1>()
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
-            let args = (joint_names, PyArray1::from_array(py, &action_1d));
+            let args = (
+                PyArray1::from_array(py, &action_1d),
+                PyModelMetadata::from(metadata.clone()),
+            );
             obj.call_method(py, "take_action", args, None)?;
             Ok(())
         })
@@ -236,10 +292,10 @@ struct PyModelRunner {
 #[pymethods]
 impl PyModelRunner {
     #[new]
-    fn new(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
-        let input_provider = Arc::new(PyModelProvider::new(provider));
+    fn __new__(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
+        let input_provider = Arc::new(PyModelProvider::__new__(provider));
 
-        let runner = tokio::runtime::Runtime::new().unwrap().block_on(async {
+        let runner = tokio::runtime::Runtime::new()?.block_on(async {
             ModelRunner::new(model_path, input_provider)
                 .await
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
@@ -252,7 +308,7 @@ impl PyModelRunner {
 
     fn init(&self) -> PyResult<Py<PyArrayDyn<f32>>> {
         let runner = self.runner.clone();
-        let result = tokio::runtime::Runtime::new().unwrap().block_on(async {
+        let result = tokio::runtime::Runtime::new()?.block_on(async {
             runner
                 .init()
                 .await
@@ -265,17 +321,14 @@ impl PyModelRunner {
         })
     }
 
-    fn step(
-        &self,
-        carry: Py<PyArrayDyn<f32>>,
-    ) -> PyResult<(Py<PyArrayDyn<f32>>, Py<PyArrayDyn<f32>>)> {
+    fn step(&self, carry: Py<PyArrayDyn<f32>>) -> PyResult<StepResult> {
         let runner = self.runner.clone();
         let carry_array = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
             let carry_array = carry.bind(py);
             Ok(carry_array.to_owned_array())
         })?;
 
-        let result = tokio::runtime::Runtime::new().unwrap().block_on(async {
+        let result = tokio::runtime::Runtime::new()?.block_on(async {
             runner
                 .step(carry_array)
                 .await
@@ -297,7 +350,7 @@ impl PyModelRunner {
             Ok(action_array.to_owned_array())
         })?;
 
-        tokio::runtime::Runtime::new().unwrap().block_on(async {
+        tokio::runtime::Runtime::new()?.block_on(async {
             runner
                 .take_action(action_array)
                 .await
@@ -319,38 +372,56 @@ struct PyModelRuntime {
 #[pymethods]
 impl PyModelRuntime {
     #[new]
-    fn new(model_runner: PyModelRunner, dt: u64) -> PyResult<Self> {
+    fn __new__(model_runner: PyModelRunner, dt: u64) -> PyResult<Self> {
         Ok(Self {
             runtime: Arc::new(Mutex::new(ModelRuntime::new(model_runner.runner, dt))),
         })
     }
 
-    fn set_slowdown_factor(&self, slowdown_factor: i32) {
-        let mut runtime = self.runtime.lock().unwrap();
+    fn set_slowdown_factor(&self, slowdown_factor: i32) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
         runtime.set_slowdown_factor(slowdown_factor);
+        Ok(())
     }
 
-    fn set_magnitude_factor(&self, magnitude_factor: f32) {
-        let mut runtime = self.runtime.lock().unwrap();
+    fn set_magnitude_factor(&self, magnitude_factor: f32) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
         runtime.set_magnitude_factor(magnitude_factor);
+        Ok(())
     }
 
     fn start(&self) -> PyResult<()> {
-        let mut runtime = self.runtime.lock().unwrap();
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
         runtime
             .start()
             .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
     }
 
-    fn stop(&self) {
-        let mut runtime = self.runtime.lock().unwrap();
+    fn stop(&self) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
         runtime.stop();
+        Ok(())
     }
 }
 
 #[pymodule]
 fn rust_bindings(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(get_version, m)?)?;
+    m.add_class::<PyInputType>()?;
+    m.add_class::<PyModelMetadata>()?;
+    m.add_function(wrap_pyfunction!(metadata_from_json, m)?)?;
     m.add_class::<ModelProviderABC>()?;
     m.add_class::<PyModelRunner>()?;
     m.add_class::<PyModelRuntime>()?;