@1kbirds/chidori 0.1.20 → 0.1.23

package/src/translations/python.rs

@@ -5,9 +5,10 @@ use tonic::{Response, Status};
 use futures::executor;
 use futures::StreamExt;
 use pyo3::types::{PyDict, PyList, PyString};
-use pyo3::types::{PyBool, PyFloat, PyInt};
 use pyo3::prelude::*;
 use std::collections::HashMap;
+use std::sync::{Arc};
+use tokio::sync::Mutex;
 use std::time::Duration;
 use tokio::runtime::Runtime;
 use log::{debug, info};
@@ -21,6 +22,8 @@ use ::prompt_graph_core::proto2::serialized_value::Val;
 use ::prompt_graph_core::graph_definition::{create_prompt_node, create_op_map, create_code_node, create_component_node, create_vector_memory_node, create_observation_node, create_node_parameter, SourceNodeType, create_loader_node, create_custom_node};
 use ::prompt_graph_core::utils::wasm_error::CoreError;
 use ::prompt_graph_core::build_runtime_graph::graph_parse::{CleanedDefinitionGraph, CleanIndividualNode, construct_query_from_output_type, derive_for_individual_node};
+use crate::register_node_handle;
+use crate::translations::rust::{Chidori, CustomNodeCreateOpts, DenoCodeNodeCreateOpts, GraphBuilder, Handler, NodeHandle, PromptNodeCreateOpts, VectorMemoryNodeCreateOpts};
 
 #[derive(Debug)]
 pub struct CoreErrorWrapper(CoreError);
@@ -84,12 +87,27 @@ impl Into<pyo3::PyResult<()>> for PyErrWrapper {
     }
 }
 
