Fix Qdrant test (#25)

* fixing qdrant test * add back feature
ThousandBirdsInc · Oct 4, 2023 · 30363a2 · 30363a2
1 parent 7ebbeb2
commit 30363a2
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 20 deletions.
diff --git a/.idea/toolchain.iml b/.idea/toolchain.iml
diff --git a/toolchain/Cargo.lock b/toolchain/Cargo.lock
diff --git a/toolchain/prompt-graph-core/protobufs/DSL_v1.proto b/toolchain/prompt-graph-core/protobufs/DSL_v1.proto
@@ -6,6 +6,31 @@ package promptgraph;
 // This is used to provide a language agnostic interface
 // for defining PromptGraphs
 
+
+// TODO: capabilities also gives us a security model
+
+// A Node is the core primitive within a PromptGraph, it is
+// at a base level, just a function. However it includes an
+// execution_capabilities which is used to determine where in our
+// environment the function should be invoked. This allows us
+// to determine in what environment this function must run.
+message Node {
+  string handle = 1;
+  string execution_capabilities = 2;
+  repeated string argument_array = 3;
+  string inner_function_handle = 4;
+}
+
+message Triggerable {
+  string node_handle = 1;
+  Query query = 2;
+}
+
+message Subscribeable {
+  string node_handle = 1;
+  OutputType output = 2;
+}
+
 enum SupportedChatModel {
   GPT_4 = 0;
   GPT_4_0314 = 1;

diff --git a/toolchain/prompt-graph-exec/Cargo.toml b/toolchain/prompt-graph-exec/Cargo.toml
@@ -22,6 +22,7 @@ starlark = ["dep:starlark"]
 parquet = [
 #    "dep:parquet", "dep:arrow"
 ]
+qdrant = []
 default = []
 
 [dependencies]

diff --git a/toolchain/prompt-graph-exec/src/runtime_nodes/node_memory/node.rs b/toolchain/prompt-graph-exec/src/runtime_nodes/node_memory/node.rs
@@ -202,8 +202,10 @@ pub async fn initialize_node_memory_init(n: &PromptGraphNodeMemory, _core: &Item
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
     use prompt_graph_core::graph_definition::create_vector_memory_node;
-    use prompt_graph_core::proto::item;
+    use prompt_graph_core::proto::{item, NodeWillExecuteOnBranch};
+    use sled::Config as SledConfig;
     use anyhow::Result;
     use super::*;
 
@@ -217,6 +219,8 @@ mod tests {
         // -e QDRANT__SERVICE__GRPC_PORT="6334" \
         // qdrant/qdrant
 
+        let db = SledConfig::new().temporary(true).flush_every_ms(None).open().unwrap();
+        let tree = db.open_tree("test").unwrap();
         let config = QdrantClientConfig::from_url("http://localhost:6334");
         let client = QdrantClient::new(Some(config)).unwrap();
         let _ = client.delete_collection("test_exec_memory_node_qdrant").await;
@@ -235,10 +239,16 @@ mod tests {
             vec![]
         ).unwrap();
 
-        let nwe = NodeWillExecute {
-            source_node: "".to_string(),
-            change_values_used_in_execution: vec![],
-            matched_query_index: 0
+
+        let nwe = NodeWillExecuteOnBranch {
+            branch: 0,
+            counter: 0,
+            custom_node_type_name: None,
+            node: Some(NodeWillExecute {
+                source_node: "".to_string(),
+                change_values_used_in_execution: vec![],
+                matched_query_index: 0,
+            })
         };
 
         if let (core, item::Item::NodeMemory(n)) = (write.core.unwrap(), write.item.unwrap()) {
@@ -249,22 +259,28 @@ mod tests {
                 0).await.unwrap();
 
             let ctx = NodeExecutionContext {
-                node_will_execute: &nwe,
+                node_will_execute_on_branch: &nwe,
                 item_core: &core,
                 item: &item::Item::NodeMemory(n),
                 namespaces: &HashSet::from(["".to_string()]),
-                template_partials: &HashMap::new()
+                template_partials: &HashMap::new(),
+                tree: &tree
             };
             execute_node_memory(&ctx).await.unwrap();
         } else {
             assert!(false);
         }
 
 
-        let nwe = NodeWillExecute {
-            source_node: "".to_string(),
-            change_values_used_in_execution: vec![],
-            matched_query_index: 0
+        let nwe = NodeWillExecuteOnBranch {
+            branch: 0,
+            counter: 0,
+            custom_node_type_name: None,
+            node: Some(NodeWillExecute {
+                source_node: "".to_string(),
+                change_values_used_in_execution: vec![],
+                matched_query_index: 0,
+            })
         };
 
         let read = create_vector_memory_node(
@@ -281,11 +297,12 @@ mod tests {
 
         if let (core, item::Item::NodeMemory(n)) = (read.core.unwrap(), read.item.unwrap()) {
             let ctx = NodeExecutionContext {
-                node_will_execute: &nwe,
+                node_will_execute_on_branch: &nwe,
                 item_core: &core,
                 item: &item::Item::NodeMemory(n),
                 namespaces: &HashSet::from(["".to_string()]),
-                template_partials: &HashMap::new()
+                template_partials: &HashMap::new(),
+                tree: &tree
             };
             let recollection = execute_node_memory( &ctx ).await.unwrap();
             assert_eq!(recollection[0].path, Some(Path { address: vec![ "".to_string(), "key".to_string() ] }));