limithub_code_block_sdk/
lib.rs

1pub mod error;
2pub mod port;
3pub mod value;
4
5use std::{
6    collections::{BTreeSet, HashMap},
7    future::Future,
8    io,
9    net::SocketAddr,
10    sync::Arc,
11};
12
13use axum::{
14    http::StatusCode,
15    routing::{get, post},
16    Extension, Json, Router,
17};
18
19pub use error::{Error, Result};
20pub use port::Port;
21use serde::Serialize;
22use tokio::net::TcpListener;
23pub use value::{PortValue, PortValueType};
24
25#[derive(Debug, Clone, Serialize)]
26pub struct CodeBlockAppBuilder {
27    inputs: Vec<Port>,
28    outputs: Vec<Port>,
29}
30
31impl CodeBlockAppBuilder {
32    pub fn new() -> Self {
33        CodeBlockAppBuilder {
34            inputs: Vec::new(),
35            outputs: Vec::new(),
36        }
37    }
38
39    /// Append input ports of the code block, it is required to have at least one input
40    pub fn inputs(mut self, inputs: impl IntoIterator<Item = Port>) -> Self {
41        self.inputs.extend(inputs);
42        self
43    }
44
45    // Append output ports of the code block
46    pub fn outputs(mut self, outputs: impl IntoIterator<Item = Port>) -> Self {
47        self.outputs.extend(outputs);
48        self
49    }
50
51    pub fn build(self) -> Result<CodeBlockApp> {
52        if self.inputs.is_empty() {
53            return Err(Error::EmptyPorts("inputs"));
54        }
55
56        let mut inputs = BTreeSet::new();
57        let mut outputs = BTreeSet::new();
58
59        for input in self.inputs {
60            let name = input.name.clone();
61            inputs
62                .insert(input)
63                .then_some(())
64                .ok_or(Error::NameConflict(name))?;
65        }
66
67        for output in self.outputs {
68            let name = output.name.clone();
69            outputs
70                .insert(output)
71                .then_some(())
72                .ok_or(Error::NameConflict(name))?;
73        }
74
75        Ok(CodeBlockApp { inputs, outputs })
76    }
77}
78
79#[derive(Debug, Clone, Serialize)]
80pub struct CodeBlockApp {
81    inputs: BTreeSet<Port>,
82    outputs: BTreeSet<Port>,
83}
84
85impl CodeBlockApp {
86    /// Serve the code block
87    pub fn serve<Fut>(
88        self,
89        handler: fn(HashMap<String, PortValue>) -> Fut,
90    ) -> impl Future<Output = io::Result<()>>
91    where
92        Fut: Future<Output = Result<HashMap<String, PortValue>, String>> + Send + 'static,
93    {
94        let app = Router::new()
95            .route("/ports", get(get_ports))
96            .route("/run", post(move |app, req| run(app, req, handler)))
97            .layer(Extension(Arc::new(self)));
98
99        let port = std::env::var("PORT")
100            .ok()
101            .and_then(|p| p.parse().ok())
102            .unwrap_or(3000);
103
104        async move {
105            let listener = TcpListener::bind(SocketAddr::from(([0, 0, 0, 0], port))).await?;
106            axum::serve(listener, app).await
107        }
108    }
109
110    fn validate_inputs(&self, inputs: &HashMap<String, PortValue>) -> Result<()> {
111        self.inputs
112            .iter()
113            .map(|i| {
114                inputs
115                    .contains_key(&i.name)
116                    .then_some(())
117                    .ok_or(Error::MissingPort(i.name.clone()))
118            })
119            .collect::<Result<()>>()
120    }
121
122    fn validate_outputs(&self, outputs: &HashMap<String, PortValue>) -> Result<()> {
123        self.outputs
124            .iter()
125            .map(|i| {
126                outputs
127                    .contains_key(&i.name)
128                    .then_some(())
129                    .ok_or(Error::MissingPort(i.name.clone()))
130            })
131            .collect::<Result<()>>()
132    }
133}
134
135async fn get_ports(Extension(app): Extension<Arc<CodeBlockApp>>) -> Json<CodeBlockApp> {
136    Json(app.as_ref().clone())
137}
138
139async fn run<Fut>(
140    Extension(app): Extension<Arc<CodeBlockApp>>,
141    Json(req): Json<HashMap<String, PortValue>>,
142    handler: fn(HashMap<String, PortValue>) -> Fut,
143) -> Result<Json<HashMap<String, PortValue>>, (StatusCode, String)>
144where
145    Fut: Future<Output = Result<HashMap<String, PortValue>, String>> + 'static + Send,
146{
147    app.validate_inputs(&req)
148        .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
149
150    let outputs = handler(req)
151        .await
152        .map(|res| Json(res))
153        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
154
155    app.validate_outputs(&outputs)
156        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
157
158    Ok(outputs)
159}
160
161#[cfg(test)]
162mod tests {
163    use std::collections::HashMap;
164
165    use crate::{CodeBlockAppBuilder, Port, PortValue, PortValueType};
166
167    #[test]
168    fn test() {
169        _ = CodeBlockAppBuilder::new()
170            .inputs([Port::new(String::from("port1"), PortValueType::Bool)])
171            .build()
172            .unwrap()
173            .serve(handler);
174    }
175
176    async fn handler(
177        _req: HashMap<String, PortValue>,
178    ) -> Result<HashMap<String, PortValue>, String> {
179        Ok(HashMap::new())
180    }
181}
limithub_code_block_sdk/lib.rs

limithub_code_block_sdk/
lib.rs
