96 lines
2.7 KiB
Rust
96 lines
2.7 KiB
Rust
use crate::codec::{decode, encode, hash_embedding, EMBEDDING_DIM};
|
|
use engram_core::{EngramDb, MemoryTier, Node, NodeType};
|
|
use neuron_domain::{ExecutionContext, NeuronError, NeuronResult};
|
|
use tracing::debug;
|
|
use uuid::Uuid;
|
|
|
|
pub struct ContextStore {
|
|
db: EngramDb,
|
|
}
|
|
|
|
impl ContextStore {
|
|
pub fn new(db: EngramDb) -> Self {
|
|
Self { db }
|
|
}
|
|
|
|
pub fn put(&self, ctx: &ExecutionContext) -> NeuronResult<Uuid> {
|
|
debug!("ContextStore::put id={}", ctx.id);
|
|
let content = encode(ctx)?;
|
|
let text = format!("{} {}", ctx.process_name, ctx.objective);
|
|
let embedding = hash_embedding(&text, EMBEDDING_DIM);
|
|
|
|
let node = Node::new(
|
|
NodeType::Event,
|
|
embedding,
|
|
content,
|
|
MemoryTier::Working,
|
|
0.7,
|
|
)
|
|
.with_id(ctx.id);
|
|
|
|
self.db
|
|
.put_node(node)
|
|
.map_err(|e| NeuronError::Storage(e.to_string()))
|
|
}
|
|
|
|
pub fn get(&self, id: Uuid) -> NeuronResult<ExecutionContext> {
|
|
debug!("ContextStore::get id={}", id);
|
|
let node = self
|
|
.db
|
|
.get_node(id)
|
|
.map_err(|e| NeuronError::Storage(e.to_string()))?
|
|
.ok_or_else(|| NeuronError::NotFound(id.to_string()))?;
|
|
decode(&node.content)
|
|
}
|
|
|
|
pub fn list(&self) -> NeuronResult<Vec<ExecutionContext>> {
|
|
let nodes = self
|
|
.db
|
|
.scan_nodes()
|
|
.map_err(|e| NeuronError::Storage(e.to_string()))?;
|
|
let mut out = Vec::new();
|
|
for node in nodes {
|
|
if node.node_type == NodeType::Event {
|
|
if let Ok(ctx) = decode::<ExecutionContext>(&node.content) {
|
|
out.push(ctx);
|
|
}
|
|
}
|
|
}
|
|
Ok(out)
|
|
}
|
|
|
|
pub fn update(&self, ctx: &ExecutionContext) -> NeuronResult<()> {
|
|
self.put(ctx).map(|_| ())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use neuron_domain::{ContextStatus, ExecutionContext};
|
|
use tempfile::tempdir;
|
|
|
|
fn make_store() -> ContextStore {
|
|
let dir = tempdir().unwrap();
|
|
let db = engram_core::EngramDb::open(dir.path()).unwrap();
|
|
std::mem::forget(dir);
|
|
ContextStore::new(db)
|
|
}
|
|
|
|
#[test]
|
|
fn put_and_get_context() {
|
|
let store = make_store();
|
|
let ctx = ExecutionContext::new(
|
|
"write_code".into(),
|
|
"implement the store layer".into(),
|
|
"context store working".into(),
|
|
Some("neuron-rs".into()),
|
|
);
|
|
let id = ctx.id;
|
|
store.put(&ctx).unwrap();
|
|
let back = store.get(id).unwrap();
|
|
assert_eq!(back.process_name, "write_code");
|
|
assert_eq!(back.status, ContextStatus::Active);
|
|
}
|
|
}
|