use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use anyhow::{Context, Result};
use chrono::Utc;
use kuzu::{Connection, Database, SystemConfig, Value};
use uuid::Uuid;
use crate::schema::{
Conversation, Entity, Memory, MemoryType, Relation, RelationType, SyncEntry, SyncNodeType,
SyncOp, SyncState,
};
use crate::store::Store;
pub struct KuzuStore {
db: Database,
machine_id: String,
sync_seq: AtomicU64,
}
impl KuzuStore {
pub fn open(path: &Path, machine_id: String) -> Result<Self> {
let db = Database::new(path, SystemConfig::default()).context("opening KùzuDB database")?;
let store = Self {
db,
machine_id,
sync_seq: AtomicU64::new(1),
};
store.init_schema()?;
store.init_sync_seq()?;
Ok(store)
}
pub fn in_memory(machine_id: String) -> Result<Self> {
let db =
Database::in_memory(SystemConfig::default()).context("creating in-memory KùzuDB")?;
let store = Self {
db,
machine_id,
sync_seq: AtomicU64::new(1),
};
store.init_schema()?;
Ok(store)
}
pub fn machine_id(&self) -> &str {
&self.machine_id
}
fn conn(&self) -> Result<Connection<'_>> {
Connection::new(&self.db).context("creating connection")
}
pub fn diagnostic(&self) -> Result<String> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (m:Memory) RETURN count(m);")?;
let total: i64 = result
.next()
.map(|r| match &r[0] {
Value::Int64(v) => *v,
_ => -1,
})
.unwrap_or(-1);
let mut result = conn.query(
"MATCH (m:Memory) WHERE size(m.embedding) > 0 AND m.embedding[1] <> 0.0 RETURN count(m);",
)?;
let with_emb: i64 = result
.next()
.map(|r| match &r[0] {
Value::Int64(v) => *v,
_ => -1,
})
.unwrap_or(-1);
let zeros: Vec<String> = (0..384).map(|_| "0.0".to_string()).collect();
let zeros_str = format!("[{}]", zeros.join(","));
let idx_query = format!(
"CALL QUERY_VECTOR_INDEX('Memory', 'memory_emb_idx', {}, 1) YIELD node, distance RETURN distance;",
zeros_str
);
let idx_status = match conn.query(&idx_query) {
Ok(mut r) => match r.next() {
Some(row) => format!("ok (distance: {:?})", &row[0]),
None => "ok (0 results)".to_string(),
},
Err(e) => format!("error: {e}"),
};
Ok(format!(
"total memories: {total}\nwith embeddings: {with_emb}\nvector index: {idx_status}"
))
}
fn init_schema(&self) -> Result<()> {
let conn = self.conn()?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS Memory(
id STRING PRIMARY KEY,
content STRING,
embedding FLOAT[384],
memory_type STRING,
confidence FLOAT,
created_at STRING,
last_accessed STRING,
access_count INT64,
source STRING,
source_id STRING
);",
)
.context("creating Memory table")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS Entity(
id STRING PRIMARY KEY,
name STRING,
entity_type STRING,
embedding FLOAT[384],
aliases STRING[]
);",
)
.context("creating Entity table")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS Conversation(
id STRING PRIMARY KEY,
source STRING,
machine_id STRING,
started_at STRING,
project_path STRING
);",
)
.context("creating Conversation table")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS IngestLog(
file_path STRING PRIMARY KEY,
file_hash STRING,
ingested_at STRING,
memory_count INT64,
source STRING
);",
)
.context("creating IngestLog table")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS RELATES_TO(
FROM Memory TO Memory,
strength FLOAT,
context STRING
);",
)
.context("creating RELATES_TO rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS MENTIONS(
FROM Memory TO Entity,
position INT64
);",
)
.context("creating MENTIONS rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS DERIVED_FROM(
FROM Memory TO Conversation,
transformation STRING
);",
)
.context("creating DERIVED_FROM rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS DISTILLED_FROM(
FROM Memory TO Memory,
model STRING
);",
)
.context("creating DISTILLED_FROM rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS CONTRADICTS(
FROM Memory TO Memory,
resolution STRING
);",
)
.context("creating CONTRADICTS rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS REINFORCES(
FROM Memory TO Memory
);",
)
.context("creating REINFORCES rel")?;
conn.query(
"CREATE REL TABLE IF NOT EXISTS SUPERSEDES(
FROM Memory TO Memory,
reason STRING
);",
)
.context("creating SUPERSEDES rel")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS SyncLog(
seq INT64 PRIMARY KEY,
op STRING,
node_type STRING,
node_id STRING,
machine_id STRING,
timestamp STRING,
data STRING
);",
)
.context("creating SyncLog table")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS SyncState(
peer_id STRING PRIMARY KEY,
last_seq INT64,
last_sync_at STRING
);",
)
.context("creating SyncState table")?;
conn.query(
"CREATE NODE TABLE IF NOT EXISTS ConsolidationLog(
memory_id STRING PRIMARY KEY,
distilled_id STRING,
consolidated_at STRING,
model STRING
);",
)
.context("creating ConsolidationLog table")?;
conn.query(
"CALL CREATE_VECTOR_INDEX('Memory', 'memory_emb_idx', 'embedding', metric := 'cosine');",
)
.ok();
Ok(())
}
fn init_sync_seq(&self) -> Result<()> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (s:SyncLog) RETURN max(s.seq);")?;
if let Some(row) = result.next() {
if let Value::Int64(max_seq) = &row[0] {
self.sync_seq.store((*max_seq as u64) + 1, Ordering::SeqCst);
}
}
Ok(())
}
fn log_sync_entry(
&self,
conn: &Connection<'_>,
op: SyncOp,
node_type: SyncNodeType,
node_id: &str,
data: Option<&str>,
) -> Result<()> {
let seq = self.sync_seq.fetch_add(1, Ordering::SeqCst);
let timestamp = chrono::Utc::now().to_rfc3339();
let op_str = match op {
SyncOp::Create => "create",
SyncOp::Update => "update",
SyncOp::Delete => "delete",
};
let nt_str = match node_type {
SyncNodeType::Memory => "memory",
SyncNodeType::Entity => "entity",
SyncNodeType::Conversation => "conversation",
SyncNodeType::Relation => "relation",
};
let machine_escaped = escape_cypher(&self.machine_id);
let node_id_escaped = escape_cypher(node_id);
let data_escaped = data.map(|d| escape_cypher(d)).unwrap_or_default();
let query = format!(
"CREATE (:SyncLog {{
seq: {seq},
op: '{op_str}',
node_type: '{nt_str}',
node_id: '{node_id_escaped}',
machine_id: '{machine_escaped}',
timestamp: '{timestamp}',
data: '{data_escaped}'
}});"
);
conn.query(&query)?;
Ok(())
}
}
impl Store for KuzuStore {
fn store_memory(&self, memory: &Memory) -> Result<()> {
let conn = self.conn()?;
self.create_memory_node(&conn, memory)?;
let data = serde_json::to_string(memory).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Memory,
&memory.id.to_string(),
data.as_deref(),
)?;
Ok(())
}
fn get_memory(&self, id: Uuid) -> Result<Option<Memory>> {
let conn = self.conn()?;
let query = format!(
"MATCH (m:Memory {{id: '{}'}}) RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id;",
id
);
let mut result = conn.query(&query)?;
match result.next() {
Some(row) => Ok(Some(row_to_memory(&row)?)),
None => Ok(None),
}
}
fn delete_memory(&self, id: Uuid) -> Result<()> {
let conn = self.conn()?;
conn.query(&format!(
"MATCH (m:Memory {{id: '{}'}}) DETACH DELETE m;",
id
))?;
self.log_sync_entry(
&conn,
SyncOp::Delete,
SyncNodeType::Memory,
&id.to_string(),
None,
)?;
Ok(())
}
fn store_entity(&self, entity: &Entity) -> Result<()> {
let conn = self.conn()?;
self.create_entity_node(&conn, entity)?;
let data = serde_json::to_string(entity).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Entity,
&entity.id.to_string(),
data.as_deref(),
)?;
Ok(())
}
fn get_entity(&self, id: Uuid) -> Result<Option<Entity>> {
let conn = self.conn()?;
let mut result = conn.query(&format!(
"MATCH (e:Entity {{id: '{}'}}) RETURN e.id, e.name, e.entity_type;",
id
))?;
match result.next() {
Some(row) => Ok(Some(row_to_entity(&row)?)),
None => Ok(None),
}
}
fn find_entity_by_name(&self, name: &str) -> Result<Option<Entity>> {
let conn = self.conn()?;
let name_escaped = escape_cypher(name);
let query = format!(
"MATCH (e:Entity) WHERE e.name = '{}' RETURN e.id, e.name, e.entity_type;",
name_escaped
);
let mut result = conn.query(&query)?;
match result.next() {
Some(row) => Ok(Some(row_to_entity(&row)?)),
None => Ok(None),
}
}
fn store_conversation(&self, conversation: &Conversation) -> Result<()> {
let conn = self.conn()?;
self.create_conversation_node(&conn, conversation)?;
let data = serde_json::to_string(conversation).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Conversation,
&conversation.id.to_string(),
data.as_deref(),
)?;
Ok(())
}
fn store_relation(&self, relation: &Relation) -> Result<()> {
let conn = self.conn()?;
self.create_relation_edge(&conn, relation)?;
let rel_type = format!("{:?}", relation.relation_type).to_lowercase();
let node_id = format!("{}:{}:{}", relation.from_id, relation.to_id, rel_type);
let data = serde_json::to_string(relation).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Relation,
&node_id,
data.as_deref(),
)?;
Ok(())
}
fn get_relations(
&self,
node_id: Uuid,
relation_type: Option<RelationType>,
) -> Result<Vec<Relation>> {
let conn = self.conn()?;
let id = node_id.to_string();
let query = match relation_type {
Some(RelationType::RelatesTo) => format!(
"MATCH (a:Memory {{id: '{}'}})-[r:RELATES_TO]->(b:Memory) RETURN b.id, r.strength, r.context;",
id
),
Some(RelationType::Contradicts) => format!(
"MATCH (a:Memory {{id: '{}'}})-[r:CONTRADICTS]->(b:Memory) RETURN b.id, r.resolution;",
id
),
Some(RelationType::Reinforces) => format!(
"MATCH (a:Memory {{id: '{}'}})-[:REINFORCES]->(b:Memory) RETURN b.id;",
id
),
Some(RelationType::Supersedes) => format!(
"MATCH (a:Memory {{id: '{}'}})-[r:SUPERSEDES]->(b:Memory) RETURN b.id, r.reason;",
id
),
Some(RelationType::Mentions) => format!(
"MATCH (a:Memory {{id: '{}'}})-[:MENTIONS]->(b:Entity) RETURN b.id;",
id
),
Some(RelationType::DerivedFrom) => format!(
"MATCH (a:Memory {{id: '{}'}})-[:DERIVED_FROM]->(b:Conversation) RETURN b.id;",
id
),
Some(RelationType::DistilledFrom) => format!(
"MATCH (a:Memory {{id: '{}'}})-[:DISTILLED_FROM]->(b:Memory) RETURN b.id;",
id
),
None => {
let mut relations = Vec::new();
let mem_to_mem = [
("RELATES_TO", RelationType::RelatesTo),
("REINFORCES", RelationType::Reinforces),
("CONTRADICTS", RelationType::Contradicts),
("SUPERSEDES", RelationType::Supersedes),
("DISTILLED_FROM", RelationType::DistilledFrom),
];
for (label, rt) in &mem_to_mem {
let q = format!(
"MATCH (a:Memory {{id: '{}'}})-[:{}]->(b:Memory) RETURN b.id;",
id, label
);
if let Ok(mut result) = conn.query(&q) {
for row in &mut result {
let to_id_str = value_to_string(&row[0]);
let to_id = Uuid::parse_str(&to_id_str).unwrap_or_default();
relations.push(Relation {
from_id: node_id,
to_id,
relation_type: *rt,
strength: 1.0,
context: None,
});
}
}
}
for (label, target, rt) in [
("MENTIONS", "Entity", RelationType::Mentions),
("DERIVED_FROM", "Conversation", RelationType::DerivedFrom),
] {
let q = format!(
"MATCH (a:Memory {{id: '{}'}})-[:{}]->(b:{}) RETURN b.id;",
id, label, target
);
if let Ok(mut result) = conn.query(&q) {
for row in &mut result {
let to_id_str = value_to_string(&row[0]);
let to_id = Uuid::parse_str(&to_id_str).unwrap_or_default();
relations.push(Relation {
from_id: node_id,
to_id,
relation_type: rt,
strength: 1.0,
context: None,
});
}
}
}
return Ok(relations);
},
};
let mut result = conn.query(&query)?;
let mut relations = Vec::new();
for row in &mut result {
let to_id_str = value_to_string(&row[0]);
let to_id = Uuid::parse_str(&to_id_str).unwrap_or_default();
relations.push(Relation {
from_id: node_id,
to_id,
relation_type: relation_type.unwrap_or(RelationType::RelatesTo),
strength: 1.0,
context: None,
});
}
Ok(relations)
}
fn vector_search(&self, embedding: &[f32], limit: usize) -> Result<Vec<(Memory, f32)>> {
let conn = self.conn()?;
let embedding_str = format_embedding(embedding);
let query = format!(
"CALL QUERY_VECTOR_INDEX('Memory', 'memory_emb_idx', {}, {}) YIELD node, distance
RETURN node.id, node.content, node.memory_type, node.confidence,
node.created_at, node.last_accessed, node.access_count,
node.source, node.source_id, distance;",
embedding_str, limit
);
let mut result = conn.query(&query)?;
let mut results = Vec::new();
for row in &mut result {
let memory = row_to_memory(&row[..9])?;
let distance = value_to_f32(&row[9]);
let similarity = 1.0 - distance;
results.push((memory, similarity));
}
Ok(results)
}
fn traverse(&self, start_id: Uuid, depth: u32) -> Result<Vec<(Memory, Vec<Relation>)>> {
let conn = self.conn()?;
let query = format!(
"MATCH (start:Memory {{id: '{}'}})-[r:RELATES_TO|REINFORCES*1..{}]->(m:Memory)
RETURN DISTINCT m.id, m.content, m.memory_type, m.confidence,
m.created_at, m.last_accessed, m.access_count, m.source, m.source_id;",
start_id, depth
);
let mut result = conn.query(&query)?;
let mut results = Vec::new();
for row in &mut result {
let memory = row_to_memory(&row)?;
results.push((memory, Vec::new()));
}
Ok(results)
}
fn memories_by_source(&self, source: &str) -> Result<Vec<Memory>> {
let conn = self.conn()?;
let source_escaped = escape_cypher(source);
let query = format!(
"MATCH (m:Memory) WHERE m.source = '{}' RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id;",
source_escaped
);
let mut result = conn.query(&query)?;
let mut memories = Vec::new();
for row in &mut result {
memories.push(row_to_memory(&row)?);
}
Ok(memories)
}
fn memories_by_type(&self, memory_type: MemoryType) -> Result<Vec<Memory>> {
let conn = self.conn()?;
let type_str = format!("{:?}", memory_type).to_lowercase();
let query = format!(
"MATCH (m:Memory) WHERE m.memory_type = '{}' RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id;",
type_str
);
let mut result = conn.query(&query)?;
let mut memories = Vec::new();
for row in &mut result {
memories.push(row_to_memory(&row)?);
}
Ok(memories)
}
fn memories_needing_decay(&self, threshold_days: u32) -> Result<Vec<Memory>> {
let conn = self.conn()?;
let cutoff = chrono::Utc::now() - chrono::Duration::days(threshold_days as i64);
let cutoff_str = cutoff.to_rfc3339();
let query = format!(
"MATCH (m:Memory) WHERE m.last_accessed < '{}' AND m.confidence > 0.05 RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id;",
cutoff_str
);
let mut result = conn.query(&query)?;
let mut memories = Vec::new();
for row in &mut result {
memories.push(row_to_memory(&row)?);
}
Ok(memories)
}
fn update_memory(&self, memory: &Memory) -> Result<()> {
let conn = self.conn()?;
let id = memory.id.to_string();
let last_accessed = memory.last_accessed.to_rfc3339();
let query = format!(
"MATCH (m:Memory {{id: '{}'}}) SET m.confidence = {}, m.last_accessed = '{}', m.access_count = {};",
id, memory.confidence, last_accessed, memory.access_count
);
conn.query(&query)?;
let data = serde_json::to_string(memory).ok();
self.log_sync_entry(
&conn,
SyncOp::Update,
SyncNodeType::Memory,
&id,
data.as_deref(),
)?;
Ok(())
}
fn text_search(&self, query: &str, limit: usize) -> Result<Vec<Memory>> {
let conn = self.conn()?;
let query_escaped = escape_cypher(&query.to_lowercase());
let cypher = format!(
"MATCH (m:Memory) WHERE lower(m.content) CONTAINS '{}' RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id LIMIT {};",
query_escaped, limit
);
let mut result = conn.query(&cypher)?;
let mut memories = Vec::new();
for row in &mut result {
memories.push(row_to_memory(&row)?);
}
Ok(memories)
}
fn memory_count(&self) -> Result<usize> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (m:Memory) RETURN count(m);")?;
match result.next() {
Some(row) => match &row[0] {
Value::Int64(n) => Ok(*n as usize),
_ => Ok(0),
},
None => Ok(0),
}
}
}
impl KuzuStore {
pub fn find_or_create_entity(&self, name: &str, entity_type: &str) -> Result<Entity> {
if let Some(existing) = self.find_entity_by_name(name)? {
return Ok(existing);
}
let entity = Entity::new(name.to_string(), entity_type.to_string());
self.store_entity(&entity)?;
Ok(entity)
}
pub fn memory_content_exists(&self, content_prefix: &str) -> Result<bool> {
let conn = self.conn()?;
let escaped = escape_cypher(content_prefix);
let query = format!(
"MATCH (m:Memory) WHERE starts_with(m.content, '{}') RETURN m.id LIMIT 1;",
escaped
);
let mut result = conn.query(&query)?;
Ok(result.next().is_some())
}
pub fn unconsolidated_memories(&self, limit: usize) -> Result<Vec<Memory>> {
let conn = self.conn()?;
let query = format!(
"MATCH (m:Memory) WHERE NOT EXISTS {{MATCH (c:ConsolidationLog) WHERE c.memory_id = m.id}} AND m.source <> 'consolidation' RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id LIMIT {};",
limit
);
let mut result = conn.query(&query)?;
let mut memories = Vec::new();
for row in &mut result {
memories.push(row_to_memory(&row)?);
}
Ok(memories)
}
pub fn mark_consolidated(&self, raw_id: Uuid, distilled_id: Uuid, model: &str) -> Result<()> {
let conn = self.conn()?;
let now = Utc::now().to_rfc3339();
let query = format!(
"CREATE (:ConsolidationLog {{memory_id: '{}', distilled_id: '{}', consolidated_at: '{}', model: '{}'}});",
raw_id, distilled_id, now, escape_cypher(model)
);
conn.query(&query)?;
Ok(())
}
pub fn consolidation_count(&self) -> Result<usize> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (c:ConsolidationLog) RETURN count(c);")?;
match result.next() {
Some(row) => match &row[0] {
Value::Int64(n) => Ok(*n as usize),
_ => Ok(0),
},
None => Ok(0),
}
}
pub fn clear_ingest_log(&self) -> Result<usize> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (l:IngestLog) RETURN count(l);")?;
let count = match result.next() {
Some(row) => match &row[0] {
Value::Int64(n) => *n as usize,
_ => 0,
},
None => 0,
};
conn.query("MATCH (l:IngestLog) DETACH DELETE l;")?;
Ok(count)
}
pub fn delete_memories_by_source(&self, source: &str) -> Result<usize> {
let conn = self.conn()?;
let escaped = escape_cypher(source);
let mut result = conn.query(&format!(
"MATCH (m:Memory) WHERE m.source = '{}' RETURN count(m);",
escaped
))?;
let count = match result.next() {
Some(row) => match &row[0] {
Value::Int64(n) => *n as usize,
_ => 0,
},
None => 0,
};
conn.query(&format!(
"MATCH (m:Memory) WHERE m.source = '{}' DETACH DELETE m;",
escaped
))?;
Ok(count)
}
pub fn is_file_ingested(&self, file_path: &str) -> Result<bool> {
let conn = self.conn()?;
let escaped = escape_cypher(file_path);
let mut result = conn.query(&format!(
"MATCH (l:IngestLog {{file_path: '{}'}}) RETURN l.file_path;",
escaped
))?;
Ok(result.next().is_some())
}
pub fn is_file_changed(&self, file_path: &str, file_hash: &str) -> Result<bool> {
let conn = self.conn()?;
let escaped = escape_cypher(file_path);
let mut result = conn.query(&format!(
"MATCH (l:IngestLog {{file_path: '{}'}}) RETURN l.file_hash;",
escaped
))?;
match result.next() {
Some(row) => {
let stored_hash = value_to_string(&row[0]);
Ok(stored_hash != file_hash)
}
None => Ok(true),
}
}
pub fn mark_ingested(
&self,
file_path: &str,
file_hash: &str,
memory_count: usize,
source: &str,
) -> Result<()> {
let conn = self.conn()?;
let path_escaped = escape_cypher(file_path);
let hash_escaped = escape_cypher(file_hash);
let source_escaped = escape_cypher(source);
let now = chrono::Utc::now().to_rfc3339();
conn.query(&format!(
"MERGE (l:IngestLog {{file_path: '{}'}})
SET l.file_hash = '{}', l.ingested_at = '{}', l.memory_count = {}, l.source = '{}';",
path_escaped, hash_escaped, now, memory_count as i64, source_escaped
))?;
Ok(())
}
pub fn ingested_file_count(&self) -> Result<usize> {
let conn = self.conn()?;
let mut result = conn.query("MATCH (l:IngestLog) RETURN count(l);")?;
match result.next() {
Some(row) => match &row[0] {
Value::Int64(n) => Ok(*n as usize),
_ => Ok(0),
},
None => Ok(0),
}
}
}
impl KuzuStore {
fn create_memory_node(&self, conn: &Connection<'_>, memory: &Memory) -> Result<()> {
let id = memory.id.to_string();
let memory_type = format!("{:?}", memory.memory_type).to_lowercase();
let created_at = memory.created_at.to_rfc3339();
let last_accessed = memory.last_accessed.to_rfc3339();
let embedding_str = format_embedding(&memory.embedding);
let content_escaped = escape_cypher(&memory.content);
let source_escaped = escape_cypher(&memory.source);
let source_id_escaped = escape_cypher(&memory.source_id);
let query = format!(
"CREATE (:Memory {{
id: '{id}',
content: '{content_escaped}',
embedding: {embedding_str},
memory_type: '{memory_type}',
confidence: {confidence},
created_at: '{created_at}',
last_accessed: '{last_accessed}',
access_count: {access_count},
source: '{source_escaped}',
source_id: '{source_id_escaped}'
}});",
confidence = memory.confidence,
access_count = memory.access_count,
);
conn.query(&query)?;
Ok(())
}
fn create_entity_node(&self, conn: &Connection<'_>, entity: &Entity) -> Result<()> {
let id = entity.id.to_string();
let embedding_str = format_embedding(&entity.embedding);
let aliases_str = format_string_array(&entity.aliases);
let name_escaped = escape_cypher(&entity.name);
let etype_escaped = escape_cypher(&entity.entity_type);
let query = format!(
"CREATE (:Entity {{
id: '{id}',
name: '{name_escaped}',
entity_type: '{etype_escaped}',
embedding: {embedding_str},
aliases: {aliases_str}
}});"
);
conn.query(&query)?;
Ok(())
}
fn create_conversation_node(
&self,
conn: &Connection<'_>,
conversation: &Conversation,
) -> Result<()> {
let id = conversation.id.to_string();
let started_at = conversation.started_at.to_rfc3339();
let source_escaped = escape_cypher(&conversation.source);
let machine_escaped = escape_cypher(&conversation.machine_id);
let project_escaped = escape_cypher(conversation.project_path.as_deref().unwrap_or(""));
let query = format!(
"CREATE (:Conversation {{
id: '{id}',
source: '{source_escaped}',
machine_id: '{machine_escaped}',
started_at: '{started_at}',
project_path: '{project_escaped}'
}});"
);
conn.query(&query)?;
Ok(())
}
fn create_relation_edge(&self, conn: &Connection<'_>, relation: &Relation) -> Result<()> {
let from_id = relation.from_id.to_string();
let to_id = relation.to_id.to_string();
let query = match relation.relation_type {
RelationType::RelatesTo => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Memory {{id: '{}'}}) CREATE (a)-[:RELATES_TO {{strength: {}, context: '{}'}}]->(b);",
from_id, to_id, relation.strength, relation.context.as_deref().unwrap_or("")
),
RelationType::Mentions => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Entity {{id: '{}'}}) CREATE (a)-[:MENTIONS {{position: 0}}]->(b);",
from_id, to_id
),
RelationType::DerivedFrom => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Conversation {{id: '{}'}}) CREATE (a)-[:DERIVED_FROM {{transformation: '{}'}}]->(b);",
from_id, to_id, relation.context.as_deref().unwrap_or("direct")
),
RelationType::Contradicts => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Memory {{id: '{}'}}) CREATE (a)-[:CONTRADICTS {{resolution: '{}'}}]->(b);",
from_id, to_id, relation.context.as_deref().unwrap_or("")
),
RelationType::Reinforces => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Memory {{id: '{}'}}) CREATE (a)-[:REINFORCES]->(b);",
from_id, to_id
),
RelationType::DistilledFrom => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Memory {{id: '{}'}}) CREATE (a)-[:DISTILLED_FROM {{model: '{}'}}]->(b);",
from_id, to_id, relation.context.as_deref().unwrap_or("")
),
RelationType::Supersedes => format!(
"MATCH (a:Memory {{id: '{}'}}), (b:Memory {{id: '{}'}}) CREATE (a)-[:SUPERSEDES {{reason: '{}'}}]->(b);",
from_id, to_id, relation.context.as_deref().unwrap_or("")
),
};
conn.query(&query)?;
Ok(())
}
}
impl KuzuStore {
pub fn get_conversation(&self, id: Uuid) -> Result<Option<Conversation>> {
let conn = self.conn()?;
let mut result = conn.query(&format!(
"MATCH (c:Conversation {{id: '{}'}}) RETURN c.id, c.source, c.machine_id, c.started_at, c.project_path;",
id
))?;
match result.next() {
Some(row) => {
let id = Uuid::parse_str(&value_to_string(&row[0])).unwrap_or_default();
let source = value_to_string(&row[1]);
let machine_id = value_to_string(&row[2]);
let started_at_str = value_to_string(&row[3]);
let project_path = {
let s = value_to_string(&row[4]);
if s.is_empty() {
None
} else {
Some(s)
}
};
let started_at = chrono::DateTime::parse_from_rfc3339(&started_at_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
Ok(Some(Conversation {
id,
source,
machine_id,
started_at,
project_path,
}))
}
None => Ok(None),
}
}
pub fn get_memory_with_embedding(&self, id: Uuid) -> Result<Option<Memory>> {
let conn = self.conn()?;
let query = format!(
"MATCH (m:Memory {{id: '{}'}}) RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id, m.embedding;",
id
);
let mut result = conn.query(&query)?;
match result.next() {
Some(row) => {
let mut memory = row_to_memory(&row[..9])?;
memory.embedding = value_to_f32_vec(&row[9]);
Ok(Some(memory))
}
None => Ok(None),
}
}
pub fn import_memory(&self, memory: &Memory) -> Result<()> {
let conn = self.conn()?;
self.create_memory_node(&conn, memory)
}
pub fn import_conversation(&self, conversation: &Conversation) -> Result<()> {
let conn = self.conn()?;
self.create_conversation_node(&conn, conversation)
}
pub fn import_entity(&self, entity: &Entity) -> Result<()> {
let conn = self.conn()?;
self.create_entity_node(&conn, entity)
}
pub fn import_relation(&self, relation: &Relation) -> Result<()> {
let conn = self.conn()?;
self.create_relation_edge(&conn, relation)
}
pub fn import_delete_memory(&self, id: Uuid) -> Result<()> {
let conn = self.conn()?;
conn.query(&format!(
"MATCH (m:Memory {{id: '{}'}}) DETACH DELETE m;",
id
))?;
Ok(())
}
pub fn import_or_update_memory(&self, memory: &Memory) -> Result<()> {
let conn = self.conn()?;
let id = memory.id.to_string();
let last_accessed = memory.last_accessed.to_rfc3339();
let content_escaped = escape_cypher(&memory.content);
let embedding_str = format_embedding(&memory.embedding);
let memory_type = format!("{:?}", memory.memory_type).to_lowercase();
let source_escaped = escape_cypher(&memory.source);
let source_id_escaped = escape_cypher(&memory.source_id);
let created_at = memory.created_at.to_rfc3339();
let query = format!(
"MATCH (m:Memory {{id: '{id}'}}) SET m.content = '{content_escaped}', m.embedding = {embedding_str}, m.memory_type = '{memory_type}', m.confidence = {confidence}, m.created_at = '{created_at}', m.last_accessed = '{last_accessed}', m.access_count = {access_count}, m.source = '{source_escaped}', m.source_id = '{source_id_escaped}';",
confidence = memory.confidence,
access_count = memory.access_count,
);
conn.query(&query)?;
Ok(())
}
pub fn sync_log_since(&self, after_seq: u64) -> Result<Vec<SyncEntry>> {
self.sync_log_page(after_seq, None)
}
pub fn sync_log_page(&self, after_seq: u64, limit: Option<usize>) -> Result<Vec<SyncEntry>> {
let conn = self.conn()?;
let machine_escaped = escape_cypher(&self.machine_id);
let limit_clause = match limit {
Some(n) => format!(" LIMIT {}", n),
None => String::new(),
};
let query = format!(
"MATCH (s:SyncLog) WHERE s.seq > {} AND s.machine_id = '{}' RETURN s.seq, s.op, s.node_type, s.node_id, s.machine_id, s.timestamp, s.data ORDER BY s.seq{};",
after_seq, machine_escaped, limit_clause
);
let mut result = conn.query(&query)?;
let mut entries = Vec::new();
for row in &mut result {
let seq = match &row[0] {
Value::Int64(n) => *n as u64,
_ => 0,
};
let op_str = value_to_string(&row[1]);
let node_type_str = value_to_string(&row[2]);
let node_id = value_to_string(&row[3]);
let machine_id = value_to_string(&row[4]);
let timestamp_str = value_to_string(&row[5]);
let data_str = value_to_string(&row[6]);
let op = match op_str.as_str() {
"create" => SyncOp::Create,
"update" => SyncOp::Update,
"delete" => SyncOp::Delete,
_ => continue,
};
let node_type = match node_type_str.as_str() {
"memory" => SyncNodeType::Memory,
"entity" => SyncNodeType::Entity,
"conversation" => SyncNodeType::Conversation,
"relation" => SyncNodeType::Relation,
_ => continue,
};
let timestamp = chrono::DateTime::parse_from_rfc3339(×tamp_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
let data = if data_str.is_empty() {
None
} else {
Some(data_str)
};
entries.push(SyncEntry {
seq,
op,
node_type,
node_id,
machine_id,
timestamp,
data,
});
}
Ok(entries)
}
pub fn get_sync_state(&self, peer_id: &str) -> Result<Option<SyncState>> {
let conn = self.conn()?;
let escaped = escape_cypher(peer_id);
let mut result = conn.query(&format!(
"MATCH (s:SyncState {{peer_id: '{}'}}) RETURN s.peer_id, s.last_seq, s.last_sync_at;",
escaped
))?;
match result.next() {
Some(row) => {
let peer_id = value_to_string(&row[0]);
let last_seq = match &row[1] {
Value::Int64(n) => *n as u64,
_ => 0,
};
let last_sync_at_str = value_to_string(&row[2]);
let last_sync_at = chrono::DateTime::parse_from_rfc3339(&last_sync_at_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
Ok(Some(SyncState {
peer_id,
last_seq,
last_sync_at,
}))
}
None => Ok(None),
}
}
pub fn set_sync_state(&self, state: &SyncState) -> Result<()> {
let conn = self.conn()?;
let peer_escaped = escape_cypher(&state.peer_id);
let now = state.last_sync_at.to_rfc3339();
conn.query(&format!(
"MERGE (s:SyncState {{peer_id: '{}'}}) SET s.last_seq = {}, s.last_sync_at = '{}';",
peer_escaped, state.last_seq, now
))?;
Ok(())
}
pub fn get_all_sync_states(&self) -> Result<Vec<SyncState>> {
let conn = self.conn()?;
let mut result =
conn.query("MATCH (s:SyncState) RETURN s.peer_id, s.last_seq, s.last_sync_at;")?;
let mut states = Vec::new();
for row in &mut result {
let peer_id = value_to_string(&row[0]);
let last_seq = match &row[1] {
Value::Int64(n) => *n as u64,
_ => 0,
};
let last_sync_at_str = value_to_string(&row[2]);
let last_sync_at = chrono::DateTime::parse_from_rfc3339(&last_sync_at_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
states.push(SyncState {
peer_id,
last_seq,
last_sync_at,
});
}
Ok(states)
}
pub fn backfill_sync_log(&self) -> Result<u64> {
let conn = self.conn()?;
let mut count = 0u64;
let mut result = conn.query(
"MATCH (m:Memory) RETURN m.id, m.content, m.memory_type, m.confidence, m.created_at, m.last_accessed, m.access_count, m.source, m.source_id, m.embedding;",
)?;
for row in &mut result {
let mut memory = row_to_memory(&row[..9])?;
memory.embedding = value_to_f32_vec(&row[9]);
let data = serde_json::to_string(&memory).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Memory,
&memory.id.to_string(),
data.as_deref(),
)?;
count += 1;
}
let mut result = conn.query(
"MATCH (c:Conversation) RETURN c.id, c.source, c.machine_id, c.started_at, c.project_path;",
)?;
for row in &mut result {
let id = Uuid::parse_str(&value_to_string(&row[0])).unwrap_or_default();
let source = value_to_string(&row[1]);
let machine_id = value_to_string(&row[2]);
let started_at_str = value_to_string(&row[3]);
let project_path = {
let s = value_to_string(&row[4]);
if s.is_empty() {
None
} else {
Some(s)
}
};
let started_at = chrono::DateTime::parse_from_rfc3339(&started_at_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
let conv = Conversation {
id,
source,
machine_id,
started_at,
project_path,
};
let data = serde_json::to_string(&conv).ok();
self.log_sync_entry(
&conn,
SyncOp::Create,
SyncNodeType::Conversation,
&conv.id.to_string(),
data.as_deref(),
)?;
count += 1;
}
Ok(count)
}
}
fn format_embedding(embedding: &[f32]) -> String {
if embedding.is_empty() {
let zeros: Vec<String> = (0..384).map(|_| "0.0".to_string()).collect();
return format!("[{}]", zeros.join(","));
}
let parts: Vec<String> = embedding.iter().map(|v| format!("{}", v)).collect();
format!("[{}]", parts.join(","))
}
fn escape_cypher(s: &str) -> String {
s.replace('\\', "\\\\").replace('\'', "\\'")
}
fn format_string_array(items: &[String]) -> String {
let parts: Vec<String> = items
.iter()
.map(|s| format!("'{}'", s.replace('\'', "''")))
.collect();
format!("[{}]", parts.join(","))
}
fn value_to_string(val: &Value) -> String {
match val {
Value::String(s) => s.clone(),
other => format!("{:?}", other),
}
}
fn value_to_f32(val: &Value) -> f32 {
match val {
Value::Float(f) => *f,
Value::Double(d) => *d as f32,
Value::Int64(i) => *i as f32,
_ => 0.0,
}
}
fn value_to_f32_vec(val: &Value) -> Vec<f32> {
match val {
Value::Array(_, items) | Value::List(_, items) => {
items.iter().map(|v| value_to_f32(v)).collect()
}
_ => Vec::new(),
}
}
fn row_to_memory(row: &[Value]) -> Result<Memory> {
let id_str = value_to_string(&row[0]);
let id = Uuid::parse_str(&id_str).unwrap_or_default();
let content = value_to_string(&row[1]);
let memory_type_str = value_to_string(&row[2]);
let confidence = value_to_f32(&row[3]);
let created_at_str = value_to_string(&row[4]);
let last_accessed_str = value_to_string(&row[5]);
let access_count = match &row[6] {
Value::Int64(i) => *i as u32,
_ => 0,
};
let source = value_to_string(&row[7]);
let source_id = value_to_string(&row[8]);
let memory_type = match memory_type_str.as_str() {
"episodic" => MemoryType::Episodic,
"procedural" => MemoryType::Procedural,
_ => MemoryType::Semantic,
};
let created_at = chrono::DateTime::parse_from_rfc3339(&created_at_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
let last_accessed = chrono::DateTime::parse_from_rfc3339(&last_accessed_str)
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(|_| chrono::Utc::now());
Ok(Memory {
id,
content,
embedding: Vec::new(),
memory_type,
confidence,
created_at,
last_accessed,
access_count,
source,
source_id,
})
}
fn row_to_entity(row: &[Value]) -> Result<Entity> {
let id_str = value_to_string(&row[0]);
let id = Uuid::parse_str(&id_str).unwrap_or_default();
let name = value_to_string(&row[1]);
let entity_type = value_to_string(&row[2]);
Ok(Entity {
id,
name,
entity_type,
embedding: Vec::new(),
aliases: Vec::new(),
})
}
unsafe impl Send for KuzuStore {}
unsafe impl Sync for KuzuStore {}