use std::collections::HashMap;
use pyo3::{prelude::*, types::PyDict, PyObject, ToPyObject};
use crate::{
callbacks::{interface::EventHandler, PyEventCallbackInterface},
comm::{ChannelHandler, PyMessageInterface},
};
#[derive(Debug, Clone)]
#[pyclass]
pub struct Config {
pub version: f32,
pub querent_id: String,
pub querent_name: String,
pub workflow: WorkflowConfig,
pub collectors: Vec<CollectorConfig>,
pub engines: Vec<EngineConfig>,
pub resource: Option<ResourceConfig>,
}
impl ToPyObject for Config {
fn to_object(&self, py: Python) -> PyObject {
let config_dict = PyDict::new(py);
config_dict.set_item("version", self.version).unwrap();
config_dict.set_item("querent_id", &self.querent_id).unwrap();
config_dict.set_item("querent_name", &self.querent_name).unwrap();
config_dict.set_item("workflow", &self.workflow).unwrap();
config_dict.set_item("collectors", &self.collectors).unwrap();
config_dict.set_item("engines", &self.engines).unwrap();
config_dict.set_item("resource", &self.resource).unwrap();
config_dict.to_object(py)
}
}
impl Default for Config {
fn default() -> Self {
Config {
version: 0.1,
querent_id: "querent".to_string(),
querent_name: "Querent".to_string(),
workflow: WorkflowConfig {
name: "workflow".to_string(),
id: "workflow".to_string(),
config: HashMap::new(),
inner_channel: None,
channel: None,
inner_event_handler: Some(EventHandler::new(None)),
event_handler: None,
inner_tokens_feader: None,
tokens_feader: None,
},
collectors: vec![],
engines: vec![],
resource: None,
}
}
}
#[derive(Debug, Clone)]
#[pyclass]
pub struct WorkflowConfig {
pub name: String,
pub id: String,
pub config: HashMap<String, String>,
pub inner_channel: Option<ChannelHandler>,
#[pyo3(get, set)]
pub channel: Option<PyObject>,
pub inner_event_handler: Option<EventHandler>,
#[pyo3(get, set)]
pub event_handler: Option<PyObject>,
pub inner_tokens_feader: Option<ChannelHandler>,
#[pyo3(get, set)]
pub tokens_feader: Option<PyObject>,
}
impl ToPyObject for WorkflowConfig {
fn to_object(&self, py: Python) -> PyObject {
let workflow_dict = PyDict::new(py);
workflow_dict.set_item("name", &self.name).unwrap();
workflow_dict.set_item("id", &self.id).unwrap();
workflow_dict.set_item("config", &self.config).unwrap();
if let Some(inner_channel) = &self.inner_channel {
let channel_interface = PyMessageInterface::new(inner_channel.clone());
let channel: PyObject =
Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
workflow_dict.set_item("channel", channel).unwrap();
}
if let Some(inner_event_handler) = &self.inner_event_handler {
let event_interface = PyEventCallbackInterface::new(inner_event_handler.clone());
let event_handler: PyObject =
Py::new(py, event_interface).expect("Unable to create class").into_py(py);
workflow_dict.set_item("event_handler", event_handler).unwrap();
}
if let Some(inner_tokens_feader) = &self.inner_tokens_feader {
let channel_interface = PyMessageInterface::new(inner_tokens_feader.clone());
let tokens_feader: PyObject =
Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
workflow_dict.set_item("tokens_feader", tokens_feader).unwrap();
}
workflow_dict.to_object(py)
}
}
#[derive(Debug, Clone)]
#[pyclass]
pub struct CollectorConfig {
pub id: String,
pub name: String,
pub backend: String,
pub config: HashMap<String, String>,
pub inner_channel: Option<ChannelHandler>,
#[pyo3(get, set)]
pub channel: Option<PyObject>,
}
impl ToPyObject for CollectorConfig {
fn to_object(&self, py: Python) -> PyObject {
let collector_dict = PyDict::new(py);
collector_dict.set_item("id", &self.id).unwrap();
collector_dict.set_item("name", &self.name).unwrap();
collector_dict.set_item("backend", &self.backend).unwrap();
collector_dict.set_item("config", &self.config).unwrap();
if let Some(inner_channel) = &self.inner_channel {
let channel_interface = PyMessageInterface::new(inner_channel.clone());
let channel: PyObject =
Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
collector_dict.set_item("channel", channel).unwrap();
}
collector_dict.to_object(py)
}
}
#[derive(Debug, Clone)]
#[pyclass]
pub struct EngineConfig {
pub id: String,
pub name: String,
pub config: HashMap<String, String>,
pub inner_channel: Option<ChannelHandler>,
#[pyo3(get, set)]
pub channel: Option<PyObject>,
}
impl ToPyObject for EngineConfig {
fn to_object(&self, py: Python) -> PyObject {
let engine_dict = PyDict::new(py);
engine_dict.set_item("id", &self.id).unwrap();
engine_dict.set_item("name", &self.name).unwrap();
engine_dict.set_item("config", &self.config).unwrap();
if let Some(inner_channel) = &self.inner_channel {
let channel_interface = PyMessageInterface::new(inner_channel.clone());
let channel: PyObject =
Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
engine_dict.set_item("channel", channel).unwrap();
}
engine_dict.to_object(py)
}
}
#[derive(Debug, Clone)]
pub struct ResourceConfig {
pub id: String,
pub max_workers_allowed: Option<u32>,
pub max_workers_per_collector: Option<u32>,
pub max_workers_per_engine: Option<u32>,
pub max_workers_per_querent: Option<u32>,
}
impl<'a> FromPyObject<'a> for ResourceConfig {
fn extract(ob: &'a PyAny) -> PyResult<Self> {
let id = ob.getattr("id")?.extract()?;
let max_workers_allowed = ob.getattr("max_workers_allowed")?.extract()?;
let max_workers_per_collector = ob.getattr("max_workers_per_collector")?.extract()?;
let max_workers_per_engine = ob.getattr("max_workers_per_engine")?.extract()?;
let max_workers_per_querent = ob.getattr("max_workers_per_querent")?.extract()?;
Ok(ResourceConfig {
id,
max_workers_allowed,
max_workers_per_collector,
max_workers_per_engine,
max_workers_per_querent,
})
}
}
impl ToPyObject for ResourceConfig {
fn to_object(&self, py: Python) -> PyObject {
let resource_dict = PyDict::new(py);
resource_dict.set_item("id", &self.id).unwrap();
resource_dict
.set_item("max_workers_allowed", &self.max_workers_allowed)
.unwrap();
resource_dict
.set_item("max_workers_per_collector", &self.max_workers_per_collector)
.unwrap();
resource_dict
.set_item("max_workers_per_engine", &self.max_workers_per_engine)
.unwrap();
resource_dict
.set_item("max_workers_per_querent", &self.max_workers_per_querent)
.unwrap();
resource_dict.to_object(py)
}
}
#[pymethods]
impl Config {
#[new]
fn new(
version: f32,
querent_id: String,
querent_name: String,
workflow: WorkflowConfig,
collectors: Vec<CollectorConfig>,
engines: Vec<EngineConfig>,
resource: Option<ResourceConfig>,
) -> Self {
Config { version, querent_id, querent_name, workflow, collectors, engines, resource }
}
}