use crate::index::TextIndex;
use crate::{schema, text, types::*, Error, Result};
use fs4::fs_std::FileExt;
use parking_lot::{Mutex, RwLock};
use rusqlite::{params, params_from_iter, types::Value as SqlValue, Connection, OptionalExtension, Transaction, TransactionBehavior};
use serde::de::DeserializeOwned;
use serde_json::Value;
use std::{collections::{BTreeMap, BTreeSet, HashMap, HashSet}, fs::{File, OpenOptions}, path::{Path, PathBuf}, sync::{atomic::{AtomicU64, Ordering}, Arc}};
const SELF_HEAL_ATTEMPTS: usize = 1000;
pub(crate) struct Writer { pub conn: Connection, _file_lock: File }
pub(crate) struct Readers { pub idle: Vec<Connection> }
pub(crate) struct VectorCache {
generation: AtomicU64,
bumps: Mutex<HashMap<String, u64>>,
entries: Mutex<HashMap<(String, String, String), Option<Arc<crate::embeddings::Partition>>>>,
}
impl VectorCache {
fn new() -> Self {
Self { generation: AtomicU64::new(0), bumps: Mutex::new(HashMap::new()), entries: Mutex::new(HashMap::new()) }
}
pub fn epoch_of(&self, namespace: &str) -> u64 {
self.generation.load(Ordering::SeqCst) + self.bumps.lock().get(namespace).copied().unwrap_or(0)
}
pub fn invalidate(&self) {
self.generation.fetch_add(1, Ordering::SeqCst);
self.entries.lock().clear();
}
pub fn invalidate_namespaces(&self, namespaces: &HashSet<String>) {
{
let mut bumps = self.bumps.lock();
for namespace in namespaces { *bumps.entry(namespace.clone()).or_insert(0) += 1; }
}
self.entries.lock().retain(|(_, namespace, _), _| !namespaces.contains(namespace));
}
}
thread_local! {
static TOUCHED_NAMESPACES: std::cell::RefCell<Option<HashSet<String>>> = const { std::cell::RefCell::new(None) };
}
pub(crate) fn touch_namespace(namespace: &str) {
TOUCHED_NAMESPACES.with(|slot| {
if let Some(touched) = slot.borrow_mut().as_mut() { touched.insert(text::normalized_tag(namespace)); }
});
}
pub(crate) fn touch_record_namespace(conn: &Connection, record_id: i64) -> Result<()> {
if let Some(namespace) = namespace_of(conn, record_id)? { touch_namespace(&namespace); }
Ok(())
}
pub(crate) fn namespace_of(conn: &Connection, record_id: i64) -> Result<Option<String>> {
Ok(conn.query_row("SELECT s.text FROM records r JOIN strings s ON s.id=r.namespace_id WHERE r.id=?1",
[record_id], |r| r.get(0)).optional()?)
}
struct TouchLog(Option<HashSet<String>>);
impl TouchLog {
fn install() -> Self {
Self(TOUCHED_NAMESPACES.with(|slot| slot.borrow_mut().replace(HashSet::new())))
}
fn take(&self) -> HashSet<String> {
TOUCHED_NAMESPACES.with(|slot| slot.borrow_mut().take()).unwrap_or_default()
}
}
impl Drop for TouchLog {
fn drop(&mut self) {
let previous = self.0.take();
TOUCHED_NAMESPACES.with(|slot| *slot.borrow_mut() = previous);
}
}
pub(crate) struct Engine {
pub writer: Mutex<Option<Writer>>,
pub readers: Mutex<Option<Readers>>,
pub index: RwLock<Option<Arc<TextIndex>>>,
pub vectors: VectorCache,
pub embedders: crate::embeddings::EmbedderRegistry,
pub rerankers: crate::search::RerankerRegistry,
pub events: crate::events::EventRegistry,
pub vectorizer: std::sync::OnceLock<Arc<crate::embeddings::Vectorizer>>,
pub degraded: Mutex<Vec<Degrade>>,
pub root: PathBuf,
}
fn open_reader(root: &Path) -> Result<Connection> {
let conn = Connection::open(root.join("store.sqlite3"))?;
conn.execute_batch("PRAGMA busy_timeout=5000; PRAGMA synchronous=NORMAL; PRAGMA foreign_keys=ON;")?;
Ok(conn)
}
#[derive(Clone)]
pub struct KnowledgeBase { pub(crate) engine: Arc<Engine> }
pub(crate) struct ReadGuard<'a> { engine: &'a Engine, conn: Option<Connection> }
impl ReadGuard<'_> {
pub fn conn(&self) -> &Connection { self.conn.as_ref().expect("read connection lives until drop") }
}
impl Drop for ReadGuard<'_> {
fn drop(&mut self) {
let Some(conn) = self.conn.take() else { return };
if let Some(readers) = self.engine.readers.lock().as_mut() { readers.idle.push(conn); }
}
}
impl KnowledgeBase {
pub fn open(directory: impl AsRef<Path>) -> Result<Self> {
std::fs::create_dir_all(directory.as_ref())?;
let root = std::fs::canonicalize(directory.as_ref())?;
let file_lock = OpenOptions::new().create(true).truncate(false).read(true).write(true).open(root.join("writer.lock"))?;
if !file_lock.try_lock_exclusive()? { return Err(Error::Locked(root.display().to_string())); }
let mut write_conn = Connection::open(root.join("store.sqlite3"))?;
schema::initialize(&mut write_conn)?;
let index = Arc::new(TextIndex::open(&root)?);
index.recover(&write_conn)?;
let reader = open_reader(&root)?;
let engine = Arc::new(Engine {
writer: Mutex::new(Some(Writer { conn: write_conn, _file_lock: file_lock })),
readers: Mutex::new(Some(Readers { idle: vec![reader] })),
index: RwLock::new(Some(index)), vectors: VectorCache::new(),
embedders: crate::embeddings::EmbedderRegistry::new(),
rerankers: crate::search::RerankerRegistry::new(),
events: crate::events::EventRegistry::default(),
vectorizer: std::sync::OnceLock::new(),
degraded: Mutex::new(Vec::new()), root,
});
engine.vectorizer.set(crate::embeddings::Vectorizer::start(&engine)?).unwrap_or_else(|_| unreachable!("vectorizer starts once"));
Ok(Self { engine })
}
pub fn directory(&self) -> &Path { &self.engine.root }
pub(crate) fn index(&self) -> Result<Arc<TextIndex>> {
self.engine.index.read().clone().ok_or(Error::Closed)
}
pub(crate) fn index_documents(&self, docs: &[crate::index::IndexDocument]) -> Result<()> {
self.index()?.stage(docs)
}
pub fn close(&self) -> Result<()> {
if let Some(vectorizer) = self.engine.vectorizer.get() { vectorizer.stop(); }
let mut guard = self.engine.writer.lock();
let result = match guard.as_ref() {
Some(writer) => self.index()?.sync(&writer.conn),
None => Ok(()),
};
*guard = None;
*self.engine.index.write() = None;
*self.engine.readers.lock() = None;
result
}
pub(crate) fn read(&self) -> Result<ReadGuard<'_>> {
let conn = {
let mut readers = self.engine.readers.lock();
match readers.as_mut() {
Some(readers) => match readers.idle.pop() {
Some(conn) => conn,
None => open_reader(&self.engine.root)?,
},
None => return Err(Error::Closed),
}
};
Ok(ReadGuard { engine: &self.engine, conn: Some(conn) })
}
pub(crate) fn partition(&self, conn: &Connection, space: &crate::embeddings::EmbeddingSpace,
namespace: &str, scope: &str) -> Result<Option<Arc<crate::embeddings::Partition>>> {
let key = (space.id.clone(), namespace.to_string(), scope.to_string());
let epoch = self.engine.vectors.epoch_of(namespace);
let cached = self.engine.vectors.entries.lock().get(&key).cloned();
if let Some(partition) = cached { return Ok(partition); }
let loaded = crate::embeddings::Partition::load(conn, space, namespace, scope)?.map(Arc::new);
{
let mut entries = self.engine.vectors.entries.lock();
if self.engine.vectors.epoch_of(namespace) == epoch { entries.insert(key, loaded.clone()); }
}
Ok(loaded)
}
pub(crate) fn sync_index_if_behind(&self, conn: &Connection) -> Result<()> {
let pending: i64 = conn.query_row("SELECT COUNT(*) FROM index_updates", [], |r| r.get(0))?;
if pending == 0 { return Ok(()); }
let filling = self.engine.vectorizer.get().map(|vectorizer| vectorizer.clone());
for _ in 0..SELF_HEAL_ATTEMPTS {
match self.engine.writer.try_lock() {
Some(mut guard) => return match guard.as_mut() {
Some(writer) => self.index()?.sync(&writer.conn),
None => Ok(()),
},
None => {
let still: i64 = conn.query_row("SELECT COUNT(*) FROM index_updates", [], |r| r.get(0))?;
if still == 0 { return Ok(()); }
if !filling.as_ref().is_some_and(|vectorizer| vectorizer.is_filling()) { return Ok(()); }
std::thread::sleep(std::time::Duration::from_millis(1));
}
}
}
Ok(())
}
pub(crate) fn mutate<T>(&self, f: impl FnOnce(&Transaction<'_>) -> Result<T>) -> Result<WriteReceipt<T>> {
let mut guard = self.engine.writer.lock();
let writer = guard.as_mut().ok_or(Error::Closed)?;
let changed_before = writer.conn.total_changes();
let log = TouchLog::install();
let tx = writer.conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let value = f(&tx)?;
let revision = current_revision(&tx)?;
tx.commit()?;
let touched = log.take();
drop(log);
if touched.is_empty() {
if writer.conn.total_changes() > changed_before { self.engine.vectors.invalidate(); }
} else {
self.engine.vectors.invalidate_namespaces(&touched);
}
self.invalidate_readiness(&writer.conn);
if let Some(vectorizer) = self.engine.vectorizer.get() { vectorizer.notify_work(); }
Ok(WriteReceipt { value, revision })
}
pub(crate) fn mutate_meta<T>(&self, f: impl FnOnce(&Transaction<'_>) -> Result<T>) -> Result<WriteReceipt<T>> {
let mut guard = self.engine.writer.lock();
let writer = guard.as_mut().ok_or(Error::Closed)?;
let tx = writer.conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let value = f(&tx)?;
let revision = current_revision(&tx)?;
tx.commit()?;
Ok(WriteReceipt { value, revision })
}
fn invalidate_readiness(&self, conn: &Connection) {
let Ok(mut stmt) = conn.prepare("SELECT DISTINCT s.text FROM index_updates u
JOIN records r ON r.id=u.record_id JOIN strings s ON s.id=r.namespace_id") else { return };
let Ok(namespaces) = stmt.query_map([], |r| r.get::<_, String>(0)) else { return };
for namespace in namespaces.flatten() {
let _ = crate::embeddings::clear_vector_ready(conn, &namespace);
}
}
pub fn memories(&self) -> crate::memory::MemoryStore { crate::memory::MemoryStore(self.clone()) }
pub fn graph(&self) -> crate::graph::GraphStore { crate::graph::GraphStore(self.clone()) }
pub fn notes(&self) -> crate::notes::NoteStore { crate::notes::NoteStore(self.clone()) }
pub fn embeddings(&self) -> crate::embeddings::EmbeddingStore { crate::embeddings::EmbeddingStore(self.clone()) }
pub(crate) fn write<T>(&self, f: impl FnOnce(&Writer) -> Result<T>) -> Result<T> {
let mut guard = self.engine.writer.lock();
let value = f(guard.as_mut().ok_or(Error::Closed)?)?;
self.engine.vectors.invalidate();
Ok(value)
}
pub(crate) fn note_degrade(&self, degrade: Degrade) {
let mut observed = self.engine.degraded.lock();
if !observed.contains(°rade) {
observed.push(degrade);
if observed.len() > 8 { observed.remove(0); }
}
}
pub fn update_index(&self) -> Result<HealthReport> {
self.catch_up_index()?;
self.health()
}
pub(crate) fn catch_up_index(&self) -> Result<()> {
let mut guard = self.engine.writer.lock();
let writer = guard.as_mut().ok_or(Error::Closed)?;
self.index()?.sync(&writer.conn)
}
pub fn register_event_sink<F: Fn(&crate::events::LogEvent) + Send + Sync + 'static>(&self, sink: F) {
self.engine.events.set(Arc::new(sink));
}
pub fn unregister_event_sink(&self) -> bool { self.engine.events.clear() }
pub fn event_sink_registered(&self) -> bool { self.engine.events.is_registered() }
pub fn rebuild_indexes(&self) -> Result<HealthReport> {
let sink = self.engine.events.get();
let started = std::time::Instant::now();
{
let mut guard = self.engine.writer.lock();
let writer = guard.as_mut().ok_or(Error::Closed)?;
self.index()?.rebuild(&writer.conn)?;
self.engine.vectors.invalidate();
}
let ms = started.elapsed().as_millis() as u64;
let report = self.health()?;
if let Some(sink) = sink {
let mut event = crate::events::LogEvent::new("index_rebuild");
event.ms = ms;
event.documents = Some(report.index_document_count);
event.format = Some(crate::index::FORMAT.to_string());
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| sink(&event)));
}
Ok(report)
}
pub fn rebuild_progress(&self) -> Result<RebuildProgressReport> {
Ok(self.index()?.rebuild_progress())
}
pub fn health(&self) -> Result<HealthReport> {
let state = self.read()?;
let conn = state.conn();
let mut counts = BTreeMap::new();
let mut stmt = conn.prepare("SELECT kind, COUNT(*) FROM records GROUP BY kind")?;
for row in stmt.query_map([], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?)))? {
let (code, count) = row?;
let name = RecordKind::from_code(code).map(|k| k.as_str().to_string()).unwrap_or_else(|| code.to_string());
counts.insert(name, count as usize);
}
let record_count = counts.values().sum();
let mut foreign = conn.prepare("PRAGMA foreign_key_check")?;
let mut foreign_key_errors = 0;
let mut rows = foreign.query([])?;
while rows.next()?.is_some() { foreign_key_errors += 1; }
Ok(HealthReport {
schema_version: schema::SCHEMA_VERSION,
revision: current_revision(conn)?,
indexed_revision: meta(conn, "indexed_revision")?, record_count,
index_document_count: self.index()?.document_count(),
pending_index_updates: conn.query_row("SELECT COUNT(*) FROM index_updates", [], |r| r.get::<_, i64>(0))? as usize,
sqlite_integrity: conn.query_row("PRAGMA quick_check", [], |r| r.get(0))?,
foreign_key_errors, counts,
embedder_spaces: self.engine.embedders.space_ids(),
reranker_registered: self.engine.rerankers.is_registered(),
last_degraded: self.engine.degraded.lock().clone(),
})
}
pub fn backup(&self, target: impl AsRef<Path>) -> Result<()> {
let target = target.as_ref();
let state = self.read()?;
let reservation = OpenOptions::new().write(true).create_new(true).open(target)?;
drop(reservation);
if let Err(err) = state.conn().backup(rusqlite::MAIN_DB, target, None) {
let _ = std::fs::remove_file(target);
return Err(err.into());
}
Ok(())
}
pub fn restore(snapshot: impl AsRef<Path>, directory: impl AsRef<Path>) -> Result<Self> {
let source = Connection::open_with_flags(snapshot.as_ref(), rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
let app: i64 = source.pragma_query_value(None, "application_id", |r| r.get(0))?;
let version: i64 = source.pragma_query_value(None, "user_version", |r| r.get(0))?;
if app != schema::APPLICATION_ID { return Err(Error::Validation("snapshot is not a p-memory database".into())); }
if version != schema::SCHEMA_VERSION { return Err(Error::SchemaVersion { found: version, supported: schema::SCHEMA_VERSION }); }
std::fs::create_dir(directory.as_ref())?;
source.backup(rusqlite::MAIN_DB, directory.as_ref().join("store.sqlite3"), None)?;
Self::open(directory)
}
}
pub(crate) fn now_us() -> i64 { chrono::Utc::now().timestamp_micros() }
pub(crate) fn meta(conn: &Connection, key: &str) -> Result<i64> {
Ok(conn.query_row("SELECT value FROM meta WHERE key=?1", [key], |r| r.get(0))?)
}
pub(crate) fn meta_opt(conn: &Connection, key: &str) -> Result<Option<i64>> {
Ok(conn.query_row("SELECT value FROM meta WHERE key=?1", [key], |r| r.get(0)).optional()?)
}
pub(crate) fn set_meta(conn: &Connection, key: &str, value: i64) -> Result<()> {
conn.execute("INSERT INTO meta(key,value) VALUES (?1,?2) ON CONFLICT(key) DO UPDATE SET value=excluded.value", params![key, value])?;
Ok(())
}
pub(crate) fn clear_meta(conn: &Connection, key: &str) -> Result<()> {
conn.execute("DELETE FROM meta WHERE key=?1", [key])?;
Ok(())
}
pub(crate) fn current_revision(conn: &Connection) -> Result<i64> { meta(conn, "revision") }
pub(crate) fn next_revision(conn: &Connection, record_id: i64) -> Result<i64> {
conn.execute("UPDATE meta SET value=value+1 WHERE key='revision'", [])?;
let revision = current_revision(conn)?;
conn.execute("INSERT INTO index_updates(revision,record_id) VALUES (?1,?2)", params![revision, record_id])?;
Ok(revision)
}
pub(crate) fn term_id(conn: &Connection, text_value: &str) -> Result<i64> {
let normalized = text::normalized_tag(text_value);
conn.execute("INSERT OR IGNORE INTO strings(text) VALUES (?1)", [&normalized])?;
Ok(conn.query_row("SELECT id FROM strings WHERE text=?1", [&normalized], |r| r.get(0))?)
}
pub(crate) fn term_text(conn: &Connection, id: i64) -> Result<String> {
Ok(conn.query_row("SELECT text FROM strings WHERE id=?1", [id], |r| r.get(0))?)
}
pub(crate) fn record_namespaces(conn: &Connection) -> Result<Vec<String>> {
let mut stmt = conn.prepare("SELECT DISTINCT s.text FROM records r JOIN strings s ON s.id=r.namespace_id ORDER BY s.text")?;
let mut namespaces = Vec::new();
for row in stmt.query_map([], |r| r.get::<_, String>(0))? { namespaces.push(row?); }
Ok(namespaces)
}
pub(crate) fn validate_identity(label: &str, value: &str) -> Result<()> {
if value.trim().is_empty() || value != value.trim() || value.chars().any(char::is_control) {
return Err(Error::Validation(format!("{label} must be nonempty, trimmed, and contain no control characters")));
}
Ok(())
}
pub(crate) fn validate_filter(filter: &ReadFilter) -> Result<()> {
validate_identity("namespace", &filter.namespace)?;
if filter.scopes.is_empty() { return Err(Error::Validation("at least one explicit read scope is required".into())); }
for scope in &filter.scopes { validate_identity("scope", scope)?; }
Ok(())
}
pub(crate) fn validate_limit(limit: usize) -> Result<()> {
if !(1..=10_000).contains(&limit) { return Err(Error::Validation("limit must be between 1 and 10000".into())); }
Ok(())
}
pub(crate) fn normalize_tags(tags: &[String]) -> Vec<String> {
tags.iter().map(|label| text::normalized_tag(label)).filter(|tag| !tag.is_empty()).collect::<BTreeSet<_>>().into_iter().collect()
}
pub(crate) fn tags_prefix(kind: RecordKind, tags: &[String], exclude: &[String], payload: &Value) -> String {
let carries = match kind {
RecordKind::Note => false,
RecordKind::Chunk => payload.get("ordinal").and_then(Value::as_u64) == Some(0),
_ => true,
};
if !carries { return String::new(); }
tags.iter().filter(|tag| !exclude.contains(tag)).cloned().collect::<Vec<_>>().join(" ")
}
pub(crate) fn split_note_path(relative: &str) -> (Vec<String>, String) {
let segments: Vec<&str> = relative.split('/').filter(|segment| !segment.is_empty()).collect();
let Some((last, dirs)) = segments.split_last() else { return (Vec::new(), String::new()); };
let stem = last.rsplit_once('.').map(|(stem, _)| stem).unwrap_or(last).trim();
(dirs.iter().map(|segment| segment.to_string()).collect(), stem.to_string())
}
pub(crate) fn note_path_parts(conn: &Connection, note_id: i64) -> (Vec<String>, String) {
let Ok((path, name)) = conn.query_row("SELECT path,name FROM notes WHERE record_id=?1", [note_id],
|r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))) else {
return (Vec::new(), String::new());
};
if Path::new(&path).is_absolute() { return (Vec::new(), name); }
(split_note_path(&path).0, name)
}
pub(crate) fn index_columns(conn: &Connection, kind: RecordKind, payload: &Value) -> (String, String, Vec<String>) {
if kind == RecordKind::Chunk {
if payload.get("ordinal").and_then(Value::as_u64).unwrap_or(0) != 0 { return (String::new(), String::new(), Vec::new()); }
let note_id = payload.get("note_id").and_then(Value::as_i64).unwrap_or(0);
let (dirs, stem) = note_path_parts(conn, note_id);
let mut exclude = dirs.clone();
if !stem.is_empty() { exclude.push(stem.clone()); }
return (stem, dirs.join(" "), exclude);
}
(record_name(kind, payload), String::new(), Vec::new())
}
pub(crate) fn put_record(conn: &Connection, kind: RecordKind, input: &RecordInput,
payload: &Value, text: &str) -> Result<(RecordHeader, crate::index::IndexDocument)> {
validate_identity("namespace", &input.namespace)?;
validate_identity("scope", &input.scope)?;
for evidence in &input.evidence {
if evidence.source.trim().is_empty() { return Err(Error::Validation("evidence source is required".into())); }
match (evidence.offset, evidence.limit) {
(None, None) => {},
(Some(offset), Some(limit)) if offset >= 1 && limit >= 1 => {},
_ => return Err(Error::Validation("evidence offset/limit must be a 1-based start and a positive line count".into())),
}
}
let namespace_id = term_id(conn, &input.namespace)?;
touch_namespace(&input.namespace);
let scope_id = term_id(conn, &input.scope)?;
let existing = match input.id {
Some(id) => Some(conn.query_row("SELECT created_at_us,updated_at_us,revision,scope_id FROM records WHERE id=?1", [id],
|r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?, r.get::<_, i64>(2)?, r.get::<_, i64>(3)?))).optional()?
.ok_or_else(|| Error::NotFound(id.to_string()))?),
None => None,
};
if existing.as_ref().is_some_and(|v| v.3 != scope_id) {
return Err(Error::Conflict("an existing record cannot change scope; copy it to a new ID explicitly".into()));
}
if let Some(expected) = input.expected_revision {
if existing.as_ref().map(|v| v.2) != Some(expected) { return Err(Error::StaleRevision(input.id.map(|v| v.to_string()).unwrap_or_default())); }
}
let now = now_us();
let created = existing.as_ref().map(|v| v.0).unwrap_or(input.created_at_us.unwrap_or(now));
let updated = input.updated_at_us.unwrap_or_else(|| now.max(existing.as_ref().map(|v| v.1).unwrap_or(created)));
if updated < created { return Err(Error::Validation("updated_at_us precedes created_at_us".into())); }
let tags = normalize_tags(&input.tags);
let fingerprint = record_fingerprint(text, &tags);
let metadata_json = serde_json::to_string(&input.metadata)?;
let evidence_json = serde_json::to_string(&input.evidence)?;
let payload_json = serde_json::to_string(payload)?;
let (id, revision) = match input.id {
Some(id) => {
let revision = next_revision(conn, id)?;
conn.execute("UPDATE records SET namespace_id=?2,kind=?3,scope_id=?4,updated_at_us=?5,revision=?6,metadata_json=?7,
evidence_json=?8,fingerprint=?9,payload_json=?10 WHERE id=?1",
params![id, namespace_id, kind.code(), scope_id, updated, revision, metadata_json, evidence_json,
fingerprint, payload_json])?;
conn.execute("DELETE FROM embeddings WHERE record_id=?1 AND fingerprint<>?2", params![id, fingerprint])?;
(id, revision)
}
None => {
conn.execute("INSERT INTO records(namespace_id,kind,scope_id,created_at_us,updated_at_us,revision,metadata_json,evidence_json,
fingerprint,payload_json) VALUES (?1,?2,?3,?4,?5,0,?6,?7,?8,?9)",
params![namespace_id, kind.code(), scope_id, created, updated, metadata_json, evidence_json,
fingerprint, payload_json])?;
let id = conn.last_insert_rowid();
let revision = next_revision(conn, id)?;
conn.execute("UPDATE records SET revision=?2 WHERE id=?1", params![id, revision])?;
(id, revision)
}
};
let tag_ids = set_record_tags(conn, id, &tags)?;
let (name, path, exclude) = index_columns(conn, kind, payload);
let document = crate::index::IndexDocument { id, namespace_id, scope_id, kind,
text: text.to_string(), name, path,
note_id: if kind == RecordKind::Chunk { payload.get("note_id").and_then(Value::as_i64).unwrap_or(0) } else { 0 },
tags_prefix: tags_prefix(kind, &tags, &exclude, payload), tag_ids };
Ok((RecordHeader { id, namespace: input.namespace.clone(), kind, scope: input.scope.clone(),
created_at_us: created, updated_at_us: updated, revision, tags,
evidence: input.evidence.clone(), metadata: input.metadata.clone() }, document))
}
pub(crate) fn record_fingerprint(text: &str, tags: &[String]) -> String {
text::digest(&format!("text-v1\n{text}\n{}", tags.join(" ")))
}
pub(crate) fn set_record_tags(conn: &Connection, id: i64, tags: &[String]) -> Result<Vec<i64>> {
conn.execute("DELETE FROM record_tags WHERE record_id=?1", [id])?;
let mut tag_ids = Vec::with_capacity(tags.len());
for tag in tags {
let tag_id = term_id(conn, tag)?;
conn.execute("INSERT OR IGNORE INTO record_tags(record_id,tag_id) VALUES (?1,?2)", params![id, tag_id])?;
tag_ids.push(tag_id);
}
Ok(tag_ids)
}
pub(crate) fn record_tag_pairs(conn: &Connection, id: i64) -> Result<Vec<(i64, String)>> {
let mut stmt = conn.prepare("SELECT t.id,t.text FROM record_tags rt JOIN strings t ON t.id=rt.tag_id WHERE rt.record_id=?1 ORDER BY t.text")?;
let mut pairs = Vec::new();
for row in stmt.query_map([id], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))? { pairs.push(row?); }
Ok(pairs)
}
pub(crate) fn index_document(conn: &Connection, id: i64, kind: RecordKind, text: String) -> Result<crate::index::IndexDocument> {
let (namespace_id, scope_id, payload_json): (i64, i64, String) = conn.query_row("SELECT namespace_id,scope_id,payload_json FROM records WHERE id=?1",
[id], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)))?;
let payload: Value = serde_json::from_str(&payload_json)?;
let pairs = record_tag_pairs(conn, id)?;
let tags: Vec<String> = pairs.iter().map(|(_, tag)| tag.clone()).collect();
let (name, path, exclude) = index_columns(conn, kind, &payload);
Ok(crate::index::IndexDocument { id, namespace_id, scope_id, kind, text,
name, path,
note_id: if kind == RecordKind::Chunk { payload.get("note_id").and_then(Value::as_i64).unwrap_or(0) } else { 0 },
tags_prefix: tags_prefix(kind, &tags, &exclude, &payload),
tag_ids: pairs.into_iter().map(|(tag_id, _)| tag_id).collect() })
}
pub(crate) fn chunk_notes(conn: &Connection, ids: &[i64]) -> Result<BTreeMap<i64, (i64, usize)>> {
let mut out = BTreeMap::new();
if ids.is_empty() { return Ok(out); }
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
let mut stmt = conn.prepare(&format!("SELECT record_id,note_id,\"offset\" FROM chunks WHERE record_id IN ({placeholders})"))?;
for row in stmt.query_map(params_from_iter(ids.iter().copied()), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?, r.get::<_, i64>(2)?)))? {
let (id, note_id, offset) = row?;
out.insert(id, (note_id, offset.max(0) as usize));
}
Ok(out)
}
pub(crate) fn namespace_root(conn: &Connection, namespace_id: i64) -> Result<Option<String>> {
Ok(conn.query_row("SELECT root FROM namespace_roots WHERE namespace_id=?1", [namespace_id], |r| r.get(0)).optional()?)
}
pub(crate) fn absolute_note_path(conn: &Connection, namespace_id: i64, stored: &str) -> String {
match namespace_root(conn, namespace_id) {
Ok(Some(root)) => Path::new(&root).join(stored.replace('/', std::path::MAIN_SEPARATOR_STR)).to_string_lossy().into_owned(),
_ => stored.to_string(),
}
}
pub(crate) fn record_value(conn: &Connection, key: &RecordKey) -> Result<Option<Value>> {
Ok(record_values(conn, &[key.id])?.remove(&key.id))
}
pub(crate) fn record_values(conn: &Connection, ids: &[i64]) -> Result<BTreeMap<i64, Value>> {
let mut out = BTreeMap::new();
if ids.is_empty() { return Ok(out); }
let placeholders = vec!["?"; ids.len()].join(",");
let params = ids.iter().map(|id| SqlValue::Integer(*id)).collect::<Vec<_>>();
let mut stmt = conn.prepare(&format!("SELECT r.id,r.namespace_id,n.text,r.kind,s.text,r.created_at_us,r.updated_at_us,r.revision,
r.metadata_json,r.evidence_json,r.payload_json FROM records r
JOIN strings n ON n.id=r.namespace_id JOIN strings s ON s.id=r.scope_id WHERE r.id IN ({placeholders}) ORDER BY r.id"))?;
let mut rows: Vec<(i64, i64, String, i64, String, i64, i64, i64, String, String, String)> = Vec::new();
for row in stmt.query_map(params_from_iter(params.iter().cloned()), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?,
r.get::<_, String>(2)?, r.get::<_, i64>(3)?, r.get::<_, String>(4)?, r.get::<_, i64>(5)?, r.get::<_, i64>(6)?,
r.get::<_, i64>(7)?, r.get::<_, String>(8)?, r.get::<_, String>(9)?, r.get::<_, String>(10)?)))? {
rows.push(row?);
}
let mut tags_stmt = conn.prepare(&format!("SELECT rt.record_id,t.text FROM record_tags rt JOIN strings t ON t.id=rt.tag_id \
WHERE rt.record_id IN ({placeholders}) ORDER BY rt.record_id,t.text"))?;
let mut tags: BTreeMap<i64, Vec<String>> = BTreeMap::new();
for row in tags_stmt.query_map(params_from_iter(params.iter().cloned()), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))? {
let (id, tag) = row?;
tags.entry(id).or_default().push(tag);
}
let mut note_meta: BTreeMap<i64, (String, String)> = BTreeMap::new();
{
let mut stmt = conn.prepare(&format!("SELECT n.record_id,n.path,n.name FROM notes n \
WHERE n.record_id IN ({placeholders})"))?;
for row in stmt.query_map(params_from_iter(params.iter().cloned()), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?, r.get::<_, String>(2)?)))? {
let (id, path, name) = row?;
note_meta.insert(id, (path, name));
}
}
for (id, namespace_id, namespace, kind_code, scope, created, updated, revision, metadata, evidence, payload) in rows {
let kind = RecordKind::from_code(kind_code).ok_or_else(|| Error::Validation("invalid stored record kind".into()))?;
let header = RecordHeader { id, namespace, kind, scope,
created_at_us: created, updated_at_us: updated, revision, tags: tags.remove(&id).unwrap_or_default(),
metadata: serde_json::from_str(&metadata)?, evidence: serde_json::from_str(&evidence)? };
let mut value = serde_json::to_value(header)?;
let object = value.as_object_mut().ok_or_else(|| Error::Validation("invalid stored header".into()))?;
let mut payload: Metadata = serde_json::from_str(&payload)?;
if let Some(type_id) = payload.get("memory_type_id").and_then(Value::as_i64) {
payload.insert("memory_type".into(), Value::String(term_text(conn, type_id)?));
payload.remove("memory_type_id");
}
if kind == RecordKind::Note {
let (stored, name) = note_meta.remove(&id).unwrap_or_default();
let source = absolute_note_path(conn, namespace_id, &stored);
payload.insert("source".into(), Value::String(source));
payload.insert("title".into(), Value::String(name));
}
object.extend(payload);
out.insert(id, value);
}
Ok(out)
}
pub(crate) fn matches_filter(conn: &Connection, key: &RecordKey, filter: &ReadFilter) -> Result<bool> {
validate_filter(filter)?;
let row: Option<(i64, i64)> = conn.query_row("SELECT namespace_id,scope_id FROM records WHERE id=?1", [key.id],
|r| Ok((r.get(0)?, r.get(1)?))).optional()?;
let Some((namespace_id, scope_id)) = row else { return Ok(false) };
if term_text(conn, namespace_id)? != text::normalized_tag(&filter.namespace) { return Ok(false); }
let scope = term_text(conn, scope_id)?;
if !filter.scopes.iter().any(|s| text::normalized_tag(s) == scope) { return Ok(false); }
for tag in &filter.tags {
let exists: bool = conn.query_row("SELECT EXISTS(SELECT 1 FROM record_tags rt JOIN strings t ON t.id=rt.tag_id WHERE rt.record_id=?1 AND t.text=?2)",
params![key.id, text::normalized_tag(tag)], |r| r.get(0))?;
if !exists { return Ok(false); }
}
Ok(true)
}
pub(crate) fn get<T: DeserializeOwned>(conn: &Connection, key: &RecordKey, filter: &ReadFilter) -> Result<T> {
if !matches_filter(conn, key, filter)? { return Err(Error::NotFound(key.id.to_string())); }
serde_json::from_value(record_value(conn, key)?.ok_or_else(|| Error::NotFound(key.id.to_string()))?).map_err(Error::from)
}
pub(crate) fn load_many<T: DeserializeOwned>(conn: &Connection, ids: &[i64], filter: &ReadFilter) -> Result<BTreeMap<i64, T>> {
let mut out = BTreeMap::new();
if ids.is_empty() { return Ok(out); }
validate_filter(filter)?;
let (condition, values) = filter_sql(filter, &[], true)?;
let placeholders = vec!["?"; ids.len()].join(",");
let mut stmt = conn.prepare(&format!("SELECT r.id FROM records r WHERE r.id IN ({placeholders}) AND {condition} ORDER BY r.id"))?;
let params = ids.iter().map(|id| SqlValue::Integer(*id)).chain(values).collect::<Vec<_>>();
let allowed = stmt.query_map(params_from_iter(params), |r| r.get::<_, i64>(0))?.collect::<std::result::Result<Vec<_>, _>>()?;
for (id, value) in record_values(conn, &allowed)? {
out.insert(id, serde_json::from_value(value)?);
}
Ok(out)
}
pub(crate) fn filter_sql(filter: &ReadFilter, kinds: &[RecordKind], by_ids: bool) -> Result<(String, Vec<SqlValue>)> {
validate_filter(filter)?;
let mut query = if by_ids {
"+r.namespace_id=(SELECT id FROM strings WHERE text=?)".to_string()
} else {
"r.namespace_id=(SELECT id FROM strings WHERE text=?)".to_string()
};
let mut values = vec![SqlValue::Text(text::normalized_tag(&filter.namespace))];
query.push_str(" AND r.scope_id IN (SELECT id FROM strings WHERE text IN (");
query.push_str(&vec!["?"; filter.scopes.len()].join(",")); query.push_str("))");
values.extend(filter.scopes.iter().map(|s| SqlValue::Text(text::normalized_tag(s))));
if !kinds.is_empty() {
query.push_str(" AND r.kind IN ("); query.push_str(&vec!["?"; kinds.len()].join(",")); query.push(')');
values.extend(kinds.iter().map(|k| SqlValue::Integer(k.code())));
}
for tag in &filter.tags {
query.push_str(" AND EXISTS(SELECT 1 FROM record_tags rt JOIN strings t ON t.id=rt.tag_id WHERE rt.record_id=r.id AND t.text=?)");
values.push(SqlValue::Text(text::normalized_tag(tag)));
}
Ok((query, values))
}
pub(crate) fn select_ids(conn: &Connection, filter: &ReadFilter, kinds: &[RecordKind]) -> Result<Vec<i64>> {
let (condition, values) = filter_sql(filter, kinds, false)?;
let mut stmt = conn.prepare(&format!("SELECT r.id FROM records r WHERE {condition} ORDER BY r.id"))?;
let ids = stmt.query_map(params_from_iter(values), |r| r.get::<_, i64>(0))?
.collect::<std::result::Result<Vec<_>, _>>()?;
Ok(ids)
}
pub(crate) fn count_matches(conn: &Connection, filter: &ReadFilter, kinds: &[RecordKind]) -> Result<usize> {
let (condition, values) = filter_sql(filter, kinds, false)?;
let count: i64 = conn.query_row(&format!("SELECT COUNT(*) FROM records r WHERE {condition}"), params_from_iter(values), |r| r.get(0))?;
Ok(count as usize)
}
pub(crate) fn record_text(kind: RecordKind, payload: &Value) -> String {
let field = |key: &str| payload.get(key).and_then(Value::as_str).unwrap_or("").to_string();
match kind {
RecordKind::Memory => field("judgment"),
RecordKind::Entity => entity_body(payload),
RecordKind::Relation => format!("{} {} {} {}", field("subject_name"), field("predicate"), field("object_name"), field("reason")),
RecordKind::Event => format!("{} {} {} {}", field("name"), field("summary"), name_list(payload), field("reason")),
RecordKind::Note | RecordKind::Chunk => String::new(),
}
}
pub(crate) fn record_name(kind: RecordKind, payload: &Value) -> String {
match kind {
RecordKind::Entity => payload.get("name").and_then(Value::as_str).unwrap_or("").to_string(),
_ => String::new(),
}
}
pub(crate) fn event_text_lengths(conn: &Connection, ids: &[i64]) -> Result<BTreeMap<i64, usize>> {
let mut out = BTreeMap::new();
if ids.is_empty() { return Ok(out); }
let placeholders = vec!["?"; ids.len()].join(",");
let params = ids.iter().map(|id| SqlValue::Integer(*id)).collect::<Vec<_>>();
let field = |name: &str| format!(
"CASE WHEN json_type(r.payload_json,'$.{name}')='text' THEN json_extract(r.payload_json,'$.{name}') ELSE '' END");
let names = "COALESCE(CASE WHEN json_type(r.payload_json,'$.participant_names')='array' \
THEN (SELECT group_concat(j.value,' ') FROM json_each(r.payload_json,'$.participant_names') j \
WHERE j.type='text') ELSE '' END,'')";
let mut stmt = conn.prepare(&format!(
"SELECT r.id, LENGTH({} || ' ' || {} || ' ' || {names} || ' ' || {}) \
FROM records r WHERE r.id IN ({placeholders}) ORDER BY r.id",
field("name"), field("summary"), field("reason")))?;
for row in stmt.query_map(params_from_iter(params.iter().cloned()), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?)))? {
let (id, length) = row?;
out.insert(id, length.max(0) as usize);
}
Ok(out)
}
pub(crate) fn entity_names(conn: &Connection, ids: &[i64]) -> Result<BTreeMap<i64, String>> {
let mut out = BTreeMap::new();
if ids.is_empty() { return Ok(out); }
let placeholders = vec!["?"; ids.len()].join(",");
let mut stmt = conn.prepare(&format!("SELECT record_id,name FROM entities WHERE record_id IN ({placeholders})"))?;
let params = ids.iter().map(|id| SqlValue::Integer(*id)).collect::<Vec<_>>();
for row in stmt.query_map(params_from_iter(params), |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))? {
let (id, name) = row?;
out.insert(id, name);
}
Ok(out)
}
fn name_list(payload: &Value) -> String {
payload.get("participant_names").and_then(Value::as_array)
.map(|names| names.iter().filter_map(Value::as_str).collect::<Vec<_>>().join(" "))
.unwrap_or_default()
}
fn entity_body(payload: &Value) -> String {
let aliases = payload.get("aliases").and_then(Value::as_array)
.map(|a| a.iter().filter_map(Value::as_str).collect::<Vec<_>>().join(" ")).unwrap_or_default();
let summary = payload.get("summary").and_then(Value::as_str).unwrap_or("");
let attr_text = payload.get("attributes").and_then(Value::as_object).map(|attrs| {
attrs.iter().map(|(key, values)| {
let joined = values.as_array().map(|v| v.iter().filter_map(Value::as_str).collect::<Vec<_>>().join(" ")).unwrap_or_default();
format!("{key} {joined}")
}).collect::<Vec<_>>().join(" ")
}).unwrap_or_default();
format!("{aliases} {summary} {attr_text}")
}
pub(crate) fn select_keys(conn: &Connection, filter: &ReadFilter, kinds: &[RecordKind], limit: usize, after: Option<&str>) -> Result<Vec<RecordKey>> { let (mut condition, mut values) = filter_sql(filter, kinds, false)?;
if let Some(cursor) = after {
let id: i64 = cursor.parse().map_err(|_| Error::Validation("invalid page cursor".into()))?;
condition.push_str(" AND r.id>?");
values.push(SqlValue::Integer(id));
}
values.push(SqlValue::Integer(limit.min(i64::MAX as usize) as i64));
let mut stmt = conn.prepare(&format!("SELECT r.id FROM records r WHERE {condition} ORDER BY r.id LIMIT ?"))?;
let rows = stmt.query_map(params_from_iter(values), |r| r.get::<_, i64>(0))?;
let mut keys = Vec::new();
for row in rows { keys.push(RecordKey { id: row? }); }
Ok(keys)
}
pub(crate) fn list<T: DeserializeOwned>(conn: &Connection, kind: RecordKind, request: &PageRequest) -> Result<Page<T>> {
validate_limit(request.limit)?;
let mut keys = select_keys(conn, &request.filter, &[kind], request.limit + 1, request.after.as_deref())?;
let has_more = keys.len() > request.limit;
keys.truncate(request.limit);
let next_cursor = if has_more { keys.last().map(RecordKey::index_key) } else { None };
let ids: Vec<i64> = keys.iter().map(|key| key.id).collect();
let mut loaded: BTreeMap<i64, T> = load_many(conn, &ids, &request.filter)?;
let items = keys.into_iter().filter_map(|key| loaded.remove(&key.id)).collect::<Vec<_>>();
Ok(Page { items, next_cursor })
}
pub(crate) fn delete_record(conn: &Connection, key: &RecordKey) -> Result<bool> {
let namespace = namespace_of(conn, key.id)?;
let changed = conn.execute("DELETE FROM records WHERE id=?1", [key.id]);
let changed = match changed {
Err(rusqlite::Error::SqliteFailure(err, _)) if err.code == rusqlite::ErrorCode::ConstraintViolation =>
return Err(Error::Conflict(format!("record {} is still referenced", key.id))),
other => other?,
};
if changed > 0 {
next_revision(conn, key.id)?;
if let Some(namespace) = namespace { touch_namespace(&namespace); }
}
Ok(changed > 0)
}
#[cfg(test)]
mod tests {
use super::*;
fn batched_load_plan(conn: &Connection, ids: &[i64], by_ids: bool) -> String {
let (condition, values) = filter_sql(&ReadFilter::default(), &[], by_ids).unwrap();
let placeholders = vec!["?"; ids.len()].join(",");
let sql = format!("EXPLAIN QUERY PLAN SELECT r.id FROM records r WHERE r.id IN ({placeholders}) AND {condition} ORDER BY r.id");
let params: Vec<SqlValue> = ids.iter().map(|id| SqlValue::Integer(*id)).chain(values).collect();
let mut stmt = conn.prepare(&sql).unwrap();
let plans: Vec<String> = stmt.query_map(params_from_iter(params), |row| row.get::<_, String>(3))
.unwrap().map(|row| row.unwrap()).collect();
plans.join(" | ")
}
fn seed(kb: &KnowledgeBase, rows: i64) {
let inputs: Vec<crate::MemoryInput> = (1..=rows).map(|i| crate::MemoryInput::new(format!("记录 {i}"))).collect();
kb.memories().upsert_many(&inputs).unwrap();
}
fn sample_ids() -> Vec<i64> { (1..=10).collect() }
#[test]
fn batched_load_stays_on_the_primary_key() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
seed(&kb, 100);
let guard = kb.read().unwrap();
let plan = batched_load_plan(guard.conn(), &sample_ids(), true);
assert!(plan.contains("INTEGER PRIMARY KEY"), "批量取回退化为扫索引:{plan}");
}
#[test]
fn event_text_lengths_match_record_text() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
let entity = |name: &str| crate::EntityInput { record: Default::default(), name: name.into(),
entity_type: "person".into(), aliases: vec![], attributes: BTreeMap::new(), summary: String::new() };
let created = kb.graph().apply_batch(&crate::GraphBatch {
entities: vec![entity("甲"), entity("乙")], ..Default::default()
}).unwrap().value;
let (first, second) = (created.entities[0].header.id, created.entities[1].header.id);
let created = kb.graph().apply_batch(&crate::GraphBatch {
events: vec![
crate::EventInput { record: Default::default(), name: "别鹤典仪".into(), summary: "两人同去".into(),
participants: vec![first, second], confidence: 1.0, reason: "有人证".into() },
crate::EventInput { record: Default::default(), name: "堂中自语".into(), summary: String::new(),
participants: vec![first], confidence: 1.0, reason: String::new() },
], ..Default::default()
}).unwrap().value;
let ids: Vec<i64> = created.events.iter().map(|event| event.header.id).collect();
{
let raw = Connection::open(dir.path().join("store.sqlite3")).unwrap();
let payloads = [
r#"{"name":7,"summary":"只剩数字名","participant_names":"甲 乙","reason":null}"#,
r#"{"name":"正常","summary":null,"participant_names":["甲",7,"乙"],"reason":"理由"}"#,
];
for (id, payload) in ids.iter().zip(payloads) {
raw.execute("UPDATE records SET payload_json=?1 WHERE id=?2", params![payload, id]).unwrap();
}
}
let guard = kb.read().unwrap();
let conn = guard.conn();
let lengths = event_text_lengths(conn, &ids).unwrap();
assert_eq!(lengths.len(), ids.len(), "每条事件都该有长度");
for id in ids {
let payload = record_values(conn, &[id]).unwrap().remove(&id).unwrap();
assert_eq!(lengths[&id], record_text(RecordKind::Event, &payload).chars().count(),
"事件 {id} 的 SQL 长度与 record_text 不一致");
}
}
fn fixture_space() -> crate::embeddings::EmbeddingSpace {
crate::embeddings::EmbeddingSpace { id: "v".into(), model: "fixture/v1".into(),
dimension: 2, text_version: 1, encoding: "f32".into() }
}
fn cache_partition(kb: &KnowledgeBase, space: &crate::embeddings::EmbeddingSpace, namespace: &str) {
let guard = kb.read().unwrap();
kb.partition(guard.conn(), space, namespace, "public").unwrap();
}
fn cached_namespaces(kb: &KnowledgeBase) -> BTreeSet<String> {
kb.engine.vectors.entries.lock().keys().map(|(_, namespace, _)| namespace.clone()).collect()
}
fn namespace_filter(namespace: &str) -> ReadFilter {
ReadFilter { namespace: namespace.into(), scopes: vec!["public".into()], tags: vec![], note_ids: vec![] }
}
#[test]
fn invalidating_one_namespace_leaves_the_others_alone() {
let cache = VectorCache::new();
let key = |namespace: &str| ("v".to_string(), namespace.to_string(), "public".to_string());
cache.entries.lock().insert(key("a"), None);
cache.entries.lock().insert(key("b"), None);
let epoch_b = cache.epoch_of("b");
cache.invalidate_namespaces(&HashSet::from(["a".to_string()]));
assert!(cache.entries.lock().get(&key("a")).is_none(), "写过的领域要清掉条目");
assert!(cache.entries.lock().get(&key("b")).is_some(), "没写过的领域不该被牵连");
assert_eq!(cache.epoch_of("b"), epoch_b, "没写过的领域版本号不动");
assert_ne!(cache.epoch_of("a"), epoch_b, "写过的领域版本号要前进,在途载入才会作废");
let epoch_a = cache.epoch_of("a");
cache.invalidate();
assert!(cache.entries.lock().is_empty());
assert!(cache.epoch_of("a") > epoch_a && cache.epoch_of("b") > epoch_b);
}
#[test]
fn writing_one_namespace_keeps_other_vector_partitions_cached() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
let space = fixture_space();
for namespace in ["a", "b"] { cache_partition(&kb, &space, namespace); }
assert_eq!(cached_namespaces(&kb), BTreeSet::from(["a".to_string(), "b".to_string()]));
let mut input = crate::MemoryInput::new("写在 a 领域的一条");
input.record.namespace = "a".into();
kb.memories().upsert(input).unwrap();
assert_eq!(cached_namespaces(&kb), BTreeSet::from(["b".to_string()]), "只该清掉被写的那个领域");
}
#[test]
fn deleting_a_record_evicts_only_its_own_namespace() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
let mut input = crate::MemoryInput::new("要被删掉的一条");
input.record.namespace = "a".into();
let id = kb.memories().upsert(input).unwrap().value.header.id;
let space = fixture_space();
for namespace in ["a", "b"] { cache_partition(&kb, &space, namespace); }
kb.memories().delete(id, &namespace_filter("a")).unwrap();
assert_eq!(cached_namespaces(&kb), BTreeSet::from(["b".to_string()]), "删掉的领域要清,别的领域留着");
}
#[test]
fn filling_vectors_only_evicts_the_namespaces_it_wrote() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
let space = fixture_space();
kb.embeddings().register_space(space.clone()).unwrap();
kb.embeddings().register_embedder("v", |texts: &[String]| -> std::result::Result<Vec<Vec<f32>>, crate::EmbedCallbackError> {
Ok(texts.iter().map(|_| vec![1.0f32, 0.0]).collect())
}).unwrap();
for namespace in ["a", "b"] { cache_partition(&kb, &space, namespace); }
kb.memories().upsert(crate::MemoryInput::new("补齐用的一条")).unwrap();
kb.embeddings().sync("v", 32).unwrap();
let cached = cached_namespaces(&kb);
assert!(cached.contains("a") && cached.contains("b"),
"补齐只写了 default 领域,a 与 b 的分区缓存不该被牵连:{cached:?}");
}
#[test]
fn an_unregistered_write_falls_back_to_invalidating_everything() {
let dir = tempfile::tempdir().unwrap();
let kb = KnowledgeBase::open(dir.path()).unwrap();
let space = fixture_space();
for namespace in ["a", "b"] { cache_partition(&kb, &space, namespace); }
kb.mutate(|tx| Ok(tx.execute("INSERT INTO meta(key,value) VALUES ('cache_probe',1)
ON CONFLICT(key) DO UPDATE SET value=excluded.value", [])?)).unwrap();
assert!(cached_namespaces(&kb).is_empty(), "登记为空却改过行时必须整体失效");
}
}