-pub struct PyExecutionStatus(Response<ExecutionStatus>);
+
+pub struct PyExecutionStatus(ExecutionStatus);
 
 
 impl IntoPy<Py<PyAny>> for PyExecutionStatus {
     fn into_py(self, py: Python) -> Py<PyAny> {
-        let PyExecutionStatus(resp) = self;
+        let exec_status = self.0;
+        let dict = PyDict::new(py);
+        dict.set_item("id", exec_status.id).unwrap();
+        dict.set_item("monotonic_counter", exec_status.monotonic_counter).unwrap();
+        dict.set_item("branch", exec_status.branch).unwrap();
+        dict.into_py(py)
+    }
+}
+
+pub struct PyResponseExecutionStatus(Response<ExecutionStatus>);
+
+
+impl IntoPy<Py<PyAny>> for PyResponseExecutionStatus {
+    fn into_py(self, py: Python) -> Py<PyAny> {
+        let PyResponseExecutionStatus(resp) = self;
         let exec_status = resp.into_inner();
         let dict = PyDict::new(py);
         dict.set_item("id", exec_status.id).unwrap();
@@ -223,6 +241,7 @@ fn pyany_to_serialized_value(p: &PyAny) -> SerializedValue {
                 "float" => {
                     let val = p.extract::<f32>().unwrap();
                     SerializedValue {
+                        //         file: Some(File {
                         val: Some(Val::Float(val)),
                     }
                 }
@@ -267,6 +286,36 @@ fn pyany_to_serialized_value(p: &PyAny) -> SerializedValue {
 }
 
 
+use pyo3::prelude::*;
+use pyo3::types::IntoPyDict;
+use serde_json::json;
+
+fn py_to_json<'p>(py: Python<'p>, v: &PyAny) -> serde_json::Value {
+    if v.is_none() {
+        json!(null)
+    } else if let Ok(b) = v.extract::<bool>() {
+        json!(b)
+    } else if let Ok(i) = v.extract::<i64>() {
+        json!(i)
+    } else if let Ok(f) = v.extract::<f64>() {
+        json!(f)
+    } else if let Ok(dict) = v.extract::<HashMap<String, Py<PyAny>>>() {
+        let mut m = serde_json::map::Map::new();
+        for (key, value) in dict {
+            m.insert(key, py_to_json(py, value.as_ref(py)));
+        }
+        json!(m)
+    } else if let Ok(list) = v.extract::<Vec<Py<PyAny>>>() {
+        let v: Vec<serde_json::Value> = list.iter().map(|p| py_to_json(py, p.as_ref(py))).collect();
+        json!(v)
+    } else if let Ok(s) = v.extract::<String>() {
+        json!(s)
+    } else {
+        json!(null)
+    }
+}
+
+
 
 fn dict_to_paths<'p>(
     py: Python<'p>,
@@ -353,142 +402,80 @@ async fn get_client(url: String) -> Result<ExecutionRuntimeClient<tonic::transpo
 // TODO: return a handle to nodes so that we can understand and inspect them
 // TODO: include a __repr__ method on those nodes
 
-
-// TODO: add an api for wiring together a system ot nodes
-// TODO: require the description of the complete system, in order to avoid mutably handling changing edges
-// TODO: system description must be fully encapsulated in a single object and applied all at once
-
 #[pyclass]
 #[derive(Clone)]
-struct NodeHandle {
-    url: String,
-    file_id: String,
-    node: Item,
-    exec_status: ExecutionStatus,
-    indiv: CleanIndividualNode
+struct PyNodeHandle {
+    n: NodeHandle,
 }
 
-impl NodeHandle {
-    fn from(url: String, file_id: String, node: Item, exec_status: ExecutionStatus) -> anyhow::Result<NodeHandle> {
-        let indiv = derive_for_individual_node(&node)?;
-        Ok(NodeHandle {
-            url,
-            file_id,
-            node,
-            exec_status,
-            indiv
-        })
+impl PyNodeHandle {
+    fn from(node_handle: NodeHandle) -> anyhow::Result<PyNodeHandle> {
+        Ok(PyNodeHandle { n: node_handle })
     }
 }
 
 #[pymethods]
-impl NodeHandle {
+impl PyNodeHandle {
     fn get_name(&self) -> String {
-        self.node.core.as_ref().unwrap().name.clone()
+        self.n.get_name()
     }
 
     /// This updates the definition of this node to query for the target NodeHandle's output. Moving forward
     /// it will execute whenever the target node resolves.
-    #[pyo3(signature = (node_handle=None))]
-    fn run_when<'a>(mut self_: PyRefMut<'_, Self>, py: Python<'a>, node_handle: Option<NodeHandle>) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        let mut node = self_.node.clone();
+    fn run_when<'a>(mut self_: PyRefMut<'_, Self>, py: Python<'a>, graph_builder: &mut PyGraphBuilder, other_node_handle: PyNodeHandle) -> PyResult<&'a PyAny> {
+        let mut n = self_.n.clone();
+        let g = Arc::clone(&graph_builder.g);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            // TODO: scan the queries and don't insert if the query already exists
-            if let Some(node_handle) = node_handle {
-                let queries = &mut node.core.as_mut().unwrap().queries;
-                let q = construct_query_from_output_type(
-                    &node_handle.get_name(),
-                    &node_handle.get_name(),
-                    &node_handle.indiv.output_path
-                ).map_err(AnyhowErrWrapper)?;
-                queries.push(Query { query: Some(q)});
-                Ok(push_file_merge(&url, &file_id, node).await?)
-            } else {
-                Err(PyErr::new::<PyTypeError, _>("node_handle must be a NodeHandle"))
-            }
+            let mut graph_builder = g.lock().await;
+            Ok(n.run_when(&mut graph_builder, &other_node_handle.n)
+                .map_err(AnyhowErrWrapper)?)
         })
     }
 
-    #[pyo3(signature = (branch=0, frame=0))]
-    fn query<'a>(mut self_: PyRefMut<'_, Self>, py: Python<'a>, branch: u64, frame: u64) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        let name = self_.get_name();
-
-        let query = construct_query_from_output_type(&name, &name, &self_.indiv.output_path)
-            .map_err(AnyhowErrWrapper)?;
-
-        // TODO: we need this to watch for changes to the query instead
-        // TODO: or this should await until either the target counter has elapsed or the query has resulted
-
-        pyo3_asyncio::tokio::future_into_py(py, async move {
-            let mut client = get_client(url).await?;
-            Ok(PyQueryAtFrameResponse(client.run_query(QueryAtFrame {
-                id: file_id,
-                query: Some(Query {
-                    query: Some(query)
-                }),
-                frame,
-                branch,
-            }).await.map_err(PyErrWrapper::from)?))
-        })
-    }
+    // #[pyo3(signature = (branch=0, frame=0))]
+    // fn query<'a>(mut self_: PyRefMut<'_, Self>, py: Python<'a>, file_id: String, url: String, branch: u64, frame: u64) -> PyResult<&'a PyAny> {
+    //     pyo3_asyncio::tokio::future_into_py(py, async move {
+    //         Ok(PyQueryAtFrameResponse(self_.n.query(file_id, url, branch, frame)
+    //             .await.map_err(PyErrWrapper::from)?))
+    //     })
+    // }
 
     fn __str__(&self) -> PyResult<String>   {
         // TODO: best practice is that these could be used to re-construct the same object
         let name = self.get_name();
-        Ok(format!("NodeHandle(file_id={}, node={})", self.file_id, name))
+        Ok(format!("NodeHandle(node={})", name))
     }
 
     fn __repr__(&self) -> PyResult<String> {
         let name = self.get_name();
-        Ok(format!("NodeHandle(file_id={}, node={})", self.file_id, name))
+        Ok(format!("NodeHandle(node={})", name))
     }
 }
 
 
 // TODO: all operations only apply to a specific branch at a time
 // TODO: maintain an internal map of the generated change responses for node additions to the associated query necessary to get that result
-#[pyclass]
-struct Chidori {
+#[pyclass(name="Chidori")]
+struct PyChidori {
+    c: Arc<Mutex<Chidori>>,
     file_id: String,
     current_head: u64,
     current_branch: u64,
     url: String
 }
 
-
-
-async fn push_file_merge(url: &String, file_id: &String, node: Item) -> Result<NodeHandle, PyErr> {
-    let mut client = get_client(url.clone()).await?;
-    let exec_status = client.merge(RequestFileMerge {
-        id: file_id.clone(),
-        file: Some(File {
-            nodes: vec![node.clone()],
-            ..Default::default()
-        }),
-        branch: 0,
-    }).await.map_err(PyErrWrapper::from)?.into_inner();
-    Ok(NodeHandle::from(
-        url.clone(),
-        file_id.clone(),
-        node,
-        exec_status
-    ).map_err(AnyhowErrWrapper)?)
-}
-
 // TODO: internally all operations should have an assigned counter
 //       we can keep the actual target counter hidden from the host sdk
 #[pymethods]
-impl Chidori {
+impl PyChidori {
 
     #[new]
     #[pyo3(signature = (file_id=String::from("0"), url=String::from("http://127.0.0.1:9800"), api_token=None))]
     fn new(file_id: String, url: String, api_token: Option<String>) -> Self {
         debug!("Creating new Chidori instance with file_id={}, url={}, api_token={:?}", file_id, url, api_token);
-        Chidori {
+        let c = Chidori::new(file_id.clone(), url.clone());
+        PyChidori {
+            c: Arc::new(Mutex::new(c)),
             file_id,
             current_head: 0,
             current_branch: 0,
@@ -497,32 +484,11 @@ impl Chidori {
     }
 
     fn start_server<'a>(mut self_: PyRefMut<'_, Self>, py: Python<'a>, file_path: Option<String>) -> PyResult<&'a PyAny> {
-        let url_server = self_.url.clone();
-        std::thread::spawn(move || {
-            let result = run_server(url_server, file_path);
-            match result {
-                Ok(_) => { },
-                Err(e) => {
-                    println!("Error running server: {}", e);
-                },
-            }
-        });
-
+        let c = Arc::clone(&self_.c);
         let url = self_.url.clone();
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            'retry: loop {
-                let client = get_client(url.clone());
-                match client.await {
-                    Ok(connection) => {
-                        eprintln!("Connection successfully established {:?}", &url);
-                        break 'retry
-                    },
-                    Err(e) => {
-                        eprintln!("Error connecting to server: {} with Error {}. Retrying...", &url, &e.0);
-                        std::thread::sleep(std::time::Duration::from_millis(1000));
-                    }
-                }
-            }
+            let c = c.lock().await;
+            c.start_server(file_path).await.map_err(AnyhowErrWrapper)?;
             Ok(())
         })
     }
@@ -533,7 +499,7 @@ impl Chidori {
         let url = self_.url.clone();
         pyo3_asyncio::tokio::future_into_py(py, async move {
             let mut client = get_client(url).await?;
-            Ok(PyExecutionStatus(client.play(RequestAtFrame {
+            Ok(PyResponseExecutionStatus(client.play(RequestAtFrame {
                 id: file_id,
                 frame,
                 branch,
@@ -549,7 +515,7 @@ impl Chidori {
         let branch = self_.current_branch.clone();
         pyo3_asyncio::tokio::future_into_py(py, async move {
             let mut client = get_client(url).await?;
-            Ok(PyExecutionStatus(client.pause(RequestAtFrame {
+            Ok(PyResponseExecutionStatus(client.pause(RequestAtFrame {
                 id: file_id,
                 frame,
                 branch,
@@ -570,7 +536,7 @@ impl Chidori {
                 diverges_at_counter: 0,
             }).await.map_err(PyErrWrapper::from)?;
             // TODO: need to somehow handle writing to the current_branch
-            Ok(PyExecutionStatus(result_branch))
+            Ok(PyResponseExecutionStatus(result_branch))
         })
     }
 
@@ -696,213 +662,147 @@ impl Chidori {
         })
     }
 
-
-    // TODO: nodes that are added should return a clean definition of what their addition looks like
-    // TODO: adding a node should also display any errors
-    /// x = None
-    /// with open("/Users/coltonpierson/Downloads/files_and_dirs.zip", "rb") as zip_file:
-    ///     contents = zip_file.read()
-    ///     x = await p.load_zip_file("LoadZip", """ output: String """, contents)
-    /// x
-    #[pyo3(signature = (name=String::new(), output_tables=vec![], output=String::new(), bytes=vec![]))]
-    fn load_zip_file<'a>(
-        mut self_: PyRefMut<'_, Self>,
-        py: Python<'a>,
-        name: String,
-        output_tables: Vec<String>,
-        output: String,
-        bytes: Vec<u8>
-    ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        pyo3_asyncio::tokio::future_into_py(py, async move {
-            let node = create_loader_node(
-                name,
-                vec![],
-                output,
-                LoadFrom::ZipfileBytes(bytes),
-                output_tables
-            );
-            Ok(push_file_merge(&url, &file_id, node).await?)
-        })
-    }
-
-    // TODO: nodes that are added should return a clean definition of what their addition looks like
-    // TODO: adding a node should also display any errors
-    #[pyo3(signature = (name=String::new(), queries=vec![None], output_tables=vec![], template=String::new(), model=String::from("GPT_3_5_TURBO")))]
-    fn prompt_node<'a>(
+    pub fn register_custom_node_handle<'a>(
         mut self_: PyRefMut<'_, Self>,
         py: Python<'a>,
-        name: String,
-        queries: Vec<Option<String>>,
-        output_tables: Vec<String>,
-        template: String,
-        model: String
+        key: String,
+        handler: PyObject
     ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
+        let c = Arc::clone(&self_.c);
+        let handler = Arc::new(handler);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            let node = create_prompt_node(
-                name,
-                queries,
-                template,
-                model,
-                output_tables
-            ).map_err(PyErrWrapper::from)?;
-            Ok(push_file_merge(&url, &file_id, node).await?)
-        })
-    }
-
-    fn poll_local_code_node_execution<'a>(
-        mut self_: PyRefMut<'_, Self>,
-        py: Python<'a>,
-    ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-
-        pyo3_asyncio::tokio::future_into_py(py, async move {
-            let mut client = get_client(url).await?;
-            let result = client.poll_node_will_execute_events(FilteredPollNodeWillExecuteEventsRequest {
-                id: file_id.clone(),
-            }).await.map_err(PyErrWrapper::from)?;
-            debug!("poll_local_code_node_execution result = {:?}", result);
-            Ok(PyRespondPollNodeWillExecuteEvents(result))
+            let mut c = c.lock().await;
+            c.register_custom_node_handle(key, Handler::new(
+                move |n| {
+                    let handler_clone = Arc::clone(&handler);
+                    Box::pin(async move {
+                        let result = Python::with_gil(|py|  {
+                            let fut = handler_clone.as_ref().call(py, (NodeWillExecuteOnBranchWrapper(n).to_object(py), ), None)?;
+                            pyo3_asyncio::tokio::into_future(fut.as_ref(py))
+                        })?.await;
+                        match result {
+                            Ok(py_obj) => {
+                                Python::with_gil(|py|  {
+                                    let json_value = py_to_json(py, py_obj.as_ref(py));
+                                    Ok(json_value)
+                                })
+                            },
+                            Err(err) => Err(anyhow::Error::new(err)),
+                        }
+                    })
+                }
+            ));
+            Ok(())
         })
     }
 
-    #[pyo3(signature = (branch=0, counter=0))]
-    fn ack_local_code_node_execution<'a>(
+    fn run_custom_node_loop<'a>(
         mut self_: PyRefMut<'_, Self>,
         py: Python<'a>,
-        branch: u64,
-        counter: u64,
     ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
+        let c = Arc::clone(&self_.c);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            let mut client = get_client(url).await?;
-            Ok(PyExecutionStatus(client.ack_node_will_execute_event(RequestAckNodeWillExecuteEvent {
-                id: file_id.clone(),
-                branch,
-                counter,
-            }).await.map_err(PyErrWrapper::from)?))
+            let mut c = c.lock().await;
+            Ok(c.run_custom_node_loop().await.map_err(AnyhowErrWrapper)?)
         })
     }
 
-    #[pyo3(signature = (branch=0, counter=0, node_name=String::new(), response=None))]
-    fn respond_local_code_node_execution<'a>(
-        mut self_: PyRefMut<'_, Self>,
-        py: Python<'a>,
-        branch: u64,
-        counter: u64,
-        node_name: String,
-        response: Option<PyObject>
-    ) -> PyResult<&'a PyAny> {
-
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-
-        // TODO: need parent counters from the original change
-        // TODO: need source node
 
-        let response_paths = if let Some(response) = response {
-            dict_to_paths(py, response.downcast::<PyDict>(py)?)?
-        } else {
-            vec![]
-        };
+    //
+    // fn observation_node(mut self_: PyRefMut<'_, Self>, name: String, query_def: Option<String>, template: String, model: String) -> PyResult<()> {
+    //     let file_id = self_.file_id.clone();
+    //     let node = create_observation_node(
+    //         "".to_string(),
+    //         None,
+    //         "".to_string(),
+    //     );
+    //     executor::block_on(self_.client.merge(RequestFileMerge {
+    //         id: file_id,
+    //         file: Some(File {
+    //             nodes: vec![node],
+    //             ..Default::default()
+    //         }),
+    //         branch: 0,
+    //     }));
+    //     Ok(())
+    // }
+}
 
-        pyo3_asyncio::tokio::future_into_py(py, async move {
-            let mut client = get_client(url).await?;
+#[pyclass(name="GraphBuilder")]
+#[derive(Clone)]
+struct PyGraphBuilder {
+    g: Arc<Mutex<GraphBuilder>>,
+}
 
-            // TODO: need to add the output table paths to these
-            let filled_values = response_paths.into_iter().map(|path| {
-                ChangeValue {
-                    path: Some(Path {
-                        address: path.0,
-                    }),
-                    value: Some(path.1),
-                    branch,
-                }
-            });
+#[pymethods]
+impl PyGraphBuilder {
 
-            // TODO: this needs to look more like a real change
-            Ok(PyExecutionStatus(client.push_worker_event(FileAddressedChangeValueWithCounter {
-                branch,
-                counter,
-                node_name,
-                id: file_id.clone(),
-                change: Some(ChangeValueWithCounter {
-                    filled_values: filled_values.collect(),
-                    parent_monotonic_counters: vec![],
-                    monotonic_counter: counter,
-                    branch,
-                    source_node: "".to_string(),
-                })
-            }).await.map_err(PyErrWrapper::from)?))
-        })
+    #[new]
+    fn new() -> Self {
+        let g = GraphBuilder::new();
+        PyGraphBuilder {
+            g: Arc::new(Mutex::new(g)),
+        }
     }
 
-    // TODO: handle dispatch to this handler - should accept a callback
     // https://github.com/PyO3/pyo3/issues/525
-    #[pyo3(signature = (name=String::new(), queries=vec![None], output_tables=vec![], output=String::from("type O {}"), node_type_name=String::new()))]
+    #[pyo3(signature = (name=String::new(), queries=vec!["None".to_string()], output_tables=vec![], output=String::from("type O {}"), node_type_name=String::new()))]
     fn custom_node<'a>(
         mut self_: PyRefMut<'_, Self>,
         py: Python<'a>,
         name: String,
-        queries: Vec<Option<String>>,
+        queries: Option<Vec<String>>,
         output_tables: Vec<String>,
         output: String,
         node_type_name: String,
     ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        let branch = self_.current_branch;
+        let g = Arc::clone(&self_.g);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            // Register the node with the system
-            let node = create_custom_node(
+            let mut graph_builder = g.lock().await;
+            let nh = graph_builder.custom_node(CustomNodeCreateOpts {
                 name,
                 queries,
-                output,
+                output_tables: Some(output_tables),
+                output: Some(output),
                 node_type_name,
-                output_tables
-            );
-            Ok(push_file_merge(&url, &file_id, node).await?)
+            }).map_err(AnyhowErrWrapper)?;
+            Ok(PyNodeHandle::from(nh).map_err(AnyhowErrWrapper)?)
         })
     }
 
-    #[pyo3(signature = (name=String::new(), queries=vec![None], output_tables=vec![], output=String::from("type O { output: String }"), code=String::new(), is_template=false))]
+    #[pyo3(signature = (name=String::new(), queries=vec!["None".to_string()], output_tables=None, output=None, code=String::new(), is_template=None))]
     fn deno_code_node<'a>(
         mut self_: PyRefMut<'_, Self>,
         py: Python<'a>,
         name: String,
-        queries: Vec<Option<String>>,
-        output_tables: Vec<String>,
-        output: String,
+        queries: Option<Vec<String>>,
+        output_tables: Option<Vec<String>>,
+        output: Option<String>,
         code: String,
-        is_template: bool
+        is_template: Option<bool>
     ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        let branch = self_.current_branch.clone();
+        let g = Arc::clone(&self_.g);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            let node = create_code_node(
+            let mut graph_builder = g.lock().await;
+            let nh = graph_builder.deno_code_node(DenoCodeNodeCreateOpts {
                 name,
                 queries,
+                output_tables,
                 output,
-                SourceNodeType::Code("DENO".to_string(), code, is_template),
-                output_tables
-            );
-            Ok(push_file_merge(&url, &file_id, node).await?)
+                code,
+                is_template,
+            }).map_err(AnyhowErrWrapper)?;
+            Ok(PyNodeHandle::from(nh).map_err(AnyhowErrWrapper)?)
         })
     }
 
 
-    #[pyo3(signature = (name=String::new(), queries=vec![None], output_tables=vec![], output=String::from("type O { }"), template=String::new(), action="WRITE".to_string(), embedding_model="TEXT_EMBEDDING_ADA_002".to_string(), db_vendor="QDRANT".to_string(), collection_name=String::new()))]
+    #[pyo3(signature = (name=String::new(), queries=vec!["None".to_string()], output_tables=vec![], output=String::from("type O { }"), template=String::new(), action="WRITE".to_string(), embedding_model="TEXT_EMBEDDING_ADA_002".to_string(), db_vendor="QDRANT".to_string(), collection_name=String::new()))]
     fn vector_memory_node<'a>(
         mut self_: PyRefMut<'_, Self>,
         py: Python<'a>,
         name: String,
-        queries: Vec<Option<String>>,
+        queries: Option<Vec<String>>,
         output_tables: Vec<String>,
         output: String,
         template: String,
@@ -911,45 +811,98 @@ impl Chidori {
         db_vendor: String,
         collection_name: String,
     ) -> PyResult<&'a PyAny> {
-        let file_id = self_.file_id.clone();
-        let url = self_.url.clone();
-        let branch = self_.current_branch.clone();
-
+        let g = Arc::clone(&self_.g);
         pyo3_asyncio::tokio::future_into_py(py, async move {
-            let node = create_vector_memory_node(
+            let mut graph_builder = g.lock().await;
+            let nh = graph_builder.vector_memory_node(VectorMemoryNodeCreateOpts {
                 name,
                 queries,
-                output,
-                action,
-                embedding_model,
-                template,
-                db_vendor,
+                output_tables: Some(output_tables),
+                output: Some(output),
+                template: Some(template),
+                action: Some(action),
+                embedding_model: Some(embedding_model),
+                db_vendor: Some(db_vendor),
                 collection_name,
-                output_tables
-            ).map_err(PyErrWrapper::from)?;
-            Ok(push_file_merge(&url, &file_id, node).await?)
+            }).map_err(AnyhowErrWrapper)?;
+            Ok(PyNodeHandle::from(nh).map_err(AnyhowErrWrapper)?)
         })
     }
 
 
-    //
-    // fn observation_node(mut self_: PyRefMut<'_, Self>, name: String, query_def: Option<String>, template: String, model: String) -> PyResult<()> {
+    // // TODO: nodes that are added should return a clean definition of what their addition looks like
+    // // TODO: adding a node should also display any errors
+    // /// x = None
+    // /// with open("/Users/coltonpierson/Downloads/files_and_dirs.zip", "rb") as zip_file:
+    // ///     contents = zip_file.read()
+    // ///     x = await p.load_zip_file("LoadZip", """ output: String """, contents)
+    // /// x
+    // #[pyo3(signature = (name=String::new(), output_tables=vec![], output=String::new(), bytes=vec![]))]
+    // fn load_zip_file<'a>(
+    //     mut self_: PyRefMut<'_, Self>,
+    //     py: Python<'a>,
+    //     name: String,
+    //     output_tables: Vec<String>,
+    //     output: String,
+    //     bytes: Vec<u8>
+    // ) -> PyResult<&'a PyAny> {
     //     let file_id = self_.file_id.clone();
-    //     let node = create_observation_node(
-    //         "".to_string(),
-    //         None,
-    //         "".to_string(),
-    //     );
-    //     executor::block_on(self_.client.merge(RequestFileMerge {
-    //         id: file_id,
-    //         file: Some(File {
-    //             nodes: vec![node],
-    //             ..Default::default()
-    //         }),
-    //         branch: 0,
-    //     }));
-    //     Ok(())
+    //     let url = self_.url.clone();
+    //     pyo3_asyncio::tokio::future_into_py(py, async move {
+    //         let node = create_loader_node(
+    //             name,
+    //             vec![],
+    //             output,
+    //             LoadFrom::ZipfileBytes(bytes),
+    //             output_tables
+    //         );
+    //         Ok(push_file_merge(&url, &file_id, node).await?)
+    //     })
     // }
+
+    // TODO: nodes that are added should return a clean definition of what their addition looks like
+    // TODO: adding a node should also display any errors
+    #[pyo3(signature = (name=String::new(), queries=vec!["None".to_string()], output_tables=None, template=String::new(), model=None))]
+    fn prompt_node<'a>(
+        mut self_: PyRefMut<'_, Self>,
+        py: Python<'a>,
+        name: String,
+        queries: Option<Vec<String>>,
+        output_tables: Option<Vec<String>>,
+        template: String,
+        model: Option<String>
+    ) -> PyResult<&'a PyAny> {
+        let g = Arc::clone(&self_.g);
+        pyo3_asyncio::tokio::future_into_py(py, async move {
+            let mut graph_builder = g.lock().await;
+            let nh = graph_builder.prompt_node(PromptNodeCreateOpts {
+                name,
+                queries,
+                output_tables,
+                template,
+                model,
+            }).map_err(AnyhowErrWrapper)?;
+            Ok(PyNodeHandle::from(nh).map_err(AnyhowErrWrapper)?)
+        })
+    }
+
+    fn commit<'a>(
+        mut self_: PyRefMut<'_, Self>,
+        py: Python<'a>,
+        c: PyRef<'_, PyChidori>,
+        branch: u64
+    ) -> PyResult<&'a PyAny> {
+        let g = Arc::clone(&self_.g);
+        let c = Arc::clone(&c.c);
+        pyo3_asyncio::tokio::future_into_py(py, async move {
+            let mut graph_builder = g.lock().await;
+            let mut chidori = c.lock().await;
+            let exec_status = graph_builder.commit(&chidori, branch).await
+                .map(PyExecutionStatus)
+                .map_err(AnyhowErrWrapper)?;
+            Ok(exec_status)
+        })
+    }
 }
 
 
@@ -959,7 +912,9 @@ impl Chidori {
 #[pymodule]
 #[pyo3(name = "chidori")]
 fn chidori(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
-    pyo3_log::init();
-    m.add_class::<Chidori>()?;
+    // pyo3_log::init();
+    m.add_class::<PyChidori>()?;
+    m.add_class::<PyGraphBuilder>()?;
     Ok(())
-}
+}
+