memoryd: the FTS5 index over session logs and notes, with labels
Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
@@ -0,0 +1,421 @@
|
||||
//! The index: SQLite FTS5 over the session logs and `notes.md` (M5a spec, section 6). Disposable:
|
||||
//! it is rebuilt from those files. What is indexed: the owner's messages, the assistant's
|
||||
//! replies, and each tool call by name and arguments; never a tool's result.
|
||||
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use proto::{DataClass, LogRecord};
|
||||
use rusqlite::{Connection, params};
|
||||
|
||||
use crate::notes::{self, Kind};
|
||||
|
||||
const SCHEMA: &str = "
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS hits USING fts5(text, tokenize = 'unicode61');
|
||||
CREATE TABLE IF NOT EXISTS sources (
|
||||
rowid INTEGER PRIMARY KEY, kind TEXT NOT NULL, session TEXT NOT NULL, epoch INTEGER NOT NULL,
|
||||
record INTEGER NOT NULL, time TEXT NOT NULL, role TEXT NOT NULL, class TEXT NOT NULL,
|
||||
untrusted INTEGER NOT NULL);
|
||||
CREATE TABLE IF NOT EXISTS progress (log TEXT PRIMARY KEY, records INTEGER NOT NULL);
|
||||
";
|
||||
|
||||
/// One indexed item.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Source {
|
||||
/// `session`, `noted`, `moved` or `dropped`.
|
||||
pub kind: String,
|
||||
pub session: String,
|
||||
pub epoch: i64,
|
||||
/// The record's line in its log, or the note's place in `notes.md` (from 0).
|
||||
pub record: i64,
|
||||
pub time: String,
|
||||
/// `owner`, `assistant` or `call` for a session record; the file (`memory.md`, `user.md`) for
|
||||
/// a moved or dropped note; empty for a noted one.
|
||||
pub role: String,
|
||||
pub class: DataClass,
|
||||
pub untrusted: bool,
|
||||
}
|
||||
|
||||
/// How much a reindex read.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub struct Counts {
|
||||
pub logs: u64,
|
||||
pub records: u64,
|
||||
pub notes: u64,
|
||||
}
|
||||
|
||||
pub struct Index {
|
||||
conn: Connection,
|
||||
max_class: DataClass,
|
||||
}
|
||||
|
||||
fn db(e: rusqlite::Error) -> String {
|
||||
format!("index: {e}")
|
||||
}
|
||||
|
||||
pub fn class_name(class: DataClass) -> &'static str {
|
||||
match class {
|
||||
DataClass::Public => "public",
|
||||
DataClass::Private => "private",
|
||||
DataClass::Secret => "secret",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn class_of(name: &str) -> DataClass {
|
||||
match name {
|
||||
"public" => DataClass::Public,
|
||||
"secret" => DataClass::Secret,
|
||||
_ => DataClass::Private,
|
||||
}
|
||||
}
|
||||
|
||||
fn int(n: usize) -> i64 {
|
||||
i64::try_from(n).unwrap_or(i64::MAX)
|
||||
}
|
||||
|
||||
impl Index {
|
||||
/// Open or create the index at `path` (its directory 0700, the file 0600).
|
||||
pub fn open(path: &Path, max_class: DataClass) -> Result<Index, String> {
|
||||
if let Some(dir) = path.parent() {
|
||||
std::fs::create_dir_all(dir).map_err(|e| format!("{}: {e}", dir.display()))?;
|
||||
std::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o700))
|
||||
.map_err(|e| format!("{}: {e}", dir.display()))?;
|
||||
}
|
||||
let conn = Connection::open(path).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
|
||||
.map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
conn.execute_batch(SCHEMA)
|
||||
.map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
Ok(Index { conn, max_class })
|
||||
}
|
||||
|
||||
/// The connection, for `recall`.
|
||||
pub fn conn(&self) -> &Connection {
|
||||
&self.conn
|
||||
}
|
||||
|
||||
fn progress(&self, log: &str) -> Result<usize, String> {
|
||||
// `self.conn.query_row("SELECT records FROM progress WHERE log = ?1", params![log], |row|
|
||||
// row.get(0))` as an Option<i64>: Err(rusqlite::Error::QueryReturnedNoRows) means None,
|
||||
// another error goes through `db`. Return it as a usize (`usize::try_from`, 0 when it does
|
||||
// not fit), 0 for None.
|
||||
let value: Option<i64> = match self.conn.query_row(
|
||||
"SELECT records FROM progress WHERE log = ?1",
|
||||
params![log],
|
||||
|row| row.get(0),
|
||||
) {
|
||||
Ok(v) => Some(v),
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => None,
|
||||
Err(e) => return Err(db(e)),
|
||||
};
|
||||
let n = match value {
|
||||
Some(v) => usize::try_from(v).unwrap_or(0),
|
||||
None => 0,
|
||||
};
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn add(&self, text: &str, source: &Source) -> Result<(), String> {
|
||||
// 1. `execute("INSERT INTO hits (text) VALUES (?1)", params![text])`.
|
||||
// 2. `let rowid = self.conn.last_insert_rowid();`
|
||||
// 3. `execute("INSERT INTO sources (rowid, kind, session, epoch, record, time, role, class,
|
||||
// untrusted) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", params![rowid, kind, session,
|
||||
// epoch, record, time, role, class_name(class), untrusted])`. Errors through `db`.
|
||||
self.conn
|
||||
.execute("INSERT INTO hits (text) VALUES (?1)", params![text])
|
||||
.map_err(db)?;
|
||||
let rowid = self.conn.last_insert_rowid();
|
||||
self.conn
|
||||
.execute(
|
||||
"INSERT INTO sources (rowid, kind, session, epoch, record, time, role, class, untrusted) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
|
||||
params![
|
||||
rowid,
|
||||
source.kind.clone(),
|
||||
source.session.clone(),
|
||||
source.epoch,
|
||||
source.record,
|
||||
source.time.clone(),
|
||||
source.role.clone(),
|
||||
class_name(source.class),
|
||||
source.untrusted,
|
||||
],
|
||||
)
|
||||
.map_err(db)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_progress(&self, log: &str, records: usize) -> Result<(), String> {
|
||||
// `execute("INSERT INTO progress (log, records) VALUES (?1, ?2) ON CONFLICT(log) DO UPDATE
|
||||
// SET records = excluded.records", params![log, int(records)])`, errors through `db`.
|
||||
self.conn
|
||||
.execute(
|
||||
"INSERT INTO progress (log, records) VALUES (?1, ?2) ON CONFLICT(log) DO UPDATE SET records = excluded.records",
|
||||
params![log, int(records)],
|
||||
)
|
||||
.map_err(db)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Index the records of one session log past what was indexed before; how many were read.
|
||||
/// The label is walked from the log's first record: the highest class of any tool result
|
||||
/// before a record (at least `private`), and whether any was untrusted.
|
||||
pub fn index_log(&self, path: &Path, session: &str, epoch: i64) -> Result<usize, String> {
|
||||
// 1. `read_to_string(path)` (an error names the path). `key` =
|
||||
// format!("{session}/{epoch}"); `done = self.progress(&key)?`.
|
||||
// 2. The label starts at class Private, untrusted false; `read` at 0.
|
||||
// 3. For each line with its place n (from 0): parse it as a LogRecord (an error ->
|
||||
// Err(format!("{}: line {}: {e}", path.display(), n + 1))). When n >= done: read += 1,
|
||||
// and `self.index_record(&record, &base)?` where base is a Source with kind "session",
|
||||
// session, epoch, record int(n), time and role empty, and the label as it stands. Then,
|
||||
// for every record (indexed or not): a ToolResult raises the label: class =
|
||||
// class.max(its class), untrusted |= its untrusted.
|
||||
// 4. `self.set_progress(&key, done + read)?`; Ok(read).
|
||||
let text = std::fs::read_to_string(path).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
let key = format!("{session}/{epoch}");
|
||||
let done = self.progress(&key)?;
|
||||
let mut class = DataClass::Private;
|
||||
let mut untrusted = false;
|
||||
let mut read = 0;
|
||||
for (n, line) in text.lines().enumerate() {
|
||||
let record: LogRecord = serde_json::from_str(line)
|
||||
.map_err(|e| format!("{}: line {}: {e}", path.display(), n + 1))?;
|
||||
if n >= done {
|
||||
read += 1;
|
||||
let base = Source {
|
||||
kind: "session".to_string(),
|
||||
session: session.to_string(),
|
||||
epoch,
|
||||
record: int(n),
|
||||
time: String::new(),
|
||||
role: String::new(),
|
||||
class,
|
||||
untrusted,
|
||||
};
|
||||
self.index_record(&record, &base)?;
|
||||
}
|
||||
if let LogRecord::ToolResult {
|
||||
class: c,
|
||||
untrusted: u,
|
||||
..
|
||||
} = &record
|
||||
{
|
||||
class = class.max(*c);
|
||||
untrusted |= *u;
|
||||
}
|
||||
}
|
||||
self.set_progress(&key, done + read)?;
|
||||
Ok(read)
|
||||
}
|
||||
|
||||
/// Index what one log record holds: the owner's message, the assistant's reply, and each tool
|
||||
/// call by name and arguments. `base` carries the kind, session, epoch, line and label; each item
|
||||
/// gets its own role and time. Nothing when the label's class is above `max_class`, and nothing
|
||||
/// for an item with no text.
|
||||
fn index_record(&self, record: &LogRecord, base: &Source) -> Result<(), String> {
|
||||
// When base.class > self.max_class: Ok(()) and nothing indexed. Otherwise each item is
|
||||
// `self.add(text, &Source { time: time.to_rfc3339(), role, ..base.clone() })`, skipped when
|
||||
// text.trim() is empty. User -> its content, role "owner". Assistant -> its content when
|
||||
// Some, role "assistant"; then each tool call as format!("{} {}", name, arguments), role
|
||||
// "call"; all with the record's time. Every other record: nothing.
|
||||
if base.class > self.max_class {
|
||||
return Ok(());
|
||||
}
|
||||
match record {
|
||||
LogRecord::User { time, content } => {
|
||||
if !content.trim().is_empty() {
|
||||
let src = Source {
|
||||
time: time.to_rfc3339(),
|
||||
role: "owner".to_string(),
|
||||
..base.clone()
|
||||
};
|
||||
self.add(content, &src)?;
|
||||
}
|
||||
}
|
||||
LogRecord::Assistant {
|
||||
time,
|
||||
content,
|
||||
tool_calls,
|
||||
..
|
||||
} => {
|
||||
if let Some(c) = content
|
||||
&& !c.trim().is_empty()
|
||||
{
|
||||
let src = Source {
|
||||
time: time.to_rfc3339(),
|
||||
role: "assistant".to_string(),
|
||||
..base.clone()
|
||||
};
|
||||
self.add(c, &src)?;
|
||||
}
|
||||
for call in tool_calls {
|
||||
let text = format!("{} {}", call.name, call.arguments);
|
||||
if !text.trim().is_empty() {
|
||||
let src = Source {
|
||||
time: time.to_rfc3339(),
|
||||
role: "call".to_string(),
|
||||
..base.clone()
|
||||
};
|
||||
self.add(&text, &src)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Index the entries of `notes.md` past what was indexed before; how many were read.
|
||||
pub fn index_notes(&self, text: &str) -> Result<usize, String> {
|
||||
// `done = self.progress("notes.md")?`. Parse `text` with notes::parse and skip the first
|
||||
// `done`. For each remaining note (with its place n): read += 1; its kind, role, class and
|
||||
// untrusted: Noted -> ("noted", "", its class, its untrusted); Moved(file) -> ("moved",
|
||||
// file.name(), Private, false); Dropped(file) -> ("dropped", file.name(), Private, false).
|
||||
// Skip it when class > self.max_class. Otherwise `self.add(¬e.text, &Source { kind,
|
||||
// session: note.session, epoch: 0, record: int(n), time: note.time, role, class, untrusted
|
||||
// })?`. Then `self.set_progress("notes.md", done + read)?`; Ok(read).
|
||||
let done = self.progress("notes.md")?;
|
||||
let notes = notes::parse(text);
|
||||
let mut read = 0;
|
||||
for (n, note) in notes.iter().enumerate() {
|
||||
if n < done {
|
||||
continue;
|
||||
}
|
||||
read += 1;
|
||||
let (kind, role, class, untrusted) = match ¬e.kind {
|
||||
Kind::Noted { class, untrusted } => {
|
||||
("noted".to_string(), String::new(), *class, *untrusted)
|
||||
}
|
||||
Kind::Moved(file) => (
|
||||
"moved".to_string(),
|
||||
file.name().to_string(),
|
||||
DataClass::Private,
|
||||
false,
|
||||
),
|
||||
Kind::Dropped(file) => (
|
||||
"dropped".to_string(),
|
||||
file.name().to_string(),
|
||||
DataClass::Private,
|
||||
false,
|
||||
),
|
||||
};
|
||||
if class > self.max_class {
|
||||
continue;
|
||||
}
|
||||
self.add(
|
||||
¬e.text,
|
||||
&Source {
|
||||
kind,
|
||||
session: note.session.clone(),
|
||||
epoch: 0,
|
||||
record: int(n),
|
||||
time: note.time.clone(),
|
||||
role,
|
||||
class,
|
||||
untrusted,
|
||||
},
|
||||
)?;
|
||||
}
|
||||
self.set_progress("notes.md", done + read)?;
|
||||
Ok(read)
|
||||
}
|
||||
|
||||
/// Index everything new: every `sessions/<id>/<epoch>.jsonl`, then `notes.md`. A log that
|
||||
/// cannot be read or parsed is skipped and named in the returned problems; it is `loopd`'s.
|
||||
pub fn catch_up(
|
||||
&self,
|
||||
sessions: &Path,
|
||||
notes_md: &Path,
|
||||
) -> Result<(Counts, Vec<String>), String> {
|
||||
// For each (session, epoch, path) of `logs(sessions)?`: `self.index_log(&path, &session,
|
||||
// epoch)`; Ok(n) -> counts.logs += 1 and counts.records += n (u64::try_from, u64::MAX when
|
||||
// it does not fit); Err(e) -> push e onto `problems` and go on. Then read `notes_md`:
|
||||
// NotFound is an empty text, another error -> Err(format!("{}: {e}", notes_md.display()));
|
||||
// counts.notes = `self.index_notes(&text)?` as u64 (u64::try_from). Ok((counts, problems)).
|
||||
let mut counts = Counts::default();
|
||||
let mut problems: Vec<String> = Vec::new();
|
||||
for (session, epoch, path) in logs(sessions)? {
|
||||
match self.index_log(&path, &session, epoch) {
|
||||
Ok(n) => {
|
||||
counts.logs += 1;
|
||||
counts.records += u64::try_from(n).unwrap_or(u64::MAX);
|
||||
}
|
||||
Err(e) => problems.push(e),
|
||||
}
|
||||
}
|
||||
let text = match std::fs::read_to_string(notes_md) {
|
||||
Ok(t) => t,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(),
|
||||
Err(e) => return Err(format!("{}: {e}", notes_md.display())),
|
||||
};
|
||||
counts.notes = u64::try_from(self.index_notes(&text)?).unwrap_or(u64::MAX);
|
||||
Ok((counts, problems))
|
||||
}
|
||||
}
|
||||
|
||||
/// Every session log under `sessions`: `(session, epoch, path)`, sorted.
|
||||
pub fn logs(sessions: &Path) -> Result<Vec<(String, i64, PathBuf)>, String> {
|
||||
// Read the directory `sessions` (NotFound -> Ok(empty); another error -> Err naming it). For
|
||||
// each entry, its name is the session; read that directory (skip it when that fails); every
|
||||
// file named `<n>.jsonl` where n parses as an i64 gives (session, n, the file's path). Sort the
|
||||
// list and return it.
|
||||
let entries = match std::fs::read_dir(sessions) {
|
||||
Ok(e) => e,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
|
||||
Err(e) => return Err(format!("{}: {e}", sessions.display())),
|
||||
};
|
||||
let mut out: Vec<(String, i64, PathBuf)> = Vec::new();
|
||||
for entry in entries {
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let name = entry.file_name();
|
||||
let session = match name.to_str() {
|
||||
Some(s) => s,
|
||||
None => continue,
|
||||
};
|
||||
let dir = entry.path();
|
||||
let sub = match std::fs::read_dir(&dir) {
|
||||
Ok(d) => d,
|
||||
Err(_) => continue,
|
||||
};
|
||||
for sub_entry in sub {
|
||||
let sub_entry = match sub_entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let sub_name = sub_entry.file_name();
|
||||
let file_name = match sub_name.to_str() {
|
||||
Some(s) => s,
|
||||
None => continue,
|
||||
};
|
||||
if let Some(stem) = file_name.strip_suffix(".jsonl")
|
||||
&& let Ok(epoch) = stem.parse::<i64>()
|
||||
{
|
||||
out.push((session.to_string(), epoch, sub_entry.path()));
|
||||
}
|
||||
}
|
||||
}
|
||||
out.sort();
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Build a fresh index at `path` from `sessions` and `notes.md`: into a new file beside it,
|
||||
/// renamed over the old one, so a running `recall` keeps its index until the new one is whole.
|
||||
pub fn rebuild(
|
||||
path: &Path,
|
||||
max_class: DataClass,
|
||||
sessions: &Path,
|
||||
notes_md: &Path,
|
||||
) -> Result<(Counts, Vec<String>), String> {
|
||||
// `fresh = path.with_extension("rebuilding")`; remove it if it is there (ignore the result).
|
||||
// `Index::open(&fresh, max_class)?`, then `catch_up(sessions, notes_md)?` on it, and drop that
|
||||
// Index. `std::fs::rename(&fresh, path)` (an error names path). Return the result of catch_up.
|
||||
let fresh = path.with_extension("rebuilding");
|
||||
let _ = std::fs::remove_file(&fresh);
|
||||
let index = Index::open(&fresh, max_class)?;
|
||||
let result = index.catch_up(sessions, notes_md)?;
|
||||
drop(index);
|
||||
std::fs::rename(&fresh, path).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
Ok(result)
|
||||
}
|
||||
Reference in New Issue
Block a user