chore: decomplect — delete worthless stub crates, unblock hanzo-vm/consensus/l2
User mandate: "if valid fix and publish, if worthless delete — one and one way only". Deleted (no consumers, broken / incomplete / aspirational): - hanzo-baml — git-only baml-runtime dep; BAML integration never wired - hanzo-db — unused alongside hanzo-database/hanzo-db-sqlite - hanzo-kbs — incomplete, missing types, no in-tree consumer - hanzo-sheet — abandoned spreadsheet UI, depends on non-existent module - hanzo-simulation — research stub, not used by node - hanzo-sovereign — 170-LoC sketch with broken ml-dsa = "0.6" (no stable release) - hanzo-tests — test-only crate with compile errors Fixed and unblocked: - hanzo-vm — uses crates.io luxprecompile-sys 0.1, hanzo-engine 0.6 - hanzo-consensus — uses crates.io lux-consensus 1.22 - hanzo-l2 — uses crates.io lux-consensus 1.22 Removed `publish = false` from the three so they can ship to crates.io. Cargo.lock will regenerate on next build.
This commit is contained in:
Hanzo Dev
Generated
+6
-2
@@ -5913,7 +5913,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-consensus"
|
||||
version = "1.1.20"
|
||||
version = "1.1.21"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
@@ -6179,7 +6179,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-l2"
|
||||
version = "1.1.20"
|
||||
version = "1.1.21"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
@@ -9308,6 +9308,8 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
[[package]]
|
||||
name = "lux-consensus"
|
||||
version = "1.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73f3e103de37e286e68c58e5188a49fcc94b2ac23c3afa2aaa1775f45cf97c68"
|
||||
dependencies = [
|
||||
"blst",
|
||||
"cc",
|
||||
@@ -9318,6 +9320,8 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "luxprecompile-sys"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f152fe85e7080ed60e68da4348079167e0329b9da964d3dba604efd9ca0f891"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
||||
+12
-25
@@ -1,10 +1,13 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"hanzo-agentic",
|
||||
"hanzo-ai-format",
|
||||
"hanzo-api",
|
||||
"hanzo-brain",
|
||||
"hanzo-compute",
|
||||
"hanzo-config",
|
||||
"hanzo-consensus",
|
||||
"hanzo-database",
|
||||
# "hanzo-db", # Excluded: not in publish list, has feature-gated dependencies
|
||||
"hanzo-db-sqlite",
|
||||
"hanzo-did",
|
||||
"hanzo-embed",
|
||||
@@ -14,37 +17,25 @@ members = [
|
||||
"hanzo-identity",
|
||||
"hanzo-job-queue-manager",
|
||||
"hanzo-jobs",
|
||||
# "hanzo-kbs", # Excluded: incomplete implementation with missing types
|
||||
"hanzo-l2",
|
||||
"hanzo-libp2p",
|
||||
"hanzo-libp2p-relayer",
|
||||
"hanzo-machine",
|
||||
"hanzo-mcp",
|
||||
"hanzo-messages",
|
||||
"hanzo-mining",
|
||||
"hanzo-vm",
|
||||
"hanzo-machine",
|
||||
"hanzo-consensus",
|
||||
"hanzo-compute",
|
||||
"hanzo-ai-format",
|
||||
"hanzo-model-discovery",
|
||||
"hanzo-models",
|
||||
"hanzo-pqc",
|
||||
"hanzo-runner",
|
||||
"hanzo-runtime",
|
||||
"hanzo-runtime-tests",
|
||||
# "hanzo-sheet", # Excluded: depends on non-existent sheet module
|
||||
# "hanzo-simulation",
|
||||
# "hanzo-tests", # Excluded: test-only crate with compilation errors
|
||||
"hanzo-tools",
|
||||
"hanzo-tools-runner",
|
||||
"hanzo-vm",
|
||||
"hanzo-wasm",
|
||||
"hanzo-wasm-runtime",
|
||||
"hanzo-zap",
|
||||
"hanzo-agentic",
|
||||
# "hanzo-sovereign", # Excluded: ml-dsa 0.6 doesn't exist on crates.io (only prerelease 0.1.0-rc.x)
|
||||
]
|
||||
exclude = [
|
||||
"hanzo-sovereign",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
@@ -55,10 +46,13 @@ authors = ["Nico Arqueros <nico@hanzo.ai>"]
|
||||
repository = "https://github.com/hanzonet/network"
|
||||
|
||||
[workspace.dependencies]
|
||||
hanzo-ai-format = { path = "./hanzo-ai-format", version = "1.1.21" }
|
||||
hanzo-api = { path = "./hanzo-api", version = "1.1.13" }
|
||||
hanzo-brain = { path = "./hanzo-brain", version = "0.1.0" }
|
||||
hanzo-compute = { path = "./hanzo-compute", version = "1.1.21" }
|
||||
hanzo-config = { path = "./hanzo-config", version = "1.1.21", package = "hanzonet-config" }
|
||||
hanzo-consensus = { path = "./hanzo-consensus", version = "1.1.20" }
|
||||
hanzo-database = { path = "./hanzo-database", version = "1.1.13" }
|
||||
hanzo-db = { path = "./hanzo-db", version = "0.1.1" }
|
||||
hanzo-db-sqlite = { path = "./hanzo-db-sqlite", version = "1.1.13" }
|
||||
hanzo-did = { path = "./hanzo-did", version = "1.1.21", package = "hanzonet-did" }
|
||||
hanzo-embed = { path = "./hanzo-embed", version = "1.1.13" }
|
||||
@@ -68,32 +62,25 @@ hanzo-http-api = { path = "./hanzo-http-api", version = "1.1.21" }
|
||||
hanzo-identity = { path = "./hanzo-identity", version = "1.1.13" }
|
||||
hanzo-job-queue-manager = { path = "./hanzo-job-queue-manager", version = "1.1.21" }
|
||||
hanzo-jobs = { path = "./hanzo-jobs", version = "1.1.13" }
|
||||
hanzo-kbs = { path = "./hanzo-kbs", version = "1.1.11" }
|
||||
hanzo-l2 = { path = "./hanzo-l2", version = "1.1.20" }
|
||||
hanzo-libp2p = { path = "./hanzo-libp2p", version = "1.1.13" }
|
||||
hanzo-libp2p-relayer = { path = "./hanzo-libp2p-relayer", version = "1.1.13" }
|
||||
hanzo-machine = { path = "./hanzo-machine", version = "1.1.21" }
|
||||
hanzo-mcp = { path = "./hanzo-mcp", version = "1.1.21", package = "hanzonet-mcp" }
|
||||
hanzo-messages = { path = "./hanzo-messages", version = "1.1.13" }
|
||||
hanzo-mining = { path = "./hanzo-mining", version = "1.1.21" }
|
||||
hanzo-consensus = { path = "./hanzo-consensus", version = "1.1.20" }
|
||||
hanzo-compute = { path = "./hanzo-compute", version = "1.1.21" }
|
||||
hanzo-ai-format = { path = "./hanzo-ai-format", version = "1.1.21" }
|
||||
hanzo-model-discovery = { path = "./hanzo-model-discovery", version = "1.1.13" }
|
||||
hanzo-models = { path = "./hanzo-models", version = "1.1.13" }
|
||||
hanzo-pqc = { path = "./hanzo-pqc", version = "1.1.21", package = "hanzonet-pqc" }
|
||||
hanzo-runner = { path = "./hanzo-runner", version = "1.1.13" }
|
||||
hanzo-runtime = { path = "./hanzo-runtime", version = "1.1.13" }
|
||||
hanzo-runtime-tests = { path = "./hanzo-runtime-tests", version = "0.1.2" }
|
||||
hanzo-sheet = { path = "./hanzo-sheet", version = "0.1.1" }
|
||||
hanzo-tests = { path = "./hanzo-tests", version = "0.1.1" }
|
||||
hanzo-tools = { path = "./hanzo-tools", version = "1.1.21" }
|
||||
hanzo-tools-runner = { path = "./hanzo-tools-runner", version = "1.0.2" }
|
||||
hanzo-wasm = { path = "./hanzo-wasm", version = "0.1.2" }
|
||||
hanzo-vm = { path = "./hanzo-vm", version = "1.1.21" }
|
||||
hanzo-machine = { path = "./hanzo-machine", version = "1.1.21" }
|
||||
hanzo-wasm = { path = "./hanzo-wasm", version = "0.1.2" }
|
||||
hanzo-wasm-runtime = { path = "./hanzo-wasm-runtime", version = "1.1.21" }
|
||||
hanzo-zap = { path = "./hanzo-zap", version = "0.6.75" }
|
||||
hanzo-brain = { path = "./hanzo-brain", version = "0.1.0" }
|
||||
|
||||
rmcp = { version = "0.8" }
|
||||
futures = "0.3.30"
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-baml"
|
||||
description = "Hanzo AI - BAML Integration Library"
|
||||
license = "MIT"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
authors = { workspace = true }
|
||||
repository = "https://github.com/hanzoai/node"
|
||||
homepage = "https://hanzo.ai"
|
||||
|
||||
[dependencies]
|
||||
baml-runtime = { git = "https://github.com/BoundaryML/baml.git", tag = "0.56.1", package = "baml-runtime", features = [
|
||||
"internal",
|
||||
] }
|
||||
baml-types = { git = "https://github.com/BoundaryML/baml.git", tag = "0.56.1", package = "baml-types" }
|
||||
env_logger = { workspace = true }
|
||||
log = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
indexmap = { version = "2.1.0", features = ["serde"] }
|
||||
regex = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
@@ -1,331 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use baml_runtime::{BamlRuntime, InternalRuntimeInterface};
|
||||
use baml_types::BamlValue;
|
||||
use indexmap::IndexMap;
|
||||
use regex::Regex;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct GeneratorConfig {
|
||||
pub output_type: String,
|
||||
pub output_dir: String,
|
||||
pub version: String,
|
||||
pub default_client_mode: String,
|
||||
}
|
||||
|
||||
impl Default for GeneratorConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
output_type: "typescript".to_string(),
|
||||
output_dir: "../src/".to_string(),
|
||||
version: "0.55.3".to_string(),
|
||||
default_client_mode: "async".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ClientConfig {
|
||||
pub provider: String,
|
||||
pub base_url: Option<String>,
|
||||
pub model: String,
|
||||
pub default_role: String,
|
||||
pub api_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BamlConfig {
|
||||
pub generator: GeneratorConfig,
|
||||
pub client: ClientConfig,
|
||||
pub dsl_class_file: Option<String>,
|
||||
pub input: Option<String>,
|
||||
pub function_name: Option<String>,
|
||||
pub param_name: Option<String>,
|
||||
}
|
||||
|
||||
impl BamlConfig {
|
||||
pub fn builder(generator: GeneratorConfig, client: ClientConfig) -> BamlConfigBuilder {
|
||||
BamlConfigBuilder {
|
||||
generator,
|
||||
client,
|
||||
dsl_class_file: None,
|
||||
input: None,
|
||||
function_name: None,
|
||||
param_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn initialize_runtime(&self, env_vars: HashMap<&str, &str>) -> Result<BamlRuntime> {
|
||||
let mut files = HashMap::new();
|
||||
files.insert(
|
||||
"generator.baml".to_string(),
|
||||
format!(
|
||||
r##"
|
||||
generator lang_ts {{
|
||||
output_type "{}"
|
||||
output_dir "{}"
|
||||
version "{}"
|
||||
default_client_mode "{}"
|
||||
}}
|
||||
"##,
|
||||
self.generator.output_type,
|
||||
self.generator.output_dir,
|
||||
self.generator.version,
|
||||
self.generator.default_client_mode
|
||||
),
|
||||
);
|
||||
|
||||
let base_url_option = if let Some(base_url) = &self.client.base_url {
|
||||
format!(r#"base_url "{}""#, base_url)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
files.insert(
|
||||
"client.baml".to_string(),
|
||||
format!(
|
||||
r##"
|
||||
client<llm> {} {{
|
||||
provider {}
|
||||
options {{
|
||||
{}
|
||||
model "{}"
|
||||
default_role "{}"
|
||||
api_key "{}"
|
||||
}}
|
||||
}}
|
||||
"##,
|
||||
"HanzoProvider",
|
||||
self.client.provider,
|
||||
base_url_option,
|
||||
self.client.model,
|
||||
self.client.default_role,
|
||||
self.client.api_key.clone().unwrap_or_default()
|
||||
),
|
||||
);
|
||||
|
||||
if let Some(dsl_class_file) = &self.dsl_class_file {
|
||||
files.insert("dsl_class.baml".to_string(), dsl_class_file.clone());
|
||||
}
|
||||
|
||||
let runtime = BamlRuntime::from_file_content("baml_src", &files, env_vars)?;
|
||||
let diagnostics = runtime.internal().diagnostics();
|
||||
eprintln!("BAML diagnostics: {:?}", diagnostics);
|
||||
|
||||
if diagnostics.has_errors() {
|
||||
let error_message = diagnostics.to_pretty_string();
|
||||
return Err(anyhow::anyhow!("BAML diagnostics errors: {}", error_message));
|
||||
}
|
||||
|
||||
Ok(runtime)
|
||||
}
|
||||
|
||||
pub fn execute(&self, runtime: &BamlRuntime, extract_data: bool) -> Result<String> {
|
||||
let ctx_manager = runtime.create_ctx_manager(BamlValue::String("none".to_string()), None);
|
||||
|
||||
let mut params = IndexMap::new();
|
||||
if let (Some(param_name), Some(input)) = (&self.param_name, &self.input) {
|
||||
let trimmed_input = input.trim();
|
||||
let context_value = if trimmed_input.starts_with('{') && trimmed_input.ends_with('}') {
|
||||
eprintln!("input is a json string: {}", trimmed_input);
|
||||
match serde_json::from_str(&trimmed_input) {
|
||||
Ok(parsed_json) => BamlConfig::from_serde_value(parsed_json),
|
||||
Err(_) => {
|
||||
eprintln!("Failed to parse JSON, attempting to unescape");
|
||||
let unescaped_input = BamlConfig::unescape_json_string(trimmed_input);
|
||||
eprintln!("Unescaped input: {}", unescaped_input);
|
||||
match serde_json::from_str(&unescaped_input) {
|
||||
Ok(parsed_json) => BamlConfig::from_serde_value(parsed_json),
|
||||
Err(e) => return Err(anyhow::anyhow!("Failed to parse JSON after unescaping: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
BamlValue::String(trimmed_input.to_string())
|
||||
};
|
||||
params.insert(param_name.clone(), context_value);
|
||||
}
|
||||
|
||||
if let Some(function_name) = &self.function_name {
|
||||
eprintln!("\n\n Params string keys {:?}\n\n", params.keys());
|
||||
eprintln!("\n\n Params string values {:?}\n\n", params.values());
|
||||
let (result, _uuid) = runtime.call_function_sync(function_name.clone(), ¶ms, &ctx_manager, None, None);
|
||||
|
||||
match result {
|
||||
Ok(response) => match response.content() {
|
||||
Ok(content) => {
|
||||
if extract_data {
|
||||
eprintln!("Extracting data from response: {}", content);
|
||||
if content.starts_with('{') && content.ends_with('}') {
|
||||
return Ok(content.to_string());
|
||||
}
|
||||
let re = Regex::new(r"```(?:json)?\s*([\s\S]*?)\s*```").unwrap();
|
||||
if let Some(captures) = re.captures(&content) {
|
||||
if let Some(matched) = captures.get(1) {
|
||||
return Ok(matched.as_str().to_string());
|
||||
}
|
||||
}
|
||||
return Ok(content.to_string());
|
||||
} else {
|
||||
return Ok(content.to_string());
|
||||
}
|
||||
}
|
||||
Err(e) => return Err(anyhow::anyhow!("Error getting content: {}", e)),
|
||||
},
|
||||
Err(e) => return Err(anyhow::anyhow!("Error: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
Err(anyhow::anyhow!("Function name not provided"))
|
||||
}
|
||||
|
||||
pub fn from_serde_value(value: serde_json::Value) -> BamlValue {
|
||||
match value {
|
||||
serde_json::Value::Null => BamlValue::Null,
|
||||
serde_json::Value::Bool(b) => BamlValue::Bool(b),
|
||||
serde_json::Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
BamlValue::Int(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
BamlValue::Float(f)
|
||||
} else {
|
||||
panic!("Unexpected number type")
|
||||
}
|
||||
}
|
||||
serde_json::Value::String(s) => BamlValue::String(s),
|
||||
serde_json::Value::Array(arr) => {
|
||||
let baml_values = arr.into_iter().map(BamlConfig::from_serde_value).collect();
|
||||
BamlValue::List(baml_values)
|
||||
}
|
||||
serde_json::Value::Object(obj) => {
|
||||
if let Some(class_name) = obj.clone().get("class_name").and_then(|v| v.as_str()) {
|
||||
let class_fields = obj
|
||||
.into_iter()
|
||||
.filter(|(k, _)| k != "class_name")
|
||||
.map(|(k, v)| (k, BamlConfig::from_serde_value(v)))
|
||||
.collect();
|
||||
BamlValue::Class(class_name.to_string(), class_fields)
|
||||
} else {
|
||||
let baml_map = obj
|
||||
.into_iter()
|
||||
.map(|(k, v)| (k, BamlConfig::from_serde_value(v)))
|
||||
.collect();
|
||||
BamlValue::Map(baml_map)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn unescape_json_string(json_str: &str) -> String {
|
||||
let re_backslash_quote = Regex::new(r#"\\""#).unwrap(); // Matches \"
|
||||
let re_backslash = Regex::new(r#"\\\\"#).unwrap(); // Matches \\
|
||||
let re_newline = Regex::new(r#"\\n"#).unwrap(); // Matches \n
|
||||
let re_tab = Regex::new(r#"\\t"#).unwrap(); // Matches \t
|
||||
|
||||
let intermediate = re_backslash_quote.replace_all(json_str, "\"");
|
||||
let intermediate = re_backslash.replace_all(&intermediate, "\\");
|
||||
let intermediate = re_newline.replace_all(&intermediate, "\n");
|
||||
let intermediate = re_tab.replace_all(&intermediate, "\t");
|
||||
|
||||
intermediate.to_string()
|
||||
}
|
||||
|
||||
/// Converts the existing DSL string to the format expected by Baml.
|
||||
pub fn convert_dsl_class_file(old_dsl: &str) -> String {
|
||||
// Define regex patterns for different escape sequences
|
||||
let re_triple_backslash_quote = Regex::new(r#"\\\\\\""#).unwrap(); // Matches \\\"
|
||||
let re_newline = Regex::new(r#"\\n"#).unwrap(); // Matches \n
|
||||
let re_quote = Regex::new(r#"\\""#).unwrap(); // Matches \"
|
||||
let re_backslash = Regex::new(r#"\\\\"#).unwrap(); // Matches \\
|
||||
|
||||
// Perform replacements using regex in the correct order
|
||||
// 1. Replace triple backslashes followed by a quote (\\\\\") with an escaped quote (\\")
|
||||
let intermediate = re_triple_backslash_quote.replace_all(old_dsl, "\\\"");
|
||||
// 2. Replace escaped newlines (\\n) with actual newlines (\n)
|
||||
let intermediate = re_newline.replace_all(&intermediate, "\n");
|
||||
// 3. Replace escaped quotes (\\") with actual quotes (")
|
||||
let intermediate = re_quote.replace_all(&intermediate, "\"");
|
||||
// 4. Replace escaped backslashes (\\\\) with a single backslash (\\)
|
||||
let intermediate = re_backslash.replace_all(&intermediate, "\\");
|
||||
|
||||
// Optionally, adjust other parts of the DSL as needed
|
||||
// For example, change client provider from Ollama to HanzoProvider
|
||||
let re_client = Regex::new(r#"client\s+\w+"#).unwrap();
|
||||
let adjusted = re_client.replace_all(&intermediate, "client HanzoProvider");
|
||||
|
||||
adjusted.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct BamlConfigBuilder {
|
||||
generator: GeneratorConfig,
|
||||
client: ClientConfig,
|
||||
dsl_class_file: Option<String>,
|
||||
input: Option<String>,
|
||||
function_name: Option<String>,
|
||||
param_name: Option<String>,
|
||||
}
|
||||
|
||||
impl BamlConfigBuilder {
|
||||
pub fn dsl_class_file(mut self, dsl_class_file: &str) -> Self {
|
||||
self.dsl_class_file = Some(dsl_class_file.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn input(mut self, input: &str) -> Self {
|
||||
self.input = Some(input.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn function_name(mut self, function_name: &str) -> Self {
|
||||
self.function_name = Some(function_name.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn param_name(mut self, param_name: &str) -> Self {
|
||||
self.param_name = Some(param_name.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> BamlConfig {
|
||||
BamlConfig {
|
||||
generator: self.generator,
|
||||
client: self.client,
|
||||
dsl_class_file: self.dsl_class_file,
|
||||
input: self.input,
|
||||
function_name: self.function_name,
|
||||
param_name: self.param_name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_unescape_json_string() {
|
||||
let input = r#"{
|
||||
\"documents\": [
|
||||
{
|
||||
\"title\": \"OmniParser Abstract\",
|
||||
\"link\": \"https://arxiv.org\",
|
||||
\"text\": \"- OmniParser for Pure Vision Based GUI Agent Yadong Lu 1 , Jianwei Yang 1 , Yelong Shen 2 , Ahmed Awadallah 1 1 Microsoft Research 2 Microsoft Gen AI {yadonglu, jianwei.yang, yeshe, ahmed.awadallah}@microsoft.com Abstract (Source: 2408.00203v1.pdf, Section: )\"
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let expected_output = r#"{
|
||||
"documents": [
|
||||
{
|
||||
"title": "OmniParser Abstract",
|
||||
"link": "https://arxiv.org",
|
||||
"text": "- OmniParser for Pure Vision Based GUI Agent Yadong Lu 1 , Jianwei Yang 1 , Yelong Shen 2 , Ahmed Awadallah 1 1 Microsoft Research 2 Microsoft Gen AI {yadonglu, jianwei.yang, yeshe, ahmed.awadallah}@microsoft.com Abstract (Source: 2408.00203v1.pdf, Section: )"
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let unescaped = BamlConfig::unescape_json_string(input);
|
||||
assert_eq!(unescaped, expected_output);
|
||||
}
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
pub mod baml_builder;
|
||||
@@ -1,85 +0,0 @@
|
||||
mod baml_builder;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Once;
|
||||
|
||||
use anyhow::Result;
|
||||
use baml_builder::{BamlConfig, ClientConfig, GeneratorConfig};
|
||||
use log::info;
|
||||
|
||||
static INIT: Once = Once::new();
|
||||
|
||||
fn main() -> Result<()> {
|
||||
INIT.call_once(|| {
|
||||
env_logger::init();
|
||||
});
|
||||
info!("Starting Baml Runtime");
|
||||
|
||||
let generator_config = GeneratorConfig {
|
||||
output_type: "typescript".to_string(),
|
||||
output_dir: "../src/".to_string(),
|
||||
version: "0.55.3".to_string(),
|
||||
default_client_mode: "async".to_string(),
|
||||
};
|
||||
|
||||
let client_config = ClientConfig {
|
||||
provider: "ollama".to_string(),
|
||||
base_url: Some("http://localhost:11434/v1".to_string()),
|
||||
model: "llama3.1:8b-instruct-q4_1".to_string(),
|
||||
default_role: "user".to_string(),
|
||||
api_key: None,
|
||||
};
|
||||
|
||||
let baml_config = BamlConfig::builder(generator_config, client_config)
|
||||
.dsl_class_file(
|
||||
r##"
|
||||
class Resume {
|
||||
name string
|
||||
email string
|
||||
experience string[]
|
||||
skills string[]
|
||||
}
|
||||
|
||||
function ExtractResume(resume: string) -> Resume {
|
||||
client HanzoProvider
|
||||
prompt #"
|
||||
Extract from this content:
|
||||
{{ resume }}
|
||||
|
||||
{{ ctx.output_format }}
|
||||
"#
|
||||
}
|
||||
"##,
|
||||
)
|
||||
.input(
|
||||
r#"
|
||||
Name: Nico Arqueros
|
||||
123 Main St, Anytown, USA
|
||||
Email: john.doe@email.com
|
||||
Phone: (555) 123-4567
|
||||
|
||||
Education:
|
||||
Bachelor of Science in Computer Science
|
||||
University of Technology, 2015-2019
|
||||
|
||||
Work Experience:
|
||||
Software Engineer, Tech Corp
|
||||
June 2019 - Present
|
||||
- Developed and maintained web applications
|
||||
- Collaborated with cross-functional teams
|
||||
|
||||
Skills:
|
||||
JavaScript, TypeScript, React, Node.js, Git
|
||||
"#,
|
||||
)
|
||||
.function_name("ExtractResume")
|
||||
.param_name("resume")
|
||||
.build();
|
||||
|
||||
let env_vars = HashMap::new();
|
||||
let runtime = baml_config.initialize_runtime(env_vars)?;
|
||||
let result = baml_config.execute(&runtime, true)?;
|
||||
println!("Execution result: {}", result);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,17 +1,14 @@
|
||||
[package]
|
||||
name = "hanzo-consensus"
|
||||
version = "1.1.20"
|
||||
version = "1.1.21"
|
||||
edition = "2021"
|
||||
description = "Native Quasar consensus engine for Hanzo Network L2"
|
||||
description = "Hanzo Network — Quasar BFT consensus engine for L2 (wraps lux-consensus)"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/hanzonet/network"
|
||||
# Cannot publish to crates.io until lux-consensus is published.
|
||||
# Also depends on the local net hanzo-pqc (publish = false).
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
# Lux consensus SDK (Quasar protocol)
|
||||
lux-consensus = { path = "../../../lux/consensus/pkg/rust" }
|
||||
lux-consensus = "1.22"
|
||||
|
||||
# PQC for dual certificates
|
||||
hanzo-pqc = { workspace = true }
|
||||
|
||||
@@ -1,79 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-db"
|
||||
version = "0.1.1"
|
||||
edition = { workspace = true }
|
||||
publish = false # Not ready for crates.io
|
||||
|
||||
# Standalone crate - not part of parent workspace
|
||||
[workspace]
|
||||
authors = { workspace = true }
|
||||
description = "Multi-backend database abstraction for Hanzo Node - Supports LanceDB, DuckDB, PostgreSQL, and more"
|
||||
license = "Apache-2.0"
|
||||
keywords = ["database", "vector", "multimodal", "lancedb", "duckdb"]
|
||||
categories = ["database", "science"]
|
||||
|
||||
[dependencies]
|
||||
# Core LanceDB with latest version (uses arrow 55.2)
|
||||
lancedb = "0.22.0"
|
||||
arrow-array = "55.2"
|
||||
arrow-schema = "55.2"
|
||||
arrow-ipc = "55.2"
|
||||
parquet = "55.2"
|
||||
|
||||
# Async runtime
|
||||
tokio = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
|
||||
# Serialization
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
bincode = { workspace = true }
|
||||
|
||||
# Error handling
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
# Utilities
|
||||
log = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
uuid = { workspace = true, features = ["v4", "serde"] }
|
||||
base64 = { workspace = true }
|
||||
blake3 = { workspace = true }
|
||||
|
||||
# Connection pooling
|
||||
deadpool = "0.12"
|
||||
deadpool-runtime = "0.1"
|
||||
|
||||
# Image processing for multimodal
|
||||
image = { version = "0.25", optional = true }
|
||||
|
||||
# Migration support
|
||||
rusqlite = { workspace = true, optional = true }
|
||||
|
||||
# Backend-specific dependencies
|
||||
duckdb = { version = "0.10", optional = true }
|
||||
sqlx = { version = "0.8", features = ["postgres", "runtime-tokio-rustls", "sqlite"], optional = true }
|
||||
redis = { version = "0.25", features = ["aio", "tokio-comp"], optional = true }
|
||||
|
||||
# Hanzo dependencies
|
||||
hanzo-messages = { version = "1.1.12", path = "../hanzo-messages" }
|
||||
hanzo-embed = { version = "1.1.12", path = "../hanzo-embed" }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-test = { workspace = true }
|
||||
tempfile = { workspace = true }
|
||||
criterion = { workspace = true }
|
||||
proptest = { workspace = true }
|
||||
env_logger = { workspace = true }
|
||||
|
||||
[features]
|
||||
default = ["multimodal", "migration"]
|
||||
multimodal = ["image"]
|
||||
migration = ["rusqlite"]
|
||||
all-backends = ["backend-duckdb", "backend-postgres", "backend-redis", "backend-sqlite"]
|
||||
backend-duckdb = ["duckdb"]
|
||||
backend-postgres = ["sqlx"]
|
||||
backend-redis = ["redis"]
|
||||
backend-sqlite = ["rusqlite"]
|
||||
|
||||
@@ -1,95 +0,0 @@
|
||||
# Hanzo DB - Multi-Backend Database Abstraction
|
||||
|
||||
A unified database abstraction layer for Hanzo Node that supports multiple backend databases, each optimized for different workloads.
|
||||
|
||||
## Supported Backends
|
||||
|
||||
| Backend | Best For | Features |
|
||||
|---------|----------|----------|
|
||||
| **LanceDB** | Vector Search, AI/ML | • Native vector operations<br>• Multimodal storage<br>• Columnar format<br>• Fast similarity search |
|
||||
| **DuckDB** | Analytics, OLAP | • In-process analytical database<br>• SQL support<br>• Columnar storage<br>• Fast aggregations |
|
||||
| **PostgreSQL** | Transactional, OLTP | • ACID compliance<br>• Rich SQL features<br>• Extensions (pgvector)<br>• Battle-tested |
|
||||
| **Redis** | Caching, Real-time | • In-memory storage<br>• Pub/sub support<br>• TTL/expiration<br>• Extremely fast |
|
||||
| **SQLite** | Embedded, Lightweight | • Zero-configuration<br>• Single file<br>• Serverless<br>• Wide compatibility |
|
||||
|
||||
## Usage
|
||||
|
||||
```rust
|
||||
use hanzo_db::{connect, HanzoDbConfig, DatabaseBackend, WorkloadType};
|
||||
|
||||
// Automatically select backend based on workload
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::for_workload(WorkloadType::VectorSearch),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Or explicitly choose a backend
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::LanceDB,
|
||||
path: Some("./data/vectors".into()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let db = connect(config).await?;
|
||||
|
||||
// Use unified interface regardless of backend
|
||||
db.create_table("embeddings", schema).await?;
|
||||
db.insert("embeddings", &records).await?;
|
||||
let results = db.vector_search(query).await?;
|
||||
```
|
||||
|
||||
## Features
|
||||
|
||||
- **Unified Interface**: Same API across all backends
|
||||
- **Automatic Backend Selection**: Choose optimal backend based on workload
|
||||
- **Migration Support**: Move data between backends
|
||||
- **Connection Pooling**: Efficient resource management
|
||||
- **Transaction Support**: ACID guarantees where supported
|
||||
- **Vector Operations**: Native support in LanceDB, extension support in PostgreSQL
|
||||
|
||||
## Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
```bash
|
||||
# Select default backend
|
||||
HANZO_DB_BACKEND=lancedb # lancedb, duckdb, postgresql, redis, sqlite
|
||||
|
||||
# Backend-specific configuration
|
||||
HANZO_DB_PATH=./storage/hanzo-db
|
||||
HANZO_DB_URL=postgresql://user:pass@localhost/hanzo
|
||||
HANZO_REDIS_URL=redis://localhost:6379
|
||||
```
|
||||
|
||||
### Feature Flags
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
hanzo_db = { version = "1.0", features = ["lancedb", "duckdb", "postgres"] }
|
||||
```
|
||||
|
||||
## Migration
|
||||
|
||||
Migrate from one backend to another:
|
||||
|
||||
```bash
|
||||
# From SQLite to LanceDB
|
||||
hanzo-migrate --from sqlite://old.db --to lancedb://./vectors
|
||||
|
||||
# From LanceDB to PostgreSQL
|
||||
hanzo-migrate --from lancedb://./vectors --to postgresql://localhost/hanzo
|
||||
```
|
||||
|
||||
## Performance Characteristics
|
||||
|
||||
| Operation | LanceDB | DuckDB | PostgreSQL | Redis | SQLite |
|
||||
|-----------|---------|---------|------------|-------|---------|
|
||||
| Vector Search | ⭐⭐⭐⭐⭐ | ⭐⭐ | ⭐⭐⭐ | ⭐ | ⭐ |
|
||||
| Analytics | ⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | ⭐ | ⭐⭐ |
|
||||
| Transactions | ⭐⭐⭐ | ⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐ | ⭐⭐⭐⭐ |
|
||||
| Caching | ⭐⭐ | ⭐ | ⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ |
|
||||
| Embedded | ⭐⭐⭐⭐ | ⭐⭐⭐⭐ | ⭐ | ⭐ | ⭐⭐⭐⭐⭐ |
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
@@ -1,467 +0,0 @@
|
||||
//! DuckDB backend implementation for analytical workloads
|
||||
//!
|
||||
//! DuckDB is an embedded analytical database optimized for OLAP workloads.
|
||||
//! It provides excellent performance for analytics, aggregations, and complex queries.
|
||||
|
||||
use crate::{DatabaseBackend, HanzoDbConfig, HanzoDbError, Result};
|
||||
use async_trait::async_trait;
|
||||
use duckdb::{Connection, params};
|
||||
use serde_json::Value;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// DuckDB backend for analytical workloads
|
||||
pub struct DuckDbBackend {
|
||||
connection: Arc<Mutex<Connection>>,
|
||||
path: PathBuf,
|
||||
config: HanzoDbConfig,
|
||||
}
|
||||
|
||||
impl DuckDbBackend {
|
||||
/// Create a new DuckDB backend
|
||||
pub fn new(config: HanzoDbConfig) -> Result<Self> {
|
||||
let path = config.path.clone().unwrap_or_else(|| {
|
||||
PathBuf::from("./storage/hanzo-duckdb.db")
|
||||
});
|
||||
|
||||
// Create parent directory if needed
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
}
|
||||
|
||||
// Open DuckDB connection
|
||||
let connection = if path.to_str() == Some(":memory:") {
|
||||
Connection::open_in_memory()
|
||||
} else {
|
||||
Connection::open(&path)
|
||||
}.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Configure DuckDB for optimal performance
|
||||
connection.execute_batch("
|
||||
-- Set memory limit
|
||||
SET memory_limit='4GB';
|
||||
-- Enable parallel execution
|
||||
SET threads=4;
|
||||
-- Enable progress bar for long queries
|
||||
SET enable_progress_bar=true;
|
||||
-- Optimize for analytics
|
||||
SET enable_optimizer=true;
|
||||
").map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(Self {
|
||||
connection: Arc::new(Mutex::new(connection)),
|
||||
path,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute analytical query
|
||||
pub fn analyze(&self, query: &str) -> Result<Vec<Value>> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let mut stmt = conn.prepare(query)
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let rows = stmt.query_map(params![], |row| {
|
||||
// Convert row to JSON value
|
||||
let mut obj = serde_json::Map::new();
|
||||
|
||||
// Get column count and names
|
||||
let column_count = row.column_count();
|
||||
for i in 0..column_count {
|
||||
let column_name = row.column_name(i)
|
||||
.unwrap_or(&format!("column_{}", i))
|
||||
.to_string();
|
||||
|
||||
// Try to get value as different types
|
||||
if let Ok(val) = row.get::<_, String>(i) {
|
||||
obj.insert(column_name, Value::String(val));
|
||||
} else if let Ok(val) = row.get::<_, i64>(i) {
|
||||
obj.insert(column_name, Value::Number(val.into()));
|
||||
} else if let Ok(val) = row.get::<_, f64>(i) {
|
||||
obj.insert(column_name, Value::from(val));
|
||||
} else if let Ok(val) = row.get::<_, bool>(i) {
|
||||
obj.insert(column_name, Value::Bool(val));
|
||||
} else {
|
||||
obj.insert(column_name, Value::Null);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Value::Object(obj))
|
||||
}).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
rows.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))
|
||||
}
|
||||
|
||||
/// Import Parquet file
|
||||
pub fn import_parquet(&self, path: &str, table_name: &str) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let query = format!(
|
||||
"CREATE TABLE {} AS SELECT * FROM read_parquet('{}')",
|
||||
table_name, path
|
||||
);
|
||||
|
||||
conn.execute(&query, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Export to Parquet
|
||||
pub fn export_parquet(&self, table_name: &str, path: &str) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let query = format!(
|
||||
"COPY {} TO '{}' (FORMAT PARQUET)",
|
||||
table_name, path
|
||||
);
|
||||
|
||||
conn.execute(&query, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create aggregate view
|
||||
pub fn create_aggregate_view(&self, view_name: &str, query: &str) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let create_view = format!(
|
||||
"CREATE VIEW {} AS {}",
|
||||
view_name, query
|
||||
);
|
||||
|
||||
conn.execute(&create_view, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatabaseBackend for DuckDbBackend {
|
||||
async fn connect(&mut self) -> Result<()> {
|
||||
// Already connected in new()
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&mut self) -> Result<()> {
|
||||
// Connection will be closed when dropped
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_table(&self, name: &str, schema: &Value) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Parse schema and create SQL
|
||||
let columns = schema.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidSchema("Schema must be object".to_string()))?;
|
||||
|
||||
let column_defs: Vec<String> = columns.iter().map(|(name, type_def)| {
|
||||
let sql_type = match type_def.as_str() {
|
||||
Some("string") => "VARCHAR",
|
||||
Some("integer") => "INTEGER",
|
||||
Some("bigint") => "BIGINT",
|
||||
Some("float") => "REAL",
|
||||
Some("double") => "DOUBLE",
|
||||
Some("boolean") => "BOOLEAN",
|
||||
Some("timestamp") => "TIMESTAMP",
|
||||
Some("date") => "DATE",
|
||||
Some("json") => "JSON",
|
||||
Some("vector") => "DOUBLE[]", // DuckDB array type for vectors
|
||||
_ => "VARCHAR",
|
||||
};
|
||||
format!("{} {}", name, sql_type)
|
||||
}).collect();
|
||||
|
||||
let create_sql = format!(
|
||||
"CREATE TABLE IF NOT EXISTS {} ({})",
|
||||
name,
|
||||
column_defs.join(", ")
|
||||
);
|
||||
|
||||
conn.execute(&create_sql, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_table(&self, name: &str) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
conn.execute(&format!("DROP TABLE IF EXISTS {}", name), params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert(&self, table: &str, records: &[Value]) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Begin transaction for bulk insert
|
||||
let tx = conn.unchecked_transaction()
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))?;
|
||||
|
||||
for record in records {
|
||||
let obj = record.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Record must be object".to_string()))?;
|
||||
|
||||
let columns: Vec<String> = obj.keys().cloned().collect();
|
||||
let placeholders: Vec<String> = (1..=columns.len())
|
||||
.map(|i| format!("${}", i))
|
||||
.collect();
|
||||
|
||||
let insert_sql = format!(
|
||||
"INSERT INTO {} ({}) VALUES ({})",
|
||||
table,
|
||||
columns.join(", "),
|
||||
placeholders.join(", ")
|
||||
);
|
||||
|
||||
// Convert values to DuckDB types
|
||||
let values: Vec<duckdb::types::Value> = obj.values().map(|v| {
|
||||
match v {
|
||||
Value::String(s) => duckdb::types::Value::Text(s.clone()),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
duckdb::types::Value::BigInt(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
duckdb::types::Value::Double(f)
|
||||
} else {
|
||||
duckdb::types::Value::Null
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => duckdb::types::Value::Boolean(*b),
|
||||
Value::Null => duckdb::types::Value::Null,
|
||||
Value::Array(a) => {
|
||||
// Handle vector data
|
||||
let floats: Vec<f64> = a.iter()
|
||||
.filter_map(|v| v.as_f64())
|
||||
.collect();
|
||||
duckdb::types::Value::List(
|
||||
floats.into_iter()
|
||||
.map(duckdb::types::Value::Double)
|
||||
.collect()
|
||||
)
|
||||
},
|
||||
Value::Object(_) => {
|
||||
// Store as JSON string
|
||||
duckdb::types::Value::Text(v.to_string())
|
||||
}
|
||||
}
|
||||
}).collect();
|
||||
|
||||
tx.execute(&insert_sql, params_from_iter(values))
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query(&self, query: &str) -> Result<Vec<Value>> {
|
||||
self.analyze(query)
|
||||
}
|
||||
|
||||
async fn update(&self, table: &str, id: &str, data: &Value) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let obj = data.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Update data must be object".to_string()))?;
|
||||
|
||||
let set_clauses: Vec<String> = obj.keys()
|
||||
.enumerate()
|
||||
.map(|(i, key)| format!("{} = ${}", key, i + 2))
|
||||
.collect();
|
||||
|
||||
let update_sql = format!(
|
||||
"UPDATE {} SET {} WHERE id = $1",
|
||||
table,
|
||||
set_clauses.join(", ")
|
||||
);
|
||||
|
||||
let mut values = vec![duckdb::types::Value::Text(id.to_string())];
|
||||
for value in obj.values() {
|
||||
values.push(json_to_duckdb_value(value));
|
||||
}
|
||||
|
||||
conn.execute(&update_sql, params_from_iter(values))
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, table: &str, id: &str) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
conn.execute(
|
||||
&format!("DELETE FROM {} WHERE id = ?", table),
|
||||
params![id]
|
||||
).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn vector_search(
|
||||
&self,
|
||||
table: &str,
|
||||
query_vector: &[f32],
|
||||
limit: usize,
|
||||
_filter: Option<&str>,
|
||||
) -> Result<Vec<Value>> {
|
||||
// DuckDB doesn't have native vector similarity search
|
||||
// but we can implement it using array operations
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Convert query vector to DuckDB array literal
|
||||
let vector_str = format!(
|
||||
"[{}]",
|
||||
query_vector.iter()
|
||||
.map(|f| f.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
);
|
||||
|
||||
// Calculate cosine similarity using array operations
|
||||
let query = format!(
|
||||
"SELECT *,
|
||||
list_cosine_similarity(embedding, {}::DOUBLE[]) as similarity
|
||||
FROM {}
|
||||
ORDER BY similarity DESC
|
||||
LIMIT {}",
|
||||
vector_str, table, limit
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(&query)
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let rows = stmt.query_map(params![], row_to_json)
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
rows.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))
|
||||
}
|
||||
|
||||
fn backend_type(&self) -> DatabaseBackend {
|
||||
DatabaseBackend::DuckDB
|
||||
}
|
||||
|
||||
fn is_vector_capable(&self) -> bool {
|
||||
true // DuckDB can store vectors as arrays
|
||||
}
|
||||
|
||||
fn is_analytical(&self) -> bool {
|
||||
true // DuckDB excels at analytics
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to convert JSON to DuckDB value
|
||||
fn json_to_duckdb_value(value: &Value) -> duckdb::types::Value {
|
||||
match value {
|
||||
Value::String(s) => duckdb::types::Value::Text(s.clone()),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
duckdb::types::Value::BigInt(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
duckdb::types::Value::Double(f)
|
||||
} else {
|
||||
duckdb::types::Value::Null
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => duckdb::types::Value::Boolean(*b),
|
||||
Value::Null => duckdb::types::Value::Null,
|
||||
Value::Array(a) => {
|
||||
let values: Vec<duckdb::types::Value> = a.iter()
|
||||
.map(json_to_duckdb_value)
|
||||
.collect();
|
||||
duckdb::types::Value::List(values)
|
||||
},
|
||||
Value::Object(_) => duckdb::types::Value::Text(value.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to convert row to JSON
|
||||
fn row_to_json(row: &duckdb::Row) -> std::result::Result<Value, duckdb::Error> {
|
||||
let mut obj = serde_json::Map::new();
|
||||
|
||||
let column_count = row.column_count();
|
||||
for i in 0..column_count {
|
||||
let column_name = row.column_name(i)
|
||||
.unwrap_or(&format!("column_{}", i))
|
||||
.to_string();
|
||||
|
||||
// Try different types
|
||||
if let Ok(val) = row.get::<_, String>(i) {
|
||||
obj.insert(column_name, Value::String(val));
|
||||
} else if let Ok(val) = row.get::<_, i64>(i) {
|
||||
obj.insert(column_name, Value::Number(val.into()));
|
||||
} else if let Ok(val) = row.get::<_, f64>(i) {
|
||||
obj.insert(column_name, Value::from(val));
|
||||
} else if let Ok(val) = row.get::<_, bool>(i) {
|
||||
obj.insert(column_name, Value::Bool(val));
|
||||
} else {
|
||||
obj.insert(column_name, Value::Null);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Value::Object(obj))
|
||||
}
|
||||
|
||||
// Helper to create params from iterator
|
||||
fn params_from_iter(values: Vec<duckdb::types::Value>) -> Vec<duckdb::types::Value> {
|
||||
values
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_duckdb_analytics() {
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::DuckDB,
|
||||
path: Some(PathBuf::from(":memory:")),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let backend = DuckDbBackend::new(config).unwrap();
|
||||
|
||||
// Create table
|
||||
let schema = serde_json::json!({
|
||||
"id": "integer",
|
||||
"value": "double",
|
||||
"category": "string"
|
||||
});
|
||||
backend.create_table("metrics", &schema).await.unwrap();
|
||||
|
||||
// Insert test data
|
||||
let records = vec![
|
||||
serde_json::json!({"id": 1, "value": 100.0, "category": "A"}),
|
||||
serde_json::json!({"id": 2, "value": 200.0, "category": "B"}),
|
||||
serde_json::json!({"id": 3, "value": 150.0, "category": "A"}),
|
||||
];
|
||||
backend.insert("metrics", &records).await.unwrap();
|
||||
|
||||
// Run analytical query
|
||||
let results = backend.analyze(
|
||||
"SELECT category, AVG(value) as avg_value FROM metrics GROUP BY category"
|
||||
).unwrap();
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -1,584 +0,0 @@
|
||||
//! LanceDB backend implementation - The PRIMARY backend for Hanzo Node
|
||||
//!
|
||||
//! LanceDB is the default and recommended backend for Hanzo Node because:
|
||||
//! - Native vector search with no external dependencies
|
||||
//! - Embedded operation - no server required
|
||||
//! - Multimodal storage (text, images, embeddings, audio, video)
|
||||
//! - Columnar storage based on Apache Arrow
|
||||
//! - Production-ready with automatic versioning
|
||||
//! - Scales from edge to cloud
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use lancedb::{connect, Connection, Table};
|
||||
use arrow_array::{RecordBatch, RecordBatchIterator, Float32Array, StringArray, Int64Array};
|
||||
use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
|
||||
use std::sync::Arc;
|
||||
use std::path::PathBuf;
|
||||
use log::{debug, info};
|
||||
|
||||
use crate::{
|
||||
HanzoDatabase, HanzoDbConfig, TableSchema, Query, QueryResult,
|
||||
VectorQuery, SearchResult, Transaction, DatabaseStats, Record,
|
||||
Value, DataType, Column, Filter, OrderBy, Index, IndexType, DistanceMetric
|
||||
};
|
||||
|
||||
/// LanceDB backend implementation
|
||||
pub struct LanceDbBackend {
|
||||
connection: Connection,
|
||||
path: PathBuf,
|
||||
config: HanzoDbConfig,
|
||||
}
|
||||
|
||||
impl LanceDbBackend {
|
||||
/// Create a new LanceDB backend
|
||||
pub async fn new(config: HanzoDbConfig) -> Result<Self> {
|
||||
let path = config.path.clone()
|
||||
.unwrap_or_else(|| PathBuf::from("./storage/hanzo-db/lancedb"));
|
||||
|
||||
// Ensure directory exists
|
||||
std::fs::create_dir_all(&path)?;
|
||||
|
||||
info!("🚀 Initializing LanceDB at {:?}", path);
|
||||
info!("📦 Native vector search enabled");
|
||||
info!("🔐 Embedded operation - no external dependencies");
|
||||
info!("🎯 Multimodal storage ready (text, images, embeddings)");
|
||||
|
||||
let connection = connect(path.to_str().unwrap()).await
|
||||
.context("Failed to connect to LanceDB")?;
|
||||
|
||||
Ok(Self {
|
||||
connection,
|
||||
path,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert our DataType to Arrow DataType
|
||||
fn to_arrow_type(dt: &DataType) -> ArrowDataType {
|
||||
match dt {
|
||||
DataType::Boolean => ArrowDataType::Boolean,
|
||||
DataType::Int32 => ArrowDataType::Int32,
|
||||
DataType::Int64 => ArrowDataType::Int64,
|
||||
DataType::Float32 => ArrowDataType::Float32,
|
||||
DataType::Float64 => ArrowDataType::Float64,
|
||||
DataType::String => ArrowDataType::Utf8,
|
||||
DataType::Binary => ArrowDataType::Binary,
|
||||
DataType::Timestamp => ArrowDataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
|
||||
DataType::Vector(dim) => ArrowDataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", ArrowDataType::Float32, false)),
|
||||
*dim as i32,
|
||||
),
|
||||
DataType::Json => ArrowDataType::Utf8, // Store JSON as string
|
||||
DataType::Array(inner) => ArrowDataType::List(
|
||||
Arc::new(Field::new("item", Self::to_arrow_type(inner), true))
|
||||
),
|
||||
DataType::Struct(fields) => {
|
||||
let arrow_fields: Vec<Field> = fields.iter()
|
||||
.map(|(name, dt)| Field::new(name, Self::to_arrow_type(dt), true))
|
||||
.collect();
|
||||
ArrowDataType::Struct(arrow_fields.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl HanzoDatabase for LanceDbBackend {
|
||||
async fn init(&self) -> Result<()> {
|
||||
info!("✅ LanceDB initialized successfully");
|
||||
debug!("Storage path: {:?}", self.path);
|
||||
debug!("Cache size: {:?}", self.config.cache_size);
|
||||
debug!("Compression: {}", self.config.enable_compression);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_table(&self, name: &str, schema: TableSchema) -> Result<()> {
|
||||
// Convert schema to Arrow schema
|
||||
let fields: Vec<Field> = schema.columns.iter()
|
||||
.map(|col| Field::new(&col.name, Self::to_arrow_type(&col.data_type), col.nullable))
|
||||
.collect();
|
||||
|
||||
let arrow_schema = Arc::new(ArrowSchema::new(fields));
|
||||
|
||||
// Create empty batch for table creation
|
||||
let batch = RecordBatch::new_empty(arrow_schema.clone());
|
||||
let batches = vec![batch];
|
||||
let batch_iter = RecordBatchIterator::new(batches.into_iter().map(Ok), arrow_schema);
|
||||
|
||||
// Create table
|
||||
self.connection.create_table(name, batch_iter, None).await
|
||||
.context(format!("Failed to create table {}", name))?;
|
||||
|
||||
// Create indexes
|
||||
for index in schema.indexes {
|
||||
match index.index_type {
|
||||
IndexType::IVF_PQ { nlist, nprobe } => {
|
||||
// Create IVF_PQ index for vector columns
|
||||
if index.columns.len() == 1 {
|
||||
let table = self.connection.open_table(name).await?;
|
||||
|
||||
// LanceDB specific index creation
|
||||
debug!("Creating IVF_PQ index on {} with nlist={}, nprobe={}",
|
||||
index.columns[0], nlist, nprobe);
|
||||
|
||||
// Note: Actual index creation API depends on LanceDB version
|
||||
// This is a placeholder for the real implementation
|
||||
}
|
||||
}
|
||||
IndexType::HNSW { max_elements, m } => {
|
||||
debug!("Creating HNSW index with max_elements={}, m={}", max_elements, m);
|
||||
// HNSW index creation
|
||||
}
|
||||
_ => {
|
||||
// Other index types
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!("✅ Table '{}' created with {} columns and {} indexes",
|
||||
name, schema.columns.len(), schema.indexes.len());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert(&self, table_name: &str, data: &[Record]) -> Result<()> {
|
||||
if data.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let table = self.connection.open_table(table_name).await
|
||||
.context(format!("Failed to open table {}", table_name))?;
|
||||
|
||||
// Get table schema
|
||||
let schema = table.schema().await?;
|
||||
|
||||
// Convert records to Arrow arrays
|
||||
let mut columns: Vec<Arc<dyn arrow_array::Array>> = Vec::new();
|
||||
|
||||
for field in schema.fields() {
|
||||
let field_name = field.name();
|
||||
|
||||
// Collect values for this column from all records
|
||||
let values: Vec<Option<Value>> = data.iter()
|
||||
.map(|record| {
|
||||
record.values.iter()
|
||||
.find(|(name, _)| name == field_name)
|
||||
.map(|(_, value)| value.clone())
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Convert to appropriate Arrow array
|
||||
let array: Arc<dyn arrow_array::Array> = match field.data_type() {
|
||||
ArrowDataType::Int64 => {
|
||||
let vals: Vec<Option<i64>> = values.iter()
|
||||
.map(|v| match v {
|
||||
Some(Value::Int64(i)) => Some(*i),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
Arc::new(Int64Array::from(vals))
|
||||
}
|
||||
ArrowDataType::Float32 => {
|
||||
let vals: Vec<Option<f32>> = values.iter()
|
||||
.map(|v| match v {
|
||||
Some(Value::Float32(f)) => Some(*f),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
Arc::new(Float32Array::from(vals))
|
||||
}
|
||||
ArrowDataType::Utf8 => {
|
||||
let vals: Vec<Option<String>> = values.iter()
|
||||
.map(|v| match v {
|
||||
Some(Value::String(s)) => Some(s.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
Arc::new(StringArray::from(vals))
|
||||
}
|
||||
ArrowDataType::FixedSizeList(_, dim) => {
|
||||
// Handle vector columns
|
||||
let vectors: Vec<Option<Vec<f32>>> = values.iter()
|
||||
.map(|v| match v {
|
||||
Some(Value::Vector(vec)) => Some(vec.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Create FixedSizeList array
|
||||
let mut builder = arrow_array::builder::FixedSizeListBuilder::new(
|
||||
arrow_array::builder::Float32Builder::new(),
|
||||
*dim
|
||||
);
|
||||
|
||||
for vec_opt in vectors {
|
||||
match vec_opt {
|
||||
Some(vec) => {
|
||||
builder.values().append_slice(&vec);
|
||||
builder.append(true);
|
||||
}
|
||||
None => builder.append(false),
|
||||
}
|
||||
}
|
||||
|
||||
Arc::new(builder.finish())
|
||||
}
|
||||
_ => {
|
||||
// Handle other types
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
columns.push(array);
|
||||
}
|
||||
|
||||
// Create RecordBatch
|
||||
let batch = RecordBatch::try_new(schema.clone(), columns)?;
|
||||
|
||||
// Insert batch
|
||||
table.add(vec![batch]).await?;
|
||||
|
||||
debug!("✅ Inserted {} records into table '{}'", data.len(), table_name);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query(&self, query: Query) -> Result<QueryResult> {
|
||||
let table = self.connection.open_table(&query.table).await?;
|
||||
|
||||
// Build query (simplified - real implementation would be more complex)
|
||||
let mut results = table.query().await?;
|
||||
|
||||
// Apply filters
|
||||
if let Some(filter) = &query.filter {
|
||||
// Convert filter to LanceDB filter format
|
||||
debug!("Applying filter: {:?}", filter);
|
||||
}
|
||||
|
||||
// Apply ordering
|
||||
for order in &query.order_by {
|
||||
debug!("Ordering by {} {}", order.column, if order.ascending { "ASC" } else { "DESC" });
|
||||
}
|
||||
|
||||
// Apply limit and offset
|
||||
if let Some(limit) = query.limit {
|
||||
results = results.limit(limit);
|
||||
}
|
||||
|
||||
let batches = results.execute().await?;
|
||||
|
||||
// Convert batches to QueryResult
|
||||
let mut rows = Vec::new();
|
||||
let mut columns = Vec::new();
|
||||
|
||||
if let Some(batch) = batches.first() {
|
||||
columns = batch.schema().fields().iter()
|
||||
.map(|f| f.name().to_string())
|
||||
.collect();
|
||||
}
|
||||
|
||||
for batch in batches {
|
||||
// Convert each row
|
||||
for row_idx in 0..batch.num_rows() {
|
||||
let mut values = Vec::new();
|
||||
|
||||
for (col_idx, field) in batch.schema().fields().iter().enumerate() {
|
||||
let column = batch.column(col_idx);
|
||||
let value = Self::extract_value(column, row_idx)?;
|
||||
values.push((field.name().to_string(), value));
|
||||
}
|
||||
|
||||
rows.push(Record { values });
|
||||
}
|
||||
}
|
||||
|
||||
Ok(QueryResult {
|
||||
columns,
|
||||
row_count: rows.len(),
|
||||
rows,
|
||||
})
|
||||
}
|
||||
|
||||
async fn vector_search(&self, query: VectorQuery) -> Result<Vec<SearchResult>> {
|
||||
let table = self.connection.open_table(&query.table).await?;
|
||||
|
||||
info!("🔍 Performing vector search in table '{}' for {} nearest neighbors",
|
||||
query.table, query.k);
|
||||
|
||||
// Build vector search query
|
||||
let mut search = table.query()
|
||||
.nearest_to(&query.vector)?
|
||||
.limit(query.k);
|
||||
|
||||
// Apply filter if present
|
||||
if let Some(filter) = &query.filter {
|
||||
debug!("Applying pre-filter: {:?}", filter);
|
||||
// Convert filter to LanceDB format
|
||||
}
|
||||
|
||||
// Set distance metric
|
||||
match query.metric {
|
||||
DistanceMetric::L2 => {
|
||||
search = search.distance_type(lancedb::query::DistanceType::L2);
|
||||
}
|
||||
DistanceMetric::Cosine => {
|
||||
search = search.distance_type(lancedb::query::DistanceType::Cosine);
|
||||
}
|
||||
DistanceMetric::InnerProduct => {
|
||||
search = search.distance_type(lancedb::query::DistanceType::Dot);
|
||||
}
|
||||
}
|
||||
|
||||
let batches = search.execute().await?;
|
||||
|
||||
// Convert to SearchResult
|
||||
let mut results = Vec::new();
|
||||
|
||||
for batch in batches {
|
||||
for row_idx in 0..batch.num_rows() {
|
||||
let mut values = Vec::new();
|
||||
let mut score = 0.0f32;
|
||||
|
||||
for (col_idx, field) in batch.schema().fields().iter().enumerate() {
|
||||
let column = batch.column(col_idx);
|
||||
|
||||
// Check if this is the distance column
|
||||
if field.name() == "_distance" {
|
||||
if let Ok(Value::Float32(dist)) = Self::extract_value(column, row_idx) {
|
||||
score = dist;
|
||||
}
|
||||
} else {
|
||||
let value = Self::extract_value(column, row_idx)?;
|
||||
values.push((field.name().to_string(), value));
|
||||
}
|
||||
}
|
||||
|
||||
results.push(SearchResult {
|
||||
record: Record { values },
|
||||
score,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
info!("✅ Found {} results with scores ranging from {:.4} to {:.4}",
|
||||
results.len(),
|
||||
results.first().map(|r| r.score).unwrap_or(0.0),
|
||||
results.last().map(|r| r.score).unwrap_or(0.0)
|
||||
);
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
async fn begin_transaction(&self) -> Result<Transaction> {
|
||||
// LanceDB handles transactions automatically
|
||||
// This is a simplified implementation
|
||||
Ok(Transaction {
|
||||
inner: Arc::new(tokio::sync::RwLock::new(crate::TransactionInner {
|
||||
backend: crate::DatabaseBackend::LanceDB,
|
||||
handle: Box::new(()),
|
||||
})),
|
||||
})
|
||||
}
|
||||
|
||||
async fn optimize(&self) -> Result<()> {
|
||||
info!("🔧 Optimizing LanceDB database...");
|
||||
|
||||
// List all tables
|
||||
let tables = self.connection.table_names().await?;
|
||||
|
||||
for table_name in tables {
|
||||
let table = self.connection.open_table(&table_name).await?;
|
||||
|
||||
// Compact table (merge small files)
|
||||
table.compact().await?;
|
||||
|
||||
// Clean up old versions
|
||||
table.cleanup_old_versions().await?;
|
||||
|
||||
debug!("✅ Optimized table '{}'", table_name);
|
||||
}
|
||||
|
||||
info!("✅ Database optimization complete");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stats(&self) -> Result<DatabaseStats> {
|
||||
let tables = self.connection.table_names().await?;
|
||||
let table_count = tables.len();
|
||||
|
||||
let mut total_rows = 0;
|
||||
let mut index_count = 0;
|
||||
|
||||
for table_name in &tables {
|
||||
let table = self.connection.open_table(table_name).await?;
|
||||
total_rows += table.count_rows().await?;
|
||||
// Count indexes (simplified)
|
||||
index_count += 1; // Assume at least one index per table
|
||||
}
|
||||
|
||||
// Calculate storage size
|
||||
let mut total_size_bytes = 0;
|
||||
if let Ok(entries) = std::fs::read_dir(&self.path) {
|
||||
for entry in entries.flatten() {
|
||||
if let Ok(metadata) = entry.metadata() {
|
||||
total_size_bytes += metadata.len() as usize;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(DatabaseStats {
|
||||
backend: crate::DatabaseBackend::LanceDB,
|
||||
table_count,
|
||||
total_rows,
|
||||
total_size_bytes,
|
||||
index_count,
|
||||
cache_hit_rate: 0.95, // LanceDB has excellent caching
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl LanceDbBackend {
|
||||
/// Extract value from Arrow column
|
||||
fn extract_value(column: &dyn arrow_array::Array, row_idx: usize) -> Result<Value> {
|
||||
use arrow_array::cast::as_primitive_array;
|
||||
|
||||
match column.data_type() {
|
||||
ArrowDataType::Int64 => {
|
||||
let array = as_primitive_array::<arrow_array::types::Int64Type>(column);
|
||||
Ok(if array.is_null(row_idx) {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::Int64(array.value(row_idx))
|
||||
})
|
||||
}
|
||||
ArrowDataType::Float32 => {
|
||||
let array = as_primitive_array::<arrow_array::types::Float32Type>(column);
|
||||
Ok(if array.is_null(row_idx) {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::Float32(array.value(row_idx))
|
||||
})
|
||||
}
|
||||
ArrowDataType::Utf8 => {
|
||||
let array = column.as_any().downcast_ref::<StringArray>().unwrap();
|
||||
Ok(if array.is_null(row_idx) {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::String(array.value(row_idx).to_string())
|
||||
})
|
||||
}
|
||||
_ => Ok(Value::Null),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lancedb_init() {
|
||||
let config = HanzoDbConfig::default();
|
||||
let backend = LanceDbBackend::new(config).await.unwrap();
|
||||
backend.init().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_table() {
|
||||
let config = HanzoDbConfig::default();
|
||||
let backend = LanceDbBackend::new(config).await.unwrap();
|
||||
|
||||
let schema = TableSchema {
|
||||
columns: vec![
|
||||
Column {
|
||||
name: "id".to_string(),
|
||||
data_type: DataType::Int64,
|
||||
nullable: false,
|
||||
default: None,
|
||||
},
|
||||
Column {
|
||||
name: "embedding".to_string(),
|
||||
data_type: DataType::Vector(384),
|
||||
nullable: false,
|
||||
default: None,
|
||||
},
|
||||
Column {
|
||||
name: "text".to_string(),
|
||||
data_type: DataType::String,
|
||||
nullable: true,
|
||||
default: None,
|
||||
},
|
||||
],
|
||||
indexes: vec![
|
||||
Index {
|
||||
name: "embedding_idx".to_string(),
|
||||
columns: vec!["embedding".to_string()],
|
||||
index_type: IndexType::IVF_PQ { nlist: 100, nprobe: 10 },
|
||||
},
|
||||
],
|
||||
constraints: vec![],
|
||||
};
|
||||
|
||||
backend.create_table("test_embeddings", schema).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_vector_search() {
|
||||
let config = HanzoDbConfig::default();
|
||||
let backend = LanceDbBackend::new(config).await.unwrap();
|
||||
|
||||
// Create table first
|
||||
let schema = TableSchema {
|
||||
columns: vec![
|
||||
Column {
|
||||
name: "id".to_string(),
|
||||
data_type: DataType::Int64,
|
||||
nullable: false,
|
||||
default: None,
|
||||
},
|
||||
Column {
|
||||
name: "embedding".to_string(),
|
||||
data_type: DataType::Vector(3),
|
||||
nullable: false,
|
||||
default: None,
|
||||
},
|
||||
],
|
||||
indexes: vec![],
|
||||
constraints: vec![],
|
||||
};
|
||||
|
||||
backend.create_table("vectors", schema).await.unwrap();
|
||||
|
||||
// Insert test data
|
||||
let records = vec![
|
||||
Record {
|
||||
values: vec![
|
||||
("id".to_string(), Value::Int64(1)),
|
||||
("embedding".to_string(), Value::Vector(vec![0.1, 0.2, 0.3])),
|
||||
],
|
||||
},
|
||||
Record {
|
||||
values: vec![
|
||||
("id".to_string(), Value::Int64(2)),
|
||||
("embedding".to_string(), Value::Vector(vec![0.4, 0.5, 0.6])),
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
backend.insert("vectors", &records).await.unwrap();
|
||||
|
||||
// Search
|
||||
let query = VectorQuery {
|
||||
table: "vectors".to_string(),
|
||||
vector: vec![0.15, 0.25, 0.35],
|
||||
k: 2,
|
||||
filter: None,
|
||||
metric: DistanceMetric::L2,
|
||||
};
|
||||
|
||||
let results = backend.vector_search(query).await.unwrap();
|
||||
assert_eq!(results.len(), 2);
|
||||
assert!(results[0].score < results[1].score); // Closer match has lower L2 distance
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
//! Database backend implementations
|
||||
|
||||
pub mod lancedb;
|
||||
pub mod duckdb;
|
||||
pub mod postgres;
|
||||
pub mod redis;
|
||||
pub mod sqlite;
|
||||
@@ -1,518 +0,0 @@
|
||||
//! PostgreSQL backend implementation for transactional workloads
|
||||
//!
|
||||
//! PostgreSQL is a powerful, open-source relational database with excellent
|
||||
//! ACID compliance, rich SQL features, and extension support (including pgvector).
|
||||
|
||||
use crate::{DatabaseBackend, HanzoDbConfig, HanzoDbError, Result};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use sqlx::{postgres::{PgPoolOptions, PgPool}, Row};
|
||||
use std::time::Duration;
|
||||
|
||||
/// PostgreSQL backend for transactional workloads
|
||||
pub struct PostgresBackend {
|
||||
pool: PgPool,
|
||||
config: HanzoDbConfig,
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
/// Create a new PostgreSQL backend
|
||||
pub async fn new(config: HanzoDbConfig) -> Result<Self> {
|
||||
let database_url = config.url.as_ref()
|
||||
.ok_or_else(|| HanzoDbError::ConfigError(
|
||||
"PostgreSQL requires database URL".to_string()
|
||||
))?;
|
||||
|
||||
// Create connection pool
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(config.max_connections.unwrap_or(10))
|
||||
.min_connections(config.min_connections.unwrap_or(2))
|
||||
.connect_timeout(Duration::from_secs(config.connect_timeout.unwrap_or(10)))
|
||||
.idle_timeout(Duration::from_secs(config.idle_timeout.unwrap_or(300)))
|
||||
.connect(database_url)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Test connection
|
||||
sqlx::query("SELECT 1")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
Ok(Self { pool, config })
|
||||
}
|
||||
|
||||
/// Enable pgvector extension
|
||||
pub async fn enable_vector_extension(&self) -> Result<()> {
|
||||
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create vector index for similarity search
|
||||
pub async fn create_vector_index(
|
||||
&self,
|
||||
table: &str,
|
||||
column: &str,
|
||||
dimensions: usize,
|
||||
index_type: &str, // "ivfflat" or "hnsw"
|
||||
) -> Result<()> {
|
||||
let index_name = format!("{}_{}_idx", table, column);
|
||||
|
||||
let query = match index_type {
|
||||
"ivfflat" => {
|
||||
format!(
|
||||
"CREATE INDEX {} ON {} USING ivfflat ({} vector_l2_ops) WITH (lists = {})",
|
||||
index_name, table, column, dimensions / 10
|
||||
)
|
||||
},
|
||||
"hnsw" => {
|
||||
format!(
|
||||
"CREATE INDEX {} ON {} USING hnsw ({} vector_l2_ops)",
|
||||
index_name, table, column
|
||||
)
|
||||
},
|
||||
_ => {
|
||||
return Err(HanzoDbError::InvalidSchema(
|
||||
format!("Unknown index type: {}", index_type)
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
sqlx::query(&query)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Begin transaction
|
||||
pub async fn begin_transaction(&self) -> Result<sqlx::Transaction<'_, sqlx::Postgres>> {
|
||||
self.pool.begin()
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))
|
||||
}
|
||||
|
||||
/// Execute raw SQL with parameters
|
||||
pub async fn execute_raw(&self, query: &str, params: Vec<Value>) -> Result<u64> {
|
||||
let mut q = sqlx::query(query);
|
||||
|
||||
for param in params {
|
||||
q = match param {
|
||||
Value::String(s) => q.bind(s),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
q.bind(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
q.bind(f)
|
||||
} else {
|
||||
q.bind(Option::<i64>::None)
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => q.bind(b),
|
||||
Value::Null => q.bind(Option::<String>::None),
|
||||
_ => q.bind(param.to_string()),
|
||||
};
|
||||
}
|
||||
|
||||
let result = q.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatabaseBackend for PostgresBackend {
|
||||
async fn connect(&mut self) -> Result<()> {
|
||||
// Already connected in new()
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&mut self) -> Result<()> {
|
||||
self.pool.close().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_table(&self, name: &str, schema: &Value) -> Result<()> {
|
||||
let columns = schema.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidSchema("Schema must be object".to_string()))?;
|
||||
|
||||
let mut column_defs = Vec::new();
|
||||
let mut has_vector = false;
|
||||
|
||||
for (col_name, type_def) in columns {
|
||||
let sql_type = match type_def {
|
||||
Value::String(s) => match s.as_str() {
|
||||
"string" => "TEXT",
|
||||
"integer" => "INTEGER",
|
||||
"bigint" => "BIGINT",
|
||||
"float" => "REAL",
|
||||
"double" => "DOUBLE PRECISION",
|
||||
"boolean" => "BOOLEAN",
|
||||
"timestamp" => "TIMESTAMP WITH TIME ZONE",
|
||||
"date" => "DATE",
|
||||
"json" | "jsonb" => "JSONB",
|
||||
"uuid" => "UUID",
|
||||
"vector" => {
|
||||
has_vector = true;
|
||||
"vector(1536)" // Default to OpenAI embedding size
|
||||
},
|
||||
t if t.starts_with("vector(") => {
|
||||
has_vector = true;
|
||||
t
|
||||
},
|
||||
_ => "TEXT",
|
||||
},
|
||||
Value::Object(obj) => {
|
||||
// Handle complex type definitions
|
||||
if let Some(Value::String(t)) = obj.get("type") {
|
||||
match t.as_str() {
|
||||
"vector" => {
|
||||
has_vector = true;
|
||||
if let Some(Value::Number(dim)) = obj.get("dimensions") {
|
||||
&format!("vector({})", dim)
|
||||
} else {
|
||||
"vector(1536)"
|
||||
}
|
||||
},
|
||||
_ => "TEXT",
|
||||
}
|
||||
} else {
|
||||
"TEXT"
|
||||
}
|
||||
},
|
||||
_ => "TEXT",
|
||||
};
|
||||
|
||||
column_defs.push(format!("{} {}", col_name, sql_type));
|
||||
}
|
||||
|
||||
// Enable pgvector if needed
|
||||
if has_vector {
|
||||
self.enable_vector_extension().await?;
|
||||
}
|
||||
|
||||
// Add standard columns if not present
|
||||
if !columns.contains_key("id") {
|
||||
column_defs.insert(0, "id SERIAL PRIMARY KEY".to_string());
|
||||
}
|
||||
if !columns.contains_key("created_at") {
|
||||
column_defs.push("created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()".to_string());
|
||||
}
|
||||
if !columns.contains_key("updated_at") {
|
||||
column_defs.push("updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()".to_string());
|
||||
}
|
||||
|
||||
let create_sql = format!(
|
||||
"CREATE TABLE IF NOT EXISTS {} ({})",
|
||||
name,
|
||||
column_defs.join(", ")
|
||||
);
|
||||
|
||||
sqlx::query(&create_sql)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Create update trigger for updated_at
|
||||
let trigger_sql = format!(
|
||||
"CREATE OR REPLACE TRIGGER update_{}_updated_at
|
||||
BEFORE UPDATE ON {}
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION update_updated_at_column()",
|
||||
name, name
|
||||
);
|
||||
|
||||
// First create the function if it doesn't exist
|
||||
sqlx::query(
|
||||
"CREATE OR REPLACE FUNCTION update_updated_at_column()
|
||||
RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
NEW.updated_at = NOW();
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ language 'plpgsql'"
|
||||
).execute(&self.pool).await.ok(); // Ignore if already exists
|
||||
|
||||
sqlx::query(&trigger_sql)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.ok(); // Ignore if trigger already exists
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_table(&self, name: &str) -> Result<()> {
|
||||
sqlx::query(&format!("DROP TABLE IF EXISTS {} CASCADE", name))
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert(&self, table: &str, records: &[Value]) -> Result<()> {
|
||||
let mut tx = self.begin_transaction().await?;
|
||||
|
||||
for record in records {
|
||||
let obj = record.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Record must be object".to_string()))?;
|
||||
|
||||
let columns: Vec<String> = obj.keys().cloned().collect();
|
||||
let placeholders: Vec<String> = (1..=columns.len())
|
||||
.map(|i| format!("${}", i))
|
||||
.collect();
|
||||
|
||||
let insert_sql = format!(
|
||||
"INSERT INTO {} ({}) VALUES ({})",
|
||||
table,
|
||||
columns.join(", "),
|
||||
placeholders.join(", ")
|
||||
);
|
||||
|
||||
let mut query = sqlx::query(&insert_sql);
|
||||
|
||||
for value in obj.values() {
|
||||
query = match value {
|
||||
Value::String(s) => query.bind(s),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
query.bind(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
query.bind(f)
|
||||
} else {
|
||||
query.bind(Option::<i64>::None)
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => query.bind(b),
|
||||
Value::Null => query.bind(Option::<String>::None),
|
||||
Value::Array(a) => {
|
||||
// Handle vector data
|
||||
if a.iter().all(|v| v.is_number()) {
|
||||
let floats: Vec<f32> = a.iter()
|
||||
.filter_map(|v| v.as_f64().map(|f| f as f32))
|
||||
.collect();
|
||||
query.bind(floats)
|
||||
} else {
|
||||
query.bind(value.to_string())
|
||||
}
|
||||
},
|
||||
Value::Object(_) => query.bind(value.to_string()),
|
||||
};
|
||||
}
|
||||
|
||||
query.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query(&self, query: &str) -> Result<Vec<Value>> {
|
||||
let rows = sqlx::query(query)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let mut results = Vec::new();
|
||||
|
||||
for row in rows {
|
||||
let mut obj = serde_json::Map::new();
|
||||
|
||||
// Get column information
|
||||
for (i, column) in row.columns().iter().enumerate() {
|
||||
let name = column.name();
|
||||
|
||||
// Try to get value as different types
|
||||
if let Ok(val) = row.try_get::<String, _>(i) {
|
||||
obj.insert(name.to_string(), Value::String(val));
|
||||
} else if let Ok(val) = row.try_get::<i64, _>(i) {
|
||||
obj.insert(name.to_string(), Value::Number(val.into()));
|
||||
} else if let Ok(val) = row.try_get::<i32, _>(i) {
|
||||
obj.insert(name.to_string(), Value::Number(val.into()));
|
||||
} else if let Ok(val) = row.try_get::<f64, _>(i) {
|
||||
obj.insert(name.to_string(), Value::from(val));
|
||||
} else if let Ok(val) = row.try_get::<f32, _>(i) {
|
||||
obj.insert(name.to_string(), Value::from(val as f64));
|
||||
} else if let Ok(val) = row.try_get::<bool, _>(i) {
|
||||
obj.insert(name.to_string(), Value::Bool(val));
|
||||
} else if let Ok(val) = row.try_get::<serde_json::Value, _>(i) {
|
||||
obj.insert(name.to_string(), val);
|
||||
} else if let Ok(val) = row.try_get::<Vec<f32>, _>(i) {
|
||||
// Handle vector data
|
||||
let array: Vec<Value> = val.iter().map(|f| Value::from(*f as f64)).collect();
|
||||
obj.insert(name.to_string(), Value::Array(array));
|
||||
} else {
|
||||
obj.insert(name.to_string(), Value::Null);
|
||||
}
|
||||
}
|
||||
|
||||
results.push(Value::Object(obj));
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
async fn update(&self, table: &str, id: &str, data: &Value) -> Result<()> {
|
||||
let obj = data.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Update data must be object".to_string()))?;
|
||||
|
||||
let set_clauses: Vec<String> = obj.keys()
|
||||
.enumerate()
|
||||
.map(|(i, key)| format!("{} = ${}", key, i + 2))
|
||||
.collect();
|
||||
|
||||
let update_sql = format!(
|
||||
"UPDATE {} SET {}, updated_at = NOW() WHERE id = $1",
|
||||
table,
|
||||
set_clauses.join(", ")
|
||||
);
|
||||
|
||||
let mut query = sqlx::query(&update_sql).bind(id);
|
||||
|
||||
for value in obj.values() {
|
||||
query = match value {
|
||||
Value::String(s) => query.bind(s),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
query.bind(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
query.bind(f)
|
||||
} else {
|
||||
query.bind(Option::<i64>::None)
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => query.bind(b),
|
||||
Value::Null => query.bind(Option::<String>::None),
|
||||
_ => query.bind(value.to_string()),
|
||||
};
|
||||
}
|
||||
|
||||
query.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, table: &str, id: &str) -> Result<()> {
|
||||
sqlx::query(&format!("DELETE FROM {} WHERE id = $1", table))
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn vector_search(
|
||||
&self,
|
||||
table: &str,
|
||||
query_vector: &[f32],
|
||||
limit: usize,
|
||||
filter: Option<&str>,
|
||||
) -> Result<Vec<Value>> {
|
||||
// Convert vector to PostgreSQL array format
|
||||
let vector_str = format!(
|
||||
"[{}]",
|
||||
query_vector.iter()
|
||||
.map(|f| f.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
|
||||
let where_clause = if let Some(f) = filter {
|
||||
format!("WHERE {}", f)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
// Use pgvector's <-> operator for L2 distance
|
||||
let query = format!(
|
||||
"SELECT *,
|
||||
embedding <-> '{}' as distance
|
||||
FROM {}
|
||||
{}
|
||||
ORDER BY embedding <-> '{}'
|
||||
LIMIT {}",
|
||||
vector_str, table, where_clause, vector_str, limit
|
||||
);
|
||||
|
||||
self.query(&query).await
|
||||
}
|
||||
|
||||
fn backend_type(&self) -> DatabaseBackend {
|
||||
DatabaseBackend::PostgreSQL
|
||||
}
|
||||
|
||||
fn is_vector_capable(&self) -> bool {
|
||||
true // With pgvector extension
|
||||
}
|
||||
|
||||
fn is_transactional(&self) -> bool {
|
||||
true // PostgreSQL has excellent ACID compliance
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_postgres_transactions() {
|
||||
// This test requires a PostgreSQL instance
|
||||
// Skip if not available
|
||||
let database_url = std::env::var("TEST_POSTGRES_URL")
|
||||
.unwrap_or_else(|_| "postgresql://postgres:postgres@localhost/test".to_string());
|
||||
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::PostgreSQL,
|
||||
url: Some(database_url),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let backend = match PostgresBackend::new(config).await {
|
||||
Ok(b) => b,
|
||||
Err(_) => {
|
||||
println!("PostgreSQL not available, skipping test");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Create table with vector support
|
||||
let schema = serde_json::json!({
|
||||
"name": "string",
|
||||
"embedding": {
|
||||
"type": "vector",
|
||||
"dimensions": 384
|
||||
}
|
||||
});
|
||||
|
||||
backend.create_table("test_items", &schema).await.unwrap();
|
||||
|
||||
// Test transaction
|
||||
let mut tx = backend.begin_transaction().await.unwrap();
|
||||
|
||||
sqlx::query("INSERT INTO test_items (name) VALUES ($1)")
|
||||
.bind("test")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
tx.commit().await.unwrap();
|
||||
|
||||
// Clean up
|
||||
backend.drop_table("test_items").await.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -1,497 +0,0 @@
|
||||
//! Redis backend implementation for caching and real-time operations
|
||||
//!
|
||||
//! Redis is an in-memory data structure store, perfect for caching,
|
||||
//! session management, real-time analytics, and pub/sub messaging.
|
||||
|
||||
use crate::{DatabaseBackend, HanzoDbConfig, HanzoDbError, Result};
|
||||
use async_trait::async_trait;
|
||||
use redis::{aio::ConnectionManager, AsyncCommands, Client};
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Redis backend for caching and real-time operations
|
||||
pub struct RedisBackend {
|
||||
client: Client,
|
||||
connection: ConnectionManager,
|
||||
config: HanzoDbConfig,
|
||||
}
|
||||
|
||||
impl RedisBackend {
|
||||
/// Create a new Redis backend
|
||||
pub async fn new(config: HanzoDbConfig) -> Result<Self> {
|
||||
let redis_url = config.url.as_ref()
|
||||
.unwrap_or(&"redis://127.0.0.1:6379".to_string())
|
||||
.clone();
|
||||
|
||||
// Create Redis client
|
||||
let client = Client::open(redis_url.as_str())
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Create connection manager for connection pooling
|
||||
let connection = ConnectionManager::new(client.clone()).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
Ok(Self {
|
||||
client,
|
||||
connection,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
/// Set value with optional TTL
|
||||
pub async fn set_with_ttl(&mut self, key: &str, value: &Value, ttl_seconds: Option<u64>) -> Result<()> {
|
||||
let serialized = serde_json::to_string(value)
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))?;
|
||||
|
||||
if let Some(ttl) = ttl_seconds {
|
||||
self.connection
|
||||
.set_ex(key, serialized, ttl)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
} else {
|
||||
self.connection
|
||||
.set(key, serialized)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get value by key
|
||||
pub async fn get(&mut self, key: &str) -> Result<Option<Value>> {
|
||||
let value: Option<String> = self.connection
|
||||
.get(key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
match value {
|
||||
Some(v) => {
|
||||
let parsed = serde_json::from_str(&v)
|
||||
.map_err(|e| HanzoDbError::DeserializationError(e.to_string()))?;
|
||||
Ok(Some(parsed))
|
||||
},
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Delete key
|
||||
pub async fn del(&mut self, key: &str) -> Result<()> {
|
||||
self.connection
|
||||
.del(key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if key exists
|
||||
pub async fn exists(&mut self, key: &str) -> Result<bool> {
|
||||
let exists: bool = self.connection
|
||||
.exists(key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
Ok(exists)
|
||||
}
|
||||
|
||||
/// Set multiple keys at once
|
||||
pub async fn mset(&mut self, items: &[(String, Value)]) -> Result<()> {
|
||||
let serialized: Vec<(String, String)> = items
|
||||
.iter()
|
||||
.map(|(k, v)| {
|
||||
serde_json::to_string(v)
|
||||
.map(|s| (k.clone(), s))
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
self.connection
|
||||
.set_multiple(&serialized)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get multiple keys at once
|
||||
pub async fn mget(&mut self, keys: &[String]) -> Result<Vec<Option<Value>>> {
|
||||
let values: Vec<Option<String>> = self.connection
|
||||
.get(keys)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
values
|
||||
.into_iter()
|
||||
.map(|opt_v| {
|
||||
match opt_v {
|
||||
Some(v) => {
|
||||
serde_json::from_str(&v)
|
||||
.map(Some)
|
||||
.map_err(|e| HanzoDbError::DeserializationError(e.to_string()))
|
||||
},
|
||||
None => Ok(None),
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Increment counter
|
||||
pub async fn incr(&mut self, key: &str, delta: i64) -> Result<i64> {
|
||||
let new_value: i64 = self.connection
|
||||
.incr(key, delta)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
Ok(new_value)
|
||||
}
|
||||
|
||||
/// Add to set
|
||||
pub async fn sadd(&mut self, key: &str, members: &[String]) -> Result<()> {
|
||||
for member in members {
|
||||
self.connection
|
||||
.sadd(key, member)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get set members
|
||||
pub async fn smembers(&mut self, key: &str) -> Result<Vec<String>> {
|
||||
let members: Vec<String> = self.connection
|
||||
.smembers(key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
Ok(members)
|
||||
}
|
||||
|
||||
/// Publish message to channel
|
||||
pub async fn publish(&mut self, channel: &str, message: &Value) -> Result<()> {
|
||||
let serialized = serde_json::to_string(message)
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))?;
|
||||
|
||||
self.connection
|
||||
.publish(channel, serialized)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Add to sorted set with score
|
||||
pub async fn zadd(&mut self, key: &str, member: &str, score: f64) -> Result<()> {
|
||||
self.connection
|
||||
.zadd(key, member, score)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get sorted set range by score
|
||||
pub async fn zrangebyscore(&mut self, key: &str, min: f64, max: f64, limit: Option<usize>) -> Result<Vec<String>> {
|
||||
let mut cmd = redis::cmd("ZRANGEBYSCORE");
|
||||
cmd.arg(key).arg(min).arg(max);
|
||||
|
||||
if let Some(l) = limit {
|
||||
cmd.arg("LIMIT").arg(0).arg(l);
|
||||
}
|
||||
|
||||
let members: Vec<String> = cmd
|
||||
.query_async(&mut self.connection)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(members)
|
||||
}
|
||||
|
||||
/// Store vector embedding with metadata
|
||||
pub async fn store_embedding(
|
||||
&mut self,
|
||||
key: &str,
|
||||
embedding: &[f32],
|
||||
metadata: Option<&Value>,
|
||||
) -> Result<()> {
|
||||
// Redis doesn't have native vector search, but we can store for caching
|
||||
let data = serde_json::json!({
|
||||
"embedding": embedding,
|
||||
"metadata": metadata,
|
||||
"timestamp": chrono::Utc::now().to_rfc3339(),
|
||||
});
|
||||
|
||||
self.set_with_ttl(key, &data, Some(3600)).await
|
||||
}
|
||||
|
||||
/// Cache query result
|
||||
pub async fn cache_result(&mut self, query_hash: &str, result: &Value, ttl_seconds: u64) -> Result<()> {
|
||||
let cache_key = format!("cache:{}", query_hash);
|
||||
self.set_with_ttl(&cache_key, result, Some(ttl_seconds)).await
|
||||
}
|
||||
|
||||
/// Get cached result
|
||||
pub async fn get_cached(&mut self, query_hash: &str) -> Result<Option<Value>> {
|
||||
let cache_key = format!("cache:{}", query_hash);
|
||||
self.get(&cache_key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatabaseBackend for RedisBackend {
|
||||
async fn connect(&mut self) -> Result<()> {
|
||||
// Already connected in new()
|
||||
// Test connection
|
||||
redis::cmd("PING")
|
||||
.query_async::<_, String>(&mut self.connection)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&mut self) -> Result<()> {
|
||||
// Connection manager handles cleanup
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_table(&self, name: &str, _schema: &Value) -> Result<()> {
|
||||
// Redis doesn't have tables, use key prefixes instead
|
||||
// Store table metadata
|
||||
let mut conn = self.connection.clone();
|
||||
let table_meta = serde_json::json!({
|
||||
"name": name,
|
||||
"created_at": chrono::Utc::now().to_rfc3339(),
|
||||
"type": "redis_namespace"
|
||||
});
|
||||
|
||||
let meta_key = format!("_meta:table:{}", name);
|
||||
conn.set(meta_key, table_meta.to_string())
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_table(&self, name: &str) -> Result<()> {
|
||||
// Delete all keys with table prefix
|
||||
let mut conn = self.connection.clone();
|
||||
let pattern = format!("{}:*", name);
|
||||
|
||||
// Get all keys matching pattern
|
||||
let keys: Vec<String> = redis::cmd("KEYS")
|
||||
.arg(&pattern)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Delete keys in batches
|
||||
for chunk in keys.chunks(1000) {
|
||||
redis::cmd("DEL")
|
||||
.arg(chunk)
|
||||
.query_async::<_, ()>(&mut conn)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
// Delete table metadata
|
||||
let meta_key = format!("_meta:table:{}", name);
|
||||
conn.del(meta_key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert(&self, table: &str, records: &[Value]) -> Result<()> {
|
||||
let mut conn = self.connection.clone();
|
||||
|
||||
for (i, record) in records.iter().enumerate() {
|
||||
let key = if let Some(id) = record.get("id").and_then(|v| v.as_str()) {
|
||||
format!("{}:{}", table, id)
|
||||
} else {
|
||||
format!("{}:{}", table, uuid::Uuid::new_v4())
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(record)
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))?;
|
||||
|
||||
conn.set(key, serialized)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Add to table index
|
||||
let index_key = format!("_index:{}", table);
|
||||
conn.sadd(index_key, &key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query(&self, query: &str) -> Result<Vec<Value>> {
|
||||
// Redis doesn't support SQL, interpret as key pattern
|
||||
let mut conn = self.connection.clone();
|
||||
|
||||
let keys: Vec<String> = if query.starts_with("SELECT * FROM ") {
|
||||
// Extract table name from simple query
|
||||
let table = query
|
||||
.strip_prefix("SELECT * FROM ")
|
||||
.and_then(|s| s.split_whitespace().next())
|
||||
.ok_or_else(|| HanzoDbError::InvalidQuery("Invalid SELECT query".to_string()))?;
|
||||
|
||||
// Get all keys from table index
|
||||
let index_key = format!("_index:{}", table);
|
||||
conn.smembers(index_key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?
|
||||
} else {
|
||||
// Treat as key pattern
|
||||
redis::cmd("KEYS")
|
||||
.arg(query)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?
|
||||
};
|
||||
|
||||
// Get all values
|
||||
let mut results = Vec::new();
|
||||
for key in keys {
|
||||
let value: Option<String> = conn.get(&key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
if let Some(v) = value {
|
||||
if let Ok(parsed) = serde_json::from_str(&v) {
|
||||
results.push(parsed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
async fn update(&self, table: &str, id: &str, data: &Value) -> Result<()> {
|
||||
let mut conn = self.connection.clone();
|
||||
let key = format!("{}:{}", table, id);
|
||||
|
||||
// Get existing record
|
||||
let existing: Option<String> = conn.get(&key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let mut record = if let Some(e) = existing {
|
||||
serde_json::from_str(&e)
|
||||
.map_err(|e| HanzoDbError::DeserializationError(e.to_string()))?
|
||||
} else {
|
||||
serde_json::json!({})
|
||||
};
|
||||
|
||||
// Merge updates
|
||||
if let (Some(rec_obj), Some(data_obj)) = (record.as_object_mut(), data.as_object()) {
|
||||
for (k, v) in data_obj {
|
||||
rec_obj.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Save updated record
|
||||
let serialized = serde_json::to_string(&record)
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))?;
|
||||
|
||||
conn.set(key, serialized)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, table: &str, id: &str) -> Result<()> {
|
||||
let mut conn = self.connection.clone();
|
||||
let key = format!("{}:{}", table, id);
|
||||
|
||||
// Remove from index
|
||||
let index_key = format!("_index:{}", table);
|
||||
conn.srem(index_key, &key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Delete key
|
||||
conn.del(key)
|
||||
.await
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn vector_search(
|
||||
&self,
|
||||
_table: &str,
|
||||
query_vector: &[f32],
|
||||
limit: usize,
|
||||
_filter: Option<&str>,
|
||||
) -> Result<Vec<Value>> {
|
||||
// Redis doesn't have native vector search
|
||||
// For production, use RedisSearch with vector similarity
|
||||
// This is a placeholder that returns cached results if available
|
||||
|
||||
let mut conn = self.connection.clone();
|
||||
let query_hash = blake3::hash(
|
||||
&query_vector.iter()
|
||||
.flat_map(|f| f.to_le_bytes())
|
||||
.collect::<Vec<_>>()
|
||||
).to_hex();
|
||||
|
||||
let cache_key = format!("vector_cache:{}", query_hash.as_str());
|
||||
|
||||
if let Some(cached) = self.get_cached(query_hash.as_str()).await? {
|
||||
return Ok(vec![cached]);
|
||||
}
|
||||
|
||||
// Return empty result - real implementation would use RedisSearch
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
fn backend_type(&self) -> DatabaseBackend {
|
||||
DatabaseBackend::Redis
|
||||
}
|
||||
|
||||
fn is_cache(&self) -> bool {
|
||||
true // Redis is primarily for caching
|
||||
}
|
||||
|
||||
fn is_real_time(&self) -> bool {
|
||||
true // Redis excels at real-time operations
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_redis_operations() {
|
||||
// Skip if Redis not available
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::Redis,
|
||||
url: Some("redis://localhost:6379".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut backend = match RedisBackend::new(config).await {
|
||||
Ok(b) => b,
|
||||
Err(_) => {
|
||||
println!("Redis not available, skipping test");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Test basic operations
|
||||
let key = "test_key";
|
||||
let value = serde_json::json!({"foo": "bar"});
|
||||
|
||||
backend.set_with_ttl(key, &value, Some(60)).await.unwrap();
|
||||
|
||||
let retrieved = backend.get(key).await.unwrap();
|
||||
assert_eq!(retrieved, Some(value));
|
||||
|
||||
backend.del(key).await.unwrap();
|
||||
|
||||
let deleted = backend.get(key).await.unwrap();
|
||||
assert_eq!(deleted, None);
|
||||
}
|
||||
}
|
||||
@@ -1,626 +0,0 @@
|
||||
//! SQLite backend implementation for embedded database operations
|
||||
//!
|
||||
//! SQLite is a lightweight, serverless, self-contained SQL database engine,
|
||||
//! perfect for embedded applications, testing, and local development.
|
||||
|
||||
use crate::{DatabaseBackend, HanzoDbConfig, HanzoDbError, Result};
|
||||
use async_trait::async_trait;
|
||||
use rusqlite::{params, Connection, OptionalExtension};
|
||||
use serde_json::Value;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::task;
|
||||
|
||||
/// SQLite backend for embedded database operations
|
||||
pub struct SqliteBackend {
|
||||
connection: Arc<Mutex<Connection>>,
|
||||
path: PathBuf,
|
||||
config: HanzoDbConfig,
|
||||
}
|
||||
|
||||
impl SqliteBackend {
|
||||
/// Create a new SQLite backend
|
||||
pub fn new(config: HanzoDbConfig) -> Result<Self> {
|
||||
let path = config.path.clone().unwrap_or_else(|| {
|
||||
PathBuf::from("./storage/hanzo-sqlite.db")
|
||||
});
|
||||
|
||||
// Create parent directory if needed
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
}
|
||||
|
||||
// Open SQLite connection
|
||||
let connection = if path.to_str() == Some(":memory:") {
|
||||
Connection::open_in_memory()
|
||||
} else {
|
||||
Connection::open(&path)
|
||||
}.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Configure SQLite for optimal performance
|
||||
connection.execute_batch("
|
||||
PRAGMA journal_mode = WAL;
|
||||
PRAGMA synchronous = NORMAL;
|
||||
PRAGMA cache_size = -64000;
|
||||
PRAGMA mmap_size = 30000000000;
|
||||
PRAGMA temp_store = MEMORY;
|
||||
PRAGMA foreign_keys = ON;
|
||||
").map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Enable JSON1 extension
|
||||
connection.execute_batch("
|
||||
CREATE TABLE IF NOT EXISTS _meta (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
").map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(Self {
|
||||
connection: Arc::new(Mutex::new(connection)),
|
||||
path,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute a query and return results as JSON
|
||||
fn execute_query(&self, query: &str) -> Result<Vec<Value>> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let mut stmt = conn.prepare(query)
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let column_names: Vec<String> = stmt
|
||||
.column_names()
|
||||
.into_iter()
|
||||
.map(|s| s.to_string())
|
||||
.collect();
|
||||
|
||||
let rows = stmt.query_map(params![], |row| {
|
||||
let mut obj = serde_json::Map::new();
|
||||
|
||||
for (i, name) in column_names.iter().enumerate() {
|
||||
let value = match row.get_ref(i) {
|
||||
Ok(val) => {
|
||||
use rusqlite::types::ValueRef;
|
||||
match val {
|
||||
ValueRef::Null => Value::Null,
|
||||
ValueRef::Integer(i) => Value::Number(i.into()),
|
||||
ValueRef::Real(f) => Value::from(f),
|
||||
ValueRef::Text(s) => {
|
||||
// Try to parse as JSON first
|
||||
if let Ok(json) = serde_json::from_str(std::str::from_utf8(s).unwrap_or("")) {
|
||||
json
|
||||
} else {
|
||||
Value::String(std::str::from_utf8(s).unwrap_or("").to_string())
|
||||
}
|
||||
},
|
||||
ValueRef::Blob(b) => {
|
||||
// Convert blob to base64
|
||||
Value::String(base64::encode(b))
|
||||
},
|
||||
}
|
||||
},
|
||||
Err(_) => Value::Null,
|
||||
};
|
||||
|
||||
obj.insert(name.clone(), value);
|
||||
}
|
||||
|
||||
Ok(Value::Object(obj))
|
||||
}).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
rows.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))
|
||||
}
|
||||
|
||||
/// Create vector extension tables
|
||||
pub fn enable_vector_support(&self) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Create vector storage table
|
||||
conn.execute_batch("
|
||||
CREATE TABLE IF NOT EXISTS _vectors (
|
||||
id TEXT PRIMARY KEY,
|
||||
table_name TEXT NOT NULL,
|
||||
record_id TEXT NOT NULL,
|
||||
vector BLOB NOT NULL,
|
||||
metadata TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(table_name, record_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_vectors_table ON _vectors(table_name);
|
||||
CREATE INDEX IF NOT EXISTS idx_vectors_record ON _vectors(table_name, record_id);
|
||||
").map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Store vector embedding
|
||||
pub fn store_vector(
|
||||
&self,
|
||||
table: &str,
|
||||
record_id: &str,
|
||||
vector: &[f32],
|
||||
metadata: Option<&Value>,
|
||||
) -> Result<()> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Convert vector to bytes
|
||||
let vector_bytes: Vec<u8> = vector.iter()
|
||||
.flat_map(|f| f.to_le_bytes())
|
||||
.collect();
|
||||
|
||||
let metadata_str = metadata
|
||||
.map(|m| serde_json::to_string(m))
|
||||
.transpose()
|
||||
.map_err(|e| HanzoDbError::SerializationError(e.to_string()))?;
|
||||
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO _vectors (id, table_name, record_id, vector, metadata)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)",
|
||||
params![
|
||||
format!("{}:{}", table, record_id),
|
||||
table,
|
||||
record_id,
|
||||
vector_bytes,
|
||||
metadata_str
|
||||
]
|
||||
).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Simple cosine similarity search (not optimized)
|
||||
pub fn vector_similarity_search(
|
||||
&self,
|
||||
table: &str,
|
||||
query_vector: &[f32],
|
||||
limit: usize,
|
||||
) -> Result<Vec<Value>> {
|
||||
let conn = self.connection.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Get all vectors from table
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT record_id, vector, metadata FROM _vectors WHERE table_name = ?1"
|
||||
).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
let vectors = stmt.query_map(params![table], |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, Vec<u8>>(1)?,
|
||||
row.get::<_, Option<String>>(2)?,
|
||||
))
|
||||
}).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Calculate cosine similarity for each vector
|
||||
let mut results: Vec<(String, f32, Option<String>)> = Vec::new();
|
||||
|
||||
for vector_result in vectors {
|
||||
let (record_id, vector_bytes, metadata) = vector_result
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Convert bytes back to f32 vector
|
||||
let vector: Vec<f32> = vector_bytes
|
||||
.chunks_exact(4)
|
||||
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
||||
.collect();
|
||||
|
||||
// Calculate cosine similarity
|
||||
let similarity = cosine_similarity(query_vector, &vector);
|
||||
results.push((record_id, similarity, metadata));
|
||||
}
|
||||
|
||||
// Sort by similarity (descending)
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
|
||||
// Take top N results
|
||||
let top_results: Vec<Value> = results
|
||||
.into_iter()
|
||||
.take(limit)
|
||||
.map(|(id, score, metadata)| {
|
||||
let mut obj = serde_json::Map::new();
|
||||
obj.insert("record_id".to_string(), Value::String(id));
|
||||
obj.insert("similarity".to_string(), Value::from(score));
|
||||
if let Some(meta) = metadata {
|
||||
if let Ok(meta_json) = serde_json::from_str::<Value>(&meta) {
|
||||
obj.insert("metadata".to_string(), meta_json);
|
||||
}
|
||||
}
|
||||
Value::Object(obj)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(top_results)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatabaseBackend for SqliteBackend {
|
||||
async fn connect(&mut self) -> Result<()> {
|
||||
// Already connected in new()
|
||||
// Test connection
|
||||
let conn = self.connection.clone();
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
c.execute("SELECT 1", params![])
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&mut self) -> Result<()> {
|
||||
// Connection will be closed when dropped
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_table(&self, name: &str, schema: &Value) -> Result<()> {
|
||||
let conn = self.connection.clone();
|
||||
let name = name.to_string();
|
||||
let schema = schema.clone();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let columns = schema.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidSchema("Schema must be object".to_string()))?;
|
||||
|
||||
let mut column_defs = Vec::new();
|
||||
|
||||
// Add ID column if not present
|
||||
if !columns.contains_key("id") {
|
||||
column_defs.push("id INTEGER PRIMARY KEY AUTOINCREMENT".to_string());
|
||||
}
|
||||
|
||||
for (col_name, type_def) in columns {
|
||||
let sql_type = match type_def.as_str() {
|
||||
Some("string") => "TEXT",
|
||||
Some("integer") => "INTEGER",
|
||||
Some("bigint") => "INTEGER",
|
||||
Some("float") | Some("double") => "REAL",
|
||||
Some("boolean") => "INTEGER", // SQLite uses 0/1 for boolean
|
||||
Some("timestamp") => "DATETIME",
|
||||
Some("date") => "DATE",
|
||||
Some("json") | Some("jsonb") => "TEXT", // Store JSON as TEXT
|
||||
Some("vector") => "BLOB", // Store vectors as BLOB
|
||||
_ => "TEXT",
|
||||
};
|
||||
|
||||
column_defs.push(format!("{} {}", col_name, sql_type));
|
||||
}
|
||||
|
||||
// Add timestamps if not present
|
||||
if !columns.contains_key("created_at") {
|
||||
column_defs.push("created_at DATETIME DEFAULT CURRENT_TIMESTAMP".to_string());
|
||||
}
|
||||
if !columns.contains_key("updated_at") {
|
||||
column_defs.push("updated_at DATETIME DEFAULT CURRENT_TIMESTAMP".to_string());
|
||||
}
|
||||
|
||||
let create_sql = format!(
|
||||
"CREATE TABLE IF NOT EXISTS {} ({})",
|
||||
name,
|
||||
column_defs.join(", ")
|
||||
);
|
||||
|
||||
c.execute(&create_sql, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
// Create update trigger for updated_at
|
||||
let trigger_sql = format!(
|
||||
"CREATE TRIGGER IF NOT EXISTS update_{}_timestamp
|
||||
AFTER UPDATE ON {}
|
||||
FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE {} SET updated_at = CURRENT_TIMESTAMP WHERE id = NEW.id;
|
||||
END",
|
||||
name, name, name
|
||||
);
|
||||
|
||||
c.execute(&trigger_sql, params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
// Enable vector support if needed
|
||||
if schema.as_object()
|
||||
.and_then(|o| o.keys().find(|k| k.contains("vector") || k.contains("embedding")))
|
||||
.is_some()
|
||||
{
|
||||
self.enable_vector_support()?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_table(&self, name: &str) -> Result<()> {
|
||||
let conn = self.connection.clone();
|
||||
let name = name.to_string();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
c.execute(&format!("DROP TABLE IF EXISTS {}", name), params![])
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert(&self, table: &str, records: &[Value]) -> Result<()> {
|
||||
let conn = self.connection.clone();
|
||||
let table = table.to_string();
|
||||
let records = records.to_vec();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let tx = c.transaction()
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))?;
|
||||
|
||||
for record in records {
|
||||
let obj = record.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Record must be object".to_string()))?;
|
||||
|
||||
let columns: Vec<String> = obj.keys().cloned().collect();
|
||||
let placeholders: Vec<String> = (1..=columns.len())
|
||||
.map(|i| format!("?{}", i))
|
||||
.collect();
|
||||
|
||||
let insert_sql = format!(
|
||||
"INSERT INTO {} ({}) VALUES ({})",
|
||||
table,
|
||||
columns.join(", "),
|
||||
placeholders.join(", ")
|
||||
);
|
||||
|
||||
let params: Vec<rusqlite::types::Value> = obj.values().map(|v| {
|
||||
match v {
|
||||
Value::String(s) => rusqlite::types::Value::Text(s.clone()),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
rusqlite::types::Value::Integer(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
rusqlite::types::Value::Real(f)
|
||||
} else {
|
||||
rusqlite::types::Value::Null
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => rusqlite::types::Value::Integer(if *b { 1 } else { 0 }),
|
||||
Value::Null => rusqlite::types::Value::Null,
|
||||
Value::Array(_) | Value::Object(_) => {
|
||||
// Store complex types as JSON
|
||||
rusqlite::types::Value::Text(v.to_string())
|
||||
}
|
||||
}
|
||||
}).collect();
|
||||
|
||||
tx.execute(&insert_sql, params.as_slice())
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.map_err(|e| HanzoDbError::TransactionError(e.to_string()))?;
|
||||
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query(&self, query: &str) -> Result<Vec<Value>> {
|
||||
let conn = self.connection.clone();
|
||||
let query = query.to_string();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
// Use the helper method to execute query and get JSON results
|
||||
drop(c); // Release lock before calling execute_query
|
||||
|
||||
// Re-acquire for execute_query (this is safe as it's in the same thread)
|
||||
let backend = SqliteBackend {
|
||||
connection: conn,
|
||||
path: PathBuf::new(), // Not used in execute_query
|
||||
config: HanzoDbConfig::default(), // Not used in execute_query
|
||||
};
|
||||
|
||||
backend.execute_query(&query)
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?
|
||||
}
|
||||
|
||||
async fn update(&self, table: &str, id: &str, data: &Value) -> Result<()> {
|
||||
let conn = self.connection.clone();
|
||||
let table = table.to_string();
|
||||
let id = id.to_string();
|
||||
let data = data.clone();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
let obj = data.as_object()
|
||||
.ok_or_else(|| HanzoDbError::InvalidData("Update data must be object".to_string()))?;
|
||||
|
||||
let set_clauses: Vec<String> = obj.keys()
|
||||
.enumerate()
|
||||
.map(|(i, key)| format!("{} = ?{}", key, i + 2))
|
||||
.collect();
|
||||
|
||||
let update_sql = format!(
|
||||
"UPDATE {} SET {}, updated_at = CURRENT_TIMESTAMP WHERE id = ?1",
|
||||
table,
|
||||
set_clauses.join(", ")
|
||||
);
|
||||
|
||||
let mut params: Vec<rusqlite::types::Value> = vec![rusqlite::types::Value::Text(id)];
|
||||
|
||||
for value in obj.values() {
|
||||
params.push(match value {
|
||||
Value::String(s) => rusqlite::types::Value::Text(s.clone()),
|
||||
Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
rusqlite::types::Value::Integer(i)
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
rusqlite::types::Value::Real(f)
|
||||
} else {
|
||||
rusqlite::types::Value::Null
|
||||
}
|
||||
},
|
||||
Value::Bool(b) => rusqlite::types::Value::Integer(if *b { 1 } else { 0 }),
|
||||
Value::Null => rusqlite::types::Value::Null,
|
||||
_ => rusqlite::types::Value::Text(value.to_string()),
|
||||
});
|
||||
}
|
||||
|
||||
c.execute(&update_sql, params.as_slice())
|
||||
.map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, table: &str, id: &str) -> Result<()> {
|
||||
let conn = self.connection.clone();
|
||||
let table = table.to_string();
|
||||
let id = id.to_string();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let c = conn.lock()
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?;
|
||||
|
||||
c.execute(
|
||||
&format!("DELETE FROM {} WHERE id = ?1", table),
|
||||
params![id]
|
||||
).map_err(|e| HanzoDbError::QueryError(e.to_string()))?;
|
||||
|
||||
Ok::<(), HanzoDbError>(())
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn vector_search(
|
||||
&self,
|
||||
table: &str,
|
||||
query_vector: &[f32],
|
||||
limit: usize,
|
||||
_filter: Option<&str>,
|
||||
) -> Result<Vec<Value>> {
|
||||
let table = table.to_string();
|
||||
let query_vector = query_vector.to_vec();
|
||||
let conn = self.connection.clone();
|
||||
let path = self.path.clone();
|
||||
let config = self.config.clone();
|
||||
|
||||
task::spawn_blocking(move || {
|
||||
let backend = SqliteBackend {
|
||||
connection: conn,
|
||||
path,
|
||||
config,
|
||||
};
|
||||
|
||||
backend.vector_similarity_search(&table, &query_vector, limit)
|
||||
}).await
|
||||
.map_err(|e| HanzoDbError::ConnectionError(e.to_string()))?
|
||||
}
|
||||
|
||||
fn backend_type(&self) -> DatabaseBackend {
|
||||
DatabaseBackend::SQLite
|
||||
}
|
||||
|
||||
fn is_embedded(&self) -> bool {
|
||||
true // SQLite is embedded
|
||||
}
|
||||
|
||||
fn is_lightweight(&self) -> bool {
|
||||
true // SQLite is very lightweight
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function for cosine similarity
|
||||
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
if a.len() != b.len() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
||||
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
|
||||
if norm_a == 0.0 || norm_b == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
dot_product / (norm_a * norm_b)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sqlite_operations() {
|
||||
let config = HanzoDbConfig {
|
||||
backend: DatabaseBackend::SQLite,
|
||||
path: Some(PathBuf::from(":memory:")),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let backend = SqliteBackend::new(config).unwrap();
|
||||
|
||||
// Create table
|
||||
let schema = serde_json::json!({
|
||||
"name": "string",
|
||||
"age": "integer",
|
||||
"active": "boolean"
|
||||
});
|
||||
backend.create_table("users", &schema).await.unwrap();
|
||||
|
||||
// Insert records
|
||||
let records = vec![
|
||||
serde_json::json!({"name": "Alice", "age": 30, "active": true}),
|
||||
serde_json::json!({"name": "Bob", "age": 25, "active": false}),
|
||||
];
|
||||
backend.insert("users", &records).await.unwrap();
|
||||
|
||||
// Query records
|
||||
let results = backend.query("SELECT * FROM users WHERE age > 20").await.unwrap();
|
||||
assert_eq!(results.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_similarity() {
|
||||
let a = vec![1.0, 2.0, 3.0];
|
||||
let b = vec![1.0, 2.0, 3.0];
|
||||
assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
|
||||
|
||||
let c = vec![-1.0, -2.0, -3.0];
|
||||
assert!((cosine_similarity(&a, &c) + 1.0).abs() < 0.001);
|
||||
}
|
||||
}
|
||||
@@ -1,384 +0,0 @@
|
||||
//! # Hanzo DB - Multi-Backend Database Abstraction
|
||||
//!
|
||||
//! Production-ready database abstraction for Hanzo Node, supporting:
|
||||
//! - LanceDB for vector search and multimodal storage
|
||||
//! - DuckDB for analytics and OLAP queries
|
||||
//! - PostgreSQL for relational data
|
||||
//! - Redis for caching
|
||||
//! - SQLite for lightweight deployments
|
||||
//!
|
||||
//! Features:
|
||||
//! - Unified interface across all backends
|
||||
//! - Automatic backend selection based on workload
|
||||
//! - Connection pooling and transaction support
|
||||
//! - Migration between backends
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use log::{debug, error, info, warn};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
pub mod models;
|
||||
pub mod vector_search;
|
||||
|
||||
#[cfg(feature = "migration")]
|
||||
pub mod migration;
|
||||
|
||||
// Backend modules
|
||||
pub mod backends;
|
||||
|
||||
// Re-exports
|
||||
pub use models::*;
|
||||
pub use vector_search::*;
|
||||
|
||||
/// Database backend type
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DatabaseBackend {
|
||||
/// LanceDB for vector operations and multimodal data
|
||||
LanceDB,
|
||||
/// DuckDB for analytics and OLAP
|
||||
DuckDB,
|
||||
/// PostgreSQL for relational data
|
||||
PostgreSQL,
|
||||
/// Redis for caching
|
||||
Redis,
|
||||
/// SQLite for lightweight deployments
|
||||
SQLite,
|
||||
}
|
||||
|
||||
impl DatabaseBackend {
|
||||
/// Select optimal backend for workload type
|
||||
pub fn for_workload(workload: WorkloadType) -> Self {
|
||||
match workload {
|
||||
WorkloadType::VectorSearch => Self::LanceDB,
|
||||
WorkloadType::Analytics => Self::DuckDB,
|
||||
WorkloadType::Transactional => Self::PostgreSQL,
|
||||
WorkloadType::Cache => Self::Redis,
|
||||
WorkloadType::Embedded => Self::SQLite,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Workload type for backend selection
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum WorkloadType {
|
||||
/// Vector similarity search
|
||||
VectorSearch,
|
||||
/// Analytical queries (OLAP)
|
||||
Analytics,
|
||||
/// Transactional operations (OLTP)
|
||||
Transactional,
|
||||
/// High-speed caching
|
||||
Cache,
|
||||
/// Embedded/lightweight operations
|
||||
Embedded,
|
||||
}
|
||||
|
||||
/// Unified database configuration
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HanzoDbConfig {
|
||||
/// Selected backend
|
||||
pub backend: DatabaseBackend,
|
||||
/// Database path (for file-based backends)
|
||||
pub path: Option<PathBuf>,
|
||||
/// Connection URL (for network backends)
|
||||
pub url: Option<String>,
|
||||
/// Connection pool size
|
||||
pub pool_size: usize,
|
||||
/// Enable write-ahead logging
|
||||
pub enable_wal: bool,
|
||||
/// Cache size in bytes
|
||||
pub cache_size: Option<usize>,
|
||||
/// Enable compression
|
||||
pub enable_compression: bool,
|
||||
}
|
||||
|
||||
impl Default for HanzoDbConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
backend: DatabaseBackend::LanceDB,
|
||||
path: Some(PathBuf::from("./storage/hanzo-db")),
|
||||
url: None,
|
||||
pool_size: 16,
|
||||
enable_wal: true,
|
||||
cache_size: Some(64 * 1024 * 1024), // 64MB
|
||||
enable_compression: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Unified database trait
|
||||
#[async_trait]
|
||||
pub trait HanzoDatabase: Send + Sync {
|
||||
/// Initialize the database
|
||||
async fn init(&self) -> Result<()>;
|
||||
|
||||
/// Create a table
|
||||
async fn create_table(&self, name: &str, schema: TableSchema) -> Result<()>;
|
||||
|
||||
/// Insert data
|
||||
async fn insert(&self, table: &str, data: &[Record]) -> Result<()>;
|
||||
|
||||
/// Query data
|
||||
async fn query(&self, query: Query) -> Result<QueryResult>;
|
||||
|
||||
/// Vector search
|
||||
async fn vector_search(&self, query: VectorQuery) -> Result<Vec<SearchResult>>;
|
||||
|
||||
/// Begin transaction
|
||||
async fn begin_transaction(&self) -> Result<Transaction>;
|
||||
|
||||
/// Optimize database
|
||||
async fn optimize(&self) -> Result<()>;
|
||||
|
||||
/// Get database statistics
|
||||
async fn stats(&self) -> Result<DatabaseStats>;
|
||||
}
|
||||
|
||||
/// Table schema definition
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TableSchema {
|
||||
pub columns: Vec<Column>,
|
||||
pub indexes: Vec<Index>,
|
||||
pub constraints: Vec<Constraint>,
|
||||
}
|
||||
|
||||
/// Column definition
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Column {
|
||||
pub name: String,
|
||||
pub data_type: DataType,
|
||||
pub nullable: bool,
|
||||
pub default: Option<Value>,
|
||||
}
|
||||
|
||||
/// Data type enumeration
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DataType {
|
||||
// Scalar types
|
||||
Boolean,
|
||||
Int32,
|
||||
Int64,
|
||||
Float32,
|
||||
Float64,
|
||||
String,
|
||||
Binary,
|
||||
Timestamp,
|
||||
|
||||
// Vector types
|
||||
Vector(usize), // dimension
|
||||
|
||||
// Complex types
|
||||
Json,
|
||||
Array(Box<DataType>),
|
||||
Struct(Vec<(String, DataType)>),
|
||||
}
|
||||
|
||||
/// Index definition
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Index {
|
||||
pub name: String,
|
||||
pub columns: Vec<String>,
|
||||
pub index_type: IndexType,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum IndexType {
|
||||
BTree,
|
||||
Hash,
|
||||
IVF_PQ { nlist: usize, nprobe: usize },
|
||||
HNSW { max_elements: usize, m: usize },
|
||||
}
|
||||
|
||||
/// Database constraint
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Constraint {
|
||||
PrimaryKey(Vec<String>),
|
||||
ForeignKey { columns: Vec<String>, references: String },
|
||||
Unique(Vec<String>),
|
||||
Check(String),
|
||||
}
|
||||
|
||||
/// Query structure
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Query {
|
||||
pub table: String,
|
||||
pub select: Vec<String>,
|
||||
pub filter: Option<Filter>,
|
||||
pub order_by: Vec<OrderBy>,
|
||||
pub limit: Option<usize>,
|
||||
pub offset: Option<usize>,
|
||||
}
|
||||
|
||||
/// Filter expression
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Filter {
|
||||
Eq(String, Value),
|
||||
Ne(String, Value),
|
||||
Gt(String, Value),
|
||||
Gte(String, Value),
|
||||
Lt(String, Value),
|
||||
Lte(String, Value),
|
||||
In(String, Vec<Value>),
|
||||
Like(String, String),
|
||||
And(Box<Filter>, Box<Filter>),
|
||||
Or(Box<Filter>, Box<Filter>),
|
||||
Not(Box<Filter>),
|
||||
}
|
||||
|
||||
/// Order by clause
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OrderBy {
|
||||
pub column: String,
|
||||
pub ascending: bool,
|
||||
}
|
||||
|
||||
/// Value type
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Value {
|
||||
Null,
|
||||
Bool(bool),
|
||||
Int32(i32),
|
||||
Int64(i64),
|
||||
Float32(f32),
|
||||
Float64(f64),
|
||||
String(String),
|
||||
Binary(Vec<u8>),
|
||||
Timestamp(i64),
|
||||
Vector(Vec<f32>),
|
||||
Json(serde_json::Value),
|
||||
}
|
||||
|
||||
/// Record type
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Record {
|
||||
pub values: Vec<(String, Value)>,
|
||||
}
|
||||
|
||||
/// Query result
|
||||
#[derive(Debug)]
|
||||
pub struct QueryResult {
|
||||
pub columns: Vec<String>,
|
||||
pub rows: Vec<Record>,
|
||||
pub row_count: usize,
|
||||
}
|
||||
|
||||
/// Vector query
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct VectorQuery {
|
||||
pub table: String,
|
||||
pub vector: Vec<f32>,
|
||||
pub k: usize,
|
||||
pub filter: Option<Filter>,
|
||||
pub metric: DistanceMetric,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum DistanceMetric {
|
||||
L2,
|
||||
Cosine,
|
||||
InnerProduct,
|
||||
}
|
||||
|
||||
/// Search result
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchResult {
|
||||
pub record: Record,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
/// Transaction handle
|
||||
pub struct Transaction {
|
||||
inner: Arc<RwLock<TransactionInner>>,
|
||||
}
|
||||
|
||||
struct TransactionInner {
|
||||
backend: DatabaseBackend,
|
||||
// Backend-specific transaction handle
|
||||
handle: Box<dyn std::any::Any + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
/// Commit the transaction
|
||||
pub async fn commit(self) -> Result<()> {
|
||||
let inner = self.inner.write().await;
|
||||
// Backend-specific commit logic
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rollback the transaction
|
||||
pub async fn rollback(self) -> Result<()> {
|
||||
let inner = self.inner.write().await;
|
||||
// Backend-specific rollback logic
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Database statistics
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DatabaseStats {
|
||||
pub backend: DatabaseBackend,
|
||||
pub table_count: usize,
|
||||
pub total_rows: usize,
|
||||
pub total_size_bytes: usize,
|
||||
pub index_count: usize,
|
||||
pub cache_hit_rate: f64,
|
||||
}
|
||||
|
||||
/// Create a Hanzo database instance
|
||||
pub async fn connect(config: HanzoDbConfig) -> Result<Arc<dyn HanzoDatabase>> {
|
||||
match config.backend {
|
||||
DatabaseBackend::LanceDB => {
|
||||
let db = backends::lancedb::LanceDbBackend::new(config).await?;
|
||||
Ok(Arc::new(db))
|
||||
}
|
||||
DatabaseBackend::DuckDB => {
|
||||
let db = backends::duckdb::DuckDbBackend::new(config).await?;
|
||||
Ok(Arc::new(db))
|
||||
}
|
||||
DatabaseBackend::PostgreSQL => {
|
||||
let db = backends::postgres::PostgresBackend::new(config).await?;
|
||||
Ok(Arc::new(db))
|
||||
}
|
||||
DatabaseBackend::Redis => {
|
||||
let db = backends::redis::RedisBackend::new(config).await?;
|
||||
Ok(Arc::new(db))
|
||||
}
|
||||
DatabaseBackend::SQLite => {
|
||||
let db = backends::sqlite::SqliteBackend::new(config).await?;
|
||||
Ok(Arc::new(db))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Backend implementations module
|
||||
pub mod backends {
|
||||
pub mod lancedb;
|
||||
pub mod duckdb;
|
||||
pub mod postgres;
|
||||
pub mod redis;
|
||||
pub mod sqlite;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_backend_selection() {
|
||||
assert_eq!(
|
||||
DatabaseBackend::for_workload(WorkloadType::VectorSearch),
|
||||
DatabaseBackend::LanceDB
|
||||
);
|
||||
assert_eq!(
|
||||
DatabaseBackend::for_workload(WorkloadType::Analytics),
|
||||
DatabaseBackend::DuckDB
|
||||
);
|
||||
assert_eq!(
|
||||
DatabaseBackend::for_workload(WorkloadType::Transactional),
|
||||
DatabaseBackend::PostgreSQL
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,461 +0,0 @@
|
||||
//! # SQLite to LanceDB Migration Module
|
||||
//!
|
||||
//! Provides tools to migrate existing SQLite databases to LanceDB.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use arrow_array::{
|
||||
BooleanArray, Float32Array, Int64Array, RecordBatch, StringArray, TimestampMillisecondArray,
|
||||
ArrayRef, FixedSizeListArray,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
|
||||
use chrono::{DateTime, NaiveDateTime, Utc};
|
||||
use lancedb::Connection;
|
||||
use log::{debug, error, info, warn};
|
||||
use rusqlite::{params, Connection as SqliteConnection, Row};
|
||||
use serde_json;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::models::*;
|
||||
use crate::LanceDb;
|
||||
|
||||
/// Migration configuration
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MigrationConfig {
|
||||
/// Source SQLite database path
|
||||
pub sqlite_path: String,
|
||||
/// Target LanceDB instance
|
||||
pub target_db: Arc<LanceDb>,
|
||||
/// Batch size for migration
|
||||
pub batch_size: usize,
|
||||
/// Skip existing tables
|
||||
pub skip_existing: bool,
|
||||
/// Verify data after migration
|
||||
pub verify: bool,
|
||||
}
|
||||
|
||||
impl Default for MigrationConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
sqlite_path: "./storage/db.sqlite".to_string(),
|
||||
target_db: Arc::new(LanceDb::new().expect("Failed to create LanceDB")),
|
||||
batch_size: 1000,
|
||||
skip_existing: false,
|
||||
verify: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SQLite to LanceDB migrator
|
||||
pub struct Migrator {
|
||||
config: MigrationConfig,
|
||||
sqlite_conn: SqliteConnection,
|
||||
}
|
||||
|
||||
impl Migrator {
|
||||
/// Create new migrator instance
|
||||
pub fn new(config: MigrationConfig) -> Result<Self> {
|
||||
let sqlite_conn = SqliteConnection::open(&config.sqlite_path)
|
||||
.context("Failed to open SQLite database")?;
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
sqlite_conn,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run full migration
|
||||
pub async fn migrate_all(&mut self) -> Result<MigrationStats> {
|
||||
info!("Starting SQLite to LanceDB migration");
|
||||
let mut stats = MigrationStats::default();
|
||||
|
||||
// Migrate users
|
||||
match self.migrate_users().await {
|
||||
Ok(count) => {
|
||||
info!("Migrated {} users", count);
|
||||
stats.users_migrated = count;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to migrate users: {}", e);
|
||||
stats.errors.push(format!("Users: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Migrate tools
|
||||
match self.migrate_tools().await {
|
||||
Ok(count) => {
|
||||
info!("Migrated {} tools", count);
|
||||
stats.tools_migrated = count;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to migrate tools: {}", e);
|
||||
stats.errors.push(format!("Tools: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Migrate agents
|
||||
match self.migrate_agents().await {
|
||||
Ok(count) => {
|
||||
info!("Migrated {} agents", count);
|
||||
stats.agents_migrated = count;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to migrate agents: {}", e);
|
||||
stats.errors.push(format!("Agents: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Migrate embeddings
|
||||
match self.migrate_embeddings().await {
|
||||
Ok(count) => {
|
||||
info!("Migrated {} embeddings", count);
|
||||
stats.embeddings_migrated = count;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to migrate embeddings: {}", e);
|
||||
stats.errors.push(format!("Embeddings: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Migrate jobs
|
||||
match self.migrate_jobs().await {
|
||||
Ok(count) => {
|
||||
info!("Migrated {} jobs", count);
|
||||
stats.jobs_migrated = count;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to migrate jobs: {}", e);
|
||||
stats.errors.push(format!("Jobs: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Verify if requested
|
||||
if self.config.verify {
|
||||
self.verify_migration(&stats).await?;
|
||||
}
|
||||
|
||||
info!("Migration completed: {:?}", stats);
|
||||
Ok(stats)
|
||||
}
|
||||
|
||||
/// Migrate users table
|
||||
async fn migrate_users(&mut self) -> Result<usize> {
|
||||
info!("Migrating users table");
|
||||
|
||||
let mut stmt = self.sqlite_conn.prepare(
|
||||
"SELECT id, profile_name, identity_type, identity_public_key,
|
||||
encryption_public_key, signature_public_key,
|
||||
node_signature_public_key, node_encryption_public_key,
|
||||
permission_type, wallet_id, metadata,
|
||||
created_at, updated_at, is_active
|
||||
FROM users"
|
||||
)?;
|
||||
|
||||
let users = stmt.query_map([], |row| {
|
||||
Ok(UserRow {
|
||||
id: row.get(0)?,
|
||||
profile_name: row.get(1)?,
|
||||
identity_type: row.get(2)?,
|
||||
identity_public_key: row.get(3)?,
|
||||
encryption_public_key: row.get(4)?,
|
||||
signature_public_key: row.get(5)?,
|
||||
node_signature_public_key: row.get(6)?,
|
||||
node_encryption_public_key: row.get(7)?,
|
||||
permission_type: row.get(8)?,
|
||||
wallet_id: row.get(9)?,
|
||||
metadata: row.get(10)?,
|
||||
created_at: row.get(11)?,
|
||||
updated_at: row.get(12)?,
|
||||
is_active: row.get(13)?,
|
||||
})
|
||||
})?;
|
||||
|
||||
let mut count = 0;
|
||||
let mut batch = Vec::new();
|
||||
|
||||
for user in users {
|
||||
let user = user?;
|
||||
batch.push(user);
|
||||
|
||||
if batch.len() >= self.config.batch_size {
|
||||
self.insert_users_batch(&batch).await?;
|
||||
count += batch.len();
|
||||
batch.clear();
|
||||
}
|
||||
}
|
||||
|
||||
// Insert remaining batch
|
||||
if !batch.is_empty() {
|
||||
self.insert_users_batch(&batch).await?;
|
||||
count += batch.len();
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Insert batch of users into LanceDB
|
||||
async fn insert_users_batch(&self, users: &[UserRow]) -> Result<()> {
|
||||
let schema = user_schema();
|
||||
|
||||
// Build arrays for each column
|
||||
let mut id_builder = StringArray::builder(users.len());
|
||||
let mut profile_name_builder = StringArray::builder(users.len());
|
||||
let mut identity_type_builder = StringArray::builder(users.len());
|
||||
let mut identity_pk_builder = StringArray::builder(users.len());
|
||||
let mut encryption_pk_builder = StringArray::builder(users.len());
|
||||
let mut signature_pk_builder = StringArray::builder(users.len());
|
||||
let mut node_sig_pk_builder = StringArray::builder(users.len());
|
||||
let mut node_enc_pk_builder = StringArray::builder(users.len());
|
||||
let mut permission_builder = StringArray::builder(users.len());
|
||||
let mut wallet_builder = StringArray::builder(users.len());
|
||||
let mut metadata_builder = StringArray::builder(users.len());
|
||||
let mut created_at_builder = TimestampMillisecondArray::builder(users.len());
|
||||
let mut updated_at_builder = TimestampMillisecondArray::builder(users.len());
|
||||
let mut is_active_builder = BooleanArray::builder(users.len());
|
||||
|
||||
for user in users {
|
||||
id_builder.append_value(&user.id);
|
||||
profile_name_builder.append_value(&user.profile_name);
|
||||
identity_type_builder.append_value(&user.identity_type);
|
||||
|
||||
if let Some(v) = &user.identity_public_key {
|
||||
identity_pk_builder.append_value(v);
|
||||
} else {
|
||||
identity_pk_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.encryption_public_key {
|
||||
encryption_pk_builder.append_value(v);
|
||||
} else {
|
||||
encryption_pk_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.signature_public_key {
|
||||
signature_pk_builder.append_value(v);
|
||||
} else {
|
||||
signature_pk_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.node_signature_public_key {
|
||||
node_sig_pk_builder.append_value(v);
|
||||
} else {
|
||||
node_sig_pk_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.node_encryption_public_key {
|
||||
node_enc_pk_builder.append_value(v);
|
||||
} else {
|
||||
node_enc_pk_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.permission_type {
|
||||
permission_builder.append_value(v);
|
||||
} else {
|
||||
permission_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.wallet_id {
|
||||
wallet_builder.append_value(v);
|
||||
} else {
|
||||
wallet_builder.append_null();
|
||||
}
|
||||
|
||||
if let Some(v) = &user.metadata {
|
||||
metadata_builder.append_value(v);
|
||||
} else {
|
||||
metadata_builder.append_null();
|
||||
}
|
||||
|
||||
created_at_builder.append_value(user.created_at);
|
||||
updated_at_builder.append_value(user.updated_at);
|
||||
is_active_builder.append_value(user.is_active);
|
||||
}
|
||||
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(id_builder.finish()) as ArrayRef,
|
||||
Arc::new(profile_name_builder.finish()) as ArrayRef,
|
||||
Arc::new(identity_type_builder.finish()) as ArrayRef,
|
||||
Arc::new(identity_pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(encryption_pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(signature_pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(node_sig_pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(node_enc_pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(permission_builder.finish()) as ArrayRef,
|
||||
Arc::new(wallet_builder.finish()) as ArrayRef,
|
||||
Arc::new(metadata_builder.finish()) as ArrayRef,
|
||||
Arc::new(created_at_builder.finish()) as ArrayRef,
|
||||
Arc::new(updated_at_builder.finish()) as ArrayRef,
|
||||
Arc::new(is_active_builder.finish()) as ArrayRef,
|
||||
],
|
||||
)?;
|
||||
|
||||
// Insert into LanceDB
|
||||
let table = self.config.target_db.table("users").await?;
|
||||
table.add(&[batch]).execute().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Migrate tools table
|
||||
async fn migrate_tools(&mut self) -> Result<usize> {
|
||||
info!("Migrating tools table");
|
||||
|
||||
// Check if tools table exists
|
||||
let table_exists: bool = self.sqlite_conn.query_row(
|
||||
"SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='tools'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
|
||||
if !table_exists {
|
||||
warn!("Tools table does not exist in SQLite database");
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Similar implementation as migrate_users
|
||||
// TODO: Complete implementation
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
/// Migrate agents table
|
||||
async fn migrate_agents(&mut self) -> Result<usize> {
|
||||
info!("Migrating agents table");
|
||||
// TODO: Implement agent migration
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
/// Migrate embeddings table
|
||||
async fn migrate_embeddings(&mut self) -> Result<usize> {
|
||||
info!("Migrating embeddings table");
|
||||
// TODO: Implement embeddings migration with vector data
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
/// Migrate jobs table
|
||||
async fn migrate_jobs(&mut self) -> Result<usize> {
|
||||
info!("Migrating jobs table");
|
||||
// TODO: Implement jobs migration
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
/// Verify migration results
|
||||
async fn verify_migration(&self, stats: &MigrationStats) -> Result<()> {
|
||||
info!("Verifying migration results");
|
||||
|
||||
let db_stats = self.config.target_db.stats().await?;
|
||||
|
||||
for table_stat in &db_stats.table_stats {
|
||||
match table_stat.name.as_str() {
|
||||
"users" => {
|
||||
if table_stat.row_count < stats.users_migrated {
|
||||
warn!(
|
||||
"User count mismatch: expected {}, found {}",
|
||||
stats.users_migrated, table_stat.row_count
|
||||
);
|
||||
}
|
||||
}
|
||||
"tools" => {
|
||||
if table_stat.row_count < stats.tools_migrated {
|
||||
warn!(
|
||||
"Tools count mismatch: expected {}, found {}",
|
||||
stats.tools_migrated, table_stat.row_count
|
||||
);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
info!("Verification completed");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Migration statistics
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct MigrationStats {
|
||||
pub users_migrated: usize,
|
||||
pub tools_migrated: usize,
|
||||
pub agents_migrated: usize,
|
||||
pub embeddings_migrated: usize,
|
||||
pub jobs_migrated: usize,
|
||||
pub errors: Vec<String>,
|
||||
}
|
||||
|
||||
/// Temporary struct for SQLite user row
|
||||
struct UserRow {
|
||||
id: String,
|
||||
profile_name: String,
|
||||
identity_type: String,
|
||||
identity_public_key: Option<String>,
|
||||
encryption_public_key: Option<String>,
|
||||
signature_public_key: Option<String>,
|
||||
node_signature_public_key: Option<String>,
|
||||
node_encryption_public_key: Option<String>,
|
||||
permission_type: Option<String>,
|
||||
wallet_id: Option<String>,
|
||||
metadata: Option<String>,
|
||||
created_at: i64,
|
||||
updated_at: i64,
|
||||
is_active: bool,
|
||||
}
|
||||
|
||||
/// Migration CLI tool
|
||||
pub async fn run_migration_cli(sqlite_path: &str, lance_path: &str) -> Result<()> {
|
||||
println!("SQLite to LanceDB Migration Tool");
|
||||
println!("=================================\n");
|
||||
|
||||
// Create LanceDB instance
|
||||
let lance_config = crate::LanceDbConfig {
|
||||
path: lance_path.into(),
|
||||
..Default::default()
|
||||
};
|
||||
let lance_db = Arc::new(LanceDb::with_config(lance_config).await?);
|
||||
|
||||
// Create migration config
|
||||
let config = MigrationConfig {
|
||||
sqlite_path: sqlite_path.to_string(),
|
||||
target_db: lance_db,
|
||||
batch_size: 1000,
|
||||
skip_existing: false,
|
||||
verify: true,
|
||||
};
|
||||
|
||||
// Run migration
|
||||
let mut migrator = Migrator::new(config)?;
|
||||
let stats = migrator.migrate_all().await?;
|
||||
|
||||
// Print results
|
||||
println!("\nMigration Results:");
|
||||
println!(" Users migrated: {}", stats.users_migrated);
|
||||
println!(" Tools migrated: {}", stats.tools_migrated);
|
||||
println!(" Agents migrated: {}", stats.agents_migrated);
|
||||
println!(" Embeddings migrated: {}", stats.embeddings_migrated);
|
||||
println!(" Jobs migrated: {}", stats.jobs_migrated);
|
||||
|
||||
if !stats.errors.is_empty() {
|
||||
println!("\nErrors encountered:");
|
||||
for error in &stats.errors {
|
||||
println!(" - {}", error);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_migration_config() {
|
||||
let config = MigrationConfig::default();
|
||||
assert_eq!(config.batch_size, 1000);
|
||||
assert!(config.verify);
|
||||
}
|
||||
}
|
||||
@@ -1,384 +0,0 @@
|
||||
//! # Arrow Schema Models for LanceDB Tables
|
||||
//!
|
||||
//! Defines the data models and Arrow schemas for all Hanzo Node tables.
|
||||
|
||||
use anyhow::Result;
|
||||
use arrow_array::{
|
||||
Array, BinaryArray, Float32Array, Int32Array, Int64Array, RecordBatch, StringArray,
|
||||
TimestampMillisecondArray, BooleanArray, FixedSizeListArray,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Vector dimension for embeddings
|
||||
const EMBEDDING_DIM: i32 = 1536; // OpenAI ada-002 dimension
|
||||
|
||||
/// Create Arrow schema for users table
|
||||
pub fn user_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("profile_name", DataType::Utf8, false),
|
||||
Field::new("identity_type", DataType::Utf8, false),
|
||||
Field::new("identity_public_key", DataType::Utf8, true),
|
||||
Field::new("encryption_public_key", DataType::Utf8, true),
|
||||
Field::new("signature_public_key", DataType::Utf8, true),
|
||||
Field::new("node_signature_public_key", DataType::Utf8, true),
|
||||
Field::new("node_encryption_public_key", DataType::Utf8, true),
|
||||
Field::new("permission_type", DataType::Utf8, true),
|
||||
Field::new("wallet_id", DataType::Utf8, true),
|
||||
Field::new("metadata", DataType::Utf8, true), // JSON string
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("updated_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("is_active", DataType::Boolean, false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for tools table
|
||||
pub fn tool_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("name", DataType::Utf8, false),
|
||||
Field::new("description", DataType::Utf8, true),
|
||||
Field::new("version", DataType::Utf8, false),
|
||||
Field::new("tool_type", DataType::Utf8, false), // rust, js, python, mcp
|
||||
Field::new("parameters_schema", DataType::Utf8, true), // JSON schema
|
||||
Field::new("enabled", DataType::Boolean, false),
|
||||
Field::new("config", DataType::Utf8, true), // JSON config
|
||||
Field::new("usage_count", DataType::Int64, false),
|
||||
Field::new("avg_execution_time_ms", DataType::Float32, true),
|
||||
Field::new("last_used_at", DataType::Timestamp(TimeUnit::Millisecond, None), true),
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("updated_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for agents table
|
||||
pub fn agent_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("name", DataType::Utf8, false),
|
||||
Field::new("description", DataType::Utf8, true),
|
||||
Field::new("model", DataType::Utf8, false),
|
||||
Field::new("system_prompt", DataType::Utf8, true),
|
||||
Field::new("temperature", DataType::Float32, true),
|
||||
Field::new("max_tokens", DataType::Int32, true),
|
||||
Field::new("tools", DataType::Utf8, true), // JSON array of tool IDs
|
||||
Field::new("created_by", DataType::Utf8, false),
|
||||
Field::new("is_public", DataType::Boolean, false),
|
||||
Field::new("usage_count", DataType::Int64, false),
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("updated_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for embeddings table
|
||||
pub fn embedding_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("content", DataType::Utf8, false),
|
||||
Field::new("content_hash", DataType::Utf8, false),
|
||||
Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, false)),
|
||||
EMBEDDING_DIM,
|
||||
),
|
||||
false,
|
||||
),
|
||||
Field::new("model", DataType::Utf8, false),
|
||||
Field::new("source_type", DataType::Utf8, false), // document, chat, tool_output
|
||||
Field::new("source_id", DataType::Utf8, true),
|
||||
Field::new("metadata", DataType::Utf8, true), // JSON metadata
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for jobs table
|
||||
pub fn job_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("job_type", DataType::Utf8, false),
|
||||
Field::new("status", DataType::Utf8, false), // pending, running, completed, failed
|
||||
Field::new("priority", DataType::Int32, false),
|
||||
Field::new("payload", DataType::Utf8, false), // JSON payload
|
||||
Field::new("result", DataType::Utf8, true), // JSON result
|
||||
Field::new("error", DataType::Utf8, true),
|
||||
Field::new("created_by", DataType::Utf8, false),
|
||||
Field::new("agent_id", DataType::Utf8, true),
|
||||
Field::new("started_at", DataType::Timestamp(TimeUnit::Millisecond, None), true),
|
||||
Field::new("completed_at", DataType::Timestamp(TimeUnit::Millisecond, None), true),
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("updated_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for sessions table
|
||||
pub fn session_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("user_id", DataType::Utf8, false),
|
||||
Field::new("agent_id", DataType::Utf8, true),
|
||||
Field::new("messages", DataType::Utf8, false), // JSON array of messages
|
||||
Field::new("context", DataType::Utf8, true), // JSON context
|
||||
Field::new("token_count", DataType::Int64, false),
|
||||
Field::new("cost", DataType::Float32, true),
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
Field::new("updated_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Create Arrow schema for multimodal table
|
||||
#[cfg(feature = "multimodal")]
|
||||
pub fn multimodal_schema() -> Arc<ArrowSchema> {
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Field::new("media_type", DataType::Utf8, false), // image, audio, video
|
||||
Field::new("mime_type", DataType::Utf8, false),
|
||||
Field::new("data", DataType::Binary, false), // Raw binary data
|
||||
Field::new("thumbnail", DataType::Binary, true), // Optional thumbnail
|
||||
Field::new("width", DataType::Int32, true),
|
||||
Field::new("height", DataType::Int32, true),
|
||||
Field::new("duration_ms", DataType::Int64, true), // For audio/video
|
||||
Field::new(
|
||||
"embedding",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, false)),
|
||||
EMBEDDING_DIM,
|
||||
),
|
||||
true,
|
||||
),
|
||||
Field::new("text_content", DataType::Utf8, true), // OCR or transcription
|
||||
Field::new("metadata", DataType::Utf8, true), // JSON metadata
|
||||
Field::new("created_by", DataType::Utf8, false),
|
||||
Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), false),
|
||||
]))
|
||||
}
|
||||
|
||||
/// User model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct User {
|
||||
pub id: Uuid,
|
||||
pub profile_name: String,
|
||||
pub identity_type: String,
|
||||
pub identity_public_key: Option<String>,
|
||||
pub encryption_public_key: Option<String>,
|
||||
pub signature_public_key: Option<String>,
|
||||
pub node_signature_public_key: Option<String>,
|
||||
pub node_encryption_public_key: Option<String>,
|
||||
pub permission_type: Option<String>,
|
||||
pub wallet_id: Option<String>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
/// Tool model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Tool {
|
||||
pub id: Uuid,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub version: String,
|
||||
pub tool_type: ToolType,
|
||||
pub parameters_schema: Option<serde_json::Value>,
|
||||
pub enabled: bool,
|
||||
pub config: Option<serde_json::Value>,
|
||||
pub usage_count: i64,
|
||||
pub avg_execution_time_ms: Option<f32>,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Tool type enum
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ToolType {
|
||||
Rust,
|
||||
JavaScript,
|
||||
Python,
|
||||
MCP,
|
||||
}
|
||||
|
||||
/// Agent model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Agent {
|
||||
pub id: Uuid,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub model: String,
|
||||
pub system_prompt: Option<String>,
|
||||
pub temperature: Option<f32>,
|
||||
pub max_tokens: Option<i32>,
|
||||
pub tools: Vec<String>,
|
||||
pub created_by: String,
|
||||
pub is_public: bool,
|
||||
pub usage_count: i64,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Embedding model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Embedding {
|
||||
pub id: Uuid,
|
||||
pub content: String,
|
||||
pub content_hash: String,
|
||||
pub vector: Vec<f32>,
|
||||
pub model: String,
|
||||
pub source_type: String,
|
||||
pub source_id: Option<String>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Job model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Job {
|
||||
pub id: Uuid,
|
||||
pub job_type: String,
|
||||
pub status: JobStatus,
|
||||
pub priority: i32,
|
||||
pub payload: serde_json::Value,
|
||||
pub result: Option<serde_json::Value>,
|
||||
pub error: Option<String>,
|
||||
pub created_by: String,
|
||||
pub agent_id: Option<String>,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Job status enum
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum JobStatus {
|
||||
Pending,
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// Session model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Session {
|
||||
pub id: Uuid,
|
||||
pub user_id: String,
|
||||
pub agent_id: Option<String>,
|
||||
pub messages: Vec<Message>,
|
||||
pub context: Option<serde_json::Value>,
|
||||
pub token_count: i64,
|
||||
pub cost: Option<f32>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Message model for sessions
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
}
|
||||
|
||||
/// Tool call model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: serde_json::Value,
|
||||
pub result: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Multimodal content model
|
||||
#[cfg(feature = "multimodal")]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MultimodalContent {
|
||||
pub id: Uuid,
|
||||
pub media_type: MediaType,
|
||||
pub mime_type: String,
|
||||
pub data: Vec<u8>,
|
||||
pub thumbnail: Option<Vec<u8>>,
|
||||
pub width: Option<i32>,
|
||||
pub height: Option<i32>,
|
||||
pub duration_ms: Option<i64>,
|
||||
pub embedding: Option<Vec<f32>>,
|
||||
pub text_content: Option<String>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub created_by: String,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Media type enum
|
||||
#[cfg(feature = "multimodal")]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum MediaType {
|
||||
Image,
|
||||
Audio,
|
||||
Video,
|
||||
Document,
|
||||
}
|
||||
|
||||
/// Helper to convert models to RecordBatch
|
||||
pub trait ToRecordBatch {
|
||||
fn to_record_batch(&self) -> Result<RecordBatch>;
|
||||
}
|
||||
|
||||
/// Helper to convert from RecordBatch to models
|
||||
pub trait FromRecordBatch: Sized {
|
||||
fn from_record_batch(batch: &RecordBatch, row: usize) -> Result<Self>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_user_schema() {
|
||||
let schema = user_schema();
|
||||
assert_eq!(schema.fields().len(), 14);
|
||||
assert_eq!(schema.field(0).name(), "id");
|
||||
assert_eq!(schema.field(0).data_type(), &DataType::Utf8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_embedding_schema() {
|
||||
let schema = embedding_schema();
|
||||
let vector_field = schema.field(3);
|
||||
assert_eq!(vector_field.name(), "vector");
|
||||
|
||||
if let DataType::FixedSizeList(_, size) = vector_field.data_type() {
|
||||
assert_eq!(*size, EMBEDDING_DIM);
|
||||
} else {
|
||||
panic!("Expected FixedSizeList for vector field");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_type_serialization() {
|
||||
let tool_type = ToolType::JavaScript;
|
||||
let json = serde_json::to_string(&tool_type).unwrap();
|
||||
assert_eq!(json, r#""javascript""#);
|
||||
|
||||
let deserialized: ToolType = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(deserialized, ToolType::JavaScript));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_job_status_serialization() {
|
||||
let status = JobStatus::Running;
|
||||
let json = serde_json::to_string(&status).unwrap();
|
||||
assert_eq!(json, r#""running""#);
|
||||
|
||||
let deserialized: JobStatus = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(deserialized, JobStatus::Running));
|
||||
}
|
||||
}
|
||||
@@ -1,439 +0,0 @@
|
||||
//! # Vector Search Implementation for LanceDB
|
||||
//!
|
||||
//! High-performance semantic search, KNN search, and hybrid search capabilities.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use arrow_array::{Float32Array, RecordBatch, StringArray};
|
||||
use lancedb::{query::Query, Table};
|
||||
use log::{debug, info, warn};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Search configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchConfig {
|
||||
/// Maximum number of results
|
||||
pub limit: usize,
|
||||
/// Similarity threshold (0.0 to 1.0)
|
||||
pub threshold: Option<f32>,
|
||||
/// Enable re-ranking
|
||||
pub rerank: bool,
|
||||
/// Metadata filters
|
||||
pub filters: Option<HashMap<String, serde_json::Value>>,
|
||||
/// Search mode
|
||||
pub mode: SearchMode,
|
||||
}
|
||||
|
||||
impl Default for SearchConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
limit: 10,
|
||||
threshold: Some(0.7),
|
||||
rerank: false,
|
||||
filters: None,
|
||||
mode: SearchMode::Vector,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Search mode enum
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum SearchMode {
|
||||
/// Pure vector similarity search
|
||||
Vector,
|
||||
/// Full-text keyword search
|
||||
Keyword,
|
||||
/// Hybrid search combining vector and keyword
|
||||
Hybrid,
|
||||
/// Multi-modal search
|
||||
Multimodal,
|
||||
}
|
||||
|
||||
/// Search result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchResult {
|
||||
/// Document ID
|
||||
pub id: String,
|
||||
/// Original content
|
||||
pub content: String,
|
||||
/// Similarity score (0.0 to 1.0)
|
||||
pub score: f32,
|
||||
/// Distance metric value
|
||||
pub distance: f32,
|
||||
/// Additional metadata
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
/// Source information
|
||||
pub source: Option<String>,
|
||||
}
|
||||
|
||||
/// Vector search engine
|
||||
pub struct VectorSearch {
|
||||
table: Table,
|
||||
}
|
||||
|
||||
impl VectorSearch {
|
||||
/// Create new vector search instance
|
||||
pub fn new(table: Table) -> Self {
|
||||
Self { table }
|
||||
}
|
||||
|
||||
/// Perform semantic search using vector similarity
|
||||
pub async fn search(
|
||||
&self,
|
||||
query_vector: Vec<f32>,
|
||||
config: SearchConfig,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
info!("Performing vector search with limit: {}", config.limit);
|
||||
|
||||
let mut query = self.table
|
||||
.vector_search(&query_vector)
|
||||
.limit(config.limit);
|
||||
|
||||
// Apply metadata filters if provided
|
||||
if let Some(filters) = &config.filters {
|
||||
for (key, value) in filters {
|
||||
query = apply_filter(query, key, value)?;
|
||||
}
|
||||
}
|
||||
|
||||
let results = query
|
||||
.execute()
|
||||
.await
|
||||
.context("Failed to execute vector search")?;
|
||||
|
||||
// Convert to search results
|
||||
let mut search_results = Vec::new();
|
||||
|
||||
for batch in results.try_collect::<Vec<RecordBatch>>().await? {
|
||||
search_results.extend(batch_to_results(&batch, &config)?);
|
||||
}
|
||||
|
||||
// Apply threshold filtering
|
||||
if let Some(threshold) = config.threshold {
|
||||
search_results.retain(|r| r.score >= threshold);
|
||||
}
|
||||
|
||||
// Re-rank if enabled
|
||||
if config.rerank {
|
||||
search_results = self.rerank_results(search_results).await?;
|
||||
}
|
||||
|
||||
Ok(search_results)
|
||||
}
|
||||
|
||||
/// Perform keyword-based search
|
||||
pub async fn keyword_search(
|
||||
&self,
|
||||
query: &str,
|
||||
config: SearchConfig,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
info!("Performing keyword search for: {}", query);
|
||||
|
||||
// Full-text search
|
||||
let results = self.table
|
||||
.query()
|
||||
.full_text_search("content", query)
|
||||
.limit(config.limit)
|
||||
.execute()
|
||||
.await
|
||||
.context("Failed to execute keyword search")?;
|
||||
|
||||
let mut search_results = Vec::new();
|
||||
|
||||
for batch in results.try_collect::<Vec<RecordBatch>>().await? {
|
||||
search_results.extend(batch_to_results(&batch, &config)?);
|
||||
}
|
||||
|
||||
Ok(search_results)
|
||||
}
|
||||
|
||||
/// Perform hybrid search combining vector and keyword
|
||||
pub async fn hybrid_search(
|
||||
&self,
|
||||
query: &str,
|
||||
query_vector: Vec<f32>,
|
||||
config: SearchConfig,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
info!("Performing hybrid search");
|
||||
|
||||
// Get vector search results
|
||||
let mut vector_config = config.clone();
|
||||
vector_config.limit = config.limit * 2; // Get more candidates
|
||||
let vector_results = self.search(query_vector, vector_config).await?;
|
||||
|
||||
// Get keyword search results
|
||||
let mut keyword_config = config.clone();
|
||||
keyword_config.limit = config.limit * 2;
|
||||
let keyword_results = self.keyword_search(query, keyword_config).await?;
|
||||
|
||||
// Combine and deduplicate results
|
||||
let combined = self.combine_results(vector_results, keyword_results, &config)?;
|
||||
|
||||
Ok(combined)
|
||||
}
|
||||
|
||||
/// Find K nearest neighbors
|
||||
pub async fn knn_search(
|
||||
&self,
|
||||
query_vector: Vec<f32>,
|
||||
k: usize,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
debug!("Performing KNN search with k={}", k);
|
||||
|
||||
let results = self.table
|
||||
.vector_search(&query_vector)
|
||||
.limit(k)
|
||||
.execute()
|
||||
.await
|
||||
.context("Failed to execute KNN search")?;
|
||||
|
||||
let mut search_results = Vec::new();
|
||||
let config = SearchConfig::default();
|
||||
|
||||
for batch in results.try_collect::<Vec<RecordBatch>>().await? {
|
||||
search_results.extend(batch_to_results(&batch, &config)?);
|
||||
}
|
||||
|
||||
Ok(search_results)
|
||||
}
|
||||
|
||||
/// Search with multiple query vectors (multi-query)
|
||||
pub async fn multi_query_search(
|
||||
&self,
|
||||
query_vectors: Vec<Vec<f32>>,
|
||||
config: SearchConfig,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
info!("Performing multi-query search with {} queries", query_vectors.len());
|
||||
|
||||
let mut all_results = Vec::new();
|
||||
let mut result_map: HashMap<String, (SearchResult, f32)> = HashMap::new();
|
||||
|
||||
// Search with each query vector
|
||||
for (i, vector) in query_vectors.iter().enumerate() {
|
||||
debug!("Processing query vector {}/{}", i + 1, query_vectors.len());
|
||||
|
||||
let results = self.search(vector.clone(), config.clone()).await?;
|
||||
|
||||
// Aggregate scores
|
||||
for result in results {
|
||||
let entry = result_map
|
||||
.entry(result.id.clone())
|
||||
.or_insert((result.clone(), 0.0));
|
||||
entry.1 += result.score;
|
||||
}
|
||||
}
|
||||
|
||||
// Average scores and convert to results
|
||||
let num_queries = query_vectors.len() as f32;
|
||||
for (mut result, total_score) in result_map.into_values() {
|
||||
result.score = total_score / num_queries;
|
||||
all_results.push(result);
|
||||
}
|
||||
|
||||
// Sort by score
|
||||
all_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
|
||||
all_results.truncate(config.limit);
|
||||
|
||||
Ok(all_results)
|
||||
}
|
||||
|
||||
/// Diversity-aware search to reduce redundancy
|
||||
pub async fn diverse_search(
|
||||
&self,
|
||||
query_vector: Vec<f32>,
|
||||
config: SearchConfig,
|
||||
diversity_threshold: f32,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
info!("Performing diverse search with threshold: {}", diversity_threshold);
|
||||
|
||||
// Get more candidates than needed
|
||||
let mut search_config = config.clone();
|
||||
search_config.limit = config.limit * 3;
|
||||
let candidates = self.search(query_vector, search_config).await?;
|
||||
|
||||
// Select diverse results
|
||||
let mut selected = Vec::new();
|
||||
let mut selected_vectors: Vec<Vec<f32>> = Vec::new();
|
||||
|
||||
for candidate in candidates {
|
||||
if selected.len() >= config.limit {
|
||||
break;
|
||||
}
|
||||
|
||||
// Check similarity with already selected items
|
||||
let mut is_diverse = true;
|
||||
if !selected_vectors.is_empty() {
|
||||
// This is simplified - in production you'd fetch actual vectors
|
||||
// and compute cosine similarity
|
||||
for existing in &selected {
|
||||
if candidate.content.chars().take(50).collect::<String>()
|
||||
== existing.content.chars().take(50).collect::<String>() {
|
||||
is_diverse = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if is_diverse {
|
||||
selected.push(candidate);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
/// Re-rank results using a more sophisticated model
|
||||
async fn rerank_results(
|
||||
&self,
|
||||
mut results: Vec<SearchResult>,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
// In production, this would use a cross-encoder or other reranking model
|
||||
// For now, we'll just apply a simple boost based on metadata
|
||||
|
||||
for result in &mut results {
|
||||
if let Some(metadata) = &result.metadata {
|
||||
// Boost based on recency, quality score, etc.
|
||||
if let Some(quality) = metadata.get("quality_score") {
|
||||
if let Some(q) = quality.as_f64() {
|
||||
result.score *= (1.0 + (q as f32 * 0.1));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Combine vector and keyword results with score fusion
|
||||
fn combine_results(
|
||||
&self,
|
||||
vector_results: Vec<SearchResult>,
|
||||
keyword_results: Vec<SearchResult>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
let mut combined_map: HashMap<String, SearchResult> = HashMap::new();
|
||||
|
||||
// Add vector results with weight
|
||||
let vector_weight = 0.7;
|
||||
for mut result in vector_results {
|
||||
result.score *= vector_weight;
|
||||
combined_map.insert(result.id.clone(), result);
|
||||
}
|
||||
|
||||
// Add or update with keyword results
|
||||
let keyword_weight = 0.3;
|
||||
for mut result in keyword_results {
|
||||
result.score *= keyword_weight;
|
||||
|
||||
combined_map
|
||||
.entry(result.id.clone())
|
||||
.and_modify(|r| r.score += result.score)
|
||||
.or_insert(result);
|
||||
}
|
||||
|
||||
// Convert to vector and sort
|
||||
let mut combined: Vec<SearchResult> = combined_map.into_values().collect();
|
||||
combined.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
|
||||
combined.truncate(config.limit);
|
||||
|
||||
Ok(combined)
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply filter to query
|
||||
fn apply_filter(
|
||||
mut query: Query,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<Query> {
|
||||
match value {
|
||||
serde_json::Value::String(s) => {
|
||||
query = query.filter(format!("{} = '{}'", key, s));
|
||||
}
|
||||
serde_json::Value::Number(n) => {
|
||||
query = query.filter(format!("{} = {}", key, n));
|
||||
}
|
||||
serde_json::Value::Bool(b) => {
|
||||
query = query.filter(format!("{} = {}", key, b));
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow::anyhow!("Unsupported filter value type"));
|
||||
}
|
||||
}
|
||||
Ok(query)
|
||||
}
|
||||
|
||||
/// Convert RecordBatch to SearchResults
|
||||
fn batch_to_results(batch: &RecordBatch, config: &SearchConfig) -> Result<Vec<SearchResult>> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
let id_array = batch
|
||||
.column_by_name("id")
|
||||
.context("Missing id column")?
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.context("Invalid id column type")?;
|
||||
|
||||
let content_array = batch
|
||||
.column_by_name("content")
|
||||
.context("Missing content column")?
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.context("Invalid content column type")?;
|
||||
|
||||
let distance_array = batch
|
||||
.column_by_name("_distance")
|
||||
.map(|col| {
|
||||
col.as_any()
|
||||
.downcast_ref::<Float32Array>()
|
||||
})
|
||||
.unwrap_or(None);
|
||||
|
||||
for i in 0..batch.num_rows() {
|
||||
let id = id_array.value(i).to_string();
|
||||
let content = content_array.value(i).to_string();
|
||||
|
||||
// Convert distance to similarity score (1 - normalized_distance)
|
||||
let distance = distance_array
|
||||
.and_then(|arr| Some(arr.value(i)))
|
||||
.unwrap_or(0.0);
|
||||
let score = 1.0 / (1.0 + distance); // Convert distance to similarity
|
||||
|
||||
results.push(SearchResult {
|
||||
id,
|
||||
content,
|
||||
score,
|
||||
distance,
|
||||
metadata: None,
|
||||
source: None,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_search_config_default() {
|
||||
let config = SearchConfig::default();
|
||||
assert_eq!(config.limit, 10);
|
||||
assert_eq!(config.threshold, Some(0.7));
|
||||
assert!(!config.rerank);
|
||||
assert!(matches!(config.mode, SearchMode::Vector));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_mode_serialization() {
|
||||
let mode = SearchMode::Hybrid;
|
||||
let json = serde_json::to_string(&mode).unwrap();
|
||||
assert_eq!(json, r#""hybrid""#);
|
||||
|
||||
let deserialized: SearchMode = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(deserialized, SearchMode::Hybrid));
|
||||
}
|
||||
}
|
||||
@@ -1,47 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-kbs"
|
||||
description = "Hanzo AI - Kbs Library"
|
||||
license = "MIT"
|
||||
version = "1.1.11"
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
|
||||
[features]
|
||||
default = ["pqc"]
|
||||
pqc = ["hanzo-pqc"]
|
||||
|
||||
[dependencies]
|
||||
# Core dependencies
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
log = { workspace = true }
|
||||
dashmap = { workspace = true }
|
||||
rand = { workspace = true }
|
||||
|
||||
# Cryptography
|
||||
hpke = "0.12"
|
||||
chacha20poly1305 = "0.10"
|
||||
x25519-dalek = { workspace = true }
|
||||
ed25519-dalek = { workspace = true }
|
||||
blake3 = { workspace = true }
|
||||
hex = { workspace = true }
|
||||
|
||||
# For attestation
|
||||
base64 = { workspace = true }
|
||||
bincode = { workspace = true }
|
||||
uuid = { workspace = true, features = ["v4", "serde"] }
|
||||
|
||||
# Hanzo dependencies
|
||||
hanzo-messages = { version = "1.1.12", path = "../hanzo-messages" }
|
||||
hanzo-pqc = { workspace = true, optional = true }
|
||||
|
||||
# API dependencies
|
||||
warp = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-test = "0.4"
|
||||
@@ -1,307 +0,0 @@
|
||||
# Hanzo KBS (Key Broker Service)
|
||||
|
||||
A comprehensive Key Broker Service for Hanzo Node that manages cryptographic key lifecycle with TEE (Trusted Execution Environment) attestation and privacy tiers.
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
The Hanzo auth stack consists of three integrated components:
|
||||
|
||||
### 1. **KMS (Key Management System)** - Infisical
|
||||
- External service for centralized key management
|
||||
- Provides enterprise-grade secret storage
|
||||
- Handles encryption, rotation, and access control
|
||||
- Located at: `/Users/z/work/hanzo/kms` (Infisical deployment)
|
||||
|
||||
### 2. **KBS (Key Broker Service)** - This Module
|
||||
- Attestation-based key release
|
||||
- Privacy tier enforcement (0-4)
|
||||
- TEE integration (SEV-SNP, TDX, NVIDIA H100/Blackwell)
|
||||
- Web2/Web3 authentication support
|
||||
|
||||
### 3. **IAM (Identity & Access Management)**
|
||||
- Integrated with both Web2 (OAuth, SAML) and Web3 (DIDs, wallets)
|
||||
- Role-based access control (RBAC)
|
||||
- Capability-based security tokens
|
||||
- Cross-chain identity verification
|
||||
|
||||
## Privacy Tiers
|
||||
|
||||
The KBS implements a 5-tier privacy model with graceful degradation:
|
||||
|
||||
| Tier | Name | Description | TEE Support | Use Cases |
|
||||
|------|------|-------------|-------------|-----------|
|
||||
| **0** | Open | No encryption, public data | None | Public APIs, open datasets |
|
||||
| **1** | Encrypted | Basic encryption at rest/transit | None | Standard web services |
|
||||
| **2** | Secure | Hardware security module support | HSM/TPM | Enterprise applications |
|
||||
| **3** | GPU CC | GPU Confidential Computing | NVIDIA H100 | AI model training/inference |
|
||||
| **4** | GPU TEE-I/O | Full I/O isolation | NVIDIA Blackwell | Maximum security workloads |
|
||||
|
||||
### Graceful Degradation
|
||||
|
||||
The system automatically degrades to lower tiers when higher tier hardware is unavailable:
|
||||
|
||||
```rust
|
||||
// Tier 4 (Blackwell) → Tier 3 (H100) → Tier 2 (HSM) → Tier 1 (Software)
|
||||
let vault = VaultFactory::create_with_fallback(
|
||||
PrivacyTier::Tier4GpuTeeIo,
|
||||
&kbs,
|
||||
)?;
|
||||
```
|
||||
|
||||
## Web2/Web3 Compatibility
|
||||
|
||||
### Web2 Authentication
|
||||
- OAuth 2.0 / OIDC providers
|
||||
- SAML integration
|
||||
- API keys with HMAC
|
||||
- Session-based auth
|
||||
|
||||
### Web3 Authentication
|
||||
- Ethereum wallet signatures
|
||||
- DIDs (Decentralized Identifiers)
|
||||
- Chain-based attestation
|
||||
- Zero-knowledge proofs
|
||||
|
||||
## Integration with Infisical KMS
|
||||
|
||||
The KBS integrates with Infisical for enterprise key management:
|
||||
|
||||
```rust
|
||||
// Configure Infisical connection
|
||||
let kms_config = KmsConfig {
|
||||
endpoint: "https://kms.hanzo.ai",
|
||||
api_key: std::env::var("INFISICAL_API_KEY")?,
|
||||
project_id: std::env::var("INFISICAL_PROJECT_ID")?,
|
||||
};
|
||||
|
||||
// Create KBS with Infisical backend
|
||||
let kbs = KeyBrokerService::new(
|
||||
kms_config,
|
||||
attestation_config,
|
||||
)?;
|
||||
```
|
||||
|
||||
## TEE Support
|
||||
|
||||
### NVIDIA GPU TEE Integration
|
||||
|
||||
#### H100 Confidential Computing (Tier 3)
|
||||
```rust
|
||||
// Automatic detection of H100 CC mode
|
||||
let vault = GpuCcVault::new(kbs, device_id);
|
||||
if vault.is_cc_enabled()? {
|
||||
// Use hardware-protected keys
|
||||
vault.use_key(key_id, |key| {
|
||||
// Key operations in GPU enclave
|
||||
})?;
|
||||
}
|
||||
```
|
||||
|
||||
#### Blackwell TEE-I/O (Tier 4)
|
||||
```rust
|
||||
// Maximum security with I/O isolation
|
||||
let vault = GpuTeeIoVault::new(kbs, device_id, mig_instance);
|
||||
// All I/O is encrypted and authenticated
|
||||
vault.use_key_with_io_protection(key_id, |key| {
|
||||
// Operations with full isolation
|
||||
})?;
|
||||
```
|
||||
|
||||
### CPU TEE Support
|
||||
|
||||
- **AMD SEV-SNP**: Secure Encrypted Virtualization
|
||||
- **Intel TDX**: Trust Domain Extensions
|
||||
- **ARM CCA**: Confidential Compute Architecture
|
||||
|
||||
## API Usage
|
||||
|
||||
### Basic Key Operations
|
||||
|
||||
```rust
|
||||
use hanzo_kbs::{KeyBrokerService, PrivacyTier, KeyId};
|
||||
|
||||
// Initialize KBS
|
||||
let kbs = KeyBrokerService::new(config).await?;
|
||||
|
||||
// Request key with attestation
|
||||
let attestation = generate_attestation(TeeType::SevSnp).await?;
|
||||
let key = kbs.request_key(
|
||||
&KeyId::new("api-key"),
|
||||
attestation,
|
||||
PrivacyTier::Tier2Secure,
|
||||
).await?;
|
||||
|
||||
// Use key with automatic cleanup
|
||||
kbs.use_key(&key_id, |key_material| {
|
||||
// Perform cryptographic operations
|
||||
encrypt_data(key_material, data)
|
||||
})?;
|
||||
```
|
||||
|
||||
### Capability Tokens
|
||||
|
||||
```rust
|
||||
// Create capability token for agent
|
||||
let token = kbs.create_capability_token(
|
||||
agent_id,
|
||||
vec!["gpu_compute", "tee_access"],
|
||||
Duration::hours(1),
|
||||
)?;
|
||||
|
||||
// Verify and use token
|
||||
let authorized = kbs.verify_capability(token)?;
|
||||
```
|
||||
|
||||
## Deployment
|
||||
|
||||
### Docker Compose
|
||||
|
||||
```yaml
|
||||
version: '3.8'
|
||||
services:
|
||||
infisical-kms:
|
||||
image: infisical/infisical:latest
|
||||
environment:
|
||||
- ENCRYPTION_KEY=${ENCRYPTION_KEY}
|
||||
- AUTH_SECRET=${AUTH_SECRET}
|
||||
ports:
|
||||
- "8080:8080"
|
||||
|
||||
hanzo-node:
|
||||
image: hanzo/node:latest
|
||||
environment:
|
||||
- KMS_ENDPOINT=http://infisical-kms:8080
|
||||
- ENABLE_TEE=true
|
||||
- PRIVACY_TIER=3
|
||||
depends_on:
|
||||
- infisical-kms
|
||||
```
|
||||
|
||||
### Kubernetes
|
||||
|
||||
```yaml
|
||||
apiVersion: v1
|
||||
kind: ConfigMap
|
||||
metadata:
|
||||
name: hanzo-kbs-config
|
||||
data:
|
||||
kms_endpoint: "https://kms.hanzo.ai"
|
||||
default_privacy_tier: "2"
|
||||
enable_tee_attestation: "true"
|
||||
---
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: hanzo-node
|
||||
spec:
|
||||
template:
|
||||
spec:
|
||||
containers:
|
||||
- name: hanzo-node
|
||||
image: hanzo/node:latest
|
||||
envFrom:
|
||||
- configMapRef:
|
||||
name: hanzo-kbs-config
|
||||
```
|
||||
|
||||
## Security Considerations
|
||||
|
||||
### Key Protection Hierarchy
|
||||
|
||||
1. **Master Keys**: Never leave TEE, hardware-bound
|
||||
2. **KEKs (Key Encryption Keys)**: Wrapped by master keys
|
||||
3. **DEKs (Data Encryption Keys)**: Wrapped by KEKs
|
||||
4. **Session Keys**: Ephemeral, time-bound
|
||||
|
||||
### Attestation Chain
|
||||
|
||||
```
|
||||
Hardware Root of Trust
|
||||
↓
|
||||
Platform Firmware (PSP/ME)
|
||||
↓
|
||||
TEE Attestation (SEV/TDX/GPU)
|
||||
↓
|
||||
KBS Verification
|
||||
↓
|
||||
Key Release
|
||||
```
|
||||
|
||||
### Zero Trust Architecture
|
||||
|
||||
- No implicit trust between components
|
||||
- Every request requires attestation
|
||||
- Continuous verification
|
||||
- Least privilege access
|
||||
|
||||
## Performance Optimization
|
||||
|
||||
### Caching Strategy
|
||||
- Session key caching (5 min default)
|
||||
- Attestation result caching (1 min)
|
||||
- Capability token caching
|
||||
|
||||
### Connection Pooling
|
||||
- KMS connection pool (10 connections)
|
||||
- Database connection pool (R2D2)
|
||||
- HTTP/2 multiplexing for API calls
|
||||
|
||||
## Monitoring & Observability
|
||||
|
||||
### Metrics (Prometheus)
|
||||
- Key request latency
|
||||
- Attestation success rate
|
||||
- TEE availability
|
||||
- Vault operations/sec
|
||||
|
||||
### Logging
|
||||
- Audit logs for all key operations
|
||||
- Attestation verification logs
|
||||
- Security events
|
||||
- Performance metrics
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Run unit tests
|
||||
cargo test --lib
|
||||
|
||||
# Run integration tests (requires TEE simulator)
|
||||
cargo test --test integration
|
||||
|
||||
# Run with TEE hardware
|
||||
TEE_MODE=hardware cargo test
|
||||
|
||||
# Benchmark performance
|
||||
cargo bench
|
||||
```
|
||||
|
||||
## Migration from Other Systems
|
||||
|
||||
### From HashiCorp Vault
|
||||
```rust
|
||||
// Compatible API for easy migration
|
||||
let vault_compat = VaultCompatLayer::new(kbs);
|
||||
vault_compat.read("secret/data/api-key")?;
|
||||
```
|
||||
|
||||
### From AWS KMS
|
||||
```rust
|
||||
// AWS KMS compatibility layer
|
||||
let aws_compat = AwsKmsCompat::new(kbs);
|
||||
aws_compat.decrypt(ciphertext_blob)?;
|
||||
```
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](../../CONTRIBUTING.md) for development guidelines.
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0 - See [LICENSE](../../LICENSE) for details.
|
||||
|
||||
## Support
|
||||
|
||||
- Documentation: https://docs.hanzo.ai/kbs
|
||||
- Issues: https://github.com/hanzoai/hanzo-node/issues
|
||||
- Discord: https://discord.gg/hanzo
|
||||
@@ -1,223 +0,0 @@
|
||||
//! Attestation verification module
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::{Result, SecurityError};
|
||||
use crate::types::{AttestationType, PrivacyTier};
|
||||
|
||||
/// Result of attestation verification
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AttestationResult {
|
||||
pub verified: bool,
|
||||
pub max_tier: PrivacyTier,
|
||||
pub measurements: Vec<Measurement>,
|
||||
pub platform_info: PlatformInfo,
|
||||
pub expires_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Measurement {
|
||||
pub name: String,
|
||||
pub value: Vec<u8>,
|
||||
pub pcr_index: Option<u8>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PlatformInfo {
|
||||
pub platform_type: String,
|
||||
pub tcb_version: String,
|
||||
pub security_features: Vec<String>,
|
||||
pub vendor_info: serde_json::Value,
|
||||
}
|
||||
|
||||
impl AttestationResult {
|
||||
pub fn supports_tier(&self, tier: PrivacyTier) -> bool {
|
||||
self.verified && self.max_tier >= tier
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for verifying different attestation types
|
||||
#[async_trait]
|
||||
pub trait AttestationVerifier: Send + Sync {
|
||||
/// Verify an attestation and return the result
|
||||
async fn verify_attestation(&self, attestation: &AttestationType) -> Result<AttestationResult>;
|
||||
|
||||
/// Check if a measurement is in the allowed list
|
||||
async fn is_measurement_allowed(&self, measurement: &[u8]) -> bool;
|
||||
}
|
||||
|
||||
/// Mock attestation verifier for development
|
||||
pub struct MockAttestationVerifier;
|
||||
|
||||
#[async_trait]
|
||||
impl AttestationVerifier for MockAttestationVerifier {
|
||||
async fn verify_attestation(&self, attestation: &AttestationType) -> Result<AttestationResult> {
|
||||
match attestation {
|
||||
AttestationType::SevSnp { report, .. } => {
|
||||
// Mock SEV-SNP verification
|
||||
Ok(AttestationResult {
|
||||
verified: !report.is_empty(),
|
||||
max_tier: PrivacyTier::CpuTee,
|
||||
measurements: vec![
|
||||
Measurement {
|
||||
name: "kernel".to_string(),
|
||||
value: vec![0xAA; 32],
|
||||
pcr_index: Some(0),
|
||||
},
|
||||
],
|
||||
platform_info: PlatformInfo {
|
||||
platform_type: "AMD SEV-SNP".to_string(),
|
||||
tcb_version: "1.0.0".to_string(),
|
||||
security_features: vec!["SME".to_string(), "SEV".to_string()],
|
||||
vendor_info: serde_json::json!({}),
|
||||
},
|
||||
expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
|
||||
})
|
||||
}
|
||||
AttestationType::Tdx { quote, .. } => {
|
||||
// Mock TDX verification
|
||||
Ok(AttestationResult {
|
||||
verified: !quote.is_empty(),
|
||||
max_tier: PrivacyTier::CpuTee,
|
||||
measurements: vec![],
|
||||
platform_info: PlatformInfo {
|
||||
platform_type: "Intel TDX".to_string(),
|
||||
tcb_version: "1.0.0".to_string(),
|
||||
security_features: vec!["TDX".to_string()],
|
||||
vendor_info: serde_json::json!({}),
|
||||
},
|
||||
expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
|
||||
})
|
||||
}
|
||||
AttestationType::H100Cc { gpu_attestation, cpu_attestation } => {
|
||||
// Verify CPU attestation first
|
||||
let cpu_result = self.verify_attestation(cpu_attestation).await?;
|
||||
if !cpu_result.verified {
|
||||
return Err(SecurityError::InvalidAttestation(
|
||||
"CPU attestation failed".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
// Mock GPU CC verification
|
||||
Ok(AttestationResult {
|
||||
verified: !gpu_attestation.is_empty(),
|
||||
max_tier: PrivacyTier::GpuCc,
|
||||
measurements: cpu_result.measurements,
|
||||
platform_info: PlatformInfo {
|
||||
platform_type: "NVIDIA H100 CC".to_string(),
|
||||
tcb_version: "2.0.0".to_string(),
|
||||
security_features: vec!["GPU_CC".to_string(), "MIG".to_string()],
|
||||
vendor_info: serde_json::json!({"gpu_model": "H100"}),
|
||||
},
|
||||
expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
|
||||
})
|
||||
}
|
||||
AttestationType::BlackwellTeeIo { tee_io_report, mig_config } => {
|
||||
// Mock Blackwell TEE-I/O verification
|
||||
Ok(AttestationResult {
|
||||
verified: !tee_io_report.is_empty(),
|
||||
max_tier: PrivacyTier::GpuTeeIo,
|
||||
measurements: vec![],
|
||||
platform_info: PlatformInfo {
|
||||
platform_type: "NVIDIA Blackwell TEE-I/O".to_string(),
|
||||
tcb_version: "1.0.0".to_string(),
|
||||
security_features: vec![
|
||||
"TEE_IO".to_string(),
|
||||
"SECURE_BOOT".to_string(),
|
||||
if mig_config.is_some() { "MIG_ISOLATION" } else { "FULL_GPU" }.to_string(),
|
||||
],
|
||||
vendor_info: serde_json::json!({
|
||||
"gpu_model": "Blackwell",
|
||||
"mig_enabled": mig_config.is_some()
|
||||
}),
|
||||
},
|
||||
expires_at: chrono::Utc::now() + chrono::Duration::minutes(30),
|
||||
})
|
||||
}
|
||||
AttestationType::SimEid { eid, signature } => {
|
||||
// Mock SIM EID verification
|
||||
Ok(AttestationResult {
|
||||
verified: !eid.is_empty() && !signature.is_empty(),
|
||||
max_tier: PrivacyTier::AtRest,
|
||||
measurements: vec![],
|
||||
platform_info: PlatformInfo {
|
||||
platform_type: "SIM Card".to_string(),
|
||||
tcb_version: "1.0.0".to_string(),
|
||||
security_features: vec!["EID".to_string(), "SECURE_ELEMENT".to_string()],
|
||||
vendor_info: serde_json::json!({"eid": eid}),
|
||||
},
|
||||
expires_at: chrono::Utc::now() + chrono::Duration::hours(24),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_measurement_allowed(&self, _measurement: &[u8]) -> bool {
|
||||
// Mock implementation - allow all measurements
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Production attestation verifier that calls actual verification services
|
||||
pub struct HanzoAttestationVerifier {
|
||||
config: AttestationConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AttestationConfig {
|
||||
pub sev_snp_service_url: Option<String>,
|
||||
pub tdx_service_url: Option<String>,
|
||||
pub nvidia_cc_service_url: Option<String>,
|
||||
pub allowed_measurements: Vec<AllowedMeasurement>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AllowedMeasurement {
|
||||
pub name: String,
|
||||
pub platform: String,
|
||||
pub values: Vec<String>, // Hex-encoded
|
||||
}
|
||||
|
||||
impl HanzoAttestationVerifier {
|
||||
pub fn new(config: AttestationConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AttestationVerifier for HanzoAttestationVerifier {
|
||||
async fn verify_attestation(&self, attestation: &AttestationType) -> Result<AttestationResult> {
|
||||
match attestation {
|
||||
AttestationType::SevSnp { report, vcek_cert, platform_cert_chain } => {
|
||||
// TODO: Implement actual SEV-SNP verification
|
||||
// This would call AMD's attestation service or use a local verifier
|
||||
todo!("Implement SEV-SNP verification")
|
||||
}
|
||||
AttestationType::Tdx { quote, collateral } => {
|
||||
// TODO: Implement actual TDX verification
|
||||
// This would use Intel's attestation libraries
|
||||
todo!("Implement TDX verification")
|
||||
}
|
||||
AttestationType::H100Cc { .. } => {
|
||||
// TODO: Implement NVIDIA H100 CC verification
|
||||
todo!("Implement H100 CC verification")
|
||||
}
|
||||
AttestationType::BlackwellTeeIo { .. } => {
|
||||
// TODO: Implement Blackwell TEE-I/O verification
|
||||
todo!("Implement Blackwell verification")
|
||||
}
|
||||
AttestationType::SimEid { .. } => {
|
||||
// TODO: Implement SIM EID verification
|
||||
todo!("Implement SIM verification")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_measurement_allowed(&self, measurement: &[u8]) -> bool {
|
||||
let hex_measurement = hex::encode(measurement);
|
||||
self.config.allowed_measurements.iter().any(|allowed| {
|
||||
allowed.values.contains(&hex_measurement)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,47 +0,0 @@
|
||||
//! Error types for Hanzo Security
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum SecurityError {
|
||||
#[error("Key not found: {0}")]
|
||||
KeyNotFound(String),
|
||||
|
||||
#[error("Invalid attestation: {0}")]
|
||||
InvalidAttestation(String),
|
||||
|
||||
#[error("Policy violation: {0}")]
|
||||
PolicyViolation(String),
|
||||
|
||||
#[error("Cryptographic error: {0}")]
|
||||
CryptoError(String),
|
||||
|
||||
#[error("HSM error: {0}")]
|
||||
HsmError(String),
|
||||
|
||||
#[error("KMS error: {0}")]
|
||||
KmsError(String),
|
||||
|
||||
#[error("KBS error: {0}")]
|
||||
KbsError(String),
|
||||
|
||||
#[error("Tier mismatch: requested {requested}, available {available}")]
|
||||
TierMismatch { requested: u8, available: u8 },
|
||||
|
||||
#[error("Session expired")]
|
||||
SessionExpired,
|
||||
|
||||
#[error("Rate limit exceeded")]
|
||||
RateLimitExceeded,
|
||||
|
||||
#[error("Serialization error: {0}")]
|
||||
SerializationError(#[from] serde_json::Error),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
IoError(#[from] std::io::Error),
|
||||
|
||||
#[error("Other error: {0}")]
|
||||
Other(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, SecurityError>;
|
||||
@@ -1,364 +0,0 @@
|
||||
//! Key Broker Service (KBS) trait and implementations
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::{Result, SecurityError};
|
||||
use crate::types::*;
|
||||
use crate::kms::KeyManagementService;
|
||||
use crate::attestation::AttestationVerifier;
|
||||
|
||||
/// Key Broker Service trait - handles attestation and policy-based key release
|
||||
#[async_trait]
|
||||
pub trait KeyBrokerService: Send + Sync {
|
||||
/// Authorize key release based on attestation and policy
|
||||
async fn authorize(
|
||||
&self,
|
||||
request: KeyAuthorizationRequest,
|
||||
) -> Result<KeyAuthorizationResponse>;
|
||||
|
||||
/// Renew an existing session with fresh attestation
|
||||
async fn renew(
|
||||
&self,
|
||||
session_id: Uuid,
|
||||
attestation: AttestationType,
|
||||
) -> Result<RenewResponse>;
|
||||
|
||||
/// Revoke a session (admin or chain-triggered)
|
||||
async fn revoke(&self, session_id: Uuid, reason: RevocationReason) -> Result<()>;
|
||||
|
||||
/// Get session status
|
||||
async fn get_session_status(&self, session_id: Uuid) -> Result<SessionStatus>;
|
||||
|
||||
/// Get policy for a given tier
|
||||
async fn get_tier_policy(&self, tier: PrivacyTier) -> Result<TierPolicy>;
|
||||
}
|
||||
|
||||
/// KBS implementation that connects to KMS
|
||||
pub struct HanzoKbs<K: KeyManagementService, V: AttestationVerifier> {
|
||||
kms: K,
|
||||
verifier: V,
|
||||
config: KbsConfig,
|
||||
sessions: dashmap::DashMap<Uuid, SessionInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KbsConfig {
|
||||
pub max_session_duration: Duration,
|
||||
pub attestation_cache_ttl: Duration,
|
||||
pub rate_limit_per_minute: u32,
|
||||
pub require_chain_verification: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct RenewResponse {
|
||||
pub session_id: Uuid,
|
||||
pub new_expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub enum RevocationReason {
|
||||
AdminAction { admin_id: String, reason: String },
|
||||
ChainTriggered { transaction_hash: String },
|
||||
PolicyViolation { details: String },
|
||||
Expired,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct SessionStatus {
|
||||
pub session_id: Uuid,
|
||||
pub agent_id: String,
|
||||
pub tier: PrivacyTier,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub renewals: u32,
|
||||
pub active: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SessionInfo {
|
||||
pub agent_id: String,
|
||||
pub tier: PrivacyTier,
|
||||
pub enclave_public_key: Vec<u8>,
|
||||
pub authorized_keys: Vec<KeyId>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub renewals: u32,
|
||||
pub active: bool,
|
||||
}
|
||||
|
||||
/// Policy for each privacy tier
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct TierPolicy {
|
||||
pub tier: PrivacyTier,
|
||||
pub required_attestations: Vec<AttestationRequirement>,
|
||||
pub max_session_duration: Duration,
|
||||
pub allowed_operations: Vec<String>,
|
||||
pub key_restrictions: Vec<KeyRestriction>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct AttestationRequirement {
|
||||
pub attestation_type: String,
|
||||
pub min_tcb_version: Option<String>,
|
||||
pub allowed_measurements: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyRestriction {
|
||||
pub key_type: String,
|
||||
pub max_usage_count: Option<u64>,
|
||||
pub require_audit: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<K: KeyManagementService, V: AttestationVerifier> KeyBrokerService for HanzoKbs<K, V> {
|
||||
async fn authorize(
|
||||
&self,
|
||||
request: KeyAuthorizationRequest,
|
||||
) -> Result<KeyAuthorizationResponse> {
|
||||
// Rate limiting check
|
||||
if !self.check_rate_limit(&request.capability_token.subject).await? {
|
||||
return Err(SecurityError::RateLimitExceeded);
|
||||
}
|
||||
|
||||
// Verify attestation
|
||||
let attestation_result = self.verifier
|
||||
.verify_attestation(&request.attestation)
|
||||
.await?;
|
||||
|
||||
// Check tier compatibility
|
||||
let requested_tier = request.capability_token.tier;
|
||||
if !attestation_result.supports_tier(requested_tier) {
|
||||
return Err(SecurityError::TierMismatch {
|
||||
requested: requested_tier as u8,
|
||||
available: attestation_result.max_tier as u8,
|
||||
});
|
||||
}
|
||||
|
||||
// Verify capability token (on-chain if configured)
|
||||
if self.config.require_chain_verification {
|
||||
self.verify_capability_token(&request.capability_token).await?;
|
||||
}
|
||||
|
||||
// Get tier policy
|
||||
let policy = self.get_tier_policy(requested_tier).await?;
|
||||
|
||||
// Authorize requested keys
|
||||
let mut authorized_keys = Vec::new();
|
||||
for key_request in &request.requested_keys {
|
||||
let authorized_key = self.authorize_key(
|
||||
&key_request,
|
||||
&request.session_public_key,
|
||||
requested_tier,
|
||||
&policy,
|
||||
).await?;
|
||||
authorized_keys.push(authorized_key);
|
||||
}
|
||||
|
||||
// Create session
|
||||
let session_id = Uuid::new_v4();
|
||||
let expires_at = Utc::now() + self.config.max_session_duration;
|
||||
|
||||
self.sessions.insert(session_id, SessionInfo {
|
||||
agent_id: request.capability_token.subject.clone(),
|
||||
tier: requested_tier,
|
||||
enclave_public_key: request.session_public_key.clone(),
|
||||
authorized_keys: authorized_keys.iter().map(|k| k.key_id.clone()).collect(),
|
||||
created_at: Utc::now(),
|
||||
expires_at,
|
||||
renewals: 0,
|
||||
active: true,
|
||||
});
|
||||
|
||||
Ok(KeyAuthorizationResponse {
|
||||
session_id,
|
||||
authorized_keys,
|
||||
expires_at,
|
||||
})
|
||||
}
|
||||
|
||||
async fn renew(
|
||||
&self,
|
||||
session_id: Uuid,
|
||||
attestation: AttestationType,
|
||||
) -> Result<RenewResponse> {
|
||||
let mut session = self.sessions.get_mut(&session_id)
|
||||
.ok_or_else(|| SecurityError::SessionExpired)?;
|
||||
|
||||
if !session.active || session.expires_at < Utc::now() {
|
||||
return Err(SecurityError::SessionExpired);
|
||||
}
|
||||
|
||||
// Verify fresh attestation
|
||||
let attestation_result = self.verifier.verify_attestation(&attestation).await?;
|
||||
if !attestation_result.supports_tier(session.tier) {
|
||||
return Err(SecurityError::InvalidAttestation(
|
||||
"Attestation no longer supports required tier".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
// Update session
|
||||
session.expires_at = Utc::now() + self.config.max_session_duration;
|
||||
session.renewals += 1;
|
||||
|
||||
Ok(RenewResponse {
|
||||
session_id,
|
||||
new_expires_at: session.expires_at,
|
||||
})
|
||||
}
|
||||
|
||||
async fn revoke(&self, session_id: Uuid, reason: RevocationReason) -> Result<()> {
|
||||
if let Some(mut session) = self.sessions.get_mut(&session_id) {
|
||||
session.active = false;
|
||||
|
||||
// Log revocation
|
||||
log::info!(
|
||||
"Session {} revoked for agent {}: {:?}",
|
||||
session_id, session.agent_id, reason
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_session_status(&self, session_id: Uuid) -> Result<SessionStatus> {
|
||||
let session = self.sessions.get(&session_id)
|
||||
.ok_or_else(|| SecurityError::KeyNotFound("Session not found".to_string()))?;
|
||||
|
||||
Ok(SessionStatus {
|
||||
session_id,
|
||||
agent_id: session.agent_id.clone(),
|
||||
tier: session.tier,
|
||||
created_at: session.created_at,
|
||||
expires_at: session.expires_at,
|
||||
renewals: session.renewals,
|
||||
active: session.active && session.expires_at > Utc::now(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_tier_policy(&self, tier: PrivacyTier) -> Result<TierPolicy> {
|
||||
// This would typically load from configuration or database
|
||||
Ok(match tier {
|
||||
PrivacyTier::Open => TierPolicy {
|
||||
tier,
|
||||
required_attestations: vec![],
|
||||
max_session_duration: Duration::hours(24),
|
||||
allowed_operations: vec!["*".to_string()],
|
||||
key_restrictions: vec![],
|
||||
},
|
||||
PrivacyTier::AtRest => TierPolicy {
|
||||
tier,
|
||||
required_attestations: vec![],
|
||||
max_session_duration: Duration::hours(12),
|
||||
allowed_operations: vec!["read".to_string(), "write".to_string()],
|
||||
key_restrictions: vec![],
|
||||
},
|
||||
PrivacyTier::CpuTee => TierPolicy {
|
||||
tier,
|
||||
required_attestations: vec![
|
||||
AttestationRequirement {
|
||||
attestation_type: "SevSnp".to_string(),
|
||||
min_tcb_version: Some("1.0".to_string()),
|
||||
allowed_measurements: None,
|
||||
},
|
||||
],
|
||||
max_session_duration: Duration::hours(4),
|
||||
allowed_operations: vec!["compute".to_string()],
|
||||
key_restrictions: vec![
|
||||
KeyRestriction {
|
||||
key_type: "AgentDek".to_string(),
|
||||
max_usage_count: Some(1000),
|
||||
require_audit: true,
|
||||
},
|
||||
],
|
||||
},
|
||||
PrivacyTier::GpuCc => TierPolicy {
|
||||
tier,
|
||||
required_attestations: vec![
|
||||
AttestationRequirement {
|
||||
attestation_type: "H100Cc".to_string(),
|
||||
min_tcb_version: Some("2.0".to_string()),
|
||||
allowed_measurements: None,
|
||||
},
|
||||
],
|
||||
max_session_duration: Duration::hours(2),
|
||||
allowed_operations: vec!["gpu_compute".to_string()],
|
||||
key_restrictions: vec![
|
||||
KeyRestriction {
|
||||
key_type: "AgentDek".to_string(),
|
||||
max_usage_count: Some(100),
|
||||
require_audit: true,
|
||||
},
|
||||
],
|
||||
},
|
||||
PrivacyTier::GpuTeeIo => TierPolicy {
|
||||
tier,
|
||||
required_attestations: vec![
|
||||
AttestationRequirement {
|
||||
attestation_type: "BlackwellTeeIo".to_string(),
|
||||
min_tcb_version: Some("1.0".to_string()),
|
||||
allowed_measurements: None,
|
||||
},
|
||||
],
|
||||
max_session_duration: Duration::hours(1),
|
||||
allowed_operations: vec!["secure_inference".to_string()],
|
||||
key_restrictions: vec![
|
||||
KeyRestriction {
|
||||
key_type: "AgentDek".to_string(),
|
||||
max_usage_count: Some(10),
|
||||
require_audit: true,
|
||||
},
|
||||
],
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<K: KeyManagementService, V: AttestationVerifier> HanzoKbs<K, V> {
|
||||
pub fn new(kms: K, verifier: V, config: KbsConfig) -> Self {
|
||||
Self {
|
||||
kms,
|
||||
verifier,
|
||||
config,
|
||||
sessions: dashmap::DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn check_rate_limit(&self, subject: &str) -> Result<bool> {
|
||||
// TODO: Implement proper rate limiting with Redis or in-memory cache
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn verify_capability_token(&self, token: &CapabilityToken) -> Result<()> {
|
||||
// TODO: Verify on-chain signature
|
||||
if token.expires_at.map(|exp| exp < Utc::now()).unwrap_or(false) {
|
||||
return Err(SecurityError::PolicyViolation("Token expired".to_string()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn authorize_key(
|
||||
&self,
|
||||
request: &KeyRequest,
|
||||
enclave_public_key: &[u8],
|
||||
tier: PrivacyTier,
|
||||
policy: &TierPolicy,
|
||||
) -> Result<AuthorizedKey> {
|
||||
// TODO: Implement HPKE wrapping
|
||||
// For now, return a mock authorized key
|
||||
Ok(AuthorizedKey {
|
||||
key_id: KeyId::new(),
|
||||
hpke_wrapped_key: vec![0; 32], // Mock wrapped key
|
||||
metadata: KeyMetadata {
|
||||
key_type: format!("{:?}", request.key_type),
|
||||
tier,
|
||||
restrictions: policy.key_restrictions.iter()
|
||||
.map(|r| format!("{}: max_usage={:?}", r.key_type, r.max_usage_count))
|
||||
.collect(),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,174 +0,0 @@
|
||||
//! Key Management Service (KMS) trait and implementations
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::types::*;
|
||||
|
||||
// Declare submodules
|
||||
pub mod api;
|
||||
pub mod memory_kms;
|
||||
|
||||
/// Key Management Service trait - handles key lifecycle and storage
|
||||
#[async_trait]
|
||||
pub trait KeyManagementService: Send + Sync {
|
||||
/// Initialize the KMS with root key
|
||||
async fn initialize(&mut self, config: KmsConfig) -> Result<()>;
|
||||
|
||||
/// Create a new tenant KEK
|
||||
async fn create_tenant_kek(&self, tenant_id: &str) -> Result<TenantKek>;
|
||||
|
||||
/// Create a new agent DEK
|
||||
async fn create_agent_dek(&self, agent_id: &str, tenant_id: &str) -> Result<AgentDek>;
|
||||
|
||||
/// Wrap a key under a parent key
|
||||
async fn wrap_key(&self, key_data: &[u8], parent_key_id: &KeyId) -> Result<Vec<u8>>;
|
||||
|
||||
/// Unwrap a key (internal use only, never exposed to KBS)
|
||||
async fn unwrap_key(&self, wrapped_key: &[u8], parent_key_id: &KeyId) -> Result<Vec<u8>>;
|
||||
|
||||
/// Rotate a key
|
||||
async fn rotate_key(&self, key_id: &KeyId) -> Result<KeyId>;
|
||||
|
||||
/// Destroy a key
|
||||
async fn destroy_key(&self, key_id: &KeyId) -> Result<()>;
|
||||
|
||||
/// Get key metadata
|
||||
async fn get_key_metadata(&self, key_id: &KeyId) -> Result<KeyInfo>;
|
||||
|
||||
/// Get audit logs
|
||||
async fn get_audit_logs(
|
||||
&self,
|
||||
start: DateTime<Utc>,
|
||||
end: DateTime<Utc>,
|
||||
filter: Option<AuditFilter>,
|
||||
) -> Result<Vec<KeyAuditEntry>>;
|
||||
|
||||
/// BYOK: Import customer key
|
||||
async fn import_customer_key(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
wrapped_key: &[u8],
|
||||
key_metadata: CustomerKeyMetadata,
|
||||
) -> Result<KeyId>;
|
||||
|
||||
/// HYOK: Register customer-held key reference
|
||||
async fn register_customer_held_key(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
key_reference: CustomerKeyReference,
|
||||
) -> Result<KeyId>;
|
||||
}
|
||||
|
||||
/// KMS configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KmsConfig {
|
||||
pub hsm_type: HsmType,
|
||||
pub root_key_source: RootKeySource,
|
||||
pub audit_retention_days: u32,
|
||||
pub key_rotation_days: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum HsmType {
|
||||
Software,
|
||||
AwsCloudHsm { cluster_id: String },
|
||||
AzureKeyVault { vault_url: String },
|
||||
HashicorpVault { url: String, namespace: Option<String> },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum RootKeySource {
|
||||
Generate,
|
||||
Import { wrapped_key: Vec<u8> },
|
||||
Existing { key_id: String },
|
||||
}
|
||||
|
||||
/// Extended key information
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KeyInfo {
|
||||
pub key_id: KeyId,
|
||||
pub key_type: KeyType,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
pub rotation_due: Option<DateTime<Utc>>,
|
||||
pub state: KeyState,
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum KeyType {
|
||||
Root,
|
||||
TenantKek { tenant_id: String },
|
||||
AgentDek { agent_id: String, tenant_id: String },
|
||||
CustomerKey { tenant_id: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum KeyState {
|
||||
Active,
|
||||
Rotating,
|
||||
Expired,
|
||||
Destroyed,
|
||||
}
|
||||
|
||||
/// Audit log filter
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
pub struct AuditFilter {
|
||||
pub key_ids: Option<Vec<KeyId>>,
|
||||
pub operations: Option<Vec<KeyOperation>>,
|
||||
pub actors: Option<Vec<String>>,
|
||||
pub success_only: Option<bool>,
|
||||
}
|
||||
|
||||
/// Customer key metadata for BYOK
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CustomerKeyMetadata {
|
||||
pub algorithm: String,
|
||||
pub key_size: u32,
|
||||
pub usage: Vec<String>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Customer key reference for HYOK
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CustomerKeyReference {
|
||||
pub provider: String,
|
||||
pub key_id: String,
|
||||
pub region: Option<String>,
|
||||
pub permissions: Vec<String>,
|
||||
}
|
||||
|
||||
/// Key wrapping algorithms supported
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum KeyWrapAlgorithm {
|
||||
AesKwp256,
|
||||
Rsa4096OaepSha256,
|
||||
ChaCha20Poly1305,
|
||||
}
|
||||
|
||||
/// Key hierarchy operations
|
||||
pub struct KeyHierarchy;
|
||||
|
||||
impl KeyHierarchy {
|
||||
/// Derive a KEK from root key and tenant salt
|
||||
pub fn derive_tenant_kek(root_key: &[u8], tenant_id: &str) -> Vec<u8> {
|
||||
let salt = format!("tenant_kek_{}", tenant_id);
|
||||
let mut hasher = blake3::Hasher::new_derive_key(&salt);
|
||||
hasher.update(root_key);
|
||||
let mut output = vec![0u8; 32];
|
||||
hasher.finalize_xof().fill(&mut output);
|
||||
output
|
||||
}
|
||||
|
||||
/// Generate a random DEK
|
||||
pub fn generate_dek() -> Vec<u8> {
|
||||
use rand::RngCore;
|
||||
let mut key = vec![0u8; 32];
|
||||
rand::thread_rng().fill_bytes(&mut key);
|
||||
key
|
||||
}
|
||||
}
|
||||
@@ -1,200 +0,0 @@
|
||||
//! HTTP API for KMS operations
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use warp::{Filter, Rejection, Reply};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::kms::KeyManagementService;
|
||||
use crate::types::{KeyId, AgentDek, TenantKek};
|
||||
|
||||
/// API request/response types
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CreateTenantKekRequest {
|
||||
pub tenant_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CreateTenantKekResponse {
|
||||
pub kek: TenantKek,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CreateAgentDekRequest {
|
||||
pub agent_id: String,
|
||||
pub tenant_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CreateAgentDekResponse {
|
||||
pub dek: AgentDek,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct WrapKeyRequest {
|
||||
pub key_data_base64: String,
|
||||
pub parent_key_id: KeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct WrapKeyResponse {
|
||||
pub wrapped_key_base64: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct RotateKeyRequest {
|
||||
pub key_id: KeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct RotateKeyResponse {
|
||||
pub new_key_id: KeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct DestroyKeyRequest {
|
||||
pub key_id: KeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct GetAuditLogsRequest {
|
||||
pub start_time: chrono::DateTime<chrono::Utc>,
|
||||
pub end_time: chrono::DateTime<chrono::Utc>,
|
||||
pub filter: Option<crate::kms::AuditFilter>,
|
||||
}
|
||||
|
||||
/// Create all KMS API routes
|
||||
pub fn kms_routes<K: KeyManagementService + Clone + Send + Sync + 'static>(
|
||||
kms: Arc<K>,
|
||||
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
|
||||
let create_tenant_kek = warp::path!("keys" / "tenant")
|
||||
.and(warp::post())
|
||||
.and(warp::body::json())
|
||||
.and(with_kms(kms.clone()))
|
||||
.and_then(handle_create_tenant_kek);
|
||||
|
||||
let create_agent_dek = warp::path!("keys" / "agent")
|
||||
.and(warp::post())
|
||||
.and(warp::body::json())
|
||||
.and(with_kms(kms.clone()))
|
||||
.and_then(handle_create_agent_dek);
|
||||
|
||||
let wrap_key = warp::path!("wrap")
|
||||
.and(warp::post())
|
||||
.and(warp::body::json())
|
||||
.and(with_kms(kms.clone()))
|
||||
.and_then(handle_wrap_key);
|
||||
|
||||
let rotate_key = warp::path!("rotate")
|
||||
.and(warp::post())
|
||||
.and(warp::body::json())
|
||||
.and(with_kms(kms.clone()))
|
||||
.and_then(handle_rotate_key);
|
||||
|
||||
let destroy_key = warp::path!("destroy")
|
||||
.and(warp::post())
|
||||
.and(warp::body::json())
|
||||
.and(with_kms(kms.clone()))
|
||||
.and_then(handle_destroy_key);
|
||||
|
||||
let get_audit_logs = warp::path!("audit")
|
||||
.and(warp::get())
|
||||
.and(warp::query())
|
||||
.and(with_kms(kms))
|
||||
.and_then(handle_get_audit_logs);
|
||||
|
||||
create_tenant_kek
|
||||
.or(create_agent_dek)
|
||||
.or(wrap_key)
|
||||
.or(rotate_key)
|
||||
.or(destroy_key)
|
||||
.or(get_audit_logs)
|
||||
}
|
||||
|
||||
fn with_kms<K: KeyManagementService + Clone + Send>(
|
||||
kms: Arc<K>,
|
||||
) -> impl Filter<Extract = (Arc<K>,), Error = std::convert::Infallible> + Clone {
|
||||
warp::any().map(move || kms.clone())
|
||||
}
|
||||
|
||||
async fn handle_create_tenant_kek<K: KeyManagementService>(
|
||||
req: CreateTenantKekRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
match kms.create_tenant_kek(&req.tenant_id).await {
|
||||
Ok(kek) => Ok(warp::reply::json(&CreateTenantKekResponse { kek })),
|
||||
Err(e) => {
|
||||
log::error!("Failed to create tenant KEK: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_create_agent_dek<K: KeyManagementService>(
|
||||
req: CreateAgentDekRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
match kms.create_agent_dek(&req.agent_id, &req.tenant_id).await {
|
||||
Ok(dek) => Ok(warp::reply::json(&CreateAgentDekResponse { dek })),
|
||||
Err(e) => {
|
||||
log::error!("Failed to create agent DEK: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_wrap_key<K: KeyManagementService>(
|
||||
req: WrapKeyRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
let key_data = base64::decode(&req.key_data_base64)
|
||||
.map_err(|_| warp::reject::reject())?;
|
||||
|
||||
match kms.wrap_key(&key_data, &req.parent_key_id).await {
|
||||
Ok(wrapped) => Ok(warp::reply::json(&WrapKeyResponse {
|
||||
wrapped_key_base64: base64::encode(&wrapped),
|
||||
})),
|
||||
Err(e) => {
|
||||
log::error!("Failed to wrap key: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_rotate_key<K: KeyManagementService>(
|
||||
req: RotateKeyRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
match kms.rotate_key(&req.key_id).await {
|
||||
Ok(new_key_id) => Ok(warp::reply::json(&RotateKeyResponse { new_key_id })),
|
||||
Err(e) => {
|
||||
log::error!("Failed to rotate key: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_destroy_key<K: KeyManagementService>(
|
||||
req: DestroyKeyRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
match kms.destroy_key(&req.key_id).await {
|
||||
Ok(()) => Ok(warp::reply::with_status("", warp::http::StatusCode::NO_CONTENT)),
|
||||
Err(e) => {
|
||||
log::error!("Failed to destroy key: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_get_audit_logs<K: KeyManagementService>(
|
||||
req: GetAuditLogsRequest,
|
||||
kms: Arc<K>,
|
||||
) -> Result<impl Reply, Rejection> {
|
||||
match kms.get_audit_logs(req.start_time, req.end_time, req.filter).await {
|
||||
Ok(logs) => Ok(warp::reply::json(&logs)),
|
||||
Err(e) => {
|
||||
log::error!("Failed to get audit logs: {}", e);
|
||||
Err(warp::reject::reject())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,396 +0,0 @@
|
||||
//! In-memory KMS implementation for development and testing
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use dashmap::DashMap;
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::{Result, SecurityError};
|
||||
use crate::kms::{
|
||||
AuditFilter, CustomerKeyMetadata, CustomerKeyReference, HsmType, KeyInfo, KeyManagementService,
|
||||
KeyState, KeyType, KmsConfig, RootKeySource,
|
||||
};
|
||||
use crate::types::{
|
||||
AgentDek, KeyAuditEntry, KeyId, KeyOperation, RootKey, TenantKek,
|
||||
};
|
||||
|
||||
/// In-memory KMS implementation
|
||||
pub struct MemoryKms {
|
||||
root_key: Arc<RootKey>,
|
||||
root_key_material: Vec<u8>,
|
||||
tenant_keks: DashMap<String, (TenantKek, Vec<u8>)>, // tenant_id -> (metadata, key_material)
|
||||
agent_deks: DashMap<String, (AgentDek, Vec<u8>)>, // agent_id -> (metadata, key_material)
|
||||
key_info: DashMap<KeyId, KeyInfo>,
|
||||
audit_log: DashMap<Uuid, KeyAuditEntry>,
|
||||
config: KmsConfig,
|
||||
}
|
||||
|
||||
impl MemoryKms {
|
||||
pub fn new() -> Self {
|
||||
let root_key_id = KeyId::new();
|
||||
let root_key = RootKey {
|
||||
id: root_key_id.clone(),
|
||||
created_at: Utc::now(),
|
||||
hsm_handle: None,
|
||||
};
|
||||
|
||||
// Generate root key material
|
||||
let root_key_material = crate::kms::KeyHierarchy::generate_dek();
|
||||
|
||||
let key_info = DashMap::new();
|
||||
key_info.insert(
|
||||
root_key_id.clone(),
|
||||
KeyInfo {
|
||||
key_id: root_key_id,
|
||||
key_type: KeyType::Root,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
expires_at: None,
|
||||
rotation_due: None,
|
||||
state: KeyState::Active,
|
||||
metadata: serde_json::json!({}),
|
||||
},
|
||||
);
|
||||
|
||||
Self {
|
||||
root_key: Arc::new(root_key),
|
||||
root_key_material,
|
||||
tenant_keks: DashMap::new(),
|
||||
agent_deks: DashMap::new(),
|
||||
key_info,
|
||||
audit_log: DashMap::new(),
|
||||
config: KmsConfig {
|
||||
hsm_type: HsmType::Software,
|
||||
root_key_source: RootKeySource::Generate,
|
||||
audit_retention_days: 30,
|
||||
key_rotation_days: Some(90),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn log_audit(&self, operation: KeyOperation, key_id: Option<KeyId>, success: bool, details: serde_json::Value) {
|
||||
let entry = KeyAuditEntry {
|
||||
id: Uuid::new_v4(),
|
||||
timestamp: Utc::now(),
|
||||
operation,
|
||||
actor: "system".to_string(),
|
||||
key_id,
|
||||
success,
|
||||
details,
|
||||
};
|
||||
|
||||
self.audit_log.insert(entry.id, entry);
|
||||
}
|
||||
|
||||
fn wrap_with_aes_kwp(&self, key_data: &[u8], kek: &[u8]) -> Result<Vec<u8>> {
|
||||
// Simplified wrapping - in production use proper AES-KWP
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, KeyInit},
|
||||
ChaCha20Poly1305, Nonce,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(kek)
|
||||
.map_err(|e| SecurityError::CryptoError(e.to_string()))?;
|
||||
|
||||
let nonce = Nonce::from_slice(b"unique nonce"); // In production, use random nonce
|
||||
|
||||
cipher
|
||||
.encrypt(nonce, key_data)
|
||||
.map_err(|e| SecurityError::CryptoError(e.to_string()))
|
||||
}
|
||||
|
||||
fn unwrap_with_aes_kwp(&self, wrapped: &[u8], kek: &[u8]) -> Result<Vec<u8>> {
|
||||
// Simplified unwrapping - in production use proper AES-KWP
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, KeyInit},
|
||||
ChaCha20Poly1305, Nonce,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(kek)
|
||||
.map_err(|e| SecurityError::CryptoError(e.to_string()))?;
|
||||
|
||||
let nonce = Nonce::from_slice(b"unique nonce");
|
||||
|
||||
cipher
|
||||
.decrypt(nonce, wrapped)
|
||||
.map_err(|e| SecurityError::CryptoError(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl KeyManagementService for MemoryKms {
|
||||
async fn initialize(&mut self, config: KmsConfig) -> Result<()> {
|
||||
self.config = config;
|
||||
self.log_audit(
|
||||
KeyOperation::Create,
|
||||
Some(self.root_key.id.clone()),
|
||||
true,
|
||||
serde_json::json!({ "type": "initialization" }),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_tenant_kek(&self, tenant_id: &str) -> Result<TenantKek> {
|
||||
// Derive KEK from root key
|
||||
let kek_material = crate::kms::KeyHierarchy::derive_tenant_kek(&self.root_key_material, tenant_id);
|
||||
|
||||
// Wrap KEK under root key
|
||||
let wrapped_kek = self.wrap_with_aes_kwp(&kek_material, &self.root_key_material)?;
|
||||
|
||||
let kek_id = KeyId::new();
|
||||
let kek = TenantKek {
|
||||
id: kek_id.clone(),
|
||||
tenant_id: tenant_id.to_string(),
|
||||
wrapped_key: wrapped_kek,
|
||||
created_at: Utc::now(),
|
||||
expires_at: None,
|
||||
};
|
||||
|
||||
// Store metadata
|
||||
self.tenant_keks.insert(tenant_id.to_string(), (kek.clone(), kek_material));
|
||||
|
||||
self.key_info.insert(
|
||||
kek_id.clone(),
|
||||
KeyInfo {
|
||||
key_id: kek_id.clone(),
|
||||
key_type: KeyType::TenantKek { tenant_id: tenant_id.to_string() },
|
||||
created_at: kek.created_at,
|
||||
updated_at: kek.created_at,
|
||||
expires_at: kek.expires_at,
|
||||
rotation_due: self.config.key_rotation_days
|
||||
.map(|days| kek.created_at + chrono::Duration::days(days as i64)),
|
||||
state: KeyState::Active,
|
||||
metadata: serde_json::json!({ "tenant_id": tenant_id }),
|
||||
},
|
||||
);
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Create,
|
||||
Some(kek_id),
|
||||
true,
|
||||
serde_json::json!({ "type": "tenant_kek", "tenant_id": tenant_id }),
|
||||
);
|
||||
|
||||
Ok(kek)
|
||||
}
|
||||
|
||||
async fn create_agent_dek(&self, agent_id: &str, tenant_id: &str) -> Result<AgentDek> {
|
||||
// Get tenant KEK
|
||||
let (_, tenant_kek_material) = self.tenant_keks.get(tenant_id)
|
||||
.ok_or_else(|| SecurityError::KeyNotFound(format!("Tenant KEK not found: {}", tenant_id)))?
|
||||
.clone();
|
||||
|
||||
// Generate new DEK
|
||||
let dek_material = crate::kms::KeyHierarchy::generate_dek();
|
||||
|
||||
// Wrap DEK under tenant KEK
|
||||
let wrapped_dek = self.wrap_with_aes_kwp(&dek_material, &tenant_kek_material)?;
|
||||
|
||||
let dek_id = KeyId::new();
|
||||
let dek = AgentDek {
|
||||
id: dek_id.clone(),
|
||||
agent_id: agent_id.to_string(),
|
||||
tenant_id: tenant_id.to_string(),
|
||||
wrapped_key: wrapped_dek,
|
||||
created_at: Utc::now(),
|
||||
rotation_due: self.config.key_rotation_days
|
||||
.map(|days| Utc::now() + chrono::Duration::days(days as i64)),
|
||||
};
|
||||
|
||||
// Store metadata
|
||||
self.agent_deks.insert(agent_id.to_string(), (dek.clone(), dek_material));
|
||||
|
||||
self.key_info.insert(
|
||||
dek_id.clone(),
|
||||
KeyInfo {
|
||||
key_id: dek_id.clone(),
|
||||
key_type: KeyType::AgentDek {
|
||||
agent_id: agent_id.to_string(),
|
||||
tenant_id: tenant_id.to_string(),
|
||||
},
|
||||
created_at: dek.created_at,
|
||||
updated_at: dek.created_at,
|
||||
expires_at: None,
|
||||
rotation_due: dek.rotation_due,
|
||||
state: KeyState::Active,
|
||||
metadata: serde_json::json!({
|
||||
"agent_id": agent_id,
|
||||
"tenant_id": tenant_id,
|
||||
}),
|
||||
},
|
||||
);
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Create,
|
||||
Some(dek_id),
|
||||
true,
|
||||
serde_json::json!({
|
||||
"type": "agent_dek",
|
||||
"agent_id": agent_id,
|
||||
"tenant_id": tenant_id,
|
||||
}),
|
||||
);
|
||||
|
||||
Ok(dek)
|
||||
}
|
||||
|
||||
async fn wrap_key(&self, key_data: &[u8], parent_key_id: &KeyId) -> Result<Vec<u8>> {
|
||||
// Find parent key material
|
||||
let parent_key_material = if *parent_key_id == self.root_key.id {
|
||||
self.root_key_material.clone()
|
||||
} else {
|
||||
return Err(SecurityError::KeyNotFound(format!("Parent key not found: {:?}", parent_key_id)));
|
||||
};
|
||||
|
||||
let wrapped = self.wrap_with_aes_kwp(key_data, &parent_key_material)?;
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Wrap,
|
||||
Some(parent_key_id.clone()),
|
||||
true,
|
||||
serde_json::json!({ "wrapped_size": wrapped.len() }),
|
||||
);
|
||||
|
||||
Ok(wrapped)
|
||||
}
|
||||
|
||||
async fn unwrap_key(&self, wrapped_key: &[u8], parent_key_id: &KeyId) -> Result<Vec<u8>> {
|
||||
// Find parent key material
|
||||
let parent_key_material = if *parent_key_id == self.root_key.id {
|
||||
self.root_key_material.clone()
|
||||
} else {
|
||||
return Err(SecurityError::KeyNotFound(format!("Parent key not found: {:?}", parent_key_id)));
|
||||
};
|
||||
|
||||
let unwrapped = self.unwrap_with_aes_kwp(wrapped_key, &parent_key_material)?;
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Unwrap,
|
||||
Some(parent_key_id.clone()),
|
||||
true,
|
||||
serde_json::json!({ "unwrapped_size": unwrapped.len() }),
|
||||
);
|
||||
|
||||
Ok(unwrapped)
|
||||
}
|
||||
|
||||
async fn rotate_key(&self, key_id: &KeyId) -> Result<KeyId> {
|
||||
// TODO: Implement key rotation
|
||||
self.log_audit(
|
||||
KeyOperation::Rotate,
|
||||
Some(key_id.clone()),
|
||||
false,
|
||||
serde_json::json!({ "error": "not implemented" }),
|
||||
);
|
||||
|
||||
Err(SecurityError::Other(anyhow::anyhow!("Key rotation not implemented")))
|
||||
}
|
||||
|
||||
async fn destroy_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
if let Some(mut info) = self.key_info.get_mut(key_id) {
|
||||
info.state = KeyState::Destroyed;
|
||||
info.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Destroy,
|
||||
Some(key_id.clone()),
|
||||
true,
|
||||
serde_json::json!({}),
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_key_metadata(&self, key_id: &KeyId) -> Result<KeyInfo> {
|
||||
self.key_info.get(key_id)
|
||||
.map(|info| info.value().clone())
|
||||
.ok_or_else(|| SecurityError::KeyNotFound(format!("Key not found: {:?}", key_id)))
|
||||
}
|
||||
|
||||
async fn get_audit_logs(
|
||||
&self,
|
||||
start: DateTime<Utc>,
|
||||
end: DateTime<Utc>,
|
||||
filter: Option<AuditFilter>,
|
||||
) -> Result<Vec<KeyAuditEntry>> {
|
||||
let mut logs: Vec<KeyAuditEntry> = self.audit_log.iter()
|
||||
.filter(|entry| entry.timestamp >= start && entry.timestamp <= end)
|
||||
.map(|entry| entry.value().clone())
|
||||
.collect();
|
||||
|
||||
if let Some(f) = filter {
|
||||
logs.retain(|log| {
|
||||
let mut keep = true;
|
||||
|
||||
if let Some(ref key_ids) = f.key_ids {
|
||||
keep &= log.key_id.as_ref().map(|id| key_ids.contains(id)).unwrap_or(false);
|
||||
}
|
||||
|
||||
if let Some(ref ops) = f.operations {
|
||||
keep &= ops.iter().any(|op| std::mem::discriminant(op) == std::mem::discriminant(&log.operation));
|
||||
}
|
||||
|
||||
if let Some(ref actors) = f.actors {
|
||||
keep &= actors.contains(&log.actor);
|
||||
}
|
||||
|
||||
if let Some(success_only) = f.success_only {
|
||||
keep &= log.success == success_only;
|
||||
}
|
||||
|
||||
keep
|
||||
});
|
||||
}
|
||||
|
||||
logs.sort_by_key(|log| log.timestamp);
|
||||
Ok(logs)
|
||||
}
|
||||
|
||||
async fn import_customer_key(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
wrapped_key: &[u8],
|
||||
key_metadata: CustomerKeyMetadata,
|
||||
) -> Result<KeyId> {
|
||||
// TODO: Implement customer key import
|
||||
let key_id = KeyId::new();
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Create,
|
||||
Some(key_id.clone()),
|
||||
false,
|
||||
serde_json::json!({
|
||||
"type": "customer_key_import",
|
||||
"tenant_id": tenant_id,
|
||||
"error": "not implemented",
|
||||
}),
|
||||
);
|
||||
|
||||
Err(SecurityError::Other(anyhow::anyhow!("Customer key import not implemented")))
|
||||
}
|
||||
|
||||
async fn register_customer_held_key(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
key_reference: CustomerKeyReference,
|
||||
) -> Result<KeyId> {
|
||||
// TODO: Implement HYOK registration
|
||||
let key_id = KeyId::new();
|
||||
|
||||
self.log_audit(
|
||||
KeyOperation::Create,
|
||||
Some(key_id.clone()),
|
||||
false,
|
||||
serde_json::json!({
|
||||
"type": "hyok_registration",
|
||||
"tenant_id": tenant_id,
|
||||
"provider": key_reference.provider,
|
||||
"error": "not implemented",
|
||||
}),
|
||||
);
|
||||
|
||||
Err(SecurityError::Other(anyhow::anyhow!("HYOK registration not implemented")))
|
||||
}
|
||||
}
|
||||
@@ -1,29 +0,0 @@
|
||||
//! Hanzo KBS (Key Broker Service) Library
|
||||
//!
|
||||
//! Provides Key Management Service (KMS) and Key Broker Service (KBS) functionality
|
||||
//! for confidential computing and privacy-preserving agent execution in Hanzo nodes.
|
||||
//!
|
||||
//! This crate implements the KMS/KBS split architecture where:
|
||||
//! - KMS handles key lifecycle management and storage
|
||||
//! - KBS handles attestation verification and policy-based key release
|
||||
|
||||
pub mod error;
|
||||
pub mod types;
|
||||
pub mod kms;
|
||||
pub mod kbs;
|
||||
pub mod attestation;
|
||||
pub mod vault;
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
pub mod pqc_integration;
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
pub mod pqc_vault;
|
||||
|
||||
pub use error::{SecurityError, Result};
|
||||
pub use types::*;
|
||||
pub use kms::KeyManagementService;
|
||||
pub use kbs::KeyBrokerService;
|
||||
|
||||
// Re-export submodules from kms
|
||||
pub use kms::{memory_kms, api};
|
||||
@@ -1,287 +0,0 @@
|
||||
//! Post-Quantum Cryptography integration for KBS
|
||||
//!
|
||||
//! This module provides PQC support for the KBS when the "pqc" feature is enabled.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use hanzo_pqc::{
|
||||
kem::{Kem, KemAlgorithm, MlKem, EncapsulationKey, DecapsulationKey},
|
||||
signature::{Signature, SignatureAlgorithm, MlDsa, VerifyingKey, SigningKey, DigitalSignature},
|
||||
hybrid::{HybridMode, HybridKem, HybridEncapsulationKey},
|
||||
privacy_tiers::PrivacyTier as PqcPrivacyTier,
|
||||
config::PqcConfig,
|
||||
};
|
||||
use crate::{
|
||||
error::{Result, SecurityError},
|
||||
types::PrivacyTier,
|
||||
};
|
||||
|
||||
/// PQC Key Broker Service integration
|
||||
pub struct PqcKbs {
|
||||
ml_kem: MlKem,
|
||||
ml_dsa: MlDsa,
|
||||
config: PqcConfig,
|
||||
}
|
||||
|
||||
impl PqcKbs {
|
||||
pub fn new(tier: PrivacyTier) -> Self {
|
||||
let pqc_tier = match tier {
|
||||
PrivacyTier::Open => PqcPrivacyTier::AccessOpen,
|
||||
PrivacyTier::AtRest => PqcPrivacyTier::AccessAtRest,
|
||||
PrivacyTier::CpuTee => PqcPrivacyTier::AccessCpuTee,
|
||||
PrivacyTier::GpuCc => PqcPrivacyTier::AccessCpuTeePlusGpuCc,
|
||||
PrivacyTier::GpuTeeIo => PqcPrivacyTier::AccessGpuTeeIoMax,
|
||||
};
|
||||
|
||||
let config = PqcConfig::for_privacy_tier(pqc_tier);
|
||||
|
||||
Self {
|
||||
ml_kem: MlKem::new(),
|
||||
ml_dsa: MlDsa::new(),
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate ML-KEM keypair for DEK wrapping
|
||||
pub async fn generate_kem_keypair(&self) -> Result<(EncapsulationKey, DecapsulationKey)> {
|
||||
let alg = match self.config.default_kem {
|
||||
hanzo_pqc::config::DefaultKem::MlKem512 => KemAlgorithm::MlKem512,
|
||||
hanzo_pqc::config::DefaultKem::MlKem768 => KemAlgorithm::MlKem768,
|
||||
hanzo_pqc::config::DefaultKem::MlKem1024 => KemAlgorithm::MlKem1024,
|
||||
};
|
||||
|
||||
let keypair = self.ml_kem.generate_keypair(alg).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("KEM keypair generation failed: {:?}", e)))?;
|
||||
|
||||
Ok((keypair.encap_key, keypair.decap_key))
|
||||
}
|
||||
|
||||
/// Generate ML-DSA keypair for attestation signing
|
||||
pub async fn generate_signing_keypair(&self) -> Result<(VerifyingKey, SigningKey)> {
|
||||
let alg = match self.config.default_sig {
|
||||
hanzo_pqc::config::DefaultSig::MlDsa44 => SignatureAlgorithm::MlDsa44,
|
||||
hanzo_pqc::config::DefaultSig::MlDsa65 => SignatureAlgorithm::MlDsa65,
|
||||
hanzo_pqc::config::DefaultSig::MlDsa87 => SignatureAlgorithm::MlDsa87,
|
||||
};
|
||||
|
||||
self.ml_dsa.generate_keypair(alg).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Signing keypair generation failed: {:?}", e)))
|
||||
}
|
||||
|
||||
/// Wrap a DEK with ML-KEM
|
||||
pub async fn wrap_dek_with_mlkem(
|
||||
&self,
|
||||
dek: &[u8],
|
||||
encap_key: &EncapsulationKey,
|
||||
) -> Result<Vec<u8>> {
|
||||
// Encapsulate to get shared secret
|
||||
let output = self.ml_kem.encapsulate(encap_key).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Encapsulation failed: {:?}", e)))?;
|
||||
|
||||
// Use shared secret to derive wrapping key
|
||||
use hanzo_pqc::kdf::{HkdfKdf, Kdf};
|
||||
let kdf = HkdfKdf::new(self.config.kdf);
|
||||
let wrapping_key = kdf.derive(
|
||||
None,
|
||||
&output.shared_secret,
|
||||
b"hanzo-kbs-dek-wrap-v1",
|
||||
32,
|
||||
).map_err(|e| SecurityError::CryptoError(format!("KDF failed: {:?}", e)))?;
|
||||
|
||||
// Encrypt DEK with ChaCha20Poly1305
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, AeadCore, KeyInit, OsRng},
|
||||
ChaCha20Poly1305,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(&wrapping_key)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Cipher init failed: {:?}", e)))?;
|
||||
|
||||
let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
|
||||
let ciphertext = cipher.encrypt(&nonce, dek)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("DEK encryption failed: {:?}", e)))?;
|
||||
|
||||
// Return: ciphertext_len || kem_ciphertext || nonce || encrypted_dek
|
||||
let mut wrapped = Vec::new();
|
||||
wrapped.extend_from_slice(&(output.ciphertext.len() as u32).to_be_bytes());
|
||||
wrapped.extend_from_slice(&output.ciphertext);
|
||||
wrapped.extend_from_slice(&nonce);
|
||||
wrapped.extend_from_slice(&ciphertext);
|
||||
|
||||
Ok(wrapped)
|
||||
}
|
||||
|
||||
/// Unwrap a DEK with ML-KEM
|
||||
pub async fn unwrap_dek_with_mlkem(
|
||||
&self,
|
||||
wrapped_dek: &[u8],
|
||||
decap_key: &DecapsulationKey,
|
||||
) -> Result<Vec<u8>> {
|
||||
// Parse wrapped format
|
||||
if wrapped_dek.len() < 4 {
|
||||
return Err(SecurityError::CryptoError("Invalid wrapped DEK format".into()));
|
||||
}
|
||||
|
||||
let ct_len = u32::from_be_bytes([
|
||||
wrapped_dek[0], wrapped_dek[1], wrapped_dek[2], wrapped_dek[3]
|
||||
]) as usize;
|
||||
|
||||
if wrapped_dek.len() < 4 + ct_len + 12 {
|
||||
return Err(SecurityError::CryptoError("Invalid wrapped DEK length".into()));
|
||||
}
|
||||
|
||||
let kem_ct = &wrapped_dek[4..4 + ct_len];
|
||||
let nonce = &wrapped_dek[4 + ct_len..4 + ct_len + 12];
|
||||
let encrypted_dek = &wrapped_dek[4 + ct_len + 12..];
|
||||
|
||||
// Decapsulate to recover shared secret
|
||||
let shared_secret = self.ml_kem.decapsulate(decap_key, kem_ct).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Decapsulation failed: {:?}", e)))?;
|
||||
|
||||
// Derive wrapping key
|
||||
use hanzo_pqc::kdf::{HkdfKdf, Kdf};
|
||||
let kdf = HkdfKdf::new(self.config.kdf);
|
||||
let wrapping_key = kdf.derive(
|
||||
None,
|
||||
&shared_secret,
|
||||
b"hanzo-kbs-dek-wrap-v1",
|
||||
32,
|
||||
).map_err(|e| SecurityError::CryptoError(format!("KDF failed: {:?}", e)))?;
|
||||
|
||||
// Decrypt DEK
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, KeyInit},
|
||||
ChaCha20Poly1305, Nonce,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(&wrapping_key)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Cipher init failed: {:?}", e)))?;
|
||||
|
||||
let nonce = Nonce::from_slice(nonce);
|
||||
let dek = cipher.decrypt(nonce, encrypted_dek)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("DEK decryption failed: {:?}", e)))?;
|
||||
|
||||
Ok(dek)
|
||||
}
|
||||
|
||||
/// Sign attestation report with ML-DSA
|
||||
pub async fn sign_attestation(
|
||||
&self,
|
||||
report: &[u8],
|
||||
signing_key: &SigningKey,
|
||||
) -> Result<DigitalSignature> {
|
||||
self.ml_dsa.sign(signing_key, report).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Attestation signing failed: {:?}", e)))
|
||||
}
|
||||
|
||||
/// Verify attestation signature with ML-DSA
|
||||
pub async fn verify_attestation(
|
||||
&self,
|
||||
report: &[u8],
|
||||
signature: &DigitalSignature,
|
||||
verifying_key: &VerifyingKey,
|
||||
) -> Result<bool> {
|
||||
self.ml_dsa.verify(verifying_key, report, signature).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Attestation verification failed: {:?}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
/// PQC handshake for establishing secure channels
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PqcHandshake {
|
||||
pub version: u8,
|
||||
pub mode: HandshakeMode,
|
||||
pub encap_key: Option<EncapsulationKey>,
|
||||
pub hybrid_encap_key: Option<HybridEncapsulationKey>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
||||
pub enum HandshakeMode {
|
||||
MlKemOnly(KemAlgorithm),
|
||||
Hybrid(HybridMode),
|
||||
}
|
||||
|
||||
/// Keys established after PQC handshake
|
||||
pub struct HandshakeKeys {
|
||||
pub client_write_key: [u8; 32],
|
||||
pub server_write_key: [u8; 32],
|
||||
pub client_write_iv: [u8; 12],
|
||||
pub server_write_iv: [u8; 12],
|
||||
}
|
||||
|
||||
impl PqcKbs {
|
||||
/// Create handshake initiation message
|
||||
pub async fn initiate_handshake(&self, hybrid: bool) -> Result<(PqcHandshake, DecapsulationKey)> {
|
||||
if hybrid && self.config.hybrid {
|
||||
// Use hybrid mode
|
||||
let hybrid_kem = HybridKem::new(HybridMode::default());
|
||||
let (encap_key, decap_key) = hybrid_kem.generate_keypair(HybridMode::default()).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Hybrid keypair generation failed: {:?}", e)))?;
|
||||
|
||||
Ok((
|
||||
PqcHandshake {
|
||||
version: 1,
|
||||
mode: HandshakeMode::Hybrid(HybridMode::default()),
|
||||
encap_key: None,
|
||||
hybrid_encap_key: Some(encap_key),
|
||||
},
|
||||
// We need to store the hybrid decap key - for now return a dummy
|
||||
DecapsulationKey {
|
||||
algorithm: KemAlgorithm::MlKem768,
|
||||
key_bytes: vec![],
|
||||
}
|
||||
))
|
||||
} else {
|
||||
// Use ML-KEM only
|
||||
let (encap_key, decap_key) = self.generate_kem_keypair().await?;
|
||||
|
||||
Ok((
|
||||
PqcHandshake {
|
||||
version: 1,
|
||||
mode: HandshakeMode::MlKemOnly(encap_key.algorithm),
|
||||
encap_key: Some(encap_key),
|
||||
hybrid_encap_key: None,
|
||||
},
|
||||
decap_key
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dek_wrapping() {
|
||||
// Skip test in CI environment
|
||||
if std::env::var("CI").is_ok() {
|
||||
println!("Skipping test in CI: requires specific crypto setup");
|
||||
return;
|
||||
}
|
||||
let kbs = PqcKbs::new(PrivacyTier::CpuTee);
|
||||
let (encap_key, decap_key) = kbs.generate_kem_keypair().await.unwrap();
|
||||
|
||||
let dek = b"this is a test data encryption key!";
|
||||
let wrapped = kbs.wrap_dek_with_mlkem(dek, &encap_key).await.unwrap();
|
||||
let unwrapped = kbs.unwrap_dek_with_mlkem(&wrapped, &decap_key).await.unwrap();
|
||||
|
||||
assert_eq!(dek.to_vec(), unwrapped);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_attestation_signing() {
|
||||
// Skip test in CI environment
|
||||
if std::env::var("CI").is_ok() {
|
||||
println!("Skipping test in CI: requires specific crypto setup");
|
||||
return;
|
||||
}
|
||||
let kbs = PqcKbs::new(PrivacyTier::CpuTee);
|
||||
let (verifying_key, signing_key) = kbs.generate_signing_keypair().await.unwrap();
|
||||
|
||||
let report = b"attestation report data";
|
||||
let signature = kbs.sign_attestation(report, &signing_key).await.unwrap();
|
||||
let valid = kbs.verify_attestation(report, &signature, &verifying_key).await.unwrap();
|
||||
|
||||
assert!(valid);
|
||||
}
|
||||
}
|
||||
@@ -1,418 +0,0 @@
|
||||
//! PQC-enhanced vault implementations
|
||||
//!
|
||||
//! Provides quantum-resistant key storage and operations for different privacy tiers
|
||||
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::{
|
||||
error::{Result, SecurityError},
|
||||
types::{PrivacyTier, KeyId, AttestationType},
|
||||
vault::KeyVault,
|
||||
};
|
||||
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
use hanzo_pqc::{
|
||||
kem::{Kem, KemAlgorithm, EncapsulationKey, DecapsulationKey, MlKem},
|
||||
signature::{Signature, SignatureAlgorithm, VerifyingKey, SigningKey, MlDsa},
|
||||
config::PqcConfig,
|
||||
};
|
||||
|
||||
/// PQC-enhanced vault for Tier 2+ (CPU TEE and above)
|
||||
#[cfg(feature = "pqc")]
|
||||
pub struct PqcVault {
|
||||
tier: PrivacyTier,
|
||||
config: PqcConfig,
|
||||
ml_kem: Arc<MlKem>,
|
||||
ml_dsa: Arc<MlDsa>,
|
||||
// In-memory storage for dev/testing (production would use TEE sealed storage)
|
||||
keys: Arc<RwLock<HashMap<KeyId, ProtectedKey>>>,
|
||||
attestation: Option<AttestationType>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
#[derive(Clone)]
|
||||
struct ProtectedKey {
|
||||
/// ML-KEM encrypted key material
|
||||
wrapped_key: Vec<u8>,
|
||||
/// Decapsulation key (stored in TEE in production)
|
||||
decap_key: DecapsulationKey,
|
||||
/// Metadata
|
||||
created_at: chrono::DateTime<chrono::Utc>,
|
||||
tier: PrivacyTier,
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
impl PqcVault {
|
||||
pub async fn new(tier: PrivacyTier, attestation: Option<AttestationType>) -> Result<Self> {
|
||||
if tier < PrivacyTier::CpuTee && attestation.is_some() {
|
||||
return Err(SecurityError::PolicyViolation(
|
||||
"Attestation requires CPU TEE or higher tier".into()
|
||||
));
|
||||
}
|
||||
|
||||
let pqc_tier = match tier {
|
||||
PrivacyTier::Open => hanzo_pqc::privacy_tiers::PrivacyTier::AccessOpen,
|
||||
PrivacyTier::AtRest => hanzo_pqc::privacy_tiers::PrivacyTier::AccessAtRest,
|
||||
PrivacyTier::CpuTee => hanzo_pqc::privacy_tiers::PrivacyTier::AccessCpuTee,
|
||||
PrivacyTier::GpuCc => hanzo_pqc::privacy_tiers::PrivacyTier::AccessCpuTeePlusGpuCc,
|
||||
PrivacyTier::GpuTeeIo => hanzo_pqc::privacy_tiers::PrivacyTier::AccessGpuTeeIoMax,
|
||||
};
|
||||
|
||||
let config = PqcConfig::for_privacy_tier(pqc_tier);
|
||||
|
||||
Ok(Self {
|
||||
tier,
|
||||
config,
|
||||
ml_kem: Arc::new(MlKem::new()),
|
||||
ml_dsa: Arc::new(MlDsa::new()),
|
||||
keys: Arc::new(RwLock::new(HashMap::new())),
|
||||
attestation,
|
||||
})
|
||||
}
|
||||
|
||||
/// Generate ML-KEM keypair for this vault
|
||||
async fn generate_kem_keypair(&self) -> Result<(EncapsulationKey, DecapsulationKey)> {
|
||||
let alg = self.tier.recommended_kem();
|
||||
let keypair = self.ml_kem.generate_keypair(alg).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("KEM keypair generation failed: {:?}", e)))?;
|
||||
Ok((keypair.encap_key, keypair.decap_key))
|
||||
}
|
||||
|
||||
/// Generate ML-DSA keypair for signing
|
||||
async fn generate_signing_keypair(&self) -> Result<(VerifyingKey, SigningKey)> {
|
||||
let alg = self.tier.recommended_sig();
|
||||
self.ml_dsa.generate_keypair(alg).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Signing keypair generation failed: {:?}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
// Extension methods for PrivacyTier to get recommended algorithms
|
||||
impl PrivacyTier {
|
||||
fn recommended_kem(&self) -> KemAlgorithm {
|
||||
match self {
|
||||
PrivacyTier::Open | PrivacyTier::AtRest => KemAlgorithm::MlKem512,
|
||||
PrivacyTier::CpuTee => KemAlgorithm::MlKem768,
|
||||
PrivacyTier::GpuCc | PrivacyTier::GpuTeeIo => KemAlgorithm::MlKem1024,
|
||||
}
|
||||
}
|
||||
|
||||
fn recommended_sig(&self) -> SignatureAlgorithm {
|
||||
match self {
|
||||
PrivacyTier::Open | PrivacyTier::AtRest => SignatureAlgorithm::MlDsa44,
|
||||
PrivacyTier::CpuTee => SignatureAlgorithm::MlDsa65,
|
||||
PrivacyTier::GpuCc | PrivacyTier::GpuTeeIo => SignatureAlgorithm::MlDsa87,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
#[async_trait]
|
||||
impl KeyVault for PqcVault {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
self.tier
|
||||
}
|
||||
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()> {
|
||||
// Generate vault-specific KEM keypair
|
||||
let (encap_key, decap_key) = self.generate_kem_keypair().await?;
|
||||
|
||||
// Encapsulate key data
|
||||
let output = self.ml_kem.encapsulate(&encap_key).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Key encapsulation failed: {:?}", e)))?;
|
||||
|
||||
// Derive wrapping key from shared secret
|
||||
use hanzo_pqc::kdf::{HkdfKdf, Kdf};
|
||||
let kdf = HkdfKdf::new(self.config.kdf);
|
||||
let wrapping_key = kdf.derive(
|
||||
None,
|
||||
&output.shared_secret,
|
||||
b"hanzo-vault-wrap-v1",
|
||||
32,
|
||||
).map_err(|e| SecurityError::CryptoError(format!("KDF failed: {:?}", e)))?;
|
||||
|
||||
// Encrypt key data with ChaCha20Poly1305
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, AeadCore, KeyInit, OsRng},
|
||||
ChaCha20Poly1305,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(&wrapping_key)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Cipher init failed: {:?}", e)))?;
|
||||
|
||||
let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
|
||||
let ciphertext = cipher.encrypt(&nonce, key_data)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Key encryption failed: {:?}", e)))?;
|
||||
|
||||
// Store: kem_ct || nonce || ciphertext
|
||||
let mut wrapped = Vec::new();
|
||||
wrapped.extend_from_slice(&(output.ciphertext.len() as u32).to_be_bytes());
|
||||
wrapped.extend_from_slice(&output.ciphertext);
|
||||
wrapped.extend_from_slice(&nonce);
|
||||
wrapped.extend_from_slice(&ciphertext);
|
||||
|
||||
// Store protected key
|
||||
let protected = ProtectedKey {
|
||||
wrapped_key: wrapped,
|
||||
decap_key,
|
||||
created_at: chrono::Utc::now(),
|
||||
tier: self.tier,
|
||||
};
|
||||
|
||||
self.keys.write().await.insert(key_id.clone(), protected);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
let keys = self.keys.read().await;
|
||||
let protected = keys.get(key_id)
|
||||
.ok_or_else(|| SecurityError::KeyNotFound(key_id.0.to_string()))?;
|
||||
|
||||
// Parse wrapped format
|
||||
let wrapped = &protected.wrapped_key;
|
||||
if wrapped.len() < 4 {
|
||||
return Err(SecurityError::CryptoError("Invalid wrapped key format".into()));
|
||||
}
|
||||
|
||||
let ct_len = u32::from_be_bytes([wrapped[0], wrapped[1], wrapped[2], wrapped[3]]) as usize;
|
||||
if wrapped.len() < 4 + ct_len + 12 {
|
||||
return Err(SecurityError::CryptoError("Invalid wrapped key length".into()));
|
||||
}
|
||||
|
||||
let kem_ct = &wrapped[4..4 + ct_len];
|
||||
let nonce = &wrapped[4 + ct_len..4 + ct_len + 12];
|
||||
let ciphertext = &wrapped[4 + ct_len + 12..];
|
||||
|
||||
// Decapsulate to recover shared secret
|
||||
let shared_secret = self.ml_kem.decapsulate(&protected.decap_key, kem_ct).await
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Decapsulation failed: {:?}", e)))?;
|
||||
|
||||
// Derive wrapping key
|
||||
use hanzo_pqc::kdf::{HkdfKdf, Kdf};
|
||||
let kdf = HkdfKdf::new(self.config.kdf);
|
||||
let wrapping_key = kdf.derive(
|
||||
None,
|
||||
&shared_secret,
|
||||
b"hanzo-vault-wrap-v1",
|
||||
32,
|
||||
).map_err(|e| SecurityError::CryptoError(format!("KDF failed: {:?}", e)))?;
|
||||
|
||||
// Decrypt key data
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, KeyInit},
|
||||
ChaCha20Poly1305, Nonce,
|
||||
};
|
||||
|
||||
let cipher = ChaCha20Poly1305::new_from_slice(&wrapping_key)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Cipher init failed: {:?}", e)))?;
|
||||
|
||||
let nonce = Nonce::from_slice(nonce);
|
||||
let key_data = cipher.decrypt(nonce, ciphertext)
|
||||
.map_err(|e| SecurityError::CryptoError(format!("Key decryption failed: {:?}", e)))?;
|
||||
|
||||
// Use key in protected context
|
||||
let result = operation(&key_data);
|
||||
|
||||
// Key data is automatically zeroized when dropped
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
self.keys.write().await.remove(key_id)
|
||||
.ok_or_else(|| SecurityError::KeyNotFound(key_id.0.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
// Check if we have valid attestation for TEE tiers
|
||||
if self.tier >= PrivacyTier::CpuTee {
|
||||
Ok(self.attestation.is_some())
|
||||
} else {
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GPU Confidential Computing vault (Tier 3)
|
||||
#[cfg(feature = "pqc")]
|
||||
pub struct GpuCcVault {
|
||||
inner: PqcVault,
|
||||
gpu_device_id: u32,
|
||||
cc_enabled: bool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
impl GpuCcVault {
|
||||
pub async fn new(gpu_device_id: u32, attestation: AttestationType) -> Result<Self> {
|
||||
// Verify this is H100 CC attestation
|
||||
match &attestation {
|
||||
AttestationType::H100Cc { .. } => {},
|
||||
_ => return Err(SecurityError::PolicyViolation(
|
||||
"GPU CC vault requires H100 CC attestation".into()
|
||||
)),
|
||||
}
|
||||
|
||||
let inner = PqcVault::new(PrivacyTier::GpuCc, Some(attestation)).await?;
|
||||
|
||||
Ok(Self {
|
||||
inner,
|
||||
gpu_device_id,
|
||||
cc_enabled: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
#[async_trait]
|
||||
impl KeyVault for GpuCcVault {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::GpuCc
|
||||
}
|
||||
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()> {
|
||||
// Additional GPU-specific protections could go here
|
||||
// For H100 CC, keys are protected by encrypted DMA
|
||||
self.inner.store_key(key_id, key_data).await
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
if !self.cc_enabled {
|
||||
return Err(SecurityError::PolicyViolation(
|
||||
"GPU CC must be enabled for key operations".into()
|
||||
));
|
||||
}
|
||||
|
||||
self.inner.use_key(key_id, operation).await
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
self.inner.delete_key(key_id).await
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
Ok(self.cc_enabled && self.inner.is_initialized().await?)
|
||||
}
|
||||
}
|
||||
|
||||
/// GPU TEE-I/O vault (Tier 4 - Blackwell)
|
||||
#[cfg(feature = "pqc")]
|
||||
pub struct GpuTeeIoVault {
|
||||
inner: PqcVault,
|
||||
gpu_device_id: u32,
|
||||
tee_io_enabled: bool,
|
||||
mig_instance: Option<u32>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
impl GpuTeeIoVault {
|
||||
pub async fn new(
|
||||
gpu_device_id: u32,
|
||||
mig_instance: Option<u32>,
|
||||
attestation: AttestationType,
|
||||
) -> Result<Self> {
|
||||
// Verify this is Blackwell TEE-I/O attestation
|
||||
match &attestation {
|
||||
AttestationType::BlackwellTeeIo { .. } => {},
|
||||
_ => return Err(SecurityError::PolicyViolation(
|
||||
"GPU TEE-I/O vault requires Blackwell attestation".into()
|
||||
)),
|
||||
}
|
||||
|
||||
let inner = PqcVault::new(PrivacyTier::GpuTeeIo, Some(attestation)).await?;
|
||||
|
||||
Ok(Self {
|
||||
inner,
|
||||
gpu_device_id,
|
||||
tee_io_enabled: true,
|
||||
mig_instance,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "pqc")]
|
||||
#[async_trait]
|
||||
impl KeyVault for GpuTeeIoVault {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::GpuTeeIo
|
||||
}
|
||||
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()> {
|
||||
// Blackwell TEE-I/O provides inline NVLink protection
|
||||
// Keys are protected end-to-end with near-native performance
|
||||
self.inner.store_key(key_id, key_data).await
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
if !self.tee_io_enabled {
|
||||
return Err(SecurityError::PolicyViolation(
|
||||
"TEE-I/O must be enabled for key operations".into()
|
||||
));
|
||||
}
|
||||
|
||||
// MIG isolation check for multi-tenant scenarios
|
||||
if let Some(mig) = self.mig_instance {
|
||||
log::info!("Using key in MIG instance {}", mig);
|
||||
}
|
||||
|
||||
self.inner.use_key(key_id, operation).await
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
self.inner.delete_key(key_id).await
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
Ok(self.tee_io_enabled && self.inner.is_initialized().await?)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(feature = "pqc")]
|
||||
async fn test_pqc_vault_operations() {
|
||||
// Skip test in CI environment
|
||||
if std::env::var("CI").is_ok() {
|
||||
println!("Skipping test in CI: requires specific crypto setup");
|
||||
return;
|
||||
}
|
||||
let vault = PqcVault::new(PrivacyTier::CpuTee, None).await.unwrap();
|
||||
|
||||
let key_id = KeyId::new();
|
||||
let key_data = b"secret key material";
|
||||
|
||||
// Store key
|
||||
vault.store_key(&key_id, key_data).await.unwrap();
|
||||
|
||||
// Use key
|
||||
let result = vault.use_key(&key_id, |data| {
|
||||
assert_eq!(data, key_data);
|
||||
42
|
||||
}).await.unwrap();
|
||||
|
||||
assert_eq!(result, 42);
|
||||
|
||||
// Delete key
|
||||
vault.delete_key(&key_id).await.unwrap();
|
||||
|
||||
// Verify deletion
|
||||
assert!(vault.use_key(&key_id, |_| ()).await.is_err());
|
||||
}
|
||||
}
|
||||
@@ -1,225 +0,0 @@
|
||||
//! Core types for Hanzo Security
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Privacy tiers for agent execution
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
|
||||
#[repr(u8)]
|
||||
pub enum PrivacyTier {
|
||||
/// Tier 0: Open - No privacy guarantees, suitable for public data
|
||||
Open = 0,
|
||||
|
||||
/// Tier 1: At-rest encryption with SIM/FileVault
|
||||
AtRest = 1,
|
||||
|
||||
/// Tier 2: CPU TEE (SEV-SNP, TDX)
|
||||
CpuTee = 2,
|
||||
|
||||
/// Tier 3: CPU TEE + GPU Confidential Computing (H100 CC)
|
||||
GpuCc = 3,
|
||||
|
||||
/// Tier 4: GPU TEE-I/O (Blackwell) - Maximum privacy
|
||||
GpuTeeIo = 4,
|
||||
}
|
||||
|
||||
impl PrivacyTier {
|
||||
pub fn requires_attestation(&self) -> bool {
|
||||
*self >= PrivacyTier::CpuTee
|
||||
}
|
||||
|
||||
pub fn requires_kbs(&self) -> bool {
|
||||
*self >= PrivacyTier::CpuTee
|
||||
}
|
||||
}
|
||||
|
||||
/// Supported attestation types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AttestationType {
|
||||
/// AMD SEV-SNP attestation
|
||||
SevSnp {
|
||||
report: Vec<u8>,
|
||||
vcek_cert: Vec<u8>,
|
||||
platform_cert_chain: Vec<u8>,
|
||||
},
|
||||
|
||||
/// Intel TDX attestation
|
||||
Tdx {
|
||||
quote: Vec<u8>,
|
||||
collateral: Vec<u8>,
|
||||
},
|
||||
|
||||
/// NVIDIA H100 Confidential Computing
|
||||
H100Cc {
|
||||
gpu_attestation: Vec<u8>,
|
||||
cpu_attestation: Box<AttestationType>,
|
||||
},
|
||||
|
||||
/// NVIDIA Blackwell TEE-I/O
|
||||
BlackwellTeeIo {
|
||||
tee_io_report: Vec<u8>,
|
||||
mig_config: Option<MigConfiguration>,
|
||||
},
|
||||
|
||||
/// SIM card attestation
|
||||
SimEid {
|
||||
eid: String,
|
||||
signature: Vec<u8>,
|
||||
},
|
||||
}
|
||||
|
||||
/// MIG (Multi-Instance GPU) configuration for Blackwell
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MigConfiguration {
|
||||
pub instance_id: u32,
|
||||
pub memory_size_mb: u64,
|
||||
pub compute_units: u32,
|
||||
}
|
||||
|
||||
/// Node security mode configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum NodeSecurityMode {
|
||||
/// Software only - no hardware security
|
||||
SoftwareOnly,
|
||||
|
||||
/// SIM card security only
|
||||
SimOnly,
|
||||
|
||||
/// SIM + TEE security
|
||||
SimTee,
|
||||
}
|
||||
|
||||
/// Key types in the hierarchy
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub struct KeyId(pub Uuid);
|
||||
|
||||
impl KeyId {
|
||||
pub fn new() -> Self {
|
||||
Self(Uuid::new_v4())
|
||||
}
|
||||
}
|
||||
|
||||
/// Root Key - Top of the hierarchy
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RootKey {
|
||||
pub id: KeyId,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub hsm_handle: Option<String>,
|
||||
}
|
||||
|
||||
/// Key Encryption Key (per tenant)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TenantKek {
|
||||
pub id: KeyId,
|
||||
pub tenant_id: String,
|
||||
pub wrapped_key: Vec<u8>, // Wrapped under root key
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Data Encryption Key (per agent)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentDek {
|
||||
pub id: KeyId,
|
||||
pub agent_id: String,
|
||||
pub tenant_id: String,
|
||||
pub wrapped_key: Vec<u8>, // Wrapped under tenant KEK
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub rotation_due: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Session key for enclave use
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SessionKey {
|
||||
pub id: KeyId,
|
||||
pub agent_id: String,
|
||||
pub hpke_wrapped_key: Vec<u8>,
|
||||
pub enclave_public_key: Vec<u8>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub tier: PrivacyTier,
|
||||
}
|
||||
|
||||
/// Key authorization request from KBS to KMS
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyAuthorizationRequest {
|
||||
pub attestation: AttestationType,
|
||||
pub capability_token: CapabilityToken,
|
||||
pub session_public_key: Vec<u8>, // HPKE public key
|
||||
pub requested_keys: Vec<KeyRequest>,
|
||||
pub nonce: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Individual key request
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyRequest {
|
||||
pub key_type: KeyRequestType,
|
||||
pub agent_id: String,
|
||||
pub tenant_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub enum KeyRequestType {
|
||||
TenantKek,
|
||||
AgentDek,
|
||||
SessionKey { duration_secs: u64 },
|
||||
}
|
||||
|
||||
/// On-chain capability token for policy enforcement
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CapabilityToken {
|
||||
pub id: String,
|
||||
pub subject: String, // Agent or tenant ID
|
||||
pub tier: PrivacyTier,
|
||||
pub permissions: Vec<String>,
|
||||
pub issued_at: DateTime<Utc>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
pub chain_signature: Vec<u8>,
|
||||
}
|
||||
|
||||
/// KBS authorization response
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyAuthorizationResponse {
|
||||
pub session_id: Uuid,
|
||||
pub authorized_keys: Vec<AuthorizedKey>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct AuthorizedKey {
|
||||
pub key_id: KeyId,
|
||||
pub hpke_wrapped_key: Vec<u8>,
|
||||
pub metadata: KeyMetadata,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyMetadata {
|
||||
pub key_type: String,
|
||||
pub tier: PrivacyTier,
|
||||
pub restrictions: Vec<String>,
|
||||
}
|
||||
|
||||
/// Audit log entry for key operations
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KeyAuditEntry {
|
||||
pub id: Uuid,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub operation: KeyOperation,
|
||||
pub actor: String,
|
||||
pub key_id: Option<KeyId>,
|
||||
pub success: bool,
|
||||
pub details: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum KeyOperation {
|
||||
Create,
|
||||
Wrap,
|
||||
Unwrap,
|
||||
Rotate,
|
||||
Destroy,
|
||||
Authorize,
|
||||
Revoke,
|
||||
}
|
||||
@@ -1,499 +0,0 @@
|
||||
//! Vault implementations for different privacy tiers
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::error::{Result, SecurityError};
|
||||
use crate::types::{PrivacyTier, KeyId, SessionKey, AttestationType};
|
||||
use crate::kms::KeyManagementService;
|
||||
use crate::kbs::KeyBrokerService;
|
||||
|
||||
/// Trait for secure key storage and usage
|
||||
#[async_trait]
|
||||
pub trait KeyVault: Send + Sync {
|
||||
/// Get the privacy tier this vault supports
|
||||
fn tier(&self) -> PrivacyTier;
|
||||
|
||||
/// Store a key in the vault
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()>;
|
||||
|
||||
/// Use a key for an operation (key never leaves vault in plaintext)
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send;
|
||||
|
||||
/// Delete a key from the vault
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()>;
|
||||
|
||||
/// Check if vault is properly initialized
|
||||
async fn is_initialized(&self) -> Result<bool>;
|
||||
}
|
||||
|
||||
/// Configuration for different vault types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct VaultConfig {
|
||||
pub vault_type: VaultType,
|
||||
pub kms_url: Option<String>,
|
||||
pub kbs_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum VaultType {
|
||||
File { path: String },
|
||||
Sim { eid: String },
|
||||
GpuCc { device_id: u32 },
|
||||
GpuTeeIo { device_id: u32, mig_instance: Option<u32> },
|
||||
}
|
||||
|
||||
/// Tier 1: File-based vault with at-rest encryption
|
||||
pub struct FileVault<K: KeyManagementService> {
|
||||
kms: K,
|
||||
base_path: std::path::PathBuf,
|
||||
tenant_id: String,
|
||||
}
|
||||
|
||||
impl<K: KeyManagementService> FileVault<K> {
|
||||
pub fn new(kms: K, base_path: impl AsRef<Path>, tenant_id: String) -> Self {
|
||||
Self {
|
||||
kms,
|
||||
base_path: base_path.as_ref().to_path_buf(),
|
||||
tenant_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<K: KeyManagementService> KeyVault for FileVault<K> {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::AtRest
|
||||
}
|
||||
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()> {
|
||||
// Wrap key with tenant KEK before storing
|
||||
let wrapped = self.kms.wrap_key(key_data, key_id).await?;
|
||||
|
||||
let key_path = self.base_path.join(format!("{}.key", key_id.0));
|
||||
tokio::fs::create_dir_all(&self.base_path).await?;
|
||||
tokio::fs::write(&key_path, &wrapped).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
let key_path = self.base_path.join(format!("{}.key", key_id.0));
|
||||
let wrapped = tokio::fs::read(&key_path).await?;
|
||||
|
||||
// Unwrap key for use
|
||||
let key_data = self.kms.unwrap_key(&wrapped, key_id).await?;
|
||||
|
||||
// Use key and zeroize after
|
||||
let result = operation(&key_data);
|
||||
|
||||
// Zeroize key data
|
||||
drop(key_data); // In production, use zeroize crate
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
let key_path = self.base_path.join(format!("{}.key", key_id.0));
|
||||
tokio::fs::remove_file(&key_path).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
Ok(self.base_path.exists())
|
||||
}
|
||||
}
|
||||
|
||||
/// Tier 1+: SIM-based vault with hardware-bound keys
|
||||
pub struct SimVault<K: KeyManagementService> {
|
||||
kms: K,
|
||||
sim_eid: String,
|
||||
file_vault: FileVault<K>, // Fallback storage
|
||||
}
|
||||
|
||||
impl<K: KeyManagementService + Clone> SimVault<K> {
|
||||
pub fn new(kms: K, base_path: impl AsRef<Path>, tenant_id: String, sim_eid: String) -> Self {
|
||||
Self {
|
||||
kms: kms.clone(),
|
||||
sim_eid,
|
||||
file_vault: FileVault::new(kms, base_path, tenant_id),
|
||||
}
|
||||
}
|
||||
|
||||
async fn bind_to_sim(&self, key_data: &[u8]) -> Result<Vec<u8>> {
|
||||
// TODO: Implement SIM binding using EID
|
||||
// For now, just add EID to key derivation
|
||||
let mut hasher = blake3::Hasher::new();
|
||||
hasher.update(key_data);
|
||||
hasher.update(self.sim_eid.as_bytes());
|
||||
Ok(hasher.finalize().as_bytes().to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<K: KeyManagementService + Clone> KeyVault for SimVault<K> {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::AtRest // Still Tier 1, but with SIM binding
|
||||
}
|
||||
|
||||
async fn store_key(&self, key_id: &KeyId, key_data: &[u8]) -> Result<()> {
|
||||
let sim_bound = self.bind_to_sim(key_data).await?;
|
||||
self.file_vault.store_key(key_id, &sim_bound).await
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
self.file_vault.use_key(key_id, operation).await
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &KeyId) -> Result<()> {
|
||||
self.file_vault.delete_key(key_id).await
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
// TODO: Check SIM availability
|
||||
self.file_vault.is_initialized().await
|
||||
}
|
||||
}
|
||||
|
||||
/// Tier 3: GPU Confidential Computing vault
|
||||
pub struct GpuCcVault<K: KeyBrokerService> {
|
||||
kbs: K,
|
||||
device_id: u32,
|
||||
session: Option<SessionKey>,
|
||||
}
|
||||
|
||||
impl<K: KeyBrokerService> GpuCcVault<K> {
|
||||
pub fn new(kbs: K, device_id: u32) -> Self {
|
||||
Self {
|
||||
kbs,
|
||||
device_id,
|
||||
session: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn ensure_session(&mut self, agent_id: &str) -> Result<()> {
|
||||
if let Some(ref session) = self.session {
|
||||
if session.expires_at > chrono::Utc::now() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// Need new session - create attestation
|
||||
let attestation = self.create_gpu_cc_attestation().await?;
|
||||
|
||||
// Request authorization from KBS
|
||||
let auth_request = crate::types::KeyAuthorizationRequest {
|
||||
attestation,
|
||||
capability_token: self.create_capability_token(agent_id)?,
|
||||
session_public_key: self.generate_session_key()?,
|
||||
requested_keys: vec![
|
||||
crate::types::KeyRequest {
|
||||
key_type: crate::types::KeyRequestType::SessionKey { duration_secs: 3600 },
|
||||
agent_id: agent_id.to_string(),
|
||||
tenant_id: None,
|
||||
},
|
||||
],
|
||||
nonce: self.generate_nonce(),
|
||||
};
|
||||
|
||||
let response = self.kbs.authorize(auth_request).await?;
|
||||
|
||||
// Store session
|
||||
if let Some(authorized_key) = response.authorized_keys.first() {
|
||||
self.session = Some(SessionKey {
|
||||
id: authorized_key.key_id.clone(),
|
||||
agent_id: agent_id.to_string(),
|
||||
hpke_wrapped_key: authorized_key.hpke_wrapped_key.clone(),
|
||||
enclave_public_key: self.get_enclave_public_key()?, // Get from GPU TEE
|
||||
created_at: chrono::Utc::now(),
|
||||
expires_at: response.expires_at,
|
||||
tier: PrivacyTier::GpuCc,
|
||||
});
|
||||
Ok(())
|
||||
} else {
|
||||
Err(SecurityError::KbsError("No session key received".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_gpu_cc_attestation(&self) -> Result<AttestationType> {
|
||||
// Get actual GPU attestation from NVIDIA driver
|
||||
let gpu_attestation = self.get_gpu_attestation_report().await?;
|
||||
|
||||
// Also get CPU attestation for full chain of trust
|
||||
let cpu_attestation = self.get_cpu_attestation().await?;
|
||||
|
||||
Ok(AttestationType::H100Cc {
|
||||
gpu_attestation,
|
||||
cpu_attestation: Box::new(cpu_attestation),
|
||||
})
|
||||
}
|
||||
|
||||
fn create_capability_token(&self, agent_id: &str) -> Result<crate::types::CapabilityToken> {
|
||||
// Create capability token with proper chain signature
|
||||
let token_id = uuid::Uuid::new_v4().to_string();
|
||||
let issued_at = chrono::Utc::now();
|
||||
let expires_at = Some(issued_at + chrono::Duration::hours(1));
|
||||
|
||||
// Sign token with chain authority (would connect to blockchain in production)
|
||||
let chain_signature = self.sign_with_chain_authority(&token_id, agent_id)?;
|
||||
|
||||
Ok(crate::types::CapabilityToken {
|
||||
id: token_id,
|
||||
subject: agent_id.to_string(),
|
||||
tier: PrivacyTier::GpuCc,
|
||||
permissions: vec!["gpu_compute".to_string(), "tee_access".to_string()],
|
||||
issued_at,
|
||||
expires_at,
|
||||
chain_signature,
|
||||
})
|
||||
}
|
||||
|
||||
fn generate_session_key(&self) -> Result<Vec<u8>> {
|
||||
// Generate HPKE key pair for session
|
||||
use hpke::aead::ChaCha20Poly1305;
|
||||
use hpke::kdf::HkdfSha256;
|
||||
use hpke::kem::X25519HkdfSha256;
|
||||
use hpke::Kem;
|
||||
use rand::RngCore;
|
||||
|
||||
type HpkeScheme = (X25519HkdfSha256, HkdfSha256, ChaCha20Poly1305);
|
||||
let kem = X25519HkdfSha256::default();
|
||||
|
||||
let mut rng = rand::thread_rng();
|
||||
let (secret_key, _public_key) = kem.gen_keypair(&mut rng);
|
||||
|
||||
// Return the secret key bytes
|
||||
Ok(secret_key.to_bytes().to_vec())
|
||||
}
|
||||
|
||||
fn generate_nonce(&self) -> Vec<u8> {
|
||||
use rand::RngCore;
|
||||
let mut nonce = vec![0u8; 32];
|
||||
rand::thread_rng().fill_bytes(&mut nonce);
|
||||
nonce
|
||||
}
|
||||
|
||||
async fn get_gpu_attestation_report(&self) -> Result<Vec<u8>> {
|
||||
// Get GPU attestation from NVIDIA driver
|
||||
// This would call NVIDIA DCGM or nvml library in production
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
// Check for NVIDIA GPU with CC mode
|
||||
if std::path::Path::new("/dev/nvidia-uvm").exists() {
|
||||
// Read attestation from GPU driver
|
||||
// In production, this would use NVIDIA attestation API
|
||||
let mut attestation = vec![0u8; 512];
|
||||
attestation[0..4].copy_from_slice(&self.device_id.to_le_bytes());
|
||||
attestation[4..8].copy_from_slice(b"H100");
|
||||
use rand::RngCore;
|
||||
rand::thread_rng().fill_bytes(&mut attestation[8..]);
|
||||
return Ok(attestation);
|
||||
}
|
||||
}
|
||||
|
||||
Err(SecurityError::AttestationFailure(
|
||||
"GPU attestation not available".to_string()
|
||||
))
|
||||
}
|
||||
|
||||
async fn get_cpu_attestation(&self) -> Result<AttestationType> {
|
||||
// Get CPU attestation for chain of trust
|
||||
crate::attestation::generate_attestation(TeeType::SevSnp).await
|
||||
}
|
||||
|
||||
fn get_enclave_public_key(&self) -> Result<Vec<u8>> {
|
||||
// Get public key from GPU enclave
|
||||
// This would use NVIDIA GPU TEE API in production
|
||||
use x25519_dalek::{PublicKey, StaticSecret};
|
||||
use rand::RngCore;
|
||||
|
||||
let mut secret_bytes = [0u8; 32];
|
||||
rand::thread_rng().fill_bytes(&mut secret_bytes);
|
||||
let secret = StaticSecret::from(secret_bytes);
|
||||
let public = PublicKey::from(&secret);
|
||||
|
||||
Ok(public.as_bytes().to_vec())
|
||||
}
|
||||
|
||||
fn sign_with_chain_authority(&self, token_id: &str, agent_id: &str) -> Result<Vec<u8>> {
|
||||
// Sign with chain authority (would connect to blockchain in production)
|
||||
use sha2::{Sha256, Digest};
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(token_id.as_bytes());
|
||||
hasher.update(agent_id.as_bytes());
|
||||
hasher.update(&self.device_id.to_le_bytes());
|
||||
|
||||
Ok(hasher.finalize().to_vec())
|
||||
}
|
||||
|
||||
fn unwrap_in_enclave(&self, wrapped_key: &[u8]) -> Result<Vec<u8>> {
|
||||
// Unwrap key using GPU TEE protection
|
||||
// This would use NVIDIA GPU enclave API in production
|
||||
if wrapped_key.len() < 32 {
|
||||
return Err(SecurityError::InvalidKey("Key too short".to_string()));
|
||||
}
|
||||
|
||||
// In production, this would decrypt using GPU TEE
|
||||
// For now, simulate unwrapping
|
||||
Ok(wrapped_key[0..32].to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<K: KeyBrokerService> KeyVault for GpuCcVault<K> {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::GpuCc
|
||||
}
|
||||
|
||||
async fn store_key(&self, _key_id: &KeyId, _key_data: &[u8]) -> Result<()> {
|
||||
// Keys are never stored locally in GPU CC mode
|
||||
Err(SecurityError::PolicyViolation(
|
||||
"GPU CC vault does not support local key storage".to_string()
|
||||
))
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
// Use HPKE-wrapped session key in GPU enclave
|
||||
// H100 uses confidential computing mode with enclave-protected memory
|
||||
let attestation = self.kbs.get_attestation_report().await?;
|
||||
|
||||
// Verify GPU CC mode is active
|
||||
if attestation.tee_type != TeeType::NvidiaH100 {
|
||||
return Err(SecurityError::AttestationFailure(
|
||||
"GPU not in H100 CC mode".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
// Request enclave-protected key from KBS
|
||||
let wrapped_key = self.kbs.request_key(
|
||||
key_id,
|
||||
attestation,
|
||||
PrivacyTier::Tier3GpuCc,
|
||||
).await?;
|
||||
|
||||
// Unwrap using GPU TEE protection
|
||||
let key = self.unwrap_in_enclave(&wrapped_key)?;
|
||||
Ok(operation(&key))
|
||||
}
|
||||
|
||||
async fn delete_key(&self, _key_id: &KeyId) -> Result<()> {
|
||||
// No local storage
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
// Check GPU availability and CC mode
|
||||
match self.kbs.get_attestation_report().await {
|
||||
Ok(report) => Ok(report.tee_type == TeeType::NvidiaH100),
|
||||
Err(_) => Ok(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tier 4: GPU TEE-I/O vault (Blackwell)
|
||||
pub struct GpuTeeIoVault<K: KeyBrokerService> {
|
||||
kbs: K,
|
||||
device_id: u32,
|
||||
mig_instance: Option<u32>,
|
||||
session: Option<SessionKey>,
|
||||
}
|
||||
|
||||
impl<K: KeyBrokerService> GpuTeeIoVault<K> {
|
||||
pub fn new(kbs: K, device_id: u32, mig_instance: Option<u32>) -> Self {
|
||||
Self {
|
||||
kbs,
|
||||
device_id,
|
||||
mig_instance,
|
||||
session: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn unwrap_in_tee_io_enclave(&self, wrapped_key: &[u8]) -> Result<Vec<u8>> {
|
||||
// Unwrap key using Blackwell TEE-I/O protection
|
||||
// This provides full I/O isolation with encrypted memory and I/O channels
|
||||
// In production, this would use NVIDIA Blackwell TEE-I/O API
|
||||
|
||||
if wrapped_key.len() < 32 {
|
||||
return Err(SecurityError::InvalidKey("Key too short".to_string()));
|
||||
}
|
||||
|
||||
// Blackwell TEE-I/O provides:
|
||||
// - Encrypted memory pages
|
||||
// - Authenticated I/O channels
|
||||
// - Hardware-enforced isolation
|
||||
// - Protection against all physical attacks
|
||||
|
||||
// In production, decrypt using Blackwell TEE-I/O enclave
|
||||
Ok(wrapped_key[0..32].to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
// Implementation similar to GpuCcVault but with Blackwell-specific attestation
|
||||
#[async_trait]
|
||||
impl<K: KeyBrokerService> KeyVault for GpuTeeIoVault<K> {
|
||||
fn tier(&self) -> PrivacyTier {
|
||||
PrivacyTier::GpuTeeIo
|
||||
}
|
||||
|
||||
// Similar implementation to GpuCcVault...
|
||||
async fn store_key(&self, _key_id: &KeyId, _key_data: &[u8]) -> Result<()> {
|
||||
Err(SecurityError::PolicyViolation(
|
||||
"GPU TEE-I/O vault does not support local key storage".to_string()
|
||||
))
|
||||
}
|
||||
|
||||
async fn use_key<F, R>(&self, key_id: &KeyId, operation: F) -> Result<R>
|
||||
where
|
||||
F: FnOnce(&[u8]) -> R + Send,
|
||||
R: Send,
|
||||
{
|
||||
// Blackwell-specific key usage with TEE-I/O protection
|
||||
// This provides the highest level of security with full I/O isolation
|
||||
let attestation = self.kbs.get_attestation_report().await?;
|
||||
|
||||
// Verify we're on Blackwell with TEE-I/O
|
||||
if attestation.tee_type != TeeType::Blackwell {
|
||||
return Err(SecurityError::AttestationFailure(
|
||||
"GPU not in Blackwell TEE-I/O mode".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
// Request key with TEE-I/O protection
|
||||
let wrapped_key = self.kbs.request_key(
|
||||
key_id,
|
||||
attestation,
|
||||
PrivacyTier::Tier4GpuTeeIo,
|
||||
).await?;
|
||||
|
||||
// Unwrap in TEE-I/O protected enclave
|
||||
// All I/O is encrypted and authenticated
|
||||
let key = self.unwrap_in_tee_io_enclave(&wrapped_key)?;
|
||||
Ok(operation(&key))
|
||||
}
|
||||
|
||||
async fn delete_key(&self, _key_id: &KeyId) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_initialized(&self) -> Result<bool> {
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
+3
-6
@@ -1,13 +1,10 @@
|
||||
[package]
|
||||
name = "hanzo-l2"
|
||||
version = "1.1.20"
|
||||
version = "1.1.21"
|
||||
edition = "2021"
|
||||
description = "Hanzo L2 bridge and sequencing on Lux Network"
|
||||
description = "Hanzo Network — L2 bridge and sequencing on Lux Network"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/hanzonet/network"
|
||||
# Cannot publish to crates.io until lux-consensus is published.
|
||||
# Also depends on the local net hanzo-pqc (publish = false).
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true }
|
||||
@@ -23,7 +20,7 @@ chrono = { workspace = true }
|
||||
hex = { workspace = true }
|
||||
|
||||
# Consensus
|
||||
lux-consensus = { path = "../../../lux/consensus/pkg/rust" }
|
||||
lux-consensus = "1.22"
|
||||
|
||||
# PQC for cross-chain signatures
|
||||
hanzo-pqc = { workspace = true }
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-sheet"
|
||||
description = "Hanzo AI - Sheet Library"
|
||||
license = "MIT"
|
||||
version = "0.1.1"
|
||||
edition = { workspace = true }
|
||||
publish = false # Not ready for crates.io
|
||||
|
||||
# Standalone crate - not part of parent workspace
|
||||
[workspace]
|
||||
authors = { workspace = true }
|
||||
|
||||
[dependencies]
|
||||
serde_json = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
uuid = { workspace = true, features = ["v4"] }
|
||||
regex = { workspace = true }
|
||||
hanzo-messages = { version = "1.1.12", path = "../hanzo-messages" }
|
||||
async-channel = { workspace = true }
|
||||
|
||||
[dependencies.serde]
|
||||
workspace = true
|
||||
features = ["derive"]
|
||||
|
||||
[[test]]
|
||||
name = "sheet_common_tests"
|
||||
path = "tests/sheet_common_tests.rs"
|
||||
|
||||
[[test]]
|
||||
name = "sheet_one_row_advanced_tests"
|
||||
path = "tests/sheet_one_row_advanced_tests.rs"
|
||||
@@ -1,19 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Fix line 755 (col.clone() is the key)
|
||||
sed -i '' '755,761s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
|
||||
# Fix line 800 (col.clone() is the key)
|
||||
sed -i '' '800,806s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
|
||||
# Fix line 812 (col.clone() is the key)
|
||||
sed -i '' '812,818s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
|
||||
# Fix line 1022 (col.clone() is the key)
|
||||
sed -i '' '1022,1028s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
|
||||
# Fix line 1049 (col.clone() is the key)
|
||||
sed -i '' '1049,1055s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
|
||||
# Fix line 1114 (col.clone() is the key)
|
||||
sed -i '' '1114,1120s/Cell {/Cell {\n id: CellId::from(col.clone()),/' src/sheet.rs
|
||||
@@ -1,76 +0,0 @@
|
||||
use regex::Regex;
|
||||
use hanzo_messages::schemas::sheet::CellId;
|
||||
|
||||
pub struct CellNameConverter;
|
||||
|
||||
impl CellNameConverter {
|
||||
pub fn column_name_to_index(name: &str) -> usize {
|
||||
let mut index = 0;
|
||||
for (i, c) in name.chars().rev().enumerate() {
|
||||
index += (c as usize - 'A' as usize + 1) * 26_usize.pow(i as u32);
|
||||
}
|
||||
index - 1
|
||||
}
|
||||
|
||||
pub fn column_index_to_name(index: usize) -> String {
|
||||
let mut index = index + 1;
|
||||
let mut name = String::new();
|
||||
while index > 0 {
|
||||
let rem = (index - 1) % 26;
|
||||
name.insert(0, (rem as u8 + b'A') as char);
|
||||
index = (index - rem - 1) / 26;
|
||||
}
|
||||
name
|
||||
}
|
||||
|
||||
pub fn cell_name_to_indices(name: &str) -> (usize, usize) {
|
||||
eprintln!("cell_name_to_indices: {}", name);
|
||||
let re = Regex::new(r"([A-Z]+)(\d+)").unwrap();
|
||||
let caps = re.captures(name).unwrap();
|
||||
let col_name = &caps[1];
|
||||
let row_index: usize = caps[2].parse().unwrap();
|
||||
(row_index - 1, Self::column_name_to_index(col_name))
|
||||
}
|
||||
|
||||
pub fn cell_indices_to_name(row: usize, col: usize) -> String {
|
||||
format!("{}{}", Self::column_index_to_name(col), row + 1)
|
||||
}
|
||||
|
||||
pub fn cell_id_to_indices(cell_id: &CellId) -> (usize, usize) {
|
||||
let parts: Vec<&str> = cell_id.0.split(':').collect();
|
||||
(parts[0].parse().unwrap(), parts[1].parse().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_column_name_conversion() {
|
||||
assert_eq!(CellNameConverter::column_name_to_index("A"), 0);
|
||||
assert_eq!(CellNameConverter::column_name_to_index("Z"), 25);
|
||||
assert_eq!(CellNameConverter::column_name_to_index("AA"), 26);
|
||||
assert_eq!(CellNameConverter::column_name_to_index("AB"), 27);
|
||||
assert_eq!(CellNameConverter::column_index_to_name(0), "A");
|
||||
assert_eq!(CellNameConverter::column_index_to_name(25), "Z");
|
||||
assert_eq!(CellNameConverter::column_index_to_name(26), "AA");
|
||||
assert_eq!(CellNameConverter::column_index_to_name(27), "AB");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cell_name_conversion() {
|
||||
assert_eq!(CellNameConverter::cell_name_to_indices("A1"), (0, 0));
|
||||
assert_eq!(CellNameConverter::cell_name_to_indices("B2"), (1, 1));
|
||||
assert_eq!(CellNameConverter::cell_name_to_indices("AA10"), (9, 26));
|
||||
assert_eq!(CellNameConverter::cell_indices_to_name(0, 0), "A1");
|
||||
assert_eq!(CellNameConverter::cell_indices_to_name(1, 1), "B2");
|
||||
assert_eq!(CellNameConverter::cell_indices_to_name(9, 26), "AA10");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cell_id_to_indices() {
|
||||
let cell_id = CellId("0:1".to_string());
|
||||
assert_eq!(CellNameConverter::cell_id_to_indices(&cell_id), (0, 1));
|
||||
}
|
||||
}
|
||||
@@ -1,72 +0,0 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use hanzo_messages::schemas::sheet::UuidString;
|
||||
|
||||
#[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ColumnDependencyManager {
|
||||
// Column -> Columns it depends on
|
||||
pub dependencies: HashMap<UuidString, HashSet<UuidString>>,
|
||||
// Column -> Columns that depend on it
|
||||
pub reverse_dependencies: HashMap<UuidString, HashSet<UuidString>>,
|
||||
}
|
||||
|
||||
impl ColumnDependencyManager {
|
||||
pub fn add_dependency(&mut self, from: UuidString, to: UuidString) {
|
||||
self.dependencies.entry(from.clone()).or_default().insert(to.clone());
|
||||
self.reverse_dependencies.entry(to).or_default().insert(from);
|
||||
}
|
||||
|
||||
pub fn remove_dependency(&mut self, from: UuidString, to: UuidString) {
|
||||
if let Some(deps) = self.dependencies.get_mut(&from) {
|
||||
deps.remove(&to);
|
||||
}
|
||||
if let Some(rev_deps) = self.reverse_dependencies.get_mut(&to) {
|
||||
rev_deps.remove(&from);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn remove_column(&mut self, col: UuidString) {
|
||||
// Remove all dependencies where the column is a key
|
||||
if let Some(deps) = self.dependencies.remove(&col) {
|
||||
for dep in deps {
|
||||
if let Some(rev_deps) = self.reverse_dependencies.get_mut(&dep) {
|
||||
rev_deps.remove(&col);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove all reverse dependencies where the column is a value
|
||||
if let Some(rev_deps) = self.reverse_dependencies.remove(&col) {
|
||||
for rev_dep in rev_deps {
|
||||
if let Some(deps) = self.dependencies.get_mut(&rev_dep) {
|
||||
deps.remove(&col);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn update_dependencies(&mut self, col: UuidString, dependencies: HashSet<UuidString>) {
|
||||
// Remove existing dependencies for the column without affecting reverse dependencies
|
||||
if let Some(existing_deps) = self.dependencies.remove(&col) {
|
||||
for dep in &existing_deps {
|
||||
if let Some(rev_deps) = self.reverse_dependencies.get_mut(dep) {
|
||||
rev_deps.remove(&col);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add new dependencies
|
||||
for dep in dependencies {
|
||||
self.add_dependency(col.clone(), dep);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_dependents(&self, column: UuidString) -> HashSet<UuidString> {
|
||||
self.dependencies.get(&column).cloned().unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn get_reverse_dependents(&self, column: UuidString) -> HashSet<UuidString> {
|
||||
self.reverse_dependencies.get(&column).cloned().unwrap_or_default()
|
||||
}
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
pub mod sheet;
|
||||
pub mod cell_name_converter;
|
||||
pub mod column_dependency_manager;
|
||||
@@ -1,6 +0,0 @@
|
||||
extern crate serde;
|
||||
extern crate serde_json;
|
||||
|
||||
fn main() {
|
||||
println!("Hello, world!");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,108 +0,0 @@
|
||||
// Local type definitions for sheet types that were previously in hanzo-message-primitives
|
||||
use serde::{Deserialize, Serialize};
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
pub type UuidString = String;
|
||||
pub type ColumnUuid = String;
|
||||
pub type RowUuid = String;
|
||||
|
||||
// CellId is a newtype wrapper around String
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
||||
pub struct CellId(pub String);
|
||||
|
||||
impl CellId {
|
||||
pub fn new(id: String) -> Self {
|
||||
CellId(id)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for CellId {
|
||||
fn from(s: String) -> Self {
|
||||
CellId(s)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for CellId {
|
||||
fn from(s: &str) -> Self {
|
||||
CellId(s.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct Cell {
|
||||
pub id: CellId,
|
||||
pub value: Option<String>,
|
||||
pub status: CellStatus,
|
||||
pub input_hash: Option<String>,
|
||||
pub last_updated: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum CellStatus {
|
||||
Empty,
|
||||
Filled,
|
||||
Computing,
|
||||
Ready,
|
||||
Pending,
|
||||
Waiting,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CellUpdateData {
|
||||
pub cell_id: CellId,
|
||||
pub value: Option<String>,
|
||||
pub row_id: RowUuid,
|
||||
pub column_id: ColumnUuid,
|
||||
pub input_hash: Option<String>,
|
||||
pub status: CellStatus,
|
||||
pub last_updated: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct CellUpdateInfo {
|
||||
pub cell_id: CellId,
|
||||
pub old_value: Option<String>,
|
||||
pub new_value: Option<String>,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub sheet_id: String,
|
||||
pub update_type: String,
|
||||
pub data: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct ColumnDefinition {
|
||||
pub id: ColumnUuid, // Alias for uuid
|
||||
pub uuid: ColumnUuid,
|
||||
pub name: String,
|
||||
pub behavior: ColumnBehavior,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum ColumnBehavior {
|
||||
Text,
|
||||
Number,
|
||||
Formula(String),
|
||||
LLM {
|
||||
prompt_template: String,
|
||||
model: Option<String>,
|
||||
},
|
||||
LLMCall {
|
||||
input: String,
|
||||
llm_provider_name: Option<String>,
|
||||
input_hash: Option<String>,
|
||||
},
|
||||
MultipleVRFiles,
|
||||
UploadedFiles,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WorkflowSheetJobData {
|
||||
pub sheet_id: String,
|
||||
pub cell_updates: Vec<CellUpdateData>,
|
||||
pub row: RowUuid,
|
||||
pub col: ColumnUuid,
|
||||
pub col_definition: ColumnDefinition,
|
||||
pub input_cells: Vec<Cell>,
|
||||
pub llm_provider_name: Option<String>,
|
||||
}
|
||||
@@ -1,534 +0,0 @@
|
||||
use chrono::Utc;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use hanzo_messages::schemas::sheet::{ColumnBehavior, ColumnDefinition, UuidString};
|
||||
use hanzo_sheet::sheet::Sheet;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_id = Uuid::new_v4().to_string();
|
||||
let column = ColumnDefinition {
|
||||
id: column_id.clone(),
|
||||
name: "Column 1".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column.clone()).await;
|
||||
assert_eq!(sheet.columns.len(), 1);
|
||||
assert_eq!(sheet.columns[&column_id], column);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_id = Uuid::new_v4().to_string();
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_id.clone(),
|
||||
name: "Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let _ = sheet.set_column(column_text.clone()).await;
|
||||
assert_eq!(sheet.columns[&column_id].name, "Text Column");
|
||||
|
||||
let updated_column_text = ColumnDefinition {
|
||||
id: column_id.clone(),
|
||||
name: "Updated Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let _ = sheet.set_column(updated_column_text.clone()).await;
|
||||
assert_eq!(sheet.columns[&column_id].name, "Updated Text Column");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_set_cell_value() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column = ColumnDefinition {
|
||||
id: column_id.clone(),
|
||||
name: "Column 1".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column).await;
|
||||
|
||||
// Add a new row before setting the cell value
|
||||
let _ = sheet.add_row(row_id.clone()).await;
|
||||
|
||||
let result = sheet
|
||||
.set_cell_value(row_id.clone(), column_id.clone(), "Test Value".to_string())
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("Test Value".to_string()));
|
||||
assert!(cell.last_updated <= Utc::now());
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_a_new_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
let _ = sheet.set_column(column_c.clone()).await;
|
||||
|
||||
assert_eq!(sheet.columns.len(), 2);
|
||||
assert_eq!(sheet.columns[&column_a_id], column_a);
|
||||
assert_eq!(sheet.columns[&column_c_id], column_c);
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_formula_evaluation() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+B".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
|
||||
// Ensure the row is created before setting the cell values
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("HelloWorld".to_string()));
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Bye".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("ByeWorld".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_formula_evaluation_multiple_references() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let column_d_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_d = ColumnDefinition {
|
||||
id: column_d_id.clone(),
|
||||
name: "Column D".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+B+C".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
let _ = sheet.set_column(column_d).await;
|
||||
|
||||
// Ensure the row is created before setting the cell values
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_c_id.clone(), "Again".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_d_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("HelloWorldAgain".to_string()));
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Bye".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_d_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("ByeWorldAgain".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_formula_evaluation_with_literals() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+\" space \"+B".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
|
||||
// Ensure the row is created before setting the cell values
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("Hello space World".to_string()));
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Bye".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell = sheet.get_cell(row_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell.value, Some("Bye space World".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dependencies_update() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+B+\"hey\"".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
|
||||
// Ensure the row is created before setting the cell values
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let dependents = sheet.column_dependency_manager.get_dependents(column_c_id.clone());
|
||||
assert!(dependents.contains(&column_a_id));
|
||||
assert!(dependents.contains(&column_b_id));
|
||||
|
||||
let reverse_dependents_a = sheet
|
||||
.column_dependency_manager
|
||||
.get_reverse_dependents(column_a_id.clone());
|
||||
let reverse_dependents_b = sheet
|
||||
.column_dependency_manager
|
||||
.get_reverse_dependents(column_b_id.clone());
|
||||
assert!(reverse_dependents_a.contains(&column_c_id));
|
||||
assert!(reverse_dependents_b.contains(&column_c_id));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_formula_reads_column_a() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+\" Copy\"".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
|
||||
// Ensure the row is created before setting the cell values
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Not Empty".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell_b_updated = sheet.get_cell(row_id.clone(), column_b_id.clone()).unwrap();
|
||||
assert_eq!(cell_b_updated.value, Some("Not Empty Copy".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_formula_evaluation_with_literals_and_copy_for_multiple_rows() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let column_d_id = Uuid::new_v4().to_string();
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+\" \"+B".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
|
||||
// Ensure the rows are created before setting the cell values
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_a_id.clone(), "Foo".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_b_id.clone(), "Bar".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell_c_0 = sheet.get_cell(row_0_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_0.value, Some("Hello World".to_string()));
|
||||
|
||||
let cell_c_1 = sheet.get_cell(row_1_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_1.value, Some("Foo Bar".to_string()));
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Bye".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cell_c_updated_0 = sheet.get_cell(row_0_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_updated_0.value, Some("Bye World".to_string()));
|
||||
|
||||
let column_d = ColumnDefinition {
|
||||
id: column_d_id.clone(),
|
||||
name: "Column D".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=C+\" Copy\"".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_d).await;
|
||||
|
||||
let cell_d_0 = sheet.get_cell(row_0_id.clone(), column_d_id.clone()).unwrap();
|
||||
assert_eq!(cell_d_0.value, Some("Bye World Copy".to_string()));
|
||||
|
||||
let cell_d_1 = sheet.get_cell(row_1_id.clone(), column_d_id.clone()).unwrap();
|
||||
sheet.print_as_ascii_table();
|
||||
assert_eq!(cell_d_1.value, Some("Foo Bar Copy".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_columns_with_formula() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let column_d_id = Uuid::new_v4().to_string();
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A+\" \"+B".to_string()),
|
||||
};
|
||||
|
||||
// Add columns
|
||||
let _ = sheet.set_column(column_a).await;
|
||||
let _ = sheet.set_column(column_b).await;
|
||||
let _ = sheet.set_column(column_c).await;
|
||||
|
||||
// Ensure the rows are created before setting the cell values
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
|
||||
// Set cell values
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_b_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_a_id.clone(), "Foo".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_b_id.clone(), "Bar".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check initial formula results
|
||||
let cell_c_0 = sheet.get_cell(row_0_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_0.value, Some("Hello World".to_string()));
|
||||
|
||||
let cell_c_1 = sheet.get_cell(row_1_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_1.value, Some("Foo Bar".to_string()));
|
||||
|
||||
// Update a cell value and check formula result
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Bye".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
let cell_c_updated_0 = sheet.get_cell(row_0_id.clone(), column_c_id.clone()).unwrap();
|
||||
assert_eq!(cell_c_updated_0.value, Some("Bye World".to_string()));
|
||||
|
||||
// Add another formula column
|
||||
let column_d = ColumnDefinition {
|
||||
id: column_d_id.clone(),
|
||||
name: "Column D".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=C+\" Copy\"".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(column_d).await;
|
||||
|
||||
let cell_d_0 = sheet.get_cell(row_0_id.clone(), column_d_id.clone()).unwrap();
|
||||
assert_eq!(cell_d_0.value, Some("Bye World Copy".to_string()));
|
||||
|
||||
let cell_d_1 = sheet.get_cell(row_1_id.clone(), column_d_id.clone()).unwrap();
|
||||
assert_eq!(cell_d_1.value, Some("Foo Bar Copy".to_string()));
|
||||
|
||||
// Remove the middle column
|
||||
sheet.remove_column(column_b_id.clone()).await.unwrap();
|
||||
sheet.print_as_ascii_table();
|
||||
|
||||
assert_eq!(sheet.columns.len(), 3);
|
||||
|
||||
// Remove the first column
|
||||
sheet.remove_column(column_a_id.clone()).await.unwrap();
|
||||
assert_eq!(sheet.columns.len(), 2);
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_formula_dependencies_text_input() {
|
||||
let sheet = Sheet::new();
|
||||
let formula = "Say Hello World";
|
||||
let dependencies = sheet.parse_formula_dependencies(formula);
|
||||
let expected: HashSet<UuidString> = HashSet::new();
|
||||
assert_eq!(dependencies, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_formula_dependencies_with_column_reference() {
|
||||
let mut sheet = Sheet::new();
|
||||
sheet.display_columns.push(Uuid::new_v4().to_string()); // Add a column to display_columns
|
||||
let formula = "=A + \" And Space\"";
|
||||
let dependencies = sheet.parse_formula_dependencies(formula);
|
||||
let mut expected = HashSet::new();
|
||||
expected.insert(sheet.display_columns[0].clone()); // Use the UUID from display_columns
|
||||
assert_eq!(dependencies, expected);
|
||||
}
|
||||
}
|
||||
@@ -1,889 +0,0 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use hanzo_messages::schemas::sheet::{CellStatus, ColumnBehavior, ColumnDefinition};
|
||||
use hanzo_sheet::sheet::Sheet;
|
||||
use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_call_column() {
|
||||
let sheet = Arc::new(Mutex::new(Sheet::new()));
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let column_text_id = Uuid::new_v4().to_string();
|
||||
let column_llm_id = Uuid::new_v4().to_string();
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_text_id.clone(),
|
||||
name: "Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let column_llm = ColumnDefinition {
|
||||
id: column_llm_id.clone(),
|
||||
name: "LLM Call Column".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "Say Hello World".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
{
|
||||
let mut sheet = sheet.lock().await;
|
||||
let _ = sheet.set_column(column_text.clone()).await;
|
||||
let _ = sheet.set_column(column_llm.clone()).await;
|
||||
}
|
||||
|
||||
assert_eq!(sheet.lock().await.columns.len(), 2);
|
||||
assert_eq!(sheet.lock().await.columns[&column_text_id], column_text);
|
||||
assert_eq!(sheet.lock().await.columns[&column_llm_id], column_llm);
|
||||
|
||||
sheet.lock().await.add_row(row_0_id.clone()).await.unwrap();
|
||||
|
||||
// Check the value of the cell after the update
|
||||
let cell_value = sheet
|
||||
.lock()
|
||||
.await
|
||||
.get_cell_value(Uuid::new_v4().to_string(), column_llm_id.clone());
|
||||
assert_eq!(cell_value, None);
|
||||
|
||||
// Simulate job completion
|
||||
let jobs = sheet
|
||||
.lock()
|
||||
.await
|
||||
.set_cell_value(row_0_id.clone(), column_llm_id.clone(), "Hello World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(jobs.len(), 0);
|
||||
|
||||
// Check the value of the cell after the job completion
|
||||
let cell_value = sheet
|
||||
.lock()
|
||||
.await
|
||||
.get_cell_value(row_0_id.clone(), column_llm_id.clone());
|
||||
assert_eq!(cell_value, Some("Hello World".to_string()));
|
||||
|
||||
{
|
||||
let sheet_locked = sheet.lock().await;
|
||||
let cell_status = sheet_locked
|
||||
.get_cell(row_0_id.clone(), column_llm_id)
|
||||
.map(|cell| &cell.status);
|
||||
assert_eq!(cell_status, Some(&CellStatus::Ready));
|
||||
sheet_locked.print_as_ascii_table();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_auto_populate_new_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A + \" Copy\"".to_string()),
|
||||
};
|
||||
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=B + \" Second Copy\"".to_string()),
|
||||
};
|
||||
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
let _ = sheet.set_column(column_b.clone()).await;
|
||||
let _ = sheet.set_column(column_c.clone()).await;
|
||||
|
||||
sheet.print_as_ascii_table();
|
||||
|
||||
// Ensure the row is created before setting the cell value
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_id.clone(), column_a_id.clone()),
|
||||
Some("Hello".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_id.clone(), column_b_id.clone()),
|
||||
Some("Hello Copy".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_id.clone(), column_c_id.clone()),
|
||||
Some("Hello Copy Second Copy".to_string())
|
||||
);
|
||||
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_call_with_dependent_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_text_id = Uuid::new_v4().to_string();
|
||||
let column_llm_id = Uuid::new_v4().to_string();
|
||||
let column_formula_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_text_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let workflow_str = r#"
|
||||
workflow WorkflowTest v0.1 {
|
||||
step Main {
|
||||
$RESULT = call opinionated_inference($INPUT)
|
||||
}
|
||||
}
|
||||
"#;
|
||||
|
||||
let column_llm = ColumnDefinition {
|
||||
id: column_llm_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "Say Hello World".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let column_formula = ColumnDefinition {
|
||||
id: column_formula_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A + \" And Space\"".to_string()),
|
||||
};
|
||||
|
||||
let _text_jobs = sheet.set_column(column_text.clone()).await;
|
||||
let _llm_jobs = sheet.set_column(column_llm.clone()).await;
|
||||
let _formula_jobs = sheet.set_column(column_formula.clone()).await;
|
||||
|
||||
// Ensure the row is created before setting the cell value
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
// Set value in Column A
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_text_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check initial state of Column C (formula depending on Column A)
|
||||
let cell_value_formula = sheet.get_cell_value(row_id.clone(), column_formula_id.clone());
|
||||
assert_eq!(cell_value_formula, Some("Hello And Space".to_string()));
|
||||
|
||||
// Change Column C formula to depend on Column B instead of Column A
|
||||
let new_column_formula = ColumnDefinition {
|
||||
id: column_formula_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=B + \" Updated\"".to_string()),
|
||||
};
|
||||
let _ = sheet.set_column(new_column_formula).await;
|
||||
|
||||
// Check Column C value before updating Column B (should be empty or default value)
|
||||
let cell_value_formula = sheet.get_cell_value(row_id.clone(), column_formula_id.clone());
|
||||
assert_eq!(cell_value_formula, Some(" Updated".to_string()));
|
||||
|
||||
// Simulate LLM call completion for Column B
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_llm_id.clone(), "Hola Mundo".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check the value of the LLM call cell (Column B) after the update
|
||||
let cell_value_llm = sheet.get_cell_value(row_id.clone(), column_llm_id.clone());
|
||||
assert_eq!(cell_value_llm, Some("Hola Mundo".to_string()));
|
||||
|
||||
// Check if Column C reflects the updated value of Column B
|
||||
let cell_value_formula = sheet.get_cell_value(row_id.clone(), column_formula_id.clone());
|
||||
assert_eq!(cell_value_formula, Some("Hola Mundo Updated".to_string()));
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_rows() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let column_c_id = Uuid::new_v4().to_string();
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
let row_2_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A + \" Processed\"".to_string()),
|
||||
};
|
||||
|
||||
let column_c = ColumnDefinition {
|
||||
id: column_c_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "Summarize: $INPUT".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
let _ = sheet.set_column(column_b.clone()).await;
|
||||
let _ = sheet.set_column(column_c.clone()).await;
|
||||
|
||||
// Ensure the rows are created before setting the cell values
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_2_id.clone()).await.unwrap();
|
||||
|
||||
// Add data to multiple rows
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_a_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_2_id.clone(), column_a_id.clone(), "Test".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check values in Column A and B
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_a_id.clone()),
|
||||
Some("Hello".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_b_id.clone()),
|
||||
Some("Hello Processed".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_1_id.clone(), column_a_id.clone()),
|
||||
Some("World".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_1_id.clone(), column_b_id.clone()),
|
||||
Some("World Processed".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_a_id.clone()),
|
||||
Some("Test".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_b_id.clone()),
|
||||
Some("Test Processed".to_string())
|
||||
);
|
||||
|
||||
// Simulate LLM calls on Column C for all rows
|
||||
for (i, row_id) in [&row_0_id, &row_1_id, &row_2_id].iter().enumerate() {
|
||||
// Simulate job completion for Column C
|
||||
sheet
|
||||
.set_cell_value(row_id.to_string(), column_c_id.clone(), format!("Summary of row {}", i))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Check final values
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_c_id.clone()),
|
||||
Some("Summary of row 0".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_1_id.clone(), column_c_id.clone()),
|
||||
Some("Summary of row 1".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_c_id.clone()),
|
||||
Some("Summary of row 2".to_string())
|
||||
);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_row() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
let row_2_id = Uuid::new_v4().to_string();
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A + \" Processed\"".to_string()),
|
||||
};
|
||||
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
let _ = sheet.set_column(column_b.clone()).await;
|
||||
|
||||
// Ensure the rows are created before setting the cell values
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_2_id.clone()).await.unwrap();
|
||||
|
||||
// Add data to multiple rows
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_a_id.clone(), "World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_2_id.clone(), column_a_id.clone(), "Test".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check initial values
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_a_id.clone()),
|
||||
Some("Hello".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_1_id.clone(), column_a_id.clone()),
|
||||
Some("World".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_a_id.clone()),
|
||||
Some("Test".to_string())
|
||||
);
|
||||
|
||||
// Remove a row
|
||||
let _ = sheet.remove_row(row_1_id.clone()).await.unwrap();
|
||||
|
||||
// Check values after row removal
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_a_id.clone()),
|
||||
Some("Hello".to_string())
|
||||
);
|
||||
assert_eq!(sheet.get_cell_value(row_1_id.clone(), column_a_id.clone()), None);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_a_id.clone()),
|
||||
Some("Test".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_b_id.clone()),
|
||||
Some("Hello Processed".to_string())
|
||||
);
|
||||
assert_eq!(sheet.get_cell_value(row_1_id.clone(), column_b_id.clone()), None);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_b_id.clone()),
|
||||
Some("Test Processed".to_string())
|
||||
);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
|
||||
// Add a new row
|
||||
let new_row_id = Uuid::new_v4().to_string();
|
||||
let _ = sheet.add_row(new_row_id.clone()).await.unwrap();
|
||||
|
||||
// Add a value to the first column of the new row
|
||||
sheet
|
||||
.set_cell_value(new_row_id.clone(), column_a_id.clone(), "New Value".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check values after adding a new row
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_a_id.clone()),
|
||||
Some("Hello".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(new_row_id.clone(), column_a_id.clone()),
|
||||
Some("New Value".to_string())
|
||||
); // New row should have the new value
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_a_id.clone()),
|
||||
Some("Test".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_0_id.clone(), column_b_id.clone()),
|
||||
Some("Hello Processed".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(new_row_id.clone(), column_b_id.clone()),
|
||||
Some("New Value Processed".to_string())
|
||||
); // Formula should be recalculated
|
||||
assert_eq!(
|
||||
sheet.get_cell_value(row_2_id.clone(), column_b_id.clone()),
|
||||
Some("Test Processed".to_string())
|
||||
);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dependency_manager() {
|
||||
let sheet = Arc::new(Mutex::new(Sheet::new()));
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_text_id = Uuid::new_v4().to_string();
|
||||
let column_llm_id = Uuid::new_v4().to_string();
|
||||
let column_formula_id = Uuid::new_v4().to_string();
|
||||
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_text_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
let column_llm = ColumnDefinition {
|
||||
id: column_llm_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "=A".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let column_formula = ColumnDefinition {
|
||||
id: column_formula_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=B + \" And Space\"".to_string()),
|
||||
};
|
||||
|
||||
{
|
||||
let mut sheet = sheet.lock().await;
|
||||
let _ = sheet.set_column(column_text.clone()).await;
|
||||
let _ = sheet.set_column(column_llm.clone()).await;
|
||||
let _ = sheet.set_column(column_formula.clone()).await;
|
||||
}
|
||||
|
||||
eprintln!("Checking initial state of the sheet");
|
||||
|
||||
assert_eq!(sheet.lock().await.columns.len(), 3);
|
||||
assert_eq!(sheet.lock().await.columns[&column_text_id], column_text);
|
||||
assert_eq!(sheet.lock().await.columns[&column_llm_id], column_llm);
|
||||
assert_eq!(sheet.lock().await.columns[&column_formula_id], column_formula);
|
||||
|
||||
sheet.lock().await.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
// Set value in Column A
|
||||
sheet
|
||||
.lock()
|
||||
.await
|
||||
.set_cell_value(row_id.clone(), column_text_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check initial state of Column C (formula depending on Column B)
|
||||
let cell_value_formula = sheet
|
||||
.lock()
|
||||
.await
|
||||
.get_cell_value(row_id.clone(), column_formula_id.clone());
|
||||
assert_eq!(cell_value_formula, Some(" And Space".to_string()));
|
||||
|
||||
eprintln!("Checking initial state of the dependency manager");
|
||||
|
||||
// Simulate LLM call completion for Column B
|
||||
sheet
|
||||
.lock()
|
||||
.await
|
||||
.set_cell_value(row_id.clone(), column_llm_id.clone(), "Hola Mundo".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check the value of the LLM call cell (Column B) after the update
|
||||
let cell_value_llm = sheet.lock().await.get_cell_value(row_id.clone(), column_llm_id.clone());
|
||||
assert_eq!(cell_value_llm, Some("Hola Mundo".to_string()));
|
||||
|
||||
// Check if Column C reflects the updated value of Column B
|
||||
let cell_value_formula = sheet
|
||||
.lock()
|
||||
.await
|
||||
.get_cell_value(row_id.clone(), column_formula_id.clone());
|
||||
assert_eq!(cell_value_formula, Some("Hola Mundo And Space".to_string()));
|
||||
|
||||
// Check dependency manager
|
||||
{
|
||||
let sheet_locked = sheet.lock().await;
|
||||
let reverse_dependents_a = sheet_locked
|
||||
.column_dependency_manager
|
||||
.get_reverse_dependents(column_text_id.clone());
|
||||
let reverse_dependents_b = sheet_locked
|
||||
.column_dependency_manager
|
||||
.get_reverse_dependents(column_llm_id.clone());
|
||||
assert!(reverse_dependents_a.contains(&column_llm_id));
|
||||
assert!(reverse_dependents_b.contains(&column_formula_id));
|
||||
}
|
||||
|
||||
// Update Column C to be a text column
|
||||
let new_column_text = ColumnDefinition {
|
||||
id: column_formula_id.clone(),
|
||||
name: "Column C".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
eprintln!("Updating Column C to be a text column");
|
||||
let _ = sheet.lock().await.set_column(new_column_text.clone()).await;
|
||||
|
||||
// Check the dependencies after updating Column C
|
||||
let sheet_locked = sheet.lock().await;
|
||||
let dependencies_c = sheet_locked
|
||||
.column_dependency_manager
|
||||
.dependencies
|
||||
.get(&column_formula_id);
|
||||
let reverse_dependents_b = sheet_locked
|
||||
.column_dependency_manager
|
||||
.reverse_dependencies
|
||||
.get(&column_llm_id);
|
||||
|
||||
eprintln!("Checking dependencies after updating Column C: {:?}", dependencies_c);
|
||||
assert!(dependencies_c.is_none());
|
||||
assert!(!reverse_dependents_b.unwrap().contains(&column_formula_id));
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet_locked.print_as_ascii_table();
|
||||
}
|
||||
|
||||
// TODO: add a new test and feature to throw an error that gets scaled up to the user
|
||||
//eg behavior: ColumnBehavior::Formula("=LLM Call Column + \" Copy\"".to_string()),
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_call_and_formula_cell_status() {
|
||||
let sheet = Arc::new(Mutex::new(Sheet::new()));
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
let column_llm_id = Uuid::new_v4().to_string();
|
||||
let column_formula_id = Uuid::new_v4().to_string();
|
||||
|
||||
let column_llm = ColumnDefinition {
|
||||
id: column_llm_id.clone(),
|
||||
name: "LLM Call Column".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "Say Hello World".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let column_formula = ColumnDefinition {
|
||||
id: column_formula_id.clone(),
|
||||
name: "Formula Column".to_string(),
|
||||
behavior: ColumnBehavior::Formula("=A + \" Copy\"".to_string()),
|
||||
};
|
||||
|
||||
{
|
||||
let mut sheet = sheet.lock().await;
|
||||
let _ = sheet.set_column(column_llm.clone()).await;
|
||||
let _ = sheet.set_column(column_formula.clone()).await;
|
||||
}
|
||||
|
||||
sheet.lock().await.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
// Check initial cell statuses
|
||||
{
|
||||
let sheet_locked = sheet.lock().await;
|
||||
let llm_cell_status = sheet_locked
|
||||
.get_cell(row_id.clone(), column_llm_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
let formula_cell_status = sheet_locked
|
||||
.get_cell(row_id.clone(), column_formula_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
|
||||
assert_eq!(llm_cell_status, Some(&CellStatus::Pending));
|
||||
assert_eq!(formula_cell_status, Some(&CellStatus::Ready)); // TODO: eventually it should be "linked" so Pending
|
||||
}
|
||||
|
||||
// Simulate LLM call completion
|
||||
sheet
|
||||
.lock()
|
||||
.await
|
||||
.set_cell_value(row_id.clone(), column_llm_id.clone(), "Hello World".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Check cell statuses after LLM call completion
|
||||
{
|
||||
let sheet_locked = sheet.lock().await;
|
||||
let llm_cell_status = sheet_locked
|
||||
.get_cell(row_id.clone(), column_llm_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
let formula_cell_status = sheet_locked
|
||||
.get_cell(row_id.clone(), column_formula_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
|
||||
assert_eq!(llm_cell_status, Some(&CellStatus::Ready));
|
||||
assert_eq!(formula_cell_status, Some(&CellStatus::Ready));
|
||||
}
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.lock().await.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_jobs_created_for_new_column() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Column A".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
// Add Column A and a row
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
|
||||
// Add multiple rows
|
||||
let mut row_ids = Vec::new();
|
||||
for _ in 0..3 {
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
row_ids.push(row_id.clone());
|
||||
// Set value in Column A
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_a_id.clone(), "Hello".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Add Column B as LLMCall and check if jobs are created
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "Column B".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "Say Hello World".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let jobs = sheet.set_column(column_b.clone()).await.unwrap();
|
||||
assert_eq!(
|
||||
jobs.len(),
|
||||
3,
|
||||
"The number of jobs should be equal to the number of rows created"
|
||||
);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_call_shouldnt_include_empty_dependencies() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_a_id = Uuid::new_v4().to_string();
|
||||
let column_b_id = Uuid::new_v4().to_string();
|
||||
|
||||
let column_a = ColumnDefinition {
|
||||
id: column_a_id.clone(),
|
||||
name: "Cities".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
|
||||
// Add Column A and create 3 rows
|
||||
let _ = sheet.set_column(column_a.clone()).await;
|
||||
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
let row_2_id = Uuid::new_v4().to_string();
|
||||
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_2_id.clone()).await.unwrap();
|
||||
|
||||
// Set values in Column A
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_a_id.clone(), "Santiago, Chile".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_a_id.clone(), "Austin, Texas".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Add Column B as LLMCall and check if jobs are created
|
||||
let column_b = ColumnDefinition {
|
||||
id: column_b_id.clone(),
|
||||
name: "New Column".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "=\"Say Hello World in \" + A".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
|
||||
let jobs = sheet.set_column(column_b.clone()).await.unwrap();
|
||||
assert_eq!(
|
||||
jobs.len(),
|
||||
2,
|
||||
"The number of jobs should be equal to the number of rows with values set"
|
||||
);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_text_column_and_get_value() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_text_id = Uuid::new_v4().to_string();
|
||||
let column_text_2_id = Uuid::new_v4().to_string();
|
||||
let row_id = Uuid::new_v4().to_string();
|
||||
|
||||
// Create a text column
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_text_id.clone(),
|
||||
name: "Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column_text.clone()).await;
|
||||
|
||||
// Add a row
|
||||
sheet.add_row(row_id.clone()).await.unwrap();
|
||||
|
||||
// Set the first item to some text
|
||||
sheet
|
||||
.set_cell_value(row_id.clone(), column_text_id.clone(), "Sample Text".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Create a second text column
|
||||
let column_text_2 = ColumnDefinition {
|
||||
id: column_text_2_id.clone(),
|
||||
name: "Second Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column_text_2.clone()).await;
|
||||
|
||||
// Try to get the value of the cell from the second column
|
||||
let cell_value = sheet.get_cell_value(row_id.clone(), column_text_2_id.clone());
|
||||
assert_eq!(cell_value, None);
|
||||
|
||||
// Retrieve rows from the sheet and ensure there are two cells
|
||||
let rows = sheet.rows.get(&row_id).unwrap();
|
||||
assert_eq!(rows.len(), 2);
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_call_column_with_two_rows() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_text_id = Uuid::new_v4().to_string();
|
||||
let column_llm_id = Uuid::new_v4().to_string();
|
||||
let row_0_id = Uuid::new_v4().to_string();
|
||||
let row_1_id = Uuid::new_v4().to_string();
|
||||
|
||||
// Create a text column
|
||||
let column_text = ColumnDefinition {
|
||||
id: column_text_id.clone(),
|
||||
name: "Text Column".to_string(),
|
||||
behavior: ColumnBehavior::Text,
|
||||
};
|
||||
let _ = sheet.set_column(column_text.clone()).await;
|
||||
|
||||
// Add two rows
|
||||
sheet.add_row(row_0_id.clone()).await.unwrap();
|
||||
sheet.add_row(row_1_id.clone()).await.unwrap();
|
||||
|
||||
// Set the first item in both rows to some text
|
||||
sheet
|
||||
.set_cell_value(row_0_id.clone(), column_text_id.clone(), "Sample Text 1".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
sheet
|
||||
.set_cell_value(row_1_id.clone(), column_text_id.clone(), "Sample Text 2".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Create an LLMCall column
|
||||
let column_llm = ColumnDefinition {
|
||||
id: column_llm_id.clone(),
|
||||
name: "LLM Call Column".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "=A".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
let jobs = sheet.set_column(column_llm.clone()).await.unwrap();
|
||||
|
||||
// Check that two jobs are created
|
||||
assert_eq!(
|
||||
jobs.len(),
|
||||
2,
|
||||
"The number of jobs should be equal to the number of rows created"
|
||||
);
|
||||
|
||||
// Check that the two cells of the second column are in pending status
|
||||
{
|
||||
let llm_cell_status_row_0 = sheet
|
||||
.get_cell(row_0_id.clone(), column_llm_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
let llm_cell_status_row_1 = sheet
|
||||
.get_cell(row_1_id.clone(), column_llm_id.clone())
|
||||
.map(|cell| &cell.status);
|
||||
|
||||
assert_eq!(llm_cell_status_row_0, Some(&CellStatus::Pending));
|
||||
assert_eq!(llm_cell_status_row_1, Some(&CellStatus::Pending));
|
||||
}
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_new_rows_should_create_jobs() {
|
||||
let mut sheet = Sheet::new();
|
||||
let column_ai_id = Uuid::new_v4().to_string();
|
||||
|
||||
// Create an AI column
|
||||
let column_ai = ColumnDefinition {
|
||||
id: column_ai_id.clone(),
|
||||
name: "AI Column".to_string(),
|
||||
behavior: ColumnBehavior::LLMCall {
|
||||
input: "=\"Hello World\"".to_string(),
|
||||
llm_provider_name: "MockProvider".to_string(),
|
||||
input_hash: None,
|
||||
},
|
||||
};
|
||||
let _ = sheet.set_column(column_ai.clone()).await;
|
||||
|
||||
// Add an extra row and check for job creation
|
||||
let row_2_id = Uuid::new_v4().to_string();
|
||||
let new_jobs = sheet.add_row(row_2_id.clone()).await.unwrap();
|
||||
assert_eq!(new_jobs.len(), 1, "A new job should be created for the newly added row");
|
||||
|
||||
// We add a new row and check for job creation
|
||||
let row_3_id = Uuid::new_v4().to_string();
|
||||
let new_jobs = sheet.add_row(row_3_id.clone()).await.unwrap();
|
||||
assert_eq!(new_jobs.len(), 1, "A new job should be created for the newly added row");
|
||||
|
||||
// Print final state of the sheet
|
||||
sheet.print_as_ascii_table();
|
||||
}
|
||||
}
|
||||
|
||||
// // TODO: add test that A (text missing) -> B (workflow depending on A) -> C (workflo depending on B)
|
||||
@@ -1,60 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-simulation"
|
||||
description = "Hanzo AI - Simulation Library"
|
||||
license = "MIT"
|
||||
version = "0.1.1"
|
||||
edition = "2021"
|
||||
publish = false # Not ready for crates.io
|
||||
|
||||
# Standalone crate - not part of parent workspace
|
||||
[workspace]
|
||||
|
||||
[dependencies]
|
||||
# Core
|
||||
tokio = { version = "1.38", features = ["full"] }
|
||||
anyhow = "1.0"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
tracing = "0.1"
|
||||
dashmap = "5.5"
|
||||
|
||||
# Database
|
||||
sqlx = { version = "0.8", features = ["runtime-tokio", "sqlite", "macros", "json"] }
|
||||
sqlite-vec = "0.1" # Vector search in SQLite
|
||||
|
||||
# USD support
|
||||
openusd = "0.1" # OpenUSD bindings
|
||||
usd-rs = "0.1" # Alternative USD library
|
||||
|
||||
# Math & Physics
|
||||
nalgebra = "0.33"
|
||||
rapier3d = "0.22" # Rust physics engine
|
||||
parry3d = "0.17" # Collision detection
|
||||
|
||||
# Graphics formats
|
||||
gltf = "1.4"
|
||||
image = "0.25"
|
||||
basis-universal = "0.3" # KTX2 texture compression
|
||||
|
||||
# Web/Networking
|
||||
warp = "0.3"
|
||||
tower = "0.4"
|
||||
tower-http = { version = "0.5", features = ["cors", "compression"] }
|
||||
|
||||
# ML Integration
|
||||
candle-core = "0.7"
|
||||
candle-nn = "0.7"
|
||||
ort = "2.0.0-rc.10" # ONNX Runtime for V-JEPA2
|
||||
|
||||
# Compute orchestration
|
||||
bollard = "0.17" # Docker API
|
||||
kube = { version = "0.95", features = ["runtime", "derive"] }
|
||||
k8s-openapi = { version = "0.23", features = ["latest"] }
|
||||
|
||||
# Security
|
||||
firecracker-sdk = "0.0.3" # MicroVM isolation
|
||||
hanzo-pqc = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = "0.5"
|
||||
tempfile = "3.10"
|
||||
@@ -1,256 +0,0 @@
|
||||
//! Zoo-1 Multi-tier Simulation System
|
||||
//!
|
||||
//! Provides three tiers of simulation:
|
||||
//! - Tier A: V-JEPA2 latent world model (cheapest)
|
||||
//! - Tier B: WebXR/WebGPU lightweight spatial (browser)
|
||||
//! - Tier C: High-fidelity physics (MuJoCo/Brax)
|
||||
|
||||
pub mod tiers;
|
||||
pub mod usd_layers;
|
||||
pub mod sql_brain;
|
||||
pub mod orchestrator;
|
||||
pub mod metaverse;
|
||||
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Simulation tier selection
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
||||
pub enum SimulationTier {
|
||||
/// Latent mental simulation using V-JEPA2
|
||||
Latent,
|
||||
/// Lightweight 3D with WebGPU/WebXR
|
||||
Lightweight,
|
||||
/// High-fidelity physics simulation
|
||||
HighFidelity,
|
||||
}
|
||||
|
||||
/// Expected Free Energy for tier selection
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ExpectedFreeEnergy {
|
||||
pub epistemic_value: f64, // Information gain
|
||||
pub pragmatic_value: f64, // Goal achievement
|
||||
pub cost: f64, // Computational cost
|
||||
}
|
||||
|
||||
impl ExpectedFreeEnergy {
|
||||
/// Select optimal simulation tier based on EFE
|
||||
pub fn select_tier(&self) -> SimulationTier {
|
||||
let value_of_information = self.epistemic_value + self.pragmatic_value;
|
||||
|
||||
if value_of_information < 0.3 {
|
||||
SimulationTier::Latent
|
||||
} else if value_of_information < 0.7 {
|
||||
SimulationTier::Lightweight
|
||||
} else {
|
||||
SimulationTier::HighFidelity
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simulation request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationRequest {
|
||||
pub question: String,
|
||||
pub scene_ref: String, // USD stage reference
|
||||
pub user_id: String,
|
||||
pub budget: SimulationBudget,
|
||||
pub policy: SimulationPolicy,
|
||||
}
|
||||
|
||||
/// Resource budget for simulation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationBudget {
|
||||
pub max_seconds: u64,
|
||||
pub max_gpu_seconds: u64,
|
||||
pub max_memory_gb: f32,
|
||||
}
|
||||
|
||||
/// Simulation execution policy
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationPolicy {
|
||||
pub tier: TierSelection,
|
||||
pub risk_tolerance: RiskTolerance,
|
||||
pub parallelism: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum TierSelection {
|
||||
Auto,
|
||||
Fixed(SimulationTier),
|
||||
Progressive, // Start with Latent, escalate if needed
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum RiskTolerance {
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
}
|
||||
|
||||
/// Simulation result with artifacts
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationResult {
|
||||
pub answer: String,
|
||||
pub confidence: f64,
|
||||
pub tier_used: SimulationTier,
|
||||
pub artifacts: Vec<SimulationArtifact>,
|
||||
pub metrics: SimulationMetrics,
|
||||
pub cost: SimulationCost,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationArtifact {
|
||||
pub artifact_type: ArtifactType,
|
||||
pub uri: String,
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum ArtifactType {
|
||||
Video,
|
||||
Image,
|
||||
UsdLayer,
|
||||
GltfScene,
|
||||
Trajectory,
|
||||
Metrics,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationMetrics {
|
||||
pub episodes_run: u32,
|
||||
pub success_rate: f64,
|
||||
pub avg_reward: f64,
|
||||
pub collision_count: u32,
|
||||
pub physics_violations: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SimulationCost {
|
||||
pub compute_seconds: f64,
|
||||
pub gpu_seconds: f64,
|
||||
pub memory_gb_hours: f64,
|
||||
pub estimated_usd: f64,
|
||||
}
|
||||
|
||||
/// Main simulation orchestrator
|
||||
pub struct SimulationOrchestrator {
|
||||
latent_tier: Arc<tiers::latent::LatentSimulator>,
|
||||
lightweight_tier: Arc<tiers::lightweight::LightweightSimulator>,
|
||||
highfidelity_tier: Arc<tiers::highfidelity::HighFidelitySimulator>,
|
||||
usd_manager: Arc<usd_layers::UsdLayerManager>,
|
||||
sql_brain: Arc<sql_brain::SqlBrain>,
|
||||
}
|
||||
|
||||
impl SimulationOrchestrator {
|
||||
pub async fn new(config: OrchestratorConfig) -> Result<Self> {
|
||||
Ok(Self {
|
||||
latent_tier: Arc::new(tiers::latent::LatentSimulator::new(config.latent_config).await?),
|
||||
lightweight_tier: Arc::new(tiers::lightweight::LightweightSimulator::new(config.lightweight_config).await?),
|
||||
highfidelity_tier: Arc::new(tiers::highfidelity::HighFidelitySimulator::new(config.hf_config).await?),
|
||||
usd_manager: Arc::new(usd_layers::UsdLayerManager::new(config.usd_config)?),
|
||||
sql_brain: Arc::new(sql_brain::SqlBrain::new(&config.sql_path).await?),
|
||||
})
|
||||
}
|
||||
|
||||
/// Process a simulation request
|
||||
pub async fn simulate(&self, request: SimulationRequest) -> Result<SimulationResult> {
|
||||
// 1. Load user context from SQL brain
|
||||
let user_context = self.sql_brain.load_user_context(&request.user_id).await?;
|
||||
|
||||
// 2. Create USD session layer for cloning
|
||||
let session_layer = self.usd_manager.create_session_layer(&request.scene_ref).await?;
|
||||
|
||||
// 3. Calculate Expected Free Energy
|
||||
let efe = self.calculate_efe(&request, &user_context).await?;
|
||||
|
||||
// 4. Select tier based on EFE and policy
|
||||
let tier = match request.policy.tier {
|
||||
TierSelection::Auto => efe.select_tier(),
|
||||
TierSelection::Fixed(tier) => tier,
|
||||
TierSelection::Progressive => SimulationTier::Latent, // Start with cheapest
|
||||
};
|
||||
|
||||
// 5. Execute simulation on selected tier
|
||||
let result = match tier {
|
||||
SimulationTier::Latent => {
|
||||
self.latent_tier.simulate(&request, &session_layer).await?
|
||||
}
|
||||
SimulationTier::Lightweight => {
|
||||
self.lightweight_tier.simulate(&request, &session_layer).await?
|
||||
}
|
||||
SimulationTier::HighFidelity => {
|
||||
self.highfidelity_tier.simulate(&request, &session_layer).await?
|
||||
}
|
||||
};
|
||||
|
||||
// 6. Store results in SQL brain
|
||||
self.sql_brain.store_episode(&request.user_id, &result).await?;
|
||||
|
||||
// 7. Optionally escalate if confidence too low
|
||||
if request.policy.tier == TierSelection::Progressive && result.confidence < 0.6 {
|
||||
// Escalate to next tier
|
||||
let next_tier = match tier {
|
||||
SimulationTier::Latent => SimulationTier::Lightweight,
|
||||
SimulationTier::Lightweight => SimulationTier::HighFidelity,
|
||||
SimulationTier::HighFidelity => SimulationTier::HighFidelity,
|
||||
};
|
||||
|
||||
if next_tier != tier {
|
||||
// Run again on higher tier
|
||||
return self.simulate_with_tier(request, next_tier).await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn calculate_efe(&self, request: &SimulationRequest, context: &sql_brain::UserContext) -> Result<ExpectedFreeEnergy> {
|
||||
// Calculate value of information based on question complexity and user history
|
||||
let epistemic_value = self.estimate_information_gain(&request.question, context).await?;
|
||||
let pragmatic_value = self.estimate_goal_value(&request.question, context).await?;
|
||||
let cost = self.estimate_compute_cost(SimulationTier::Latent);
|
||||
|
||||
Ok(ExpectedFreeEnergy {
|
||||
epistemic_value,
|
||||
pragmatic_value,
|
||||
cost,
|
||||
})
|
||||
}
|
||||
|
||||
async fn estimate_information_gain(&self, question: &str, context: &sql_brain::UserContext) -> Result<f64> {
|
||||
// Estimate based on question novelty and complexity
|
||||
// This would use embeddings to check similarity with past questions
|
||||
Ok(0.5) // Placeholder
|
||||
}
|
||||
|
||||
async fn estimate_goal_value(&self, question: &str, context: &sql_brain::UserContext) -> Result<f64> {
|
||||
// Estimate value for achieving user's goals
|
||||
Ok(0.3) // Placeholder
|
||||
}
|
||||
|
||||
fn estimate_compute_cost(&self, tier: SimulationTier) -> f64 {
|
||||
match tier {
|
||||
SimulationTier::Latent => 0.001, // $0.001 per simulation
|
||||
SimulationTier::Lightweight => 0.01, // $0.01 per simulation
|
||||
SimulationTier::HighFidelity => 0.10, // $0.10 per simulation
|
||||
}
|
||||
}
|
||||
|
||||
async fn simulate_with_tier(&self, request: SimulationRequest, tier: SimulationTier) -> Result<SimulationResult> {
|
||||
let mut req = request;
|
||||
req.policy.tier = TierSelection::Fixed(tier);
|
||||
self.simulate(req).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Orchestrator configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OrchestratorConfig {
|
||||
pub latent_config: tiers::latent::LatentConfig,
|
||||
pub lightweight_config: tiers::lightweight::LightweightConfig,
|
||||
pub hf_config: tiers::highfidelity::HighFidelityConfig,
|
||||
pub usd_config: usd_layers::UsdConfig,
|
||||
pub sql_path: String,
|
||||
}
|
||||
@@ -1,456 +0,0 @@
|
||||
//! SQL Brain - Local-first memory store with vector search
|
||||
//!
|
||||
//! Uses SQLite with sqlite-vec for in-database vector operations
|
||||
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::{sqlite::SqlitePoolOptions, Pool, Sqlite};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::SimulationResult;
|
||||
|
||||
/// User context loaded from SQL brain
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserContext {
|
||||
pub user_id: String,
|
||||
pub episodes: Vec<Episode>,
|
||||
pub scene_memories: Vec<SceneMemory>,
|
||||
pub preferences: UserPreferences,
|
||||
pub delta_weights: Option<DeltaWeights>,
|
||||
}
|
||||
|
||||
/// Episode memory
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Episode {
|
||||
pub id: String,
|
||||
pub timestamp: i64,
|
||||
pub question: String,
|
||||
pub answer: String,
|
||||
pub confidence: f64,
|
||||
pub trajectory: Vec<StateTransition>,
|
||||
pub reward: f64,
|
||||
}
|
||||
|
||||
/// State transition in episode
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StateTransition {
|
||||
pub state: serde_json::Value,
|
||||
pub action: serde_json::Value,
|
||||
pub next_state: serde_json::Value,
|
||||
pub reward: f64,
|
||||
}
|
||||
|
||||
/// Scene memory with embeddings
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SceneMemory {
|
||||
pub scene_id: String,
|
||||
pub description: String,
|
||||
pub embedding: Vec<f32>,
|
||||
pub objects: Vec<ObjectMemory>,
|
||||
pub affordances: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ObjectMemory {
|
||||
pub object_id: String,
|
||||
pub class: String,
|
||||
pub properties: serde_json::Value,
|
||||
pub embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
/// User preferences
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserPreferences {
|
||||
pub risk_tolerance: f64,
|
||||
pub exploration_rate: f64,
|
||||
pub preferred_tier: Option<String>,
|
||||
}
|
||||
|
||||
/// BitDelta 1-bit weight deltas for per-user adaptation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DeltaWeights {
|
||||
pub base_model: String,
|
||||
pub delta_bits: Vec<u8>, // 1-bit deltas packed
|
||||
pub scale: f32, // Learned scale factor
|
||||
pub sparsity_mask: Vec<bool>, // Which weights to update
|
||||
}
|
||||
|
||||
/// SQL Brain implementation
|
||||
pub struct SqlBrain {
|
||||
pool: Pool<Sqlite>,
|
||||
}
|
||||
|
||||
impl SqlBrain {
|
||||
pub async fn new(path: &str) -> Result<Self> {
|
||||
// Create database if it doesn't exist
|
||||
if !Path::new(path).exists() {
|
||||
tokio::fs::File::create(path).await?;
|
||||
}
|
||||
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(5)
|
||||
.connect(&format!("sqlite:{}", path))
|
||||
.await?;
|
||||
|
||||
// Initialize schema
|
||||
Self::init_schema(&pool).await?;
|
||||
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
async fn init_schema(pool: &Pool<Sqlite>) -> Result<()> {
|
||||
// Create tables
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
user_id TEXT PRIMARY KEY,
|
||||
created_at INTEGER NOT NULL,
|
||||
preferences TEXT NOT NULL
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS episodes (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
timestamp INTEGER NOT NULL,
|
||||
question TEXT NOT NULL,
|
||||
answer TEXT NOT NULL,
|
||||
confidence REAL NOT NULL,
|
||||
trajectory TEXT NOT NULL,
|
||||
reward REAL NOT NULL,
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS scenes (
|
||||
scene_id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
description TEXT NOT NULL,
|
||||
embedding BLOB NOT NULL,
|
||||
objects TEXT NOT NULL,
|
||||
affordances TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS objects (
|
||||
object_id TEXT PRIMARY KEY,
|
||||
scene_id TEXT NOT NULL,
|
||||
class TEXT NOT NULL,
|
||||
properties TEXT NOT NULL,
|
||||
embedding BLOB NOT NULL,
|
||||
FOREIGN KEY (scene_id) REFERENCES scenes(scene_id)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS delta_weights (
|
||||
user_id TEXT PRIMARY KEY,
|
||||
base_model TEXT NOT NULL,
|
||||
delta_bits BLOB NOT NULL,
|
||||
scale REAL NOT NULL,
|
||||
sparsity_mask BLOB NOT NULL,
|
||||
updated_at INTEGER NOT NULL,
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
// Create vector index using sqlite-vec virtual table
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS scene_embeddings USING vec0(
|
||||
scene_id TEXT PRIMARY KEY,
|
||||
embedding FLOAT[768]
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS object_embeddings USING vec0(
|
||||
object_id TEXT PRIMARY KEY,
|
||||
embedding FLOAT[768]
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load user context
|
||||
pub async fn load_user_context(&self, user_id: &str) -> Result<UserContext> {
|
||||
// Load user preferences
|
||||
let preferences: UserPreferences = sqlx::query_as!(
|
||||
UserPreferencesRow,
|
||||
"SELECT preferences FROM users WHERE user_id = ?",
|
||||
user_id
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
.map(|row| serde_json::from_str(&row.preferences).unwrap())
|
||||
.unwrap_or_default();
|
||||
|
||||
// Load recent episodes
|
||||
let episodes = sqlx::query!(
|
||||
r#"
|
||||
SELECT id, timestamp, question, answer, confidence, trajectory, reward
|
||||
FROM episodes
|
||||
WHERE user_id = ?
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 100
|
||||
"#,
|
||||
user_id
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| Episode {
|
||||
id: row.id,
|
||||
timestamp: row.timestamp,
|
||||
question: row.question,
|
||||
answer: row.answer,
|
||||
confidence: row.confidence,
|
||||
trajectory: serde_json::from_str(&row.trajectory).unwrap(),
|
||||
reward: row.reward,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Load scene memories
|
||||
let scene_memories = sqlx::query!(
|
||||
r#"
|
||||
SELECT scene_id, description, embedding, objects, affordances
|
||||
FROM scenes
|
||||
WHERE user_id = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 10
|
||||
"#,
|
||||
user_id
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| SceneMemory {
|
||||
scene_id: row.scene_id,
|
||||
description: row.description,
|
||||
embedding: bincode::deserialize(&row.embedding).unwrap(),
|
||||
objects: serde_json::from_str(&row.objects).unwrap(),
|
||||
affordances: serde_json::from_str(&row.affordances).unwrap(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Load delta weights if they exist
|
||||
let delta_weights = sqlx::query!(
|
||||
r#"
|
||||
SELECT base_model, delta_bits, scale, sparsity_mask
|
||||
FROM delta_weights
|
||||
WHERE user_id = ?
|
||||
"#,
|
||||
user_id
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
.map(|row| DeltaWeights {
|
||||
base_model: row.base_model,
|
||||
delta_bits: row.delta_bits,
|
||||
scale: row.scale,
|
||||
sparsity_mask: bincode::deserialize(&row.sparsity_mask).unwrap(),
|
||||
});
|
||||
|
||||
Ok(UserContext {
|
||||
user_id: user_id.to_string(),
|
||||
episodes,
|
||||
scene_memories,
|
||||
preferences,
|
||||
delta_weights,
|
||||
})
|
||||
}
|
||||
|
||||
/// Store simulation episode
|
||||
pub async fn store_episode(&self, user_id: &str, result: &SimulationResult) -> Result<()> {
|
||||
let episode_id = format!("ep_{}", uuid::Uuid::new_v4());
|
||||
let timestamp = chrono::Utc::now().timestamp();
|
||||
|
||||
let trajectory = serde_json::to_string(&Vec::<StateTransition>::new())?; // Placeholder
|
||||
|
||||
sqlx::query!(
|
||||
r#"
|
||||
INSERT INTO episodes (id, user_id, timestamp, question, answer, confidence, trajectory, reward)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
episode_id,
|
||||
user_id,
|
||||
timestamp,
|
||||
"", // Question from request
|
||||
result.answer,
|
||||
result.confidence,
|
||||
trajectory,
|
||||
result.metrics.avg_reward
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Vector search for similar scenes
|
||||
pub async fn search_similar_scenes(
|
||||
&self,
|
||||
embedding: &[f32],
|
||||
limit: usize,
|
||||
) -> Result<Vec<SceneMemory>> {
|
||||
// Use sqlite-vec for efficient vector search
|
||||
let query = format!(
|
||||
r#"
|
||||
SELECT s.scene_id, s.description, s.embedding, s.objects, s.affordances
|
||||
FROM scenes s
|
||||
JOIN scene_embeddings se ON s.scene_id = se.scene_id
|
||||
WHERE se.embedding MATCH ?
|
||||
ORDER BY distance
|
||||
LIMIT {}
|
||||
"#,
|
||||
limit
|
||||
);
|
||||
|
||||
let results = sqlx::query(&query)
|
||||
.bind(bincode::serialize(embedding)?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(results
|
||||
.into_iter()
|
||||
.map(|row| SceneMemory {
|
||||
scene_id: row.get("scene_id"),
|
||||
description: row.get("description"),
|
||||
embedding: bincode::deserialize(row.get("embedding")).unwrap(),
|
||||
objects: serde_json::from_str(row.get("objects")).unwrap(),
|
||||
affordances: serde_json::from_str(row.get("affordances")).unwrap(),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Update BitDelta weights for user
|
||||
pub async fn update_delta_weights(&self, user_id: &str, delta: DeltaWeights) -> Result<()> {
|
||||
let timestamp = chrono::Utc::now().timestamp();
|
||||
let sparsity_mask_bytes = bincode::serialize(&delta.sparsity_mask)?;
|
||||
|
||||
sqlx::query!(
|
||||
r#"
|
||||
INSERT OR REPLACE INTO delta_weights
|
||||
(user_id, base_model, delta_bits, scale, sparsity_mask, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
user_id,
|
||||
delta.base_model,
|
||||
delta.delta_bits,
|
||||
delta.scale,
|
||||
sparsity_mask_bytes,
|
||||
timestamp
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compute BitDelta from recent episodes
|
||||
pub async fn compute_bitdelta(&self, user_id: &str) -> Result<DeltaWeights> {
|
||||
// Load recent episodes
|
||||
let episodes = sqlx::query!(
|
||||
r#"
|
||||
SELECT trajectory, reward
|
||||
FROM episodes
|
||||
WHERE user_id = ?
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 50
|
||||
"#,
|
||||
user_id
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
// Convert trajectories to training data
|
||||
let mut training_data = Vec::new();
|
||||
for episode in episodes {
|
||||
let trajectory: Vec<StateTransition> = serde_json::from_str(&episode.trajectory)?;
|
||||
training_data.push((trajectory, episode.reward));
|
||||
}
|
||||
|
||||
// Compute 1-bit deltas (placeholder implementation)
|
||||
// In reality, this would:
|
||||
// 1. Fine-tune on the training data
|
||||
// 2. Compute weight differences
|
||||
// 3. Quantize to 1-bit with learned scale
|
||||
// 4. Apply sparsity mask
|
||||
|
||||
let delta_weights = DeltaWeights {
|
||||
base_model: "qwen3-30b-a3b".to_string(),
|
||||
delta_bits: vec![0xFF; 1024], // Placeholder
|
||||
scale: 0.01,
|
||||
sparsity_mask: vec![true; 8192], // Placeholder
|
||||
};
|
||||
|
||||
Ok(delta_weights)
|
||||
}
|
||||
}
|
||||
|
||||
// Helper structs for SQLx
|
||||
struct UserPreferencesRow {
|
||||
preferences: String,
|
||||
}
|
||||
|
||||
impl Default for UserPreferences {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
risk_tolerance: 0.5,
|
||||
exploration_rate: 0.1,
|
||||
preferred_tier: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// UUID placeholder
|
||||
mod uuid {
|
||||
pub struct Uuid;
|
||||
impl Uuid {
|
||||
pub fn new_v4() -> String {
|
||||
format!("{:x}", rand::random::<u128>())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Random placeholder
|
||||
mod rand {
|
||||
pub fn random<T>() -> T
|
||||
where
|
||||
T: Default,
|
||||
{
|
||||
T::default()
|
||||
}
|
||||
}
|
||||
@@ -1,444 +0,0 @@
|
||||
//! Tier A: V-JEPA2 Latent Simulation
|
||||
//!
|
||||
//! Cheapest tier using latent world models for mental simulation
|
||||
|
||||
use anyhow::Result;
|
||||
use candle_core::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::{SimulationRequest, SimulationResult, SimulationMetrics, SimulationCost, SimulationArtifact, ArtifactType};
|
||||
|
||||
/// Configuration for latent simulator
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LatentConfig {
|
||||
pub model_path: String,
|
||||
pub embedding_dim: usize,
|
||||
pub horizon: usize,
|
||||
pub num_rollouts: usize,
|
||||
pub temperature: f32,
|
||||
}
|
||||
|
||||
impl Default for LatentConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model_path: "models/v-jepa2-ac.onnx".to_string(),
|
||||
embedding_dim: 768,
|
||||
horizon: 50,
|
||||
num_rollouts: 100,
|
||||
temperature: 0.7,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// V-JEPA2-AC latent world model
|
||||
pub struct VJepa2Model {
|
||||
encoder: ort::Session,
|
||||
predictor: ort::Session,
|
||||
action_head: ort::Session,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl VJepa2Model {
|
||||
pub async fn load(model_path: &str) -> Result<Self> {
|
||||
// Load ONNX models for V-JEPA2 components
|
||||
let encoder = ort::Session::builder()?
|
||||
.with_model_from_file(format!("{}/encoder.onnx", model_path))?;
|
||||
|
||||
let predictor = ort::Session::builder()?
|
||||
.with_model_from_file(format!("{}/predictor.onnx", model_path))?;
|
||||
|
||||
let action_head = ort::Session::builder()?
|
||||
.with_model_from_file(format!("{}/action_head.onnx", model_path))?;
|
||||
|
||||
Ok(Self {
|
||||
encoder,
|
||||
predictor,
|
||||
action_head,
|
||||
device: Device::Cpu,
|
||||
})
|
||||
}
|
||||
|
||||
/// Encode scene to latent space
|
||||
pub async fn encode(&self, scene_state: &SceneState) -> Result<LatentState> {
|
||||
// Convert scene to tensor
|
||||
let input_tensor = self.scene_to_tensor(scene_state)?;
|
||||
|
||||
// Run encoder
|
||||
let outputs = self.encoder.run(ort::inputs!["input" => input_tensor]?)?;
|
||||
let latent = outputs["latent"].try_extract_tensor::<f32>()?;
|
||||
|
||||
Ok(LatentState {
|
||||
embedding: latent.to_vec(),
|
||||
timestamp: scene_state.timestamp,
|
||||
})
|
||||
}
|
||||
|
||||
/// Predict future latent states given actions
|
||||
pub async fn rollout(
|
||||
&self,
|
||||
initial_state: &LatentState,
|
||||
action_sequence: &[Action],
|
||||
horizon: usize,
|
||||
) -> Result<Vec<LatentState>> {
|
||||
let mut states = vec![initial_state.clone()];
|
||||
let mut current = Tensor::from_slice(
|
||||
&initial_state.embedding,
|
||||
&[1, initial_state.embedding.len()],
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
for (t, action) in action_sequence.iter().take(horizon).enumerate() {
|
||||
// Prepare action tensor
|
||||
let action_tensor = self.action_to_tensor(action)?;
|
||||
|
||||
// Run predictor
|
||||
let outputs = self.predictor.run(ort::inputs![
|
||||
"state" => current.clone(),
|
||||
"action" => action_tensor
|
||||
]?)?;
|
||||
|
||||
let next_latent = outputs["next_state"].try_extract_tensor::<f32>()?;
|
||||
|
||||
states.push(LatentState {
|
||||
embedding: next_latent.to_vec(),
|
||||
timestamp: initial_state.timestamp + t as f64 + 1.0,
|
||||
});
|
||||
|
||||
current = Tensor::from_slice(
|
||||
&next_latent,
|
||||
&[1, next_latent.len()],
|
||||
&self.device,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(states)
|
||||
}
|
||||
|
||||
/// Inverse model: predict action from state transitions
|
||||
pub async fn inverse_action(
|
||||
&self,
|
||||
state_t: &LatentState,
|
||||
state_t1: &LatentState,
|
||||
) -> Result<Action> {
|
||||
let s_t = Tensor::from_slice(
|
||||
&state_t.embedding,
|
||||
&[1, state_t.embedding.len()],
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
let s_t1 = Tensor::from_slice(
|
||||
&state_t1.embedding,
|
||||
&[1, state_t1.embedding.len()],
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
let outputs = self.action_head.run(ort::inputs![
|
||||
"state_t" => s_t,
|
||||
"state_t1" => s_t1
|
||||
]?)?;
|
||||
|
||||
let action_logits = outputs["action"].try_extract_tensor::<f32>()?;
|
||||
|
||||
// Decode action from logits
|
||||
self.decode_action(action_logits)
|
||||
}
|
||||
|
||||
fn scene_to_tensor(&self, scene: &SceneState) -> Result<Tensor> {
|
||||
// Convert scene graph to tensor representation
|
||||
// This would flatten object positions, velocities, etc.
|
||||
let mut features = Vec::new();
|
||||
|
||||
for object in &scene.objects {
|
||||
features.extend(&object.position);
|
||||
features.extend(&object.velocity);
|
||||
features.push(object.mass);
|
||||
}
|
||||
|
||||
Ok(Tensor::from_slice(
|
||||
&features,
|
||||
&[1, features.len()],
|
||||
&self.device,
|
||||
)?)
|
||||
}
|
||||
|
||||
fn action_to_tensor(&self, action: &Action) -> Result<Tensor> {
|
||||
let action_vec = match action {
|
||||
Action::Move(dx, dy, dz) => vec![1.0, *dx, *dy, *dz],
|
||||
Action::Rotate(rx, ry, rz) => vec![2.0, *rx, *ry, *rz],
|
||||
Action::Grasp(force) => vec![3.0, *force, 0.0, 0.0],
|
||||
Action::Release => vec![4.0, 0.0, 0.0, 0.0],
|
||||
};
|
||||
|
||||
Ok(Tensor::from_slice(&action_vec, &[1, 4], &self.device)?)
|
||||
}
|
||||
|
||||
fn decode_action(&self, logits: &[f32]) -> Result<Action> {
|
||||
// Simple argmax decoding
|
||||
let action_type = logits[0] as i32;
|
||||
|
||||
Ok(match action_type {
|
||||
1 => Action::Move(logits[1], logits[2], logits[3]),
|
||||
2 => Action::Rotate(logits[1], logits[2], logits[3]),
|
||||
3 => Action::Grasp(logits[1]),
|
||||
_ => Action::Release,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Latent state representation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LatentState {
|
||||
pub embedding: Vec<f32>,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
/// Scene state for encoding
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SceneState {
|
||||
pub objects: Vec<Object>,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Object {
|
||||
pub id: String,
|
||||
pub position: [f32; 3],
|
||||
pub velocity: [f32; 3],
|
||||
pub mass: f32,
|
||||
}
|
||||
|
||||
/// Action space
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum Action {
|
||||
Move(f32, f32, f32), // dx, dy, dz
|
||||
Rotate(f32, f32, f32), // rx, ry, rz
|
||||
Grasp(f32), // force
|
||||
Release,
|
||||
}
|
||||
|
||||
/// Latent simulator
|
||||
pub struct LatentSimulator {
|
||||
config: LatentConfig,
|
||||
model: Arc<Mutex<VJepa2Model>>,
|
||||
}
|
||||
|
||||
impl LatentSimulator {
|
||||
pub async fn new(config: LatentConfig) -> Result<Self> {
|
||||
let model = VJepa2Model::load(&config.model_path).await?;
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
model: Arc::new(Mutex::new(model)),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn simulate(
|
||||
&self,
|
||||
request: &SimulationRequest,
|
||||
session_layer: &str,
|
||||
) -> Result<SimulationResult> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Parse scene from USD
|
||||
let scene_state = self.load_scene_state(session_layer).await?;
|
||||
|
||||
// Encode to latent space
|
||||
let model = self.model.lock().await;
|
||||
let initial_latent = model.encode(&scene_state).await?;
|
||||
|
||||
// Generate action sequences for the question
|
||||
let action_sequences = self.generate_action_hypotheses(&request.question).await?;
|
||||
|
||||
// Run rollouts in latent space
|
||||
let mut trajectories = Vec::new();
|
||||
let mut success_count = 0;
|
||||
|
||||
for actions in action_sequences.iter().take(self.config.num_rollouts) {
|
||||
let trajectory = model.rollout(
|
||||
&initial_latent,
|
||||
actions,
|
||||
self.config.horizon,
|
||||
).await?;
|
||||
|
||||
// Evaluate trajectory
|
||||
if self.evaluate_trajectory(&trajectory, &request.question).await? {
|
||||
success_count += 1;
|
||||
}
|
||||
|
||||
trajectories.push(trajectory);
|
||||
}
|
||||
|
||||
let success_rate = success_count as f64 / self.config.num_rollouts as f64;
|
||||
let confidence = self.calculate_confidence(success_rate, trajectories.len());
|
||||
|
||||
// Generate answer
|
||||
let answer = self.generate_answer(&request.question, success_rate, &trajectories).await?;
|
||||
|
||||
// Create visualization artifact (latent trajectory plot)
|
||||
let artifact = self.create_latent_visualization(&trajectories).await?;
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
Ok(SimulationResult {
|
||||
answer,
|
||||
confidence,
|
||||
tier_used: crate::SimulationTier::Latent,
|
||||
artifacts: vec![artifact],
|
||||
metrics: SimulationMetrics {
|
||||
episodes_run: self.config.num_rollouts as u32,
|
||||
success_rate,
|
||||
avg_reward: 0.0, // Not applicable for latent
|
||||
collision_count: 0,
|
||||
physics_violations: 0,
|
||||
},
|
||||
cost: SimulationCost {
|
||||
compute_seconds: elapsed.as_secs_f64(),
|
||||
gpu_seconds: 0.0, // CPU only for now
|
||||
memory_gb_hours: 0.001,
|
||||
estimated_usd: 0.001,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_scene_state(&self, session_layer: &str) -> Result<SceneState> {
|
||||
// Load from USD session layer
|
||||
// This would parse the USD file and extract object states
|
||||
Ok(SceneState {
|
||||
objects: vec![
|
||||
Object {
|
||||
id: "agent_1".to_string(),
|
||||
position: [0.0, 0.0, 0.0],
|
||||
velocity: [0.0, 0.0, 0.0],
|
||||
mass: 1.0,
|
||||
},
|
||||
Object {
|
||||
id: "agent_2".to_string(),
|
||||
position: [2.0, 0.0, 0.0],
|
||||
velocity: [-0.5, 0.0, 0.0],
|
||||
mass: 1.0,
|
||||
},
|
||||
],
|
||||
timestamp: 0.0,
|
||||
})
|
||||
}
|
||||
|
||||
async fn generate_action_hypotheses(&self, question: &str) -> Result<Vec<Vec<Action>>> {
|
||||
// Generate diverse action sequences based on the question
|
||||
// This would use an LLM or planning algorithm
|
||||
|
||||
let mut sequences = Vec::new();
|
||||
|
||||
for _ in 0..self.config.num_rollouts {
|
||||
let mut seq = Vec::new();
|
||||
for _ in 0..self.config.horizon {
|
||||
// Random exploration for now
|
||||
seq.push(Action::Move(
|
||||
rand::random::<f32>() - 0.5,
|
||||
rand::random::<f32>() - 0.5,
|
||||
0.0,
|
||||
));
|
||||
}
|
||||
sequences.push(seq);
|
||||
}
|
||||
|
||||
Ok(sequences)
|
||||
}
|
||||
|
||||
async fn evaluate_trajectory(
|
||||
&self,
|
||||
trajectory: &[LatentState],
|
||||
question: &str,
|
||||
) -> Result<bool> {
|
||||
// Evaluate if trajectory answers the question
|
||||
// This would use a learned evaluator or heuristics
|
||||
|
||||
// Simple collision detection in latent space
|
||||
if question.contains("collision") {
|
||||
// Check if latent states get too similar (indicating collision)
|
||||
for window in trajectory.windows(2) {
|
||||
let dist = self.latent_distance(&window[0], &window[1]);
|
||||
if dist < 0.1 {
|
||||
return Ok(true); // Collision detected
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn latent_distance(&self, s1: &LatentState, s2: &LatentState) -> f32 {
|
||||
// L2 distance in latent space
|
||||
s1.embedding
|
||||
.iter()
|
||||
.zip(&s2.embedding)
|
||||
.map(|(a, b)| (a - b).powi(2))
|
||||
.sum::<f32>()
|
||||
.sqrt()
|
||||
}
|
||||
|
||||
fn calculate_confidence(&self, success_rate: f64, num_samples: usize) -> f64 {
|
||||
// Confidence based on success rate and sample size
|
||||
let base_confidence = success_rate;
|
||||
let sample_factor = (num_samples as f64 / 100.0).min(1.0);
|
||||
base_confidence * sample_factor
|
||||
}
|
||||
|
||||
async fn generate_answer(
|
||||
&self,
|
||||
question: &str,
|
||||
success_rate: f64,
|
||||
trajectories: &[Vec<LatentState>],
|
||||
) -> Result<String> {
|
||||
if question.contains("collision") {
|
||||
if success_rate > 0.5 {
|
||||
Ok(format!(
|
||||
"Collision likely ({:.1}% probability based on {} rollouts)",
|
||||
success_rate * 100.0,
|
||||
trajectories.len()
|
||||
))
|
||||
} else {
|
||||
Ok(format!(
|
||||
"Collision unlikely ({:.1}% probability based on {} rollouts)",
|
||||
success_rate * 100.0,
|
||||
trajectories.len()
|
||||
))
|
||||
}
|
||||
} else {
|
||||
Ok(format!(
|
||||
"Simulation complete. Success rate: {:.1}% over {} episodes",
|
||||
success_rate * 100.0,
|
||||
trajectories.len()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_latent_visualization(
|
||||
&self,
|
||||
trajectories: &[Vec<LatentState>],
|
||||
) -> Result<SimulationArtifact> {
|
||||
// Create a visualization of latent trajectories
|
||||
// This would generate a plot or animation
|
||||
|
||||
Ok(SimulationArtifact {
|
||||
artifact_type: ArtifactType::Trajectory,
|
||||
uri: "data:application/json,{}".to_string(), // Placeholder
|
||||
metadata: serde_json::json!({
|
||||
"type": "latent_trajectories",
|
||||
"num_trajectories": trajectories.len(),
|
||||
"horizon": self.config.horizon,
|
||||
}),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Placeholder rand module until we add the crate
|
||||
mod rand {
|
||||
pub fn random<T>() -> T
|
||||
where
|
||||
T: Default,
|
||||
{
|
||||
T::default()
|
||||
}
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-sovereign"
|
||||
version = "1.1.21"
|
||||
edition = "2021"
|
||||
description = "Hanzo Network — L1 sovereign blockchain with ML-KEM / ML-DSA primitives"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/hanzonet/network"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1.0", features = ["full"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
bincode = "1.3"
|
||||
sha2 = "0.10"
|
||||
log = "0.4"
|
||||
ml-kem = "0.2"
|
||||
ml-dsa = "0.6"
|
||||
chacha20poly1305 = "0.10"
|
||||
rand = "0.8"
|
||||
@@ -1,171 +0,0 @@
|
||||
// Hanzo Sovereign Chain with ML-KEM/ML-DSA FIPS standards
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::mpsc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Hanzo L1 Sovereign Blockchain Node
|
||||
pub struct HanzoSovereignNode {
|
||||
chain_id: u64,
|
||||
node_id: String,
|
||||
|
||||
// Transaction pool
|
||||
mempool: Arc<Mutex<Vec<Transaction>>>,
|
||||
|
||||
// Event channels
|
||||
tx_events: mpsc::Sender<TransactionEvent>,
|
||||
}
|
||||
|
||||
impl HanzoSovereignNode {
|
||||
/// Initialize Hanzo L1 sovereign chain
|
||||
pub async fn new(config: SovereignConfig) -> Result<Self, HanzoError> {
|
||||
// Create event channels
|
||||
let (tx_events, _) = mpsc::channel(10000);
|
||||
|
||||
Ok(Self {
|
||||
chain_id: config.chain_id,
|
||||
node_id: config.node_id,
|
||||
mempool: Arc::new(Mutex::new(Vec::new())),
|
||||
tx_events,
|
||||
})
|
||||
}
|
||||
|
||||
/// Start the sovereign chain
|
||||
pub async fn start(&mut self) -> Result<(), HanzoError> {
|
||||
log::info!("🚀 Starting Hanzo L1 Sovereign Chain");
|
||||
log::info!("✅ Hanzo L1 Sovereign Chain started successfully");
|
||||
log::info!(" Chain ID: {}", self.chain_id);
|
||||
log::info!(" Node ID: {}", self.node_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Submit transaction to mempool
|
||||
pub async fn submit_transaction(&mut self, tx: Transaction) -> Result<String, HanzoError> {
|
||||
let tx_hash = tx.hash();
|
||||
self.mempool.lock().unwrap().push(tx.clone());
|
||||
|
||||
self.tx_events.send(TransactionEvent {
|
||||
tx_hash: tx_hash.clone(),
|
||||
tx_type: tx.tx_type(),
|
||||
status: TransactionStatus::Pending,
|
||||
}).await?;
|
||||
|
||||
log::info!("📨 Transaction submitted: {}", tx_hash);
|
||||
Ok(tx_hash)
|
||||
}
|
||||
|
||||
/// Get node status
|
||||
pub async fn get_status(&self) -> NodeStatus {
|
||||
NodeStatus {
|
||||
chain_id: self.chain_id,
|
||||
node_id: self.node_id.clone(),
|
||||
block_height: 1,
|
||||
peer_count: 0,
|
||||
mempool_size: self.mempool.lock().unwrap().len(),
|
||||
pq_commerce_active: true,
|
||||
mesh_intelligence_active: true,
|
||||
lux_bridge_connected: true,
|
||||
zoo_sequencer_active: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Configuration and data structures
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SovereignConfig {
|
||||
pub chain_id: u64,
|
||||
pub node_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Transaction {
|
||||
pub nonce: u64,
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
pub value: u64,
|
||||
pub data: Vec<u8>,
|
||||
pub signature: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
pub fn hash(&self) -> String {
|
||||
use sha2::{Sha256, Digest};
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(&self.to_bytes());
|
||||
format!("0x{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
pub fn to_bytes(&self) -> Vec<u8> {
|
||||
bincode::serialize(self).unwrap()
|
||||
}
|
||||
|
||||
pub fn tx_type(&self) -> TransactionType {
|
||||
if self.data.is_empty() {
|
||||
TransactionType::Transfer
|
||||
} else {
|
||||
TransactionType::Contract
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NodeStatus {
|
||||
pub chain_id: u64,
|
||||
pub node_id: String,
|
||||
pub block_height: u64,
|
||||
pub peer_count: usize,
|
||||
pub mempool_size: usize,
|
||||
pub pq_commerce_active: bool,
|
||||
pub mesh_intelligence_active: bool,
|
||||
pub lux_bridge_connected: bool,
|
||||
pub zoo_sequencer_active: bool,
|
||||
}
|
||||
|
||||
// Event types
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TransactionEvent {
|
||||
pub tx_hash: String,
|
||||
pub tx_type: TransactionType,
|
||||
pub status: TransactionStatus,
|
||||
}
|
||||
|
||||
// Enums
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TransactionType {
|
||||
Transfer,
|
||||
Contract,
|
||||
PQCommerce,
|
||||
CrossChain,
|
||||
Intelligence,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TransactionStatus {
|
||||
Pending,
|
||||
Confirmed,
|
||||
Failed,
|
||||
}
|
||||
|
||||
// Error handling
|
||||
#[derive(Debug)]
|
||||
pub enum HanzoError {
|
||||
InvalidTransaction,
|
||||
NetworkError(String),
|
||||
}
|
||||
|
||||
impl std::error::Error for HanzoError {}
|
||||
|
||||
impl std::fmt::Display for HanzoError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
HanzoError::InvalidTransaction => write!(f, "Invalid transaction"),
|
||||
HanzoError::NetworkError(msg) => write!(f, "Network error: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<tokio::sync::mpsc::error::SendError<TransactionEvent>> for HanzoError {
|
||||
fn from(_: tokio::sync::mpsc::error::SendError<TransactionEvent>) -> Self {
|
||||
HanzoError::NetworkError("Channel send error".to_string())
|
||||
}
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
[package]
|
||||
name = "hanzo-tests"
|
||||
description = "Hanzo AI - Tests Library"
|
||||
license = "MIT"
|
||||
version = "0.1.1"
|
||||
edition = "2021"
|
||||
publish = false # Not ready for crates.io
|
||||
|
||||
# Standalone crate - not part of parent workspace
|
||||
[workspace]
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1.38", features = ["full"] }
|
||||
anyhow = "1.0"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
tempfile = "3.10"
|
||||
criterion = "0.5"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = "0.3"
|
||||
|
||||
# Runtime dependencies
|
||||
hanzo-wasm = {path = "../hanzo-wasm", version = "0.1.1"}
|
||||
hanzo-tools = { path = "../hanzo-tools" }
|
||||
hanzo-messages = {path = "../hanzo-messages", version = "1.1.12"}
|
||||
hanzo-db-sqlite = {path = "../hanzo-db-sqlite", version = "1.1.12"}
|
||||
|
||||
# For Docker tests
|
||||
bollard = "0.17"
|
||||
|
||||
# For test assets
|
||||
include_dir = "0.7"
|
||||
wat = "1.217"
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = { version = "0.5", features = ["html_reports"] }
|
||||
proptest = "1.0"
|
||||
|
||||
@@ -1,307 +0,0 @@
|
||||
//! Unified runtime testing framework for Hanzo Node
|
||||
//!
|
||||
//! Tests all execution environments:
|
||||
//! - Native Rust
|
||||
//! - WASM
|
||||
//! - JavaScript/TypeScript (Deno)
|
||||
//! - Python
|
||||
//! - Docker containers
|
||||
//! - Go plugins
|
||||
//! - MCP servers
|
||||
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{info, warn, error};
|
||||
|
||||
pub mod native_tests;
|
||||
pub mod wasm_tests;
|
||||
pub mod deno_tests;
|
||||
pub mod python_tests;
|
||||
pub mod docker_tests;
|
||||
pub mod unified_tests;
|
||||
|
||||
/// Standard test payload for all runtimes
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TestPayload {
|
||||
pub operation: String,
|
||||
pub input_a: f64,
|
||||
pub input_b: f64,
|
||||
pub expected: f64,
|
||||
}
|
||||
|
||||
/// Test result with performance metrics
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TestResult {
|
||||
pub runtime: String,
|
||||
pub success: bool,
|
||||
pub result: Option<f64>,
|
||||
pub execution_time: Duration,
|
||||
pub memory_usage: Option<u64>,
|
||||
pub gas_consumed: Option<u64>,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Runtime test trait - all runtimes must implement this
|
||||
#[async_trait::async_trait]
|
||||
pub trait RuntimeTest: Send + Sync {
|
||||
/// Get the runtime name
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Execute a simple computation
|
||||
async fn execute_compute(&self, payload: TestPayload) -> Result<TestResult>;
|
||||
|
||||
/// Test file I/O capabilities
|
||||
async fn test_file_io(&self, content: Vec<u8>) -> Result<Vec<u8>>;
|
||||
|
||||
/// Test network capabilities
|
||||
async fn test_network(&self, url: &str) -> Result<String>;
|
||||
|
||||
/// Test crypto operations
|
||||
async fn test_crypto(&self, data: Vec<u8>) -> Result<Vec<u8>>;
|
||||
|
||||
/// Benchmark performance
|
||||
async fn benchmark(&self) -> Result<BenchmarkResults>;
|
||||
}
|
||||
|
||||
/// Benchmark results for runtime comparison
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BenchmarkResults {
|
||||
pub runtime: String,
|
||||
pub ops_per_second: f64,
|
||||
pub avg_latency_ms: f64,
|
||||
pub p99_latency_ms: f64,
|
||||
pub memory_efficiency: f64,
|
||||
pub startup_time_ms: f64,
|
||||
}
|
||||
|
||||
/// Unified test orchestrator
|
||||
pub struct RuntimeOrchestrator {
|
||||
runtimes: Vec<Box<dyn RuntimeTest>>,
|
||||
results: Vec<TestResult>,
|
||||
}
|
||||
|
||||
impl RuntimeOrchestrator {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
runtimes: vec![],
|
||||
results: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_runtime(&mut self, runtime: Box<dyn RuntimeTest>) {
|
||||
self.runtimes.push(runtime);
|
||||
}
|
||||
|
||||
/// Run all tests across all runtimes
|
||||
pub async fn run_all_tests(&mut self) -> Result<TestReport> {
|
||||
info!("Starting unified runtime tests");
|
||||
|
||||
let test_cases = vec![
|
||||
TestPayload {
|
||||
operation: "add".to_string(),
|
||||
input_a: 10.0,
|
||||
input_b: 20.0,
|
||||
expected: 30.0,
|
||||
},
|
||||
TestPayload {
|
||||
operation: "multiply".to_string(),
|
||||
input_a: 5.0,
|
||||
input_b: 7.0,
|
||||
expected: 35.0,
|
||||
},
|
||||
TestPayload {
|
||||
operation: "fibonacci".to_string(),
|
||||
input_a: 10.0,
|
||||
input_b: 0.0,
|
||||
expected: 55.0,
|
||||
},
|
||||
];
|
||||
|
||||
let mut report = TestReport::new();
|
||||
|
||||
for runtime in &self.runtimes {
|
||||
info!("Testing runtime: {}", runtime.name());
|
||||
|
||||
for test_case in &test_cases {
|
||||
let start = Instant::now();
|
||||
match runtime.execute_compute(test_case.clone()).await {
|
||||
Ok(result) => {
|
||||
report.add_result(result);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Test failed for {}: {}", runtime.name(), e);
|
||||
report.add_result(TestResult {
|
||||
runtime: runtime.name().to_string(),
|
||||
success: false,
|
||||
result: None,
|
||||
execution_time: start.elapsed(),
|
||||
memory_usage: None,
|
||||
gas_consumed: None,
|
||||
error: Some(e.to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run benchmarks
|
||||
match runtime.benchmark().await {
|
||||
Ok(bench) => report.add_benchmark(bench),
|
||||
Err(e) => warn!("Benchmark failed for {}: {}", runtime.name(), e),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
/// Compare runtime performance
|
||||
pub async fn compare_performance(&self) -> Result<PerformanceComparison> {
|
||||
let mut comparison = PerformanceComparison::new();
|
||||
|
||||
for runtime in &self.runtimes {
|
||||
let bench = runtime.benchmark().await?;
|
||||
comparison.add_runtime_benchmark(bench);
|
||||
}
|
||||
|
||||
comparison.analyze();
|
||||
Ok(comparison)
|
||||
}
|
||||
}
|
||||
|
||||
/// Test report with all results
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct TestReport {
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
pub results: Vec<TestResult>,
|
||||
pub benchmarks: Vec<BenchmarkResults>,
|
||||
pub summary: TestSummary,
|
||||
}
|
||||
|
||||
impl TestReport {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
timestamp: chrono::Utc::now(),
|
||||
results: vec![],
|
||||
benchmarks: vec![],
|
||||
summary: TestSummary::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn add_result(&mut self, result: TestResult) {
|
||||
self.results.push(result);
|
||||
self.update_summary();
|
||||
}
|
||||
|
||||
fn add_benchmark(&mut self, benchmark: BenchmarkResults) {
|
||||
self.benchmarks.push(benchmark);
|
||||
}
|
||||
|
||||
fn update_summary(&mut self) {
|
||||
self.summary.total_tests = self.results.len();
|
||||
self.summary.passed = self.results.iter().filter(|r| r.success).count();
|
||||
self.summary.failed = self.results.iter().filter(|r| !r.success).count();
|
||||
|
||||
if !self.results.is_empty() {
|
||||
self.summary.avg_execution_time = Duration::from_millis(
|
||||
(self.results.iter()
|
||||
.map(|r| r.execution_time.as_millis() as u64)
|
||||
.sum::<u64>() / self.results.len() as u64) as u64
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
pub struct TestSummary {
|
||||
pub total_tests: usize,
|
||||
pub passed: usize,
|
||||
pub failed: usize,
|
||||
pub avg_execution_time: Duration,
|
||||
}
|
||||
|
||||
/// Performance comparison across runtimes
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PerformanceComparison {
|
||||
pub runtimes: Vec<BenchmarkResults>,
|
||||
pub fastest_runtime: Option<String>,
|
||||
pub most_efficient: Option<String>,
|
||||
pub recommendations: Vec<String>,
|
||||
}
|
||||
|
||||
impl PerformanceComparison {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
runtimes: vec![],
|
||||
fastest_runtime: None,
|
||||
most_efficient: None,
|
||||
recommendations: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
fn add_runtime_benchmark(&mut self, bench: BenchmarkResults) {
|
||||
self.runtimes.push(bench);
|
||||
}
|
||||
|
||||
fn analyze(&mut self) {
|
||||
if let Some(fastest) = self.runtimes.iter()
|
||||
.min_by(|a, b| a.avg_latency_ms.partial_cmp(&b.avg_latency_ms).unwrap()) {
|
||||
self.fastest_runtime = Some(fastest.runtime.clone());
|
||||
}
|
||||
|
||||
if let Some(efficient) = self.runtimes.iter()
|
||||
.max_by(|a, b| a.memory_efficiency.partial_cmp(&b.memory_efficiency).unwrap()) {
|
||||
self.most_efficient = Some(efficient.runtime.clone());
|
||||
}
|
||||
|
||||
// Generate recommendations based on analysis
|
||||
self.generate_recommendations();
|
||||
}
|
||||
|
||||
fn generate_recommendations(&mut self) {
|
||||
// Analyze and provide recommendations
|
||||
for runtime in &self.runtimes {
|
||||
if runtime.ops_per_second > 10000.0 {
|
||||
self.recommendations.push(
|
||||
format!("{} is suitable for high-throughput workloads", runtime.runtime)
|
||||
);
|
||||
}
|
||||
if runtime.avg_latency_ms < 1.0 {
|
||||
self.recommendations.push(
|
||||
format!("{} is ideal for low-latency requirements", runtime.runtime)
|
||||
);
|
||||
}
|
||||
if runtime.memory_efficiency > 0.9 {
|
||||
self.recommendations.push(
|
||||
format!("{} has excellent memory efficiency", runtime.runtime)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_orchestrator_creation() {
|
||||
let orchestrator = RuntimeOrchestrator::new();
|
||||
assert_eq!(orchestrator.runtimes.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_payload_serialization() {
|
||||
let payload = TestPayload {
|
||||
operation: "test".to_string(),
|
||||
input_a: 1.0,
|
||||
input_b: 2.0,
|
||||
expected: 3.0,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&payload).unwrap();
|
||||
let decoded: TestPayload = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(decoded.operation, payload.operation);
|
||||
}
|
||||
}
|
||||
@@ -1,231 +0,0 @@
|
||||
//! WASM runtime tests
|
||||
|
||||
use crate::{RuntimeTest, TestPayload, TestResult, BenchmarkResults};
|
||||
use anyhow::Result;
|
||||
use hanzo_wasm_runtime::{WasmRuntime, WasmConfig};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
use wat::parse_str;
|
||||
|
||||
pub struct WasmRuntimeTest {
|
||||
runtime: Arc<WasmRuntime>,
|
||||
}
|
||||
|
||||
impl WasmRuntimeTest {
|
||||
pub async fn new() -> Result<Self> {
|
||||
let config = WasmConfig {
|
||||
max_memory_bytes: 128 * 1024 * 1024, // 128MB
|
||||
max_execution_time: Duration::from_secs(10),
|
||||
enable_wasi: false,
|
||||
fuel_limit: Some(1_000_000_000),
|
||||
};
|
||||
|
||||
let runtime = Arc::new(WasmRuntime::new(config)?);
|
||||
|
||||
// Load test modules
|
||||
Self::load_test_modules(&runtime).await?;
|
||||
|
||||
Ok(Self { runtime })
|
||||
}
|
||||
|
||||
async fn load_test_modules(runtime: &WasmRuntime) -> Result<()> {
|
||||
// Simple math module
|
||||
let math_wat = r#"
|
||||
(module
|
||||
(func $add (export "add") (param f64 f64) (result f64)
|
||||
local.get 0
|
||||
local.get 1
|
||||
f64.add
|
||||
)
|
||||
|
||||
(func $multiply (export "multiply") (param f64 f64) (result f64)
|
||||
local.get 0
|
||||
local.get 1
|
||||
f64.mul
|
||||
)
|
||||
|
||||
(func $fibonacci (export "fibonacci") (param i32) (result i32)
|
||||
(local i32 i32 i32)
|
||||
local.get 0
|
||||
i32.const 2
|
||||
i32.lt_s
|
||||
if
|
||||
local.get 0
|
||||
return
|
||||
end
|
||||
|
||||
i32.const 0
|
||||
local.set 1
|
||||
i32.const 1
|
||||
local.set 2
|
||||
|
||||
loop
|
||||
local.get 1
|
||||
local.get 2
|
||||
i32.add
|
||||
local.set 3
|
||||
local.get 2
|
||||
local.set 1
|
||||
local.get 3
|
||||
local.set 2
|
||||
|
||||
local.get 0
|
||||
i32.const 1
|
||||
i32.sub
|
||||
local.tee 0
|
||||
i32.const 1
|
||||
i32.gt_s
|
||||
br_if 0
|
||||
end
|
||||
|
||||
local.get 3
|
||||
)
|
||||
)
|
||||
"#;
|
||||
|
||||
let math_bytes = parse_str(math_wat)?;
|
||||
runtime.load_module("math".to_string(), math_bytes).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl RuntimeTest for WasmRuntimeTest {
|
||||
fn name(&self) -> &str {
|
||||
"WASM"
|
||||
}
|
||||
|
||||
async fn execute_compute(&self, payload: TestPayload) -> Result<TestResult> {
|
||||
let start = Instant::now();
|
||||
let fuel_before = 1_000_000_000u64;
|
||||
|
||||
let result = match payload.operation.as_str() {
|
||||
"add" => {
|
||||
// For now, return a mock result since WASM execution isn't fully implemented
|
||||
// TODO: Implement proper WASM function calling
|
||||
Some(payload.input_a + payload.input_b)
|
||||
}
|
||||
"multiply" => {
|
||||
Some(payload.input_a * payload.input_b)
|
||||
}
|
||||
"fibonacci" => {
|
||||
// Calculate fibonacci
|
||||
let n = payload.input_a as i32;
|
||||
Some(fibonacci(n) as f64)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let execution_time = start.elapsed();
|
||||
let gas_consumed = fuel_before - 900_000_000; // Mock gas consumption
|
||||
|
||||
Ok(TestResult {
|
||||
runtime: self.name().to_string(),
|
||||
success: result.is_some() && result == Some(payload.expected),
|
||||
result,
|
||||
execution_time,
|
||||
memory_usage: Some(1024 * 1024), // 1MB mock
|
||||
gas_consumed: Some(gas_consumed),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn test_file_io(&self, content: Vec<u8>) -> Result<Vec<u8>> {
|
||||
// WASM is sandboxed - no direct file I/O
|
||||
Err(anyhow::anyhow!("File I/O not supported in WASM sandbox"))
|
||||
}
|
||||
|
||||
async fn test_network(&self, _url: &str) -> Result<String> {
|
||||
// Network access requires host functions
|
||||
Ok("Network access requires host functions".to_string())
|
||||
}
|
||||
|
||||
async fn test_crypto(&self, data: Vec<u8>) -> Result<Vec<u8>> {
|
||||
// Simple hash mock - would use WASM crypto module
|
||||
Ok(data.into_iter().map(|b| b.wrapping_add(1)).collect())
|
||||
}
|
||||
|
||||
async fn benchmark(&self) -> Result<BenchmarkResults> {
|
||||
let iterations = 1000;
|
||||
let mut latencies = Vec::new();
|
||||
let start = Instant::now();
|
||||
|
||||
for _ in 0..iterations {
|
||||
let op_start = Instant::now();
|
||||
let _ = self.execute_compute(TestPayload {
|
||||
operation: "add".to_string(),
|
||||
input_a: 42.0,
|
||||
input_b: 58.0,
|
||||
expected: 100.0,
|
||||
}).await?;
|
||||
latencies.push(op_start.elapsed().as_millis() as f64);
|
||||
}
|
||||
|
||||
let total_time = start.elapsed();
|
||||
let ops_per_second = iterations as f64 / total_time.as_secs_f64();
|
||||
|
||||
latencies.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
let avg_latency_ms = latencies.iter().sum::<f64>() / latencies.len() as f64;
|
||||
let p99_latency_ms = latencies[(latencies.len() as f64 * 0.99) as usize];
|
||||
|
||||
Ok(BenchmarkResults {
|
||||
runtime: self.name().to_string(),
|
||||
ops_per_second,
|
||||
avg_latency_ms,
|
||||
p99_latency_ms,
|
||||
memory_efficiency: 0.95, // WASM has excellent memory efficiency
|
||||
startup_time_ms: 10.0, // Fast module loading
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn fibonacci(n: i32) -> i32 {
|
||||
if n <= 1 {
|
||||
return n;
|
||||
}
|
||||
let mut a = 0;
|
||||
let mut b = 1;
|
||||
for _ in 2..=n {
|
||||
let temp = a + b;
|
||||
a = b;
|
||||
b = temp;
|
||||
}
|
||||
b
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wasm_runtime_creation() {
|
||||
if std::env::var("CI").is_ok() {
|
||||
println!("Skipping WASM test in CI");
|
||||
return;
|
||||
}
|
||||
|
||||
let runtime = WasmRuntimeTest::new().await;
|
||||
assert!(runtime.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wasm_computation() {
|
||||
if std::env::var("CI").is_ok() {
|
||||
println!("Skipping WASM test in CI");
|
||||
return;
|
||||
}
|
||||
|
||||
let runtime = WasmRuntimeTest::new().await.unwrap();
|
||||
let payload = TestPayload {
|
||||
operation: "add".to_string(),
|
||||
input_a: 10.0,
|
||||
input_b: 20.0,
|
||||
expected: 30.0,
|
||||
};
|
||||
|
||||
let result = runtime.execute_compute(payload).await.unwrap();
|
||||
assert!(result.success);
|
||||
assert_eq!(result.result, Some(30.0));
|
||||
}
|
||||
}
|
||||
+3
-6
@@ -1,14 +1,11 @@
|
||||
[package]
|
||||
name = "hanzo-vm"
|
||||
description = "Lux VM interface for Hanzo L2 - wraps EVM execution"
|
||||
description = "Hanzo Network — EVM with precompiles for PQ verify, Quasar, AI inference, AI embedding (0x0101/0x0102/0x0201/0x0202)"
|
||||
license = "MIT"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
repository = "https://github.com/hanzonet/network"
|
||||
# Cannot publish to crates.io until hanzo-engine and luxprecompile-sys are published.
|
||||
# Also depends on the local net hanzo-pqc (publish = false).
|
||||
publish = false
|
||||
|
||||
[features]
|
||||
default = []
|
||||
@@ -35,10 +32,10 @@ hanzo-pqc = { workspace = true }
|
||||
# process-global and populated at startup by `hanzo-node` (or any other
|
||||
# binary embedding hanzo-vm). If no engine is registered the precompile
|
||||
# dispatchers revert — there is no in-tree fallback.
|
||||
hanzo-engine = { path = "../../engine/hanzo-engine" }
|
||||
hanzo-engine = { version = "0.6.0", path = "../../engine/hanzo-engine" }
|
||||
|
||||
# Canonical Go-backed precompile dispatch (Quasar, ML-DSA, etc.)
|
||||
luxprecompile-sys = { path = "../../../lux/precompile/bindings/rust" }
|
||||
luxprecompile-sys = "0.1"
|
||||
|
||||
# State storage (bundled SQLite for dev; MPT later)
|
||||
rusqlite = { workspace = true }
|
||||
|
||||
Reference in New Issue
Block a user