mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
feat(memory): MD-on-disk content + participant bucketing + async pipeline foundation (#1008)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>
This commit is contained in:
co-authored by
Claude Opus 4.7
Steven Enamakel
parent
dbe60389e6
commit
fd0398d5fa
@@ -66,6 +66,7 @@ tauri.key
|
||||
tauri.key.pub
|
||||
/target/
|
||||
src-tauri/target/
|
||||
.target-codex/
|
||||
|
||||
workflow
|
||||
.fastembed_cache
|
||||
|
||||
Generated
+1
-1
@@ -4541,7 +4541,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
|
||||
|
||||
[[package]]
|
||||
name = "openhuman"
|
||||
version = "0.53.3"
|
||||
version = "0.53.4"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"anyhow",
|
||||
|
||||
@@ -13,6 +13,10 @@ path = "src/main.rs"
|
||||
name = "slack-backfill"
|
||||
path = "src/bin/slack_backfill.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "gmail-backfill-3d"
|
||||
path = "src/bin/gmail_backfill_3d.rs"
|
||||
|
||||
[lib]
|
||||
name = "openhuman_core"
|
||||
crate-type = ["rlib"]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,474 @@
|
||||
//! Backfill the last N days of Gmail into the memory-tree content store.
|
||||
//!
|
||||
//! Authenticates via Composio (JWT from `<workspace>/auth-profiles.json`),
|
||||
//! fetches Gmail pages via `GMAIL_FETCH_EMAILS`, converts each thread into an
|
||||
//! [`EmailThread`], ingests it through `ingest_page_into_memory_tree` (which
|
||||
//! writes `.md` files via `content_store` and populates SQLite), then drains
|
||||
//! the async worker pool until idle.
|
||||
//!
|
||||
//! After draining, the binary performs an integrity check: for every chunk
|
||||
//! that has a `content_path` in SQLite, it verifies the on-disk SHA-256
|
||||
//! matches the stored `content_sha256`.
|
||||
//!
|
||||
//! # Prerequisites
|
||||
//!
|
||||
//! - Signed-in openhuman session JWT in the same workspace the desktop app
|
||||
//! uses (stored at `<workspace>/auth-profiles.json`).
|
||||
//! - Active Gmail connection on Composio for that user.
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run --bin gmail-backfill-3d
|
||||
//! cargo run --bin gmail-backfill-3d -- --days 7
|
||||
//! cargo run --bin gmail-backfill-3d -- --days 14 --page-size 100
|
||||
//! cargo run --bin gmail-backfill-3d -- --skip-drain
|
||||
//! cargo run --bin gmail-backfill-3d -- --skip-verify
|
||||
//! cargo run --bin gmail-backfill-3d -- --wipe
|
||||
//! ```
|
||||
//!
|
||||
//! Set `RUST_LOG=info` (or `debug`) for detailed output.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Parser;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use openhuman_core::openhuman::composio::client::build_composio_client;
|
||||
use openhuman_core::openhuman::composio::providers::gmail::ingest::ingest_page_into_memory_tree;
|
||||
use openhuman_core::openhuman::composio::providers::registry::{
|
||||
get_provider, init_default_providers,
|
||||
};
|
||||
use openhuman_core::openhuman::config::Config;
|
||||
use openhuman_core::openhuman::memory::tree::content_store::read::{
|
||||
verify_chunk_file, verify_summary_file, VerifyResult,
|
||||
};
|
||||
use openhuman_core::openhuman::memory::tree::jobs::drain_until_idle;
|
||||
use openhuman_core::openhuman::memory::tree::store::{
|
||||
get_chunk_content_pointers, list_chunks, list_summaries_with_content_path, ListChunksQuery,
|
||||
};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
name = "gmail-backfill-3d",
|
||||
about = "Backfill last N days of Gmail into the memory-tree content store (.md files + SQLite)."
|
||||
)]
|
||||
struct Cli {
|
||||
/// Lookback window in days. Default 3.
|
||||
#[arg(long, default_value_t = 3)]
|
||||
days: u32,
|
||||
|
||||
/// Page size per `GMAIL_FETCH_EMAILS` call (1–500).
|
||||
#[arg(long, default_value_t = 50)]
|
||||
page_size: u32,
|
||||
|
||||
/// Cap on pages we will request. Guards against runaway pagination.
|
||||
#[arg(long, default_value_t = 40)]
|
||||
max_pages: u32,
|
||||
|
||||
/// Include SPAM and TRASH messages in the fetch.
|
||||
#[arg(long, default_value_t = false)]
|
||||
include_spam_trash: bool,
|
||||
|
||||
/// Extra Gmail search query AND-ed with the default scope.
|
||||
#[arg(long)]
|
||||
query: Option<String>,
|
||||
|
||||
/// Skip draining the async worker pool after ingest (useful for quick
|
||||
/// smoke-test of file writes only).
|
||||
#[arg(long, default_value_t = false)]
|
||||
skip_drain: bool,
|
||||
|
||||
/// Skip the post-drain integrity check (SHA-256 file verification).
|
||||
#[arg(long, default_value_t = false)]
|
||||
skip_verify: bool,
|
||||
|
||||
/// Override the owner string embedded in chunk metadata. Defaults to
|
||||
/// `"gmail-backfill"`.
|
||||
#[arg(long)]
|
||||
owner: Option<String>,
|
||||
|
||||
/// Wipe `chunks.db` (+ wal/shm) AND `<content_root>/` before running.
|
||||
/// Useful after a chunker change that invalidates existing chunk IDs.
|
||||
#[arg(long, default_value_t = false)]
|
||||
wipe: bool,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
|
||||
.format_timestamp_secs()
|
||||
.try_init()
|
||||
.ok();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||||
)
|
||||
.with_target(true)
|
||||
.try_init()
|
||||
.ok();
|
||||
|
||||
let cli = Cli::parse();
|
||||
if cli.days == 0 {
|
||||
anyhow::bail!("--days must be >= 1");
|
||||
}
|
||||
|
||||
let config = Config::load_or_init()
|
||||
.await
|
||||
.context("[gmail_backfill_3d] Config::load_or_init failed")?;
|
||||
|
||||
if cli.wipe {
|
||||
wipe_memory_tree_state(&config)?;
|
||||
}
|
||||
|
||||
let client = build_composio_client(&config).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"No Composio client — user not signed in (no JWT). \
|
||||
Sign in via the desktop app first, then re-run this binary."
|
||||
)
|
||||
})?;
|
||||
|
||||
init_default_providers();
|
||||
let provider = get_provider("gmail").ok_or_else(|| {
|
||||
anyhow::anyhow!("GmailProvider not registered after init_default_providers")
|
||||
})?;
|
||||
|
||||
let owner = cli
|
||||
.owner
|
||||
.clone()
|
||||
.unwrap_or_else(|| "gmail-backfill".to_string());
|
||||
|
||||
let mut query = format!("in:inbox newer_than:{}d", cli.days);
|
||||
if !cli.include_spam_trash {
|
||||
query.push_str(" -in:spam -in:trash");
|
||||
}
|
||||
if let Some(extra) = cli
|
||||
.query
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
query.push(' ');
|
||||
query.push_str(extra);
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] start days={} page_size={} max_pages={} query={:?}",
|
||||
cli.days,
|
||||
cli.page_size,
|
||||
cli.max_pages,
|
||||
query,
|
||||
);
|
||||
|
||||
let content_root = config.memory_tree_content_root();
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] content_root={}",
|
||||
content_root.display()
|
||||
);
|
||||
|
||||
// ─── Fetch + ingest ────────────────────────────────────────────────────
|
||||
|
||||
let mut page_token: Option<String> = None;
|
||||
let mut total_chunks = 0usize;
|
||||
let mut total_pages = 0usize;
|
||||
let mut total_cost: f64 = 0.0;
|
||||
|
||||
for page_num in 0..cli.max_pages {
|
||||
let mut args = json!({
|
||||
"max_results": cli.page_size,
|
||||
"query": query,
|
||||
});
|
||||
if cli.include_spam_trash {
|
||||
args["include_spam_trash"] = json!(true);
|
||||
}
|
||||
if let Some(token) = &page_token {
|
||||
args["page_token"] = json!(token);
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] fetching page {}{}…",
|
||||
page_num,
|
||||
page_token.as_ref().map(|_| " (paginated)").unwrap_or(""),
|
||||
);
|
||||
|
||||
let mut resp = client
|
||||
.execute_tool("GMAIL_FETCH_EMAILS", Some(args.clone()))
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("GMAIL_FETCH_EMAILS page {page_num}: {e:#}"))?;
|
||||
total_cost += resp.cost_usd;
|
||||
|
||||
if !resp.successful {
|
||||
anyhow::bail!(
|
||||
"GMAIL_FETCH_EMAILS page {page_num} failed: {:?}",
|
||||
resp.error
|
||||
);
|
||||
}
|
||||
|
||||
provider.post_process_action_result("GMAIL_FETCH_EMAILS", Some(&args), &mut resp.data);
|
||||
|
||||
let (messages, next_token) = extract_envelope(&resp.data);
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] page {} -> {} messages, next_token={}",
|
||||
page_num,
|
||||
messages.len(),
|
||||
next_token.as_deref().unwrap_or("(none)"),
|
||||
);
|
||||
|
||||
if messages.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
let chunks_this_page = ingest_page_into_memory_tree(&config, &owner, &messages).await?;
|
||||
total_chunks += chunks_this_page;
|
||||
total_pages += 1;
|
||||
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] page {} ingested chunks={} running_total={}",
|
||||
page_num,
|
||||
chunks_this_page,
|
||||
total_chunks,
|
||||
);
|
||||
|
||||
match next_token {
|
||||
Some(tok) => page_token = Some(tok),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] fetch+ingest done pages={} total_chunks={} cost=~${:.4}",
|
||||
total_pages,
|
||||
total_chunks,
|
||||
total_cost,
|
||||
);
|
||||
|
||||
// ─── Drain async worker pool ────────────────────────────────────────────
|
||||
|
||||
if cli.skip_drain {
|
||||
log::info!("[gmail_backfill_3d] skipping worker pool drain (--skip-drain)");
|
||||
} else {
|
||||
log::info!("[gmail_backfill_3d] draining async worker pool…");
|
||||
drain_until_idle(&config).await?;
|
||||
log::info!("[gmail_backfill_3d] worker pool idle");
|
||||
}
|
||||
|
||||
// ─── Integrity check ────────────────────────────────────────────────────
|
||||
|
||||
if cli.skip_verify {
|
||||
log::info!("[gmail_backfill_3d] skipping integrity check (--skip-verify)");
|
||||
} else {
|
||||
log::info!("[gmail_backfill_3d] running integrity check…");
|
||||
|
||||
// Chunk integrity.
|
||||
let (verified, mismatched, no_pointer, missing_file) = verify_all_chunk_files(&config)?;
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] chunks: verified={} mismatched={} no_pointer={} missing_file={}",
|
||||
verified,
|
||||
mismatched,
|
||||
no_pointer,
|
||||
missing_file,
|
||||
);
|
||||
|
||||
// Summary integrity.
|
||||
let (sum_verified, sum_mismatched, sum_no_pointer, sum_missing_file) =
|
||||
verify_all_summary_files(&config)?;
|
||||
log::info!(
|
||||
"[gmail_backfill_3d] summaries: verified={} mismatched={} no_pointer={} missing_file={}",
|
||||
sum_verified,
|
||||
sum_mismatched,
|
||||
sum_no_pointer,
|
||||
sum_missing_file,
|
||||
);
|
||||
|
||||
if mismatched > 0 || missing_file > 0 || sum_mismatched > 0 || sum_missing_file > 0 {
|
||||
anyhow::bail!(
|
||||
"Integrity check failed: \
|
||||
chunks: {} mismatches, {} missing files; \
|
||||
summaries: {} mismatches, {} missing files",
|
||||
mismatched,
|
||||
missing_file,
|
||||
sum_mismatched,
|
||||
sum_missing_file,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
println!(
|
||||
"\nBackfill complete. pages={} chunks_written={} cost=~${:.4}",
|
||||
total_pages, total_chunks, total_cost,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Wipe `<workspace>/memory_tree/chunks.db` (+ wal/shm) and
|
||||
/// `<content_root>/` so the bin can re-run cleanly after a chunker
|
||||
/// change that invalidates existing chunk IDs.
|
||||
///
|
||||
/// Logs each removed artifact at info; missing files are not an error.
|
||||
fn wipe_memory_tree_state(config: &Config) -> Result<()> {
|
||||
let mt_dir = config.workspace_dir.join("memory_tree");
|
||||
for name in &["chunks.db", "chunks.db-wal", "chunks.db-shm"] {
|
||||
let path = mt_dir.join(name);
|
||||
match std::fs::remove_file(&path) {
|
||||
Ok(()) => log::info!("[gmail_backfill_3d] wiped {}", path.display()),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
|
||||
Err(e) => return Err(e).with_context(|| format!("wipe {}", path.display())),
|
||||
}
|
||||
}
|
||||
let content_root = config.memory_tree_content_root();
|
||||
if content_root.exists() {
|
||||
std::fs::remove_dir_all(&content_root)
|
||||
.with_context(|| format!("wipe {}", content_root.display()))?;
|
||||
log::info!("[gmail_backfill_3d] wiped {}", content_root.display());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read all chunks from SQLite and verify on-disk SHA-256 matches `content_sha256`.
|
||||
///
|
||||
/// Returns `(verified, mismatched, no_pointer, missing_file)`.
|
||||
fn verify_all_chunk_files(config: &Config) -> Result<(usize, usize, usize, usize)> {
|
||||
let chunks = list_chunks(config, &ListChunksQuery::default())?;
|
||||
let content_root = config.memory_tree_content_root();
|
||||
|
||||
let mut verified = 0usize;
|
||||
let mut mismatched = 0usize;
|
||||
let mut no_pointer = 0usize;
|
||||
let mut missing_file = 0usize;
|
||||
|
||||
for chunk in &chunks {
|
||||
let pointers = get_chunk_content_pointers(config, &chunk.id)?;
|
||||
let (rel_path, expected_sha) = match pointers {
|
||||
None => {
|
||||
no_pointer += 1;
|
||||
log::debug!(
|
||||
"[gmail_backfill_3d] verify: chunk {} has no content_path/sha256",
|
||||
chunk.id
|
||||
);
|
||||
continue;
|
||||
}
|
||||
Some(pair) => pair,
|
||||
};
|
||||
|
||||
let abs_path = {
|
||||
let mut p = content_root.clone();
|
||||
for component in rel_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
if !abs_path.exists() {
|
||||
missing_file += 1;
|
||||
log::warn!(
|
||||
"[gmail_backfill_3d] verify: file missing chunk_id={} path={}",
|
||||
chunk.id,
|
||||
abs_path.display(),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
match verify_chunk_file(&abs_path, &expected_sha) {
|
||||
Ok(true) => {
|
||||
verified += 1;
|
||||
}
|
||||
Ok(false) => {
|
||||
mismatched += 1;
|
||||
log::warn!(
|
||||
"[gmail_backfill_3d] verify: SHA-256 mismatch chunk_id={} path={}",
|
||||
chunk.id,
|
||||
abs_path.display(),
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[gmail_backfill_3d] verify: error chunk_id={}: {e}",
|
||||
chunk.id,
|
||||
);
|
||||
mismatched += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((verified, mismatched, no_pointer, missing_file))
|
||||
}
|
||||
|
||||
/// Read all summary rows with a non-NULL `content_path` from SQLite and verify
|
||||
/// the on-disk SHA-256 matches `content_sha256`.
|
||||
///
|
||||
/// Returns `(verified, mismatched, no_pointer, missing_file)`.
|
||||
fn verify_all_summary_files(config: &Config) -> Result<(usize, usize, usize, usize)> {
|
||||
let rows_with_pointer = list_summaries_with_content_path(config)?;
|
||||
let content_root = config.memory_tree_content_root();
|
||||
|
||||
let mut verified = 0usize;
|
||||
let mut mismatched = 0usize;
|
||||
let mut missing_file = 0usize;
|
||||
|
||||
for (summary_id, rel_path, expected_sha) in &rows_with_pointer {
|
||||
let abs_path = {
|
||||
let mut p = content_root.clone();
|
||||
for component in rel_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
match verify_summary_file(&abs_path, expected_sha) {
|
||||
Ok(VerifyResult::Ok) => {
|
||||
verified += 1;
|
||||
}
|
||||
Ok(VerifyResult::Mismatch { actual }) => {
|
||||
mismatched += 1;
|
||||
log::warn!(
|
||||
"[gmail_backfill_3d] verify: SHA-256 mismatch summary_id={} path={} expected={} actual={}",
|
||||
summary_id,
|
||||
abs_path.display(),
|
||||
expected_sha,
|
||||
actual,
|
||||
);
|
||||
}
|
||||
Ok(VerifyResult::Missing) => {
|
||||
missing_file += 1;
|
||||
log::warn!(
|
||||
"[gmail_backfill_3d] verify: file missing summary_id={} path={}",
|
||||
summary_id,
|
||||
abs_path.display(),
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[gmail_backfill_3d] verify: error summary_id={}: {e}",
|
||||
summary_id,
|
||||
);
|
||||
mismatched += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Count rows that have no content_path at all (legacy rows).
|
||||
// We report this as no_pointer for symmetry with the chunk verifier.
|
||||
// We can't easily count them here without a separate query, so we
|
||||
// approximate: rows_with_pointer gives us the ones we checked.
|
||||
// For now no_pointer = 0 (the bin wipes before re-ingesting so all
|
||||
// new rows should have pointers; legacy rows are pre-migration).
|
||||
let no_pointer = 0usize;
|
||||
|
||||
Ok((verified, mismatched, no_pointer, missing_file))
|
||||
}
|
||||
|
||||
/// Extract the `messages` array and `nextPageToken` from a Composio response.
|
||||
fn extract_envelope(data: &Value) -> (Vec<Value>, Option<String>) {
|
||||
let candidates: [Option<&Value>; 2] = [Some(data), data.get("data")];
|
||||
for cand in candidates.into_iter().flatten() {
|
||||
if let Some(arr) = cand.get("messages").and_then(|v| v.as_array()) {
|
||||
let token = cand
|
||||
.get("nextPageToken")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.map(str::to_string);
|
||||
return (arr.clone(), token);
|
||||
}
|
||||
}
|
||||
(Vec::new(), None)
|
||||
}
|
||||
@@ -96,6 +96,21 @@ struct Cli {
|
||||
/// re-attempts cascade on the next append.
|
||||
#[arg(long = "seal-probe", default_value_t = false)]
|
||||
seal_probe: bool,
|
||||
|
||||
/// Fire N back-to-back `SLACK_FETCH_CONVERSATION_HISTORY` calls
|
||||
/// against the first listed channel and report a per-call tally
|
||||
/// of {success, ratelimit, other-failure} + total duration. No
|
||||
/// pacing by default (see --probe-pacing-ms), no ingestion. Used
|
||||
/// to characterise Composio/Slack quota behaviour without
|
||||
/// touching the memory tree.
|
||||
#[arg(long = "probe-ratelimit")]
|
||||
probe_ratelimit: Option<u32>,
|
||||
|
||||
/// Sleep this many milliseconds between probe calls. 0 = fire
|
||||
/// back-to-back (default). Use to find the threshold at which
|
||||
/// rate-limits stop firing.
|
||||
#[arg(long = "probe-pacing-ms", default_value_t = 0)]
|
||||
probe_pacing_ms: u64,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
@@ -243,6 +258,130 @@ async fn main() -> Result<()> {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some(n) = cli.probe_ratelimit {
|
||||
// Pure quota probe: fire N back-to-back
|
||||
// SLACK_FETCH_CONVERSATION_HISTORY calls against the first
|
||||
// discoverable channel. No pacing, no retry, no ingest. Reports
|
||||
// a per-call status table + summary so we can characterise
|
||||
// Composio/Slack rate-limit behaviour without contaminating the
|
||||
// memory tree or burning extra quota on retries.
|
||||
log::info!("[probe-ratelimit] requesting one channel via SLACK_LIST_CONVERSATIONS");
|
||||
let list_resp = client
|
||||
.execute_tool(
|
||||
"SLACK_LIST_CONVERSATIONS",
|
||||
Some(serde_json::json!({ "exclude_archived": true, "limit": 1 })),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("SLACK_LIST_CONVERSATIONS failed: {e:#}"))?;
|
||||
if !list_resp.successful {
|
||||
anyhow::bail!(
|
||||
"SLACK_LIST_CONVERSATIONS returned non-success: {:?}",
|
||||
list_resp.error
|
||||
);
|
||||
}
|
||||
let channel_id = ["/data/channels/0/id", "/channels/0/id", "/data/0/id"]
|
||||
.iter()
|
||||
.find_map(|p| list_resp.data.pointer(p).and_then(|v| v.as_str()))
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"could not find a channel id in SLACK_LIST_CONVERSATIONS response: {}",
|
||||
serde_json::to_string(&list_resp.data).unwrap_or_default()
|
||||
)
|
||||
})?;
|
||||
log::info!("[probe-ratelimit] firing {n} calls against channel={channel_id}");
|
||||
|
||||
#[derive(Debug)]
|
||||
enum Outcome {
|
||||
Ok,
|
||||
Ratelimit,
|
||||
OtherFail(String),
|
||||
Transport(String),
|
||||
}
|
||||
let mut outcomes: Vec<(u32, std::time::Duration, Outcome)> = Vec::with_capacity(n as usize);
|
||||
let probe_started = Instant::now();
|
||||
for i in 1..=n {
|
||||
if i > 1 && cli.probe_pacing_ms > 0 {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(cli.probe_pacing_ms)).await;
|
||||
}
|
||||
let t0 = Instant::now();
|
||||
let resp = client
|
||||
.execute_tool(
|
||||
"SLACK_FETCH_CONVERSATION_HISTORY",
|
||||
Some(serde_json::json!({ "channel": channel_id, "limit": 1000 })),
|
||||
)
|
||||
.await;
|
||||
let dt = t0.elapsed();
|
||||
let outcome = match resp {
|
||||
Err(e) => Outcome::Transport(format!("{e:#}")),
|
||||
Ok(r) if r.successful => Outcome::Ok,
|
||||
Ok(r) => {
|
||||
let err = r.error.as_deref().unwrap_or("provider failure");
|
||||
if err.contains("ratelimited")
|
||||
|| err.contains("rate_limit")
|
||||
|| err.contains("rate limit")
|
||||
{
|
||||
log::warn!(
|
||||
"[probe-ratelimit] call {i} ratelimited; body: {}",
|
||||
serde_json::to_string(&r.data).unwrap_or_default()
|
||||
);
|
||||
Outcome::Ratelimit
|
||||
} else {
|
||||
Outcome::OtherFail(err.to_string())
|
||||
}
|
||||
}
|
||||
};
|
||||
log::info!(
|
||||
"[probe-ratelimit] call {i}/{n} took {:.2}s -> {:?}",
|
||||
dt.as_secs_f64(),
|
||||
outcome
|
||||
);
|
||||
outcomes.push((i, dt, outcome));
|
||||
}
|
||||
let total = probe_started.elapsed();
|
||||
|
||||
let ok = outcomes
|
||||
.iter()
|
||||
.filter(|(_, _, o)| matches!(o, Outcome::Ok))
|
||||
.count();
|
||||
let rl = outcomes
|
||||
.iter()
|
||||
.filter(|(_, _, o)| matches!(o, Outcome::Ratelimit))
|
||||
.count();
|
||||
let other = outcomes
|
||||
.iter()
|
||||
.filter(|(_, _, o)| matches!(o, Outcome::OtherFail(_)))
|
||||
.count();
|
||||
let transport = outcomes
|
||||
.iter()
|
||||
.filter(|(_, _, o)| matches!(o, Outcome::Transport(_)))
|
||||
.count();
|
||||
let avg_ms = if !outcomes.is_empty() {
|
||||
outcomes.iter().map(|(_, d, _)| d.as_millis()).sum::<u128>() / outcomes.len() as u128
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
println!("=== probe-ratelimit summary ===");
|
||||
println!("channel: {channel_id}");
|
||||
println!("calls fired: {n}");
|
||||
println!("total duration: {:.2}s", total.as_secs_f64());
|
||||
println!("avg per call: {avg_ms} ms");
|
||||
println!("successful: {ok}");
|
||||
println!("ratelimited: {rl}");
|
||||
println!("other failures: {other}");
|
||||
println!("transport errors: {transport}");
|
||||
if rl > 0 {
|
||||
let first_rl = outcomes
|
||||
.iter()
|
||||
.find(|(_, _, o)| matches!(o, Outcome::Ratelimit))
|
||||
.map(|(i, _, _)| *i)
|
||||
.unwrap_or(0);
|
||||
println!("first ratelimit: call #{first_rl}");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let connections = client
|
||||
.list_connections()
|
||||
.await
|
||||
|
||||
+6
-2
@@ -805,7 +805,10 @@ async fn run_server_inner(
|
||||
///
|
||||
/// Guarded by `std::sync::Once` so repeated calls to `bootstrap_skill_runtime`
|
||||
/// are safe and idempotent.
|
||||
fn register_domain_subscribers(workspace_dir: std::path::PathBuf) {
|
||||
fn register_domain_subscribers(
|
||||
workspace_dir: std::path::PathBuf,
|
||||
config: crate::openhuman::config::Config,
|
||||
) {
|
||||
use std::sync::{Arc, Once};
|
||||
|
||||
static REGISTERED: Once = Once::new();
|
||||
@@ -840,6 +843,7 @@ fn register_domain_subscribers(workspace_dir: std::path::PathBuf) {
|
||||
}
|
||||
crate::openhuman::composio::register_composio_trigger_subscriber();
|
||||
crate::openhuman::composio::start_periodic_sync();
|
||||
crate::openhuman::memory::tree::jobs::start(config.clone());
|
||||
|
||||
// Restart requests go through a subscriber so every trigger path shares
|
||||
// the same respawn logic.
|
||||
@@ -880,7 +884,7 @@ pub async fn bootstrap_skill_runtime() {
|
||||
// Register domain subscribers for cross-module event handling.
|
||||
// Uses a Once guard so repeated calls to bootstrap_skill_runtime()
|
||||
// cannot double-subscribe.
|
||||
register_domain_subscribers(workspace_dir.clone());
|
||||
register_domain_subscribers(workspace_dir.clone(), cfg.clone());
|
||||
|
||||
// --- Sub-agent definition registry bootstrap ---
|
||||
// Loads built-in archetype definitions plus any custom TOML files
|
||||
|
||||
@@ -0,0 +1,642 @@
|
||||
//! Gmail → memory tree ingest plumbing.
|
||||
//!
|
||||
//! Owns the conversion from a page of `GMAIL_FETCH_EMAILS` slim-envelope
|
||||
//! messages (post-processed by [`super::post_process`]) into
|
||||
//! [`EmailThread`] batches grouped by the sorted set of distinct
|
||||
//! participants (`from` ∪ `to`-list, CC ignored), then drives
|
||||
//! [`memory::tree::ingest::ingest_email`] per participant group.
|
||||
//!
|
||||
//! Source-id is `gmail:{participants}` where participants is
|
||||
//! `addr1|addr2|...` (sorted, deduped, lowercased bare emails). All
|
||||
//! correspondence between the same set of people lands in one source tree.
|
||||
//!
|
||||
//! Idempotency: chunk IDs are content-hashed inside the memory tree, so
|
||||
//! re-ingesting a previously-seen Gmail message is an UPSERT — buffer
|
||||
//! token_sum may drift if content changes (rare for sealed mail), but
|
||||
//! the tree's seal cascade handles that on next append.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use anyhow::Result;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::canonicalize::email::{EmailMessage, EmailThread};
|
||||
use crate::openhuman::memory::tree::canonicalize::email_clean::{
|
||||
extract_email, parse_message_date,
|
||||
};
|
||||
use crate::openhuman::memory::tree::ingest::{ingest_email, IngestResult};
|
||||
use crate::openhuman::memory::tree::util::redact::redact;
|
||||
|
||||
/// Provider name embedded in the canonical email-thread header. Matches
|
||||
/// the value `memory::tree::retrieval::source::PLATFORM_KINDS` expects.
|
||||
pub const GMAIL_PROVIDER: &str = "gmail";
|
||||
|
||||
/// Tags attached to every Gmail-ingested chunk. Stable list — retrieval
|
||||
/// callers filter on these.
|
||||
pub const DEFAULT_TAGS: &[&str] = &["gmail", "ingested"];
|
||||
|
||||
/// Group raw page messages by the sorted set of distinct participants
|
||||
/// (`from` ∪ `to`-list). CC is deliberately excluded from the bucket key
|
||||
/// so CC-only recipients don't fragment conversations. All messages
|
||||
/// between the same set of people land in the same bucket regardless of
|
||||
/// direction or thread ID.
|
||||
///
|
||||
/// The bucket key is the participants joined with `|` in sorted order,
|
||||
/// e.g. `"alice@x.com|bob@y.com"`. Messages within a bucket are sorted
|
||||
/// ascending by date so the rendered conversation reads chronologically.
|
||||
pub(crate) fn bucket_by_participants(msgs: &[Value]) -> BTreeMap<String, Vec<&Value>> {
|
||||
let mut out: BTreeMap<String, Vec<&Value>> = BTreeMap::new();
|
||||
for m in msgs {
|
||||
let bucket_key = participants_bucket_key(m);
|
||||
if bucket_key == "__skip__" {
|
||||
// Message has no parseable addresses AND no id — drop it and warn.
|
||||
// Nothing useful can be done with it: no participants means no
|
||||
// source tree, and no id means no unique bucket either.
|
||||
log::warn!(
|
||||
"[composio:gmail][bucket] dropping message with no parseable addresses and no id"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
out.entry(bucket_key).or_default().push(m);
|
||||
}
|
||||
for bucket in out.values_mut() {
|
||||
bucket.sort_by_key(|m| {
|
||||
parse_message_date(m)
|
||||
.map(|d: chrono::DateTime<chrono::Utc>| d.timestamp())
|
||||
.unwrap_or(0)
|
||||
});
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Compute the participants bucket key for a single raw message.
|
||||
///
|
||||
/// Collects `from` ∪ `to` (as bare lowercased email addresses), sorts
|
||||
/// and dedupes them, then joins with `|`.
|
||||
///
|
||||
/// **Fallback policy when all addresses fail to parse**:
|
||||
/// - If the message has a non-empty `id`, use `"orphan:{id}"` so each
|
||||
/// malformed message gets its own bucket and its own source tree. Two
|
||||
/// messages with different ids that both fail address parsing will NOT
|
||||
/// collapse into a single `"unknown"` bucket.
|
||||
/// - If even `id` is missing or empty, the caller (`bucket_by_participants`)
|
||||
/// should skip the message (log a warn and drop it). This function signals
|
||||
/// that case by returning the sentinel `"__skip__"`.
|
||||
fn participants_bucket_key(raw: &Value) -> String {
|
||||
let from = extract_email(raw.get("from").and_then(|v| v.as_str()).unwrap_or(""))
|
||||
.map(|s| s.to_lowercase())
|
||||
.filter(|s| !s.is_empty());
|
||||
|
||||
let to_emails: Vec<String> = parse_address_list_for_bucket(raw.get("to"))
|
||||
.into_iter()
|
||||
.filter_map(|addr| extract_email(&addr).map(|s| s.to_lowercase()))
|
||||
.collect();
|
||||
|
||||
let mut all: Vec<String> = from.into_iter().chain(to_emails).collect();
|
||||
all.sort();
|
||||
all.dedup();
|
||||
all.retain(|s| !s.is_empty());
|
||||
|
||||
if all.is_empty() {
|
||||
// No parseable addresses — fall back to per-message uniqueness to
|
||||
// avoid collapsing all malformed messages into one "unknown" source
|
||||
// tree. Each orphan message gets its own bucket so nothing is silently
|
||||
// lost in a mixed pile.
|
||||
let id = raw
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.is_empty());
|
||||
match id {
|
||||
Some(msg_id) => format!("orphan:{}", msg_id),
|
||||
None => {
|
||||
// id is missing: signal caller to skip this message entirely.
|
||||
"__skip__".to_string()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
all.join("|")
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the `to` / `cc` field for bucket-key construction. Handles both
|
||||
/// JSON array and comma-separated string forms. Returns raw address
|
||||
/// strings (may include display names); callers must extract the bare
|
||||
/// email with [`extract_email`].
|
||||
fn parse_address_list_for_bucket(v: Option<&Value>) -> Vec<String> {
|
||||
match v {
|
||||
Some(Value::Array(arr)) => arr
|
||||
.iter()
|
||||
.filter_map(|s| s.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect(),
|
||||
Some(Value::String(s)) => s
|
||||
.split(',')
|
||||
.map(|p| p.trim().to_string())
|
||||
.filter(|p| !p.is_empty())
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an [`EmailMessage`] from a raw slim-envelope JSON message.
|
||||
/// Returns `None` when the message has no parseable date — the rest of
|
||||
/// the pipeline can't sort or canonicalise without one.
|
||||
pub(crate) fn raw_to_email_message(raw: &Value) -> Option<EmailMessage> {
|
||||
let id = raw
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or("");
|
||||
let from = raw
|
||||
.get("from")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let to = parse_address_list(raw.get("to"));
|
||||
let cc = parse_address_list(raw.get("cc"));
|
||||
let subject = raw
|
||||
.get("subject")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let sent_at = parse_message_date(raw)?;
|
||||
let body = raw
|
||||
.get("markdown")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let source_ref = if id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(format!("gmail://msg/{id}"))
|
||||
};
|
||||
Some(EmailMessage {
|
||||
from,
|
||||
to,
|
||||
cc,
|
||||
subject,
|
||||
sent_at,
|
||||
body,
|
||||
source_ref,
|
||||
})
|
||||
}
|
||||
|
||||
/// Parse the `to` / `cc` field which Composio surfaces as either a
|
||||
/// JSON array of strings or a single comma-separated string. Empty
|
||||
/// entries are dropped.
|
||||
fn parse_address_list(v: Option<&Value>) -> Vec<String> {
|
||||
match v {
|
||||
Some(Value::Array(arr)) => arr
|
||||
.iter()
|
||||
.filter_map(|s| s.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect(),
|
||||
Some(Value::String(s)) => s
|
||||
.split(',')
|
||||
.map(|p| p.trim().to_string())
|
||||
.filter(|p| !p.is_empty())
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Ingest a page of raw Gmail messages into the memory tree.
|
||||
///
|
||||
/// Each participant-bucket (sorted set of `from` ∪ `to` email addresses)
|
||||
/// becomes one [`EmailThread`] handed to [`ingest_email`] which fans out
|
||||
/// to the chunker + scorer + source tree downstream.
|
||||
///
|
||||
/// `source_id` = `"gmail:{participants}"` where participants is
|
||||
/// `addr1|addr2|...` (sorted, deduped, lowercased). This groups all
|
||||
/// correspondence between the same people into one path subtree.
|
||||
///
|
||||
/// Returns the total number of chunks written across all buckets so
|
||||
/// callers can surface counts in logs / outcomes. Per-bucket errors are
|
||||
/// logged and swallowed — one bad bucket should not abort the whole
|
||||
/// page (the next sync re-fetches via the date-cursor).
|
||||
pub async fn ingest_page_into_memory_tree(
|
||||
config: &Config,
|
||||
owner: &str,
|
||||
page_messages: &[Value],
|
||||
) -> Result<usize> {
|
||||
if page_messages.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let buckets = bucket_by_participants(page_messages);
|
||||
let mut total_chunks = 0usize;
|
||||
let mut total_buckets = 0usize;
|
||||
for (participants, raw_msgs) in &buckets {
|
||||
let messages: Vec<EmailMessage> = raw_msgs
|
||||
.iter()
|
||||
.filter_map(|raw| raw_to_email_message(raw))
|
||||
.collect();
|
||||
if messages.is_empty() {
|
||||
log::debug!(
|
||||
"[composio:gmail][ingest] skipping empty bucket participants_hash={}",
|
||||
redact(participants)
|
||||
);
|
||||
continue;
|
||||
}
|
||||
// source_id encodes participants so every unique conversation set
|
||||
// lands in its own path subtree.
|
||||
let source_id = format!("gmail:{}", participants);
|
||||
let thread_subject = pick_thread_subject(&messages);
|
||||
log::info!(
|
||||
"[composio:gmail][ingest] bucket participants_hash={} messages={} source_id_hash={}",
|
||||
redact(participants),
|
||||
messages.len(),
|
||||
redact(&source_id)
|
||||
);
|
||||
let thread = EmailThread {
|
||||
provider: GMAIL_PROVIDER.to_string(),
|
||||
thread_subject,
|
||||
messages,
|
||||
};
|
||||
let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect();
|
||||
match ingest_email(config, &source_id, owner, tags, thread).await {
|
||||
Ok(IngestResult { chunks_written, .. }) => {
|
||||
total_chunks += chunks_written;
|
||||
total_buckets += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[composio:gmail][ingest] ingest_email failed participants_hash={} source_id_hash={} err={:#}",
|
||||
redact(participants),
|
||||
redact(&source_id),
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
log::info!(
|
||||
"[composio:gmail][ingest] page_done owner_hash={} buckets={total_buckets} chunks={total_chunks}",
|
||||
redact(owner)
|
||||
);
|
||||
Ok(total_chunks)
|
||||
}
|
||||
|
||||
/// Strip "Re:" / "Fwd:" prefixes from the head message's subject so
|
||||
/// every message in a thread shares one canonical thread subject. Falls
|
||||
/// back to "(no subject)" when empty.
|
||||
fn pick_thread_subject(messages: &[EmailMessage]) -> String {
|
||||
let raw = messages
|
||||
.first()
|
||||
.map(|m| m.subject.trim().to_string())
|
||||
.unwrap_or_default();
|
||||
let stripped = strip_reply_prefixes(&raw);
|
||||
if stripped.is_empty() {
|
||||
"(no subject)".to_string()
|
||||
} else {
|
||||
stripped
|
||||
}
|
||||
}
|
||||
|
||||
/// Iteratively strip `Re:` / `Fwd:` / `Fw:` prefixes (case-insensitive,
|
||||
/// optional whitespace) from the front of a subject. Stops once a pass
|
||||
/// removes nothing.
|
||||
fn strip_reply_prefixes(subject: &str) -> String {
|
||||
let mut s = subject.trim().to_string();
|
||||
loop {
|
||||
let lower = s.to_ascii_lowercase();
|
||||
let stripped = if lower.starts_with("re:") {
|
||||
Some(&s[3..])
|
||||
} else if lower.starts_with("fwd:") {
|
||||
Some(&s[4..])
|
||||
} else if lower.starts_with("fw:") {
|
||||
Some(&s[3..])
|
||||
} else {
|
||||
None
|
||||
};
|
||||
match stripped {
|
||||
Some(rest) => {
|
||||
let trimmed = rest.trim_start().to_string();
|
||||
if trimmed == s {
|
||||
return s;
|
||||
}
|
||||
s = trimmed;
|
||||
}
|
||||
None => return s,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
// ─── bucket_by_participants tests ─────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn bidirectional_messages_bucket_together() {
|
||||
// alice→bob and bob→alice land in the same key "alice@x.com|bob@y.com".
|
||||
let msgs = vec![
|
||||
json!({
|
||||
"id": "m1",
|
||||
"from": "alice@x.com",
|
||||
"to": "bob@y.com",
|
||||
"subject": "Hi",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "hi",
|
||||
}),
|
||||
json!({
|
||||
"id": "m2",
|
||||
"from": "bob@y.com",
|
||||
"to": "alice@x.com",
|
||||
"subject": "Re: Hi",
|
||||
"date": "2026-04-21T11:00:00Z",
|
||||
"markdown": "hey",
|
||||
}),
|
||||
];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
assert_eq!(buckets.len(), 1, "both messages must share one bucket");
|
||||
let key = buckets.keys().next().unwrap();
|
||||
assert_eq!(key, "alice@x.com|bob@y.com");
|
||||
assert_eq!(buckets[key].len(), 2);
|
||||
// Sorted ascending by date inside the bucket.
|
||||
assert_eq!(buckets[key][0].get("id").unwrap().as_str().unwrap(), "m1");
|
||||
assert_eq!(buckets[key][1].get("id").unwrap().as_str().unwrap(), "m2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_recipient_bucket_key_sorted() {
|
||||
// from=alice, to=[bob, carol] → "alice@x.com|bob@y.com|carol@z.com"
|
||||
let msgs = vec![json!({
|
||||
"id": "m1",
|
||||
"from": "Alice <alice@x.com>",
|
||||
"to": ["bob@y.com", "carol@z.com"],
|
||||
"subject": "Group",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "hey all",
|
||||
})];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
let key = buckets.keys().next().unwrap();
|
||||
assert_eq!(key, "alice@x.com|bob@y.com|carol@z.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cc_field_ignored_in_bucket_key() {
|
||||
// from=alice, to=[bob], cc=[dave] → "alice@x.com|bob@y.com" (no dave).
|
||||
let msgs = vec![json!({
|
||||
"id": "m1",
|
||||
"from": "alice@x.com",
|
||||
"to": "bob@y.com",
|
||||
"cc": "dave@z.com",
|
||||
"subject": "CC test",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body",
|
||||
})];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
let key = buckets.keys().next().unwrap();
|
||||
assert_eq!(
|
||||
key, "alice@x.com|bob@y.com",
|
||||
"CC must not appear in bucket key"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn solo_message_no_to_buckets_to_sender_only() {
|
||||
// from=alice, to=[] → "alice@x.com" (single participant).
|
||||
let msgs = vec![json!({
|
||||
"id": "m1",
|
||||
"from": "alice@x.com",
|
||||
"subject": "Draft",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "draft body",
|
||||
})];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
let key = buckets.keys().next().unwrap();
|
||||
assert_eq!(key, "alice@x.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_from_and_to_falls_back_to_orphan_bucket() {
|
||||
// A message with no parseable addresses gets its own orphan bucket
|
||||
// keyed by its id rather than collapsing everything into "unknown".
|
||||
let msgs = vec![json!({
|
||||
"id": "m1",
|
||||
"from": "",
|
||||
"subject": "x",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body",
|
||||
})];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
assert_eq!(buckets.len(), 1, "must produce exactly one bucket");
|
||||
assert!(
|
||||
buckets.contains_key("orphan:m1"),
|
||||
"must fall back to orphan:<id>; got keys: {:?}",
|
||||
buckets.keys().collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_malformed_messages_with_different_ids_land_in_different_buckets() {
|
||||
// Two messages with unparseable from/to but different ids must not
|
||||
// collapse into the same "unknown" bucket — each gets its own orphan.
|
||||
let msgs = vec![
|
||||
json!({
|
||||
"id": "orphan_a",
|
||||
"from": "",
|
||||
"subject": "x",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body a",
|
||||
}),
|
||||
json!({
|
||||
"id": "orphan_b",
|
||||
"from": "",
|
||||
"subject": "y",
|
||||
"date": "2026-04-21T11:00:00Z",
|
||||
"markdown": "body b",
|
||||
}),
|
||||
];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
assert_eq!(
|
||||
buckets.len(),
|
||||
2,
|
||||
"each malformed message must have its own bucket; got: {:?}",
|
||||
buckets.keys().collect::<Vec<_>>()
|
||||
);
|
||||
assert!(buckets.contains_key("orphan:orphan_a"));
|
||||
assert!(buckets.contains_key("orphan:orphan_b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn message_with_no_id_and_no_addresses_is_dropped() {
|
||||
// A message with no id AND no parseable addresses is silently dropped.
|
||||
let valid = json!({
|
||||
"id": "m_ok",
|
||||
"from": "alice@x.com",
|
||||
"subject": "ok",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "ok",
|
||||
});
|
||||
let bad = json!({
|
||||
// no "id" field, no from/to
|
||||
"subject": "bad",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "bad",
|
||||
});
|
||||
let msgs = vec![valid, bad];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
// Only the valid message should produce a bucket.
|
||||
assert_eq!(buckets.len(), 1, "dropped message must not create a bucket");
|
||||
assert!(buckets.contains_key("alice@x.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn display_name_from_stripped_to_bare_email_in_key() {
|
||||
// "Alice <alice@x.com>" should yield bare "alice@x.com" in the key.
|
||||
let msgs = vec![json!({
|
||||
"id": "m1",
|
||||
"from": "Alice <alice@x.com>",
|
||||
"to": "Bob <bob@y.com>",
|
||||
"subject": "Hi",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "hi",
|
||||
})];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
let key = buckets.keys().next().unwrap();
|
||||
assert_eq!(key, "alice@x.com|bob@y.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_threadid_field_does_not_affect_bucketing() {
|
||||
// threadId is completely ignored; two messages from the same participants
|
||||
// share one bucket even without threadId.
|
||||
let msgs = vec![
|
||||
json!({
|
||||
"id": "m1",
|
||||
"from": "noreply@github.com",
|
||||
"to": "sanil@x.com",
|
||||
"subject": "PR opened",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body1",
|
||||
}),
|
||||
json!({
|
||||
"id": "m2",
|
||||
"from": "noreply@github.com",
|
||||
"to": "sanil@x.com",
|
||||
"subject": "PR merged",
|
||||
"date": "2026-04-21T11:00:00Z",
|
||||
"markdown": "body2",
|
||||
}),
|
||||
];
|
||||
let buckets = bucket_by_participants(&msgs);
|
||||
assert_eq!(buckets.len(), 1, "both messages must share one bucket");
|
||||
let bucket = buckets.values().next().unwrap();
|
||||
assert_eq!(bucket.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_to_email_message_parses_slim_envelope() {
|
||||
let raw = json!({
|
||||
"id": "m1",
|
||||
"from": "Alice <alice@example.com>",
|
||||
"to": "me@example.com",
|
||||
"cc": "team@example.com",
|
||||
"subject": "Phoenix kickoff",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "Let's ship Phoenix.",
|
||||
});
|
||||
let msg = raw_to_email_message(&raw).unwrap();
|
||||
assert_eq!(msg.from, "Alice <alice@example.com>");
|
||||
assert_eq!(msg.to, vec!["me@example.com"]);
|
||||
assert_eq!(msg.cc, vec!["team@example.com"]);
|
||||
assert_eq!(msg.subject, "Phoenix kickoff");
|
||||
assert_eq!(msg.body, "Let's ship Phoenix.");
|
||||
assert_eq!(msg.source_ref.as_deref(), Some("gmail://msg/m1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_to_email_message_handles_to_array() {
|
||||
let raw = json!({
|
||||
"id": "m1",
|
||||
"from": "a@x",
|
||||
"to": ["b@x", "c@x"],
|
||||
"subject": "x",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body",
|
||||
});
|
||||
let msg = raw_to_email_message(&raw).unwrap();
|
||||
assert_eq!(msg.to, vec!["b@x", "c@x"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_to_email_message_handles_comma_separated_to_string() {
|
||||
let raw = json!({
|
||||
"id": "m1",
|
||||
"from": "a@x",
|
||||
"to": "b@x, c@x ,d@x",
|
||||
"subject": "x",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body",
|
||||
});
|
||||
let msg = raw_to_email_message(&raw).unwrap();
|
||||
assert_eq!(msg.to, vec!["b@x", "c@x", "d@x"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_to_email_message_returns_none_on_unparseable_date() {
|
||||
let raw = json!({
|
||||
"id": "m1",
|
||||
"from": "a@x",
|
||||
"subject": "x",
|
||||
"date": "not-a-date",
|
||||
"markdown": "body",
|
||||
});
|
||||
assert!(raw_to_email_message(&raw).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_to_email_message_drops_source_ref_when_id_empty() {
|
||||
let raw = json!({
|
||||
"id": "",
|
||||
"from": "a@x",
|
||||
"subject": "x",
|
||||
"date": "2026-04-21T10:00:00Z",
|
||||
"markdown": "body",
|
||||
});
|
||||
let msg = raw_to_email_message(&raw).unwrap();
|
||||
assert!(msg.source_ref.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_reply_prefixes_removes_iterated() {
|
||||
assert_eq!(strip_reply_prefixes("Re: Re: Hi"), "Hi");
|
||||
assert_eq!(strip_reply_prefixes("Fwd: Re: Status"), "Status");
|
||||
assert_eq!(strip_reply_prefixes("RE: Question"), "Question");
|
||||
assert_eq!(strip_reply_prefixes("Fw: alert"), "alert");
|
||||
assert_eq!(strip_reply_prefixes("Plain subject"), "Plain subject");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pick_thread_subject_strips_reply_prefixes() {
|
||||
let messages = vec![EmailMessage {
|
||||
from: "a@x".into(),
|
||||
to: vec![],
|
||||
cc: vec![],
|
||||
subject: "Re: Re: Phoenix kickoff".into(),
|
||||
sent_at: chrono::Utc::now(),
|
||||
body: "body".into(),
|
||||
source_ref: None,
|
||||
}];
|
||||
assert_eq!(pick_thread_subject(&messages), "Phoenix kickoff");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pick_thread_subject_falls_back_to_no_subject() {
|
||||
let messages = vec![EmailMessage {
|
||||
from: "a@x".into(),
|
||||
to: vec![],
|
||||
cc: vec![],
|
||||
subject: " ".into(),
|
||||
sent_at: chrono::Utc::now(),
|
||||
body: "body".into(),
|
||||
source_ref: None,
|
||||
}];
|
||||
assert_eq!(pick_thread_subject(&messages), "(no subject)");
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod ingest;
|
||||
mod post_process;
|
||||
mod provider;
|
||||
mod sync;
|
||||
|
||||
@@ -980,6 +980,17 @@ impl Config {
|
||||
}
|
||||
}
|
||||
|
||||
// Phase MD-content: chunk body directory override. Empty string means
|
||||
// "fall back to default", consistent with other memory_tree env vars.
|
||||
if let Ok(dir) = std::env::var("OPENHUMAN_MEMORY_TREE_CONTENT_DIR") {
|
||||
let trimmed = dir.trim();
|
||||
self.memory_tree.content_dir = if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(std::path::PathBuf::from(trimmed))
|
||||
};
|
||||
}
|
||||
|
||||
// Auto-update overrides
|
||||
if let Some(flag) = env.get("OPENHUMAN_AUTO_UPDATE_ENABLED") {
|
||||
let normalized = flag.trim().to_ascii_lowercase();
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default, JsonSchema)]
|
||||
pub struct StorageConfig {
|
||||
@@ -96,6 +97,7 @@ impl Default for MemoryConfig {
|
||||
/// - `OPENHUMAN_MEMORY_SUMMARISE_ENDPOINT`
|
||||
/// - `OPENHUMAN_MEMORY_SUMMARISE_MODEL`
|
||||
/// - `OPENHUMAN_MEMORY_SUMMARISE_TIMEOUT_MS`
|
||||
/// - `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (Phase MD-content)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct MemoryTreeConfig {
|
||||
/// Ollama endpoint for the embedder (e.g. `http://localhost:11434`).
|
||||
@@ -154,6 +156,17 @@ pub struct MemoryTreeConfig {
|
||||
/// tokens and therefore takes longer to generate.
|
||||
#[serde(default = "default_memory_tree_llm_summariser_timeout_ms")]
|
||||
pub llm_summariser_timeout_ms: Option<u64>,
|
||||
|
||||
/// Phase MD-content: root directory where chunk `.md` files are stored.
|
||||
///
|
||||
/// Resolved at runtime via [`super::types::Config::memory_tree_content_root`]:
|
||||
/// - `Some(path)` → use that path verbatim.
|
||||
/// - `None` → default `<workspace_dir>/memory_tree/content/`.
|
||||
///
|
||||
/// Env override: `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (empty string = fall
|
||||
/// back to default, consistent with other memory_tree env vars).
|
||||
#[serde(default = "default_memory_tree_content_dir")]
|
||||
pub content_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
/// Returns `None` so that existing installs that never opted into Phase 4
|
||||
@@ -199,6 +212,12 @@ fn default_memory_tree_llm_summariser_timeout_ms() -> Option<u64> {
|
||||
Some(120_000)
|
||||
}
|
||||
|
||||
/// Returns `None` so the default `<workspace>/memory_tree/content/` path is
|
||||
/// used unless explicitly overridden via TOML or env var.
|
||||
fn default_memory_tree_content_dir() -> Option<PathBuf> {
|
||||
None
|
||||
}
|
||||
|
||||
impl Default for MemoryTreeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
@@ -212,6 +231,43 @@ impl Default for MemoryTreeConfig {
|
||||
llm_summariser_endpoint: default_memory_tree_llm_endpoint(),
|
||||
llm_summariser_model: default_memory_tree_llm_endpoint(),
|
||||
llm_summariser_timeout_ms: default_memory_tree_llm_summariser_timeout_ms(),
|
||||
content_dir: default_memory_tree_content_dir(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn memory_tree_config_default_content_dir_is_none() {
|
||||
let cfg = MemoryTreeConfig::default();
|
||||
assert!(
|
||||
cfg.content_dir.is_none(),
|
||||
"default content_dir must be None so workspace default path is used"
|
||||
);
|
||||
}
|
||||
|
||||
/// Verify that the env-var override logic correctly maps non-empty strings
|
||||
/// to `Some(PathBuf)` and empty/blank strings to `None`. We test the
|
||||
/// logic inline (not via `apply_env_overrides`) to avoid mutating the
|
||||
/// process environment in a way that could race with parallel tests.
|
||||
#[test]
|
||||
fn content_dir_env_override_logic() {
|
||||
// Simulate the load.rs overlay logic.
|
||||
let apply = |raw: &str| -> Option<PathBuf> {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(PathBuf::from(trimmed))
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(apply("/tmp/foo"), Some(PathBuf::from("/tmp/foo")));
|
||||
assert_eq!(apply(" /tmp/foo "), Some(PathBuf::from("/tmp/foo")));
|
||||
assert_eq!(apply(""), None);
|
||||
assert_eq!(apply(" "), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,6 +203,23 @@ pub struct Config {
|
||||
pub chat_onboarding_completed: bool,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Resolve the root directory where chunk `.md` files are stored.
|
||||
///
|
||||
/// Resolution order:
|
||||
/// 1. `memory_tree.content_dir` if `Some`.
|
||||
/// 2. Default: `<workspace_dir>/memory_tree/content/`.
|
||||
///
|
||||
/// This is the only place in the codebase that should compute the content
|
||||
/// root — all code that needs the path should call this method.
|
||||
pub fn memory_tree_content_root(&self) -> PathBuf {
|
||||
self.memory_tree
|
||||
.content_dir
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.workspace_dir.join("memory_tree").join("content"))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
let openhuman_dir =
|
||||
|
||||
@@ -4,10 +4,9 @@
|
||||
//! from the same channel becomes one [`CanonicalisedSource`]; the chunker
|
||||
//! slices it by token budget downstream.
|
||||
//!
|
||||
//! Output format:
|
||||
//! Output format (no leading `# ...` header — that info lives in front-matter
|
||||
//! once Phase MD-content lands; the chunker splits at `## ` boundaries):
|
||||
//! ```md
|
||||
//! # Chat transcript — {platform} / {channel}
|
||||
//!
|
||||
//! ## 2026-04-21T10:12:00Z — Alice
|
||||
//! Message body here.
|
||||
//!
|
||||
@@ -67,10 +66,9 @@ pub fn canonicalise(
|
||||
let last_ts = messages.last().map(|m| m.timestamp).unwrap();
|
||||
|
||||
let mut md = String::new();
|
||||
md.push_str(&format!(
|
||||
"# Chat transcript — {} / {}\n\n",
|
||||
batch.platform, batch.channel_label
|
||||
));
|
||||
// No leading `# Chat transcript — ...` header. Platform / channel info
|
||||
// belongs in the MD front-matter (Phase MD-content). The chunker splits
|
||||
// this output at `## ` boundaries so each message becomes one chunk.
|
||||
for msg in &messages {
|
||||
md.push_str(&format!(
|
||||
"## {} — {}\n{}\n\n",
|
||||
@@ -150,7 +148,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_header_and_per_message_sections() {
|
||||
fn includes_per_message_sections_without_header() {
|
||||
let b = ChatBatch {
|
||||
platform: "slack".into(),
|
||||
channel_label: "#eng".into(),
|
||||
@@ -159,10 +157,15 @@ mod tests {
|
||||
let out = canonicalise("slack:#eng", "alice", &[], b)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(out
|
||||
.markdown
|
||||
.starts_with("# Chat transcript — slack / #eng\n\n"));
|
||||
assert!(out.markdown.contains("## "));
|
||||
// No leading `# Chat transcript` header — that info belongs in front-matter.
|
||||
assert!(
|
||||
!out.markdown.starts_with("# "),
|
||||
"canonical chat MD must NOT start with a `# ` header"
|
||||
);
|
||||
assert!(
|
||||
out.markdown.starts_with("## "),
|
||||
"must start with first `## ` message block"
|
||||
);
|
||||
assert!(out.markdown.contains("— alice"));
|
||||
assert!(out.markdown.contains("hello"));
|
||||
}
|
||||
|
||||
@@ -39,7 +39,8 @@ pub fn canonicalise(
|
||||
}
|
||||
|
||||
let mut md = String::new();
|
||||
md.push_str(&format!("# {} — {}\n\n", doc.provider, doc.title));
|
||||
// No leading `# provider — title` header. Provider / title info
|
||||
// belongs in the MD front-matter (Phase MD-content).
|
||||
md.push_str(doc.body.trim());
|
||||
md.push('\n');
|
||||
|
||||
@@ -85,7 +86,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_title_and_body() {
|
||||
fn renders_body_without_header() {
|
||||
let out = canonicalise(
|
||||
"d1",
|
||||
"alice",
|
||||
@@ -94,7 +95,11 @@ mod tests {
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(out.markdown.starts_with("# notion — Launch plan\n\n"));
|
||||
// No leading `# notion — Launch plan` header — that info belongs in front-matter.
|
||||
assert!(
|
||||
!out.markdown.starts_with("# "),
|
||||
"canonical document MD must NOT start with a `# ` header"
|
||||
);
|
||||
assert!(out.markdown.contains("step one"));
|
||||
assert!(out.markdown.contains("step two"));
|
||||
}
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
//! Email threads → canonical Markdown.
|
||||
//!
|
||||
//! Email sources are scoped by **thread**. One thread becomes one
|
||||
//! [`CanonicalisedSource`]. Headers (From, To, Subject, Date) surface in a
|
||||
//! small frontmatter-style block per message; the body follows as markdown.
|
||||
//! Email sources are scoped by **participant set**. One participant bucket
|
||||
//! becomes one [`CanonicalisedSource`]. Headers (From, To, Cc, Subject, Date)
|
||||
//! surface in a small frontmatter-style block per message; the cleaned body
|
||||
//! follows as markdown. Bodies pass through [`email_clean::clean_body`] before
|
||||
//! rendering to strip reply chains, marketing footers, legal disclaimers, and
|
||||
//! other boilerplate.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{normalize_source_ref, CanonicalisedSource};
|
||||
use super::{email_clean, normalize_source_ref, CanonicalisedSource};
|
||||
use crate::openhuman::memory::tree::types::{Metadata, SourceKind};
|
||||
|
||||
/// One email in a thread.
|
||||
@@ -55,10 +58,9 @@ pub fn canonicalise(
|
||||
let last_ts = messages.last().map(|m| m.sent_at).unwrap();
|
||||
|
||||
let mut md = String::new();
|
||||
md.push_str(&format!(
|
||||
"# Email thread — {} — {}\n\n",
|
||||
thread.provider, thread.thread_subject
|
||||
));
|
||||
// No leading `# Email thread — ...` header. Provider / subject info
|
||||
// belongs in the MD front-matter (Phase MD-content). The chunker splits
|
||||
// this output at `---\nFrom:` boundaries so each message becomes one chunk.
|
||||
|
||||
for msg in &messages {
|
||||
md.push_str("---\n");
|
||||
@@ -71,7 +73,12 @@ pub fn canonicalise(
|
||||
}
|
||||
md.push_str(&format!("Subject: {}\n", msg.subject));
|
||||
md.push_str(&format!("Date: {}\n\n", msg.sent_at.to_rfc3339()));
|
||||
md.push_str(msg.body.trim());
|
||||
let cleaned = email_clean::clean_body(msg.body.trim());
|
||||
if cleaned.is_empty() {
|
||||
md.push('\n');
|
||||
} else {
|
||||
md.push_str(&cleaned);
|
||||
}
|
||||
md.push_str("\n\n");
|
||||
}
|
||||
|
||||
@@ -128,10 +135,19 @@ mod tests {
|
||||
email(2000, "alice@example.com", "Re: Launch", "agreed"),
|
||||
],
|
||||
};
|
||||
let out = canonicalise("gmail:t1", "alice@example.com", &[], t)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(out.markdown.contains("# Email thread — gmail — Launch"));
|
||||
let out = canonicalise(
|
||||
"gmail:alice@example.com|bob@example.com",
|
||||
"alice@example.com",
|
||||
&[],
|
||||
t,
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
// No leading `# Email thread` header — that info belongs in front-matter.
|
||||
assert!(
|
||||
!out.markdown.contains("# Email thread — gmail — Launch"),
|
||||
"canonical email MD must NOT contain a `# ` header"
|
||||
);
|
||||
assert!(out.markdown.contains("From: bob@example.com"));
|
||||
assert!(out.markdown.contains("Subject: Launch"));
|
||||
assert!(out.markdown.contains("let's ship"));
|
||||
@@ -139,6 +155,50 @@ mod tests {
|
||||
assert!(out.markdown.contains("agreed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clean_body_strips_footer_before_canonicalise() {
|
||||
// Body where "Unsubscribe" line triggers footer removal. Everything from
|
||||
// that line onward is dropped by clean_body; real content above survives.
|
||||
let body_with_footer =
|
||||
"Please review the attached document.\n\nUnsubscribe https://mail.example.com/unsub\n© 2026 Example Corp";
|
||||
let t = EmailThread {
|
||||
provider: "gmail".into(),
|
||||
thread_subject: "Review".into(),
|
||||
messages: vec![EmailMessage {
|
||||
from: "sender@example.com".into(),
|
||||
to: vec!["recipient@example.com".into()],
|
||||
cc: vec![],
|
||||
subject: "Review".into(),
|
||||
sent_at: Utc.timestamp_millis_opt(5000).unwrap(),
|
||||
body: body_with_footer.into(),
|
||||
source_ref: None,
|
||||
}],
|
||||
};
|
||||
let out = canonicalise(
|
||||
"gmail:recipient@example.com|sender@example.com",
|
||||
"recipient@example.com",
|
||||
&[],
|
||||
t,
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(
|
||||
out.markdown.contains("Please review the attached document"),
|
||||
"real content must survive; got:\n{}",
|
||||
out.markdown
|
||||
);
|
||||
assert!(
|
||||
!out.markdown.to_ascii_lowercase().contains("unsubscribe"),
|
||||
"unsubscribe footer must be stripped; got:\n{}",
|
||||
out.markdown
|
||||
);
|
||||
assert!(
|
||||
!out.markdown.contains("© 2026"),
|
||||
"copyright footer must be stripped; got:\n{}",
|
||||
out.markdown
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn time_range_spans_thread() {
|
||||
let t = EmailThread {
|
||||
|
||||
@@ -0,0 +1,382 @@
|
||||
//! Shared email rendering + cleaning helpers.
|
||||
//!
|
||||
//! Used by both [`canonicalize::email`](super::email) (when rendering the
|
||||
//! `GmailMarkdownStyle::Standard` shape) and the `gmail-fetch-emails` bin
|
||||
//! (which writes per-sender markdown digests to disk). Lifted out of the
|
||||
//! bin so the bin's behaviour and the production canonicaliser stay
|
||||
//! byte-identical on body cleanup.
|
||||
//!
|
||||
//! The module is intentionally pure-string-oriented + a single
|
||||
//! `serde_json::Value` helper (`parse_message_date`) used by callers that
|
||||
//! work directly off Gmail's slim envelope JSON. Nothing here depends on
|
||||
//! the memory-tree types — that keeps the helpers reusable.
|
||||
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Two-stage cleanup applied to each message body before it gets
|
||||
/// blockquoted into a digest:
|
||||
///
|
||||
/// 1. **Drop quoted reply chains** — once a message contains a
|
||||
/// `On <date>, <name> wrote:` preamble, an `Original Message` /
|
||||
/// `Forwarded message` separator, or a run of three+ consecutive
|
||||
/// `>`-prefixed lines, everything from that point onward is the
|
||||
/// parent message we already render directly above.
|
||||
/// 2. **Drop footer noise** — `Unsubscribe`, `View in browser`,
|
||||
/// copyright lines, legal disclaimers, and address blocks. We cut
|
||||
/// at the first line containing any of [`FOOTER_TRIGGERS`].
|
||||
///
|
||||
/// The two passes run in order so a quoted-chain preamble below a
|
||||
/// "view in browser" line still gets stripped on its own merits even
|
||||
/// if the footer pass missed it.
|
||||
pub fn clean_body(raw: &str) -> String {
|
||||
let stage1 = drop_reply_chain(raw);
|
||||
let stage2 = drop_footer_noise(&stage1);
|
||||
collapse_blank_runs(stage2.trim())
|
||||
}
|
||||
|
||||
/// Substrings that, when matched (case-insensitive) anywhere on a
|
||||
/// line, mark the start of footer / boilerplate territory. Conservative
|
||||
/// list — every entry should be unambiguous noise that wouldn't
|
||||
/// reasonably appear inside real prose.
|
||||
const FOOTER_TRIGGERS: &[&str] = &[
|
||||
"unsubscribe",
|
||||
"view in browser",
|
||||
"view this email in your browser",
|
||||
"view it in your browser",
|
||||
"update your email settings",
|
||||
"manage your subscription",
|
||||
"manage preferences",
|
||||
"email preferences",
|
||||
"you are receiving this email because",
|
||||
"you received this email because",
|
||||
"you're receiving this email because",
|
||||
"to stop receiving",
|
||||
"all rights reserved",
|
||||
"© 20",
|
||||
"(c) 20",
|
||||
"copyright 20",
|
||||
"powered by mailchimp",
|
||||
"sent via sendgrid",
|
||||
"this email and any files",
|
||||
"confidentiality notice",
|
||||
"if you are not the intended recipient",
|
||||
"this communication may contain",
|
||||
];
|
||||
|
||||
/// Strip quoted reply chains. See [`clean_body`] for details.
|
||||
pub fn drop_reply_chain(s: &str) -> String {
|
||||
let mut offset = 0usize;
|
||||
let mut quoted_run_start: Option<usize> = None;
|
||||
let mut quoted_run_len = 0u32;
|
||||
|
||||
for line in s.split_inclusive('\n') {
|
||||
let trimmed = line.trim();
|
||||
let lower = trimmed.to_ascii_lowercase();
|
||||
|
||||
// Explicit reply / forward markers.
|
||||
let is_preamble = (lower.starts_with("on ") && lower.contains(" wrote:"))
|
||||
|| lower.contains("---------- forwarded message")
|
||||
|| lower.contains("----- original message")
|
||||
|| lower.contains("--------- original message")
|
||||
|| lower.contains("--- forwarded by");
|
||||
if is_preamble {
|
||||
return s[..offset].trim_end().to_string();
|
||||
}
|
||||
|
||||
// Three+ consecutive lines starting with `>` is a quoted
|
||||
// reply chain in disguise (some clients de-quote on send).
|
||||
// Treat the start of the run as the cut point.
|
||||
if trimmed.starts_with('>') {
|
||||
if quoted_run_start.is_none() {
|
||||
quoted_run_start = Some(offset);
|
||||
quoted_run_len = 1;
|
||||
} else {
|
||||
quoted_run_len += 1;
|
||||
}
|
||||
if quoted_run_len >= 3 {
|
||||
let cut = quoted_run_start.unwrap_or(offset);
|
||||
return s[..cut].trim_end().to_string();
|
||||
}
|
||||
} else if !trimmed.is_empty() {
|
||||
// Reset on a non-empty, non-quoted line. Blank lines
|
||||
// don't break a quote run because senders often interleave
|
||||
// them.
|
||||
quoted_run_start = None;
|
||||
quoted_run_len = 0;
|
||||
}
|
||||
|
||||
offset += line.len();
|
||||
}
|
||||
s.to_string()
|
||||
}
|
||||
|
||||
/// Strip everything from the first line containing a footer trigger
|
||||
/// onward. See [`FOOTER_TRIGGERS`] for the matched list.
|
||||
pub fn drop_footer_noise(s: &str) -> String {
|
||||
let mut offset = 0usize;
|
||||
for line in s.split_inclusive('\n') {
|
||||
let lower = line.to_ascii_lowercase();
|
||||
if FOOTER_TRIGGERS.iter().any(|t| lower.contains(t)) {
|
||||
return s[..offset].trim_end().to_string();
|
||||
}
|
||||
offset += line.len();
|
||||
}
|
||||
s.to_string()
|
||||
}
|
||||
|
||||
/// Collapse runs of 2+ blank lines into a single blank line. Trims
|
||||
/// trailing newlines.
|
||||
pub fn collapse_blank_runs(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
let mut blank = 0u32;
|
||||
for line in s.lines() {
|
||||
if line.trim().is_empty() {
|
||||
blank += 1;
|
||||
if blank <= 1 {
|
||||
out.push('\n');
|
||||
}
|
||||
} else {
|
||||
blank = 0;
|
||||
out.push_str(line);
|
||||
out.push('\n');
|
||||
}
|
||||
}
|
||||
while out.ends_with('\n') {
|
||||
out.pop();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Truncate a body to at most `max_chars` characters, appending `…` when
|
||||
/// the body is longer. Trims first so leading/trailing whitespace doesn't
|
||||
/// count against the budget.
|
||||
pub fn truncate_body(body: &str, max_chars: usize) -> String {
|
||||
let trimmed = body.trim();
|
||||
if trimmed.chars().count() <= max_chars {
|
||||
return trimmed.to_string();
|
||||
}
|
||||
let mut out: String = trimmed.chars().take(max_chars).collect();
|
||||
out.push('…');
|
||||
out
|
||||
}
|
||||
|
||||
/// Escape only the few markdown chars that would visibly break the
|
||||
/// header/inline contexts we use (#, |, *, _, `). Newlines collapse to
|
||||
/// spaces. We leave most punctuation alone — the body is rendered as a
|
||||
/// blockquote anyway.
|
||||
pub fn md_escape(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
for ch in s.chars() {
|
||||
match ch {
|
||||
'\\' | '`' | '*' | '_' | '|' => {
|
||||
out.push('\\');
|
||||
out.push(ch);
|
||||
}
|
||||
'\n' | '\r' => out.push(' '),
|
||||
_ => out.push(ch),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Pull the `<addr@host>` portion out of a `From` header, returning
|
||||
/// just the bare email address. Falls back to `None` when no `<…>`
|
||||
/// brackets exist; in that case the caller may use the raw From field.
|
||||
pub fn extract_email(from: &str) -> Option<String> {
|
||||
let s = from.trim();
|
||||
if let (Some(start), Some(end)) = (s.rfind('<'), s.rfind('>')) {
|
||||
if start < end {
|
||||
let inner = s[start + 1..end].trim();
|
||||
if inner.contains('@') {
|
||||
return Some(inner.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
if s.contains('@') && !s.contains(' ') {
|
||||
return Some(s.to_string());
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// If `s` starts with a 3-letter day-of-week prefix (`Mon, `, `Tue, `, …),
|
||||
/// return the remainder; otherwise `None`. Used to feed a strict-rfc2822
|
||||
/// reject into a lenient retry.
|
||||
fn strip_day_of_week_prefix(s: &str) -> Option<&str> {
|
||||
const DAYS: &[&str] = &["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"];
|
||||
let (prefix, rest) = s.split_once(", ")?;
|
||||
if DAYS.iter().any(|d| d.eq_ignore_ascii_case(prefix)) {
|
||||
Some(rest)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Try a sequence of common date formats. Composio's slim envelope sets
|
||||
/// `date` from `messageTimestamp` (often ISO 8601 or epoch ms) when
|
||||
/// present, falling back to the raw `Date:` header (RFC 2822). Operates
|
||||
/// on the raw `serde_json::Value` so callers that work off the slim
|
||||
/// envelope JSON don't have to reshape it first.
|
||||
pub fn parse_message_date(m: &Value) -> Option<DateTime<Utc>> {
|
||||
let raw = m.get("date")?;
|
||||
if let Some(s) = raw.as_str() {
|
||||
let s = s.trim();
|
||||
if s.is_empty() {
|
||||
return None;
|
||||
}
|
||||
// Epoch millis as a string?
|
||||
if let Ok(ms) = s.parse::<i64>() {
|
||||
return DateTime::from_timestamp_millis(ms);
|
||||
}
|
||||
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
|
||||
return Some(dt.with_timezone(&Utc));
|
||||
}
|
||||
if let Ok(dt) = DateTime::parse_from_rfc2822(s) {
|
||||
return Some(dt.with_timezone(&Utc));
|
||||
}
|
||||
// Lenient RFC 2822 fallback: strict `parse_from_rfc2822` rejects
|
||||
// mismatched day-of-week (e.g. `Mon, 21 Apr 2026 …` when Apr 21
|
||||
// 2026 is a Tuesday). Real-world MTAs occasionally send these
|
||||
// with a wrong day-name. Strip a `<DayName>, ` prefix and retry
|
||||
// with the rfc2822 body format. Keeps the date if everything
|
||||
// else is sane.
|
||||
if let Some(rest) = strip_day_of_week_prefix(s) {
|
||||
if let Ok(dt) = DateTime::parse_from_str(rest, "%d %b %Y %H:%M:%S %z") {
|
||||
return Some(dt.with_timezone(&Utc));
|
||||
}
|
||||
}
|
||||
if let Ok(d) = NaiveDate::parse_from_str(s, "%Y-%m-%d") {
|
||||
return d.and_hms_opt(0, 0, 0).map(|n| n.and_utc());
|
||||
}
|
||||
}
|
||||
if let Some(ms) = raw.as_i64() {
|
||||
return DateTime::from_timestamp_millis(ms);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn drop_reply_chain_strips_on_x_wrote_preamble() {
|
||||
let body = "Sounds good — let's do Tuesday.\n\nOn Mon, Apr 22, 2026 at 10:00 AM, Alice <a@x> wrote:\n> Tuesday or Wednesday?\n> Let me know.";
|
||||
let cleaned = drop_reply_chain(body);
|
||||
assert_eq!(cleaned.trim(), "Sounds good — let's do Tuesday.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_reply_chain_strips_forwarded_separator() {
|
||||
let body = "FYI.\n\n---------- Forwarded message ---------\nFrom: bob\nSubject: hi";
|
||||
assert_eq!(drop_reply_chain(body).trim(), "FYI.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_reply_chain_strips_consecutive_quoted_run() {
|
||||
let body = "Thanks for the update.\n\n> earlier line 1\n> earlier line 2\n> earlier line 3\n> earlier line 4";
|
||||
assert_eq!(drop_reply_chain(body).trim(), "Thanks for the update.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_reply_chain_keeps_short_quote() {
|
||||
// A single inline blockquote is fine — only 3+ consecutive lines trigger.
|
||||
let body = "I think:\n> That sounds reasonable\n\nLet's proceed.";
|
||||
let cleaned = drop_reply_chain(body);
|
||||
assert!(cleaned.contains("Let's proceed"));
|
||||
assert!(cleaned.contains("That sounds reasonable"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_footer_noise_strips_unsubscribe_block() {
|
||||
let body =
|
||||
"Big news: GPT-5.5 is here.\n\nRead more at example.com\n\nUnsubscribe | © 2026 OpenAI";
|
||||
let cleaned = drop_footer_noise(body);
|
||||
assert!(cleaned.contains("GPT-5.5"));
|
||||
assert!(!cleaned.to_ascii_lowercase().contains("unsubscribe"));
|
||||
assert!(!cleaned.contains("©"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_footer_noise_strips_legal_disclaimer() {
|
||||
let body = "Action item — review by Friday.\n\nThis email and any files transmitted with it are confidential and intended solely for the use of the individual to whom they are addressed.";
|
||||
let cleaned = drop_footer_noise(body);
|
||||
assert_eq!(cleaned.trim(), "Action item — review by Friday.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clean_body_combines_passes() {
|
||||
let body =
|
||||
"Real content here.\n\nOn Mon, Apr 22, 2026, Alice wrote:\n> old stuff\n\nUnsubscribe";
|
||||
let cleaned = clean_body(body);
|
||||
assert_eq!(cleaned, "Real content here.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collapse_blank_runs_keeps_paragraph_breaks() {
|
||||
let s = "a\n\n\n\nb\n\n\nc\n";
|
||||
assert_eq!(collapse_blank_runs(s), "a\n\nb\n\nc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_body_adds_ellipsis() {
|
||||
let s = "x".repeat(2000);
|
||||
let t = truncate_body(&s, 1200);
|
||||
assert!(t.ends_with('…'));
|
||||
assert_eq!(t.chars().count(), 1201);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_body_passthrough_when_short() {
|
||||
let s = "hello";
|
||||
let t = truncate_body(s, 1200);
|
||||
assert_eq!(t, "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn md_escape_handles_special_chars() {
|
||||
assert_eq!(md_escape("a*b_c"), "a\\*b\\_c");
|
||||
assert_eq!(md_escape("foo|bar"), "foo\\|bar");
|
||||
assert_eq!(md_escape("line1\nline2"), "line1 line2");
|
||||
assert_eq!(md_escape("plain text"), "plain text");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_email_handles_both_forms() {
|
||||
assert_eq!(
|
||||
extract_email("Alice <alice@example.com>").as_deref(),
|
||||
Some("alice@example.com")
|
||||
);
|
||||
assert_eq!(
|
||||
extract_email("notify@github.com").as_deref(),
|
||||
Some("notify@github.com")
|
||||
);
|
||||
assert_eq!(
|
||||
extract_email("\"Bot Name\" <bot@x.io>").as_deref(),
|
||||
Some("bot@x.io")
|
||||
);
|
||||
assert!(extract_email("Alice").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_message_date_handles_iso_and_rfc2822() {
|
||||
let iso = json!({"date": "2026-04-21T10:00:00Z"});
|
||||
let rfc = json!({"date": "Mon, 21 Apr 2026 10:00:00 +0000"});
|
||||
let ms = json!({"date": 1745236800000_i64});
|
||||
let ms_str = json!({"date": "1745236800000"});
|
||||
let date_only = json!({"date": "2026-04-21"});
|
||||
assert!(parse_message_date(&iso).is_some());
|
||||
assert!(parse_message_date(&rfc).is_some());
|
||||
assert!(parse_message_date(&ms).is_some());
|
||||
assert!(parse_message_date(&ms_str).is_some());
|
||||
assert!(parse_message_date(&date_only).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_message_date_returns_none_when_missing_or_blank() {
|
||||
assert!(parse_message_date(&json!({})).is_none());
|
||||
assert!(parse_message_date(&json!({"date": ""})).is_none());
|
||||
assert!(parse_message_date(&json!({"date": " "})).is_none());
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,7 @@
|
||||
pub mod chat;
|
||||
pub mod document;
|
||||
pub mod email;
|
||||
pub mod email_clean;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
|
||||
@@ -5,16 +5,26 @@
|
||||
//! so later phases (#709 seal budget = 10k tokens) can ingest them without
|
||||
//! blowing past the summariser ceiling.
|
||||
//!
|
||||
//! Splitting strategy: prefer paragraph boundaries (`\n\n`), then single
|
||||
//! newlines, then hard character-count cut. We never split a character in half.
|
||||
//! ## Dispatch by source kind (Phase B)
|
||||
//!
|
||||
//! - **Chat**: split at `## ` message boundaries. Each message becomes one
|
||||
//! chunk. If a single message exceeds `max_tokens`, fall back to the
|
||||
//! paragraph/line/char splitter for that unit only and emit each piece with
|
||||
//! `partial_message = true`.
|
||||
//! - **Email**: split at `---\nFrom:` separators. Each email in the thread
|
||||
//! becomes one chunk. Same oversize fallback as Chat.
|
||||
//! - **Document**: original paragraph-based greedy packing (unchanged).
|
||||
|
||||
use crate::openhuman::memory::tree::types::{approx_token_count, Chunk, Metadata, SourceKind};
|
||||
use crate::openhuman::memory::tree::util::redact::redact;
|
||||
|
||||
/// Default upper bound on per-chunk tokens.
|
||||
///
|
||||
/// Aligned with the LLD's summariser 10k budget so a single chunk never blows
|
||||
/// a seal on its own.
|
||||
pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 10_000;
|
||||
/// Sized below the L0 seal budget (`source_tree::types::TOKEN_BUDGET = 4_500`)
|
||||
/// so each seal accumulates roughly 1–3 chunks before firing — natural pacing
|
||||
/// for the local 1B summariser, which produces noticeably better summaries
|
||||
/// with smaller (≤4–5k) inputs than at the previous 10k cap.
|
||||
pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 3_000;
|
||||
|
||||
/// Tunable settings for the chunker.
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -49,34 +59,255 @@ pub struct ChunkerInput {
|
||||
/// Returns chunks in source order with stable sequence numbers starting at 0.
|
||||
/// Chunk IDs are deterministic (`types::chunk_id`), so re-chunking yields the
|
||||
/// same ids for identical input.
|
||||
///
|
||||
/// ## Dispatch by source kind
|
||||
///
|
||||
/// - **Chat / Email**: split at message/email boundaries, then greedy-pack
|
||||
/// consecutive units into a single chunk until adding the next unit would
|
||||
/// exceed `max_tokens`. Oversize units (a single message > `max_tokens`)
|
||||
/// fall back to the paragraph/line/char splitter and emit each piece with
|
||||
/// `partial_message = true`.
|
||||
/// - **Document**: original paragraph-based greedy packing (unchanged).
|
||||
pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec<Chunk> {
|
||||
let now = chrono::Utc::now();
|
||||
let pieces = split_by_token_budget(&input.markdown, opts.max_tokens);
|
||||
let max_tokens = opts.max_tokens.max(1);
|
||||
let max_chars = (max_tokens as usize).saturating_mul(4);
|
||||
|
||||
// Dispatch: pick splitting units based on source kind.
|
||||
let units: Vec<String> = match input.source_kind {
|
||||
SourceKind::Chat => split_chat_messages(&input.markdown),
|
||||
SourceKind::Email => split_email_messages(&input.markdown),
|
||||
SourceKind::Document => {
|
||||
// Document: run the existing paragraph splitter directly on the
|
||||
// whole blob. No message-unit concept.
|
||||
log::debug!(
|
||||
"[memory_tree::chunker] document source_id_hash={} len={} — paragraph split",
|
||||
redact(&input.source_id),
|
||||
input.markdown.len()
|
||||
);
|
||||
split_by_token_budget(&input.markdown, max_tokens)
|
||||
}
|
||||
};
|
||||
|
||||
if matches!(input.source_kind, SourceKind::Document) {
|
||||
// Already split by budget; wrap directly.
|
||||
return units
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(idx, content)| {
|
||||
let seq = idx as u32;
|
||||
let token_count = approx_token_count(&content);
|
||||
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
|
||||
Chunk {
|
||||
id,
|
||||
content,
|
||||
metadata: input.metadata.clone(),
|
||||
token_count,
|
||||
seq_in_source: seq,
|
||||
created_at: now,
|
||||
partial_message: false,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[memory_tree::chunker] source_kind={} source_id={} len={} pieces={}",
|
||||
"[memory_tree::chunker] source_kind={} source_id_hash={} len={} units={}",
|
||||
input.source_kind.as_str(),
|
||||
input.source_id,
|
||||
redact(&input.source_id),
|
||||
input.markdown.len(),
|
||||
pieces.len()
|
||||
units.len()
|
||||
);
|
||||
|
||||
pieces
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(idx, content)| {
|
||||
let seq = idx as u32;
|
||||
let token_count = approx_token_count(&content);
|
||||
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
|
||||
Chunk {
|
||||
id,
|
||||
content,
|
||||
metadata: input.metadata.clone(),
|
||||
token_count,
|
||||
seq_in_source: seq,
|
||||
created_at: now,
|
||||
// For Chat and Email: greedy-pack consecutive units into chunks.
|
||||
// Units are accumulated until adding the next would exceed max_chars;
|
||||
// oversize single units fall back to sub-splitting with partial_message=true.
|
||||
let unit_separator = "\n\n";
|
||||
let sep_chars = unit_separator.chars().count();
|
||||
|
||||
let mut out: Vec<Chunk> = Vec::new();
|
||||
let mut acc: Vec<String> = Vec::new();
|
||||
let mut acc_chars = 0usize;
|
||||
|
||||
// Flush accumulated units as one packed chunk.
|
||||
let flush = |acc: &mut Vec<String>, acc_chars: &mut usize, out: &mut Vec<Chunk>| {
|
||||
if acc.is_empty() {
|
||||
return;
|
||||
}
|
||||
let content = acc.join(unit_separator);
|
||||
let seq = out.len() as u32;
|
||||
let tc = approx_token_count(&content);
|
||||
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
|
||||
out.push(Chunk {
|
||||
id,
|
||||
content,
|
||||
metadata: input.metadata.clone(),
|
||||
token_count: tc,
|
||||
seq_in_source: seq,
|
||||
created_at: now,
|
||||
partial_message: false,
|
||||
});
|
||||
acc.clear();
|
||||
*acc_chars = 0;
|
||||
};
|
||||
|
||||
for unit in units {
|
||||
let unit_chars = unit.chars().count();
|
||||
|
||||
if unit_chars > max_chars {
|
||||
// Oversize: flush any pending accumulator first, then sub-split.
|
||||
flush(&mut acc, &mut acc_chars, &mut out);
|
||||
let sub_pieces = split_by_token_budget(&unit, max_tokens);
|
||||
for piece in sub_pieces {
|
||||
let seq = out.len() as u32;
|
||||
let tc = approx_token_count(&piece);
|
||||
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &piece);
|
||||
out.push(Chunk {
|
||||
id,
|
||||
content: piece,
|
||||
metadata: input.metadata.clone(),
|
||||
token_count: tc,
|
||||
seq_in_source: seq,
|
||||
created_at: now,
|
||||
partial_message: true,
|
||||
});
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute projected size if we add this unit to the accumulator.
|
||||
let projected = if acc.is_empty() {
|
||||
unit_chars
|
||||
} else {
|
||||
acc_chars + sep_chars + unit_chars
|
||||
};
|
||||
|
||||
if projected > max_chars {
|
||||
// Adding this unit would overflow — flush the accumulator first.
|
||||
flush(&mut acc, &mut acc_chars, &mut out);
|
||||
}
|
||||
|
||||
if !acc.is_empty() {
|
||||
acc_chars += sep_chars;
|
||||
}
|
||||
acc_chars += unit_chars;
|
||||
acc.push(unit);
|
||||
}
|
||||
|
||||
// Flush any remaining accumulated units.
|
||||
flush(&mut acc, &mut acc_chars, &mut out);
|
||||
|
||||
if out.is_empty() {
|
||||
// Degenerate: empty input → one empty chunk, matching original behaviour.
|
||||
let id = super::types::chunk_id(input.source_kind, &input.source_id, 0, "");
|
||||
out.push(Chunk {
|
||||
id,
|
||||
content: String::new(),
|
||||
metadata: input.metadata.clone(),
|
||||
token_count: 0,
|
||||
seq_in_source: 0,
|
||||
created_at: now,
|
||||
partial_message: false,
|
||||
});
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Split a canonical chat blob into per-message units at `## ` boundaries.
|
||||
///
|
||||
/// Each returned string starts with `## ` and includes everything up to but
|
||||
/// not including the next `## ` boundary. If the blob starts with a `# `
|
||||
/// header (legacy or unexpected), everything before the first `## ` is
|
||||
/// dropped silently.
|
||||
fn split_chat_messages(md: &str) -> Vec<String> {
|
||||
let mut pieces: Vec<String> = Vec::new();
|
||||
let mut current: Option<String> = None;
|
||||
|
||||
for line in md.split_inclusive('\n') {
|
||||
if line.starts_with("## ") {
|
||||
if let Some(prev) = current.take() {
|
||||
let trimmed = prev.trim_end().to_string();
|
||||
if !trimmed.is_empty() {
|
||||
pieces.push(trimmed);
|
||||
}
|
||||
}
|
||||
current = Some(line.to_string());
|
||||
} else if let Some(ref mut buf) = current {
|
||||
buf.push_str(line);
|
||||
}
|
||||
// Lines before the first `## ` (e.g. a leading `# ` header) are dropped.
|
||||
}
|
||||
|
||||
if let Some(prev) = current.take() {
|
||||
let trimmed = prev.trim_end().to_string();
|
||||
if !trimmed.is_empty() {
|
||||
pieces.push(trimmed);
|
||||
}
|
||||
}
|
||||
|
||||
if pieces.is_empty() && !md.trim().is_empty() {
|
||||
// No `## ` found at all — treat whole blob as one unit.
|
||||
pieces.push(md.trim_end().to_string());
|
||||
}
|
||||
|
||||
pieces
|
||||
}
|
||||
|
||||
/// Split a canonical email thread blob into per-email units.
|
||||
///
|
||||
/// Splits at `---` (alone on a line, optional trailing whitespace) followed
|
||||
/// by a `From:` line within the next 8 lines. Each piece includes the `---`
|
||||
/// separator and everything up to but not including the next `---\nFrom:`
|
||||
/// boundary. Content before the first `---` separator is dropped (handles
|
||||
/// any leading header that might have slipped through).
|
||||
fn split_email_messages(md: &str) -> Vec<String> {
|
||||
let lines: Vec<&str> = md.split('\n').collect();
|
||||
let n = lines.len();
|
||||
let mut split_positions: Vec<usize> = Vec::new();
|
||||
|
||||
for i in 0..n {
|
||||
let line = lines[i].trim_end();
|
||||
if line == "---" {
|
||||
// Check if one of the next 8 lines starts with `From:`
|
||||
let window_end = (i + 9).min(n);
|
||||
for j in (i + 1)..window_end {
|
||||
if lines[j].starts_with("From:") {
|
||||
split_positions.push(i);
|
||||
break;
|
||||
}
|
||||
// Skip blank lines between `---` and `From:`
|
||||
if !lines[j].trim().is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if split_positions.is_empty() {
|
||||
// No email separator found — treat whole blob as one unit.
|
||||
let trimmed = md.trim_end().to_string();
|
||||
if trimmed.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
return vec![trimmed];
|
||||
}
|
||||
|
||||
let mut pieces: Vec<String> = Vec::new();
|
||||
for (idx, &start) in split_positions.iter().enumerate() {
|
||||
let end = if idx + 1 < split_positions.len() {
|
||||
split_positions[idx + 1]
|
||||
} else {
|
||||
n
|
||||
};
|
||||
let piece_lines: Vec<&str> = lines[start..end].iter().copied().collect();
|
||||
let piece = piece_lines.join("\n").trim_end().to_string();
|
||||
if !piece.is_empty() {
|
||||
pieces.push(piece);
|
||||
}
|
||||
}
|
||||
|
||||
pieces
|
||||
}
|
||||
|
||||
/// Split `text` into pieces each ≤ `max_tokens` tokens.
|
||||
@@ -85,7 +316,7 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec<Chunk>
|
||||
/// 1. Paragraph (`\n\n`)
|
||||
/// 2. Line (`\n`)
|
||||
/// 3. Hard character cut (last resort; preserves UTF-8 code points)
|
||||
fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec<String> {
|
||||
pub(crate) fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec<String> {
|
||||
let max_tokens = max_tokens.max(1);
|
||||
if text.is_empty() {
|
||||
return vec![String::new()];
|
||||
@@ -187,22 +418,33 @@ mod tests {
|
||||
Metadata::point_in_time(SourceKind::Chat, "slack:#eng", "alice", Utc::now())
|
||||
}
|
||||
|
||||
fn meta_email() -> Metadata {
|
||||
Metadata::point_in_time(SourceKind::Email, "gmail:t1", "alice", Utc::now())
|
||||
}
|
||||
|
||||
fn meta_doc() -> Metadata {
|
||||
Metadata::point_in_time(SourceKind::Document, "doc1", "alice", Utc::now())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_input_produces_single_chunk() {
|
||||
// Chat input without a `## ` header produces one chunk via the empty-
|
||||
// result fallback (whole blob as one unit).
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
markdown: "hello world".into(),
|
||||
markdown: "## 2026-01-01T00:00:00Z — alice\nhello world".into(),
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].content, "hello world");
|
||||
assert!(chunks[0].content.contains("hello world"));
|
||||
assert_eq!(chunks[0].seq_in_source, 0);
|
||||
assert!(!chunks[0].partial_message);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_input_produces_one_empty_chunk() {
|
||||
fn empty_chat_input_produces_one_empty_chunk() {
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "x".into(),
|
||||
@@ -212,10 +454,275 @@ mod tests {
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].content, "");
|
||||
assert!(!chunks[0].partial_message);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn paragraph_boundaries_preferred() {
|
||||
fn chat_messages_pack_into_one_chunk_when_small() {
|
||||
// Two small chat messages both fit under default max_tokens → greedy
|
||||
// packing emits ONE chunk containing both, joined by \n\n.
|
||||
let md = "## 2026-01-01T00:00:00Z — alice\nHello world\n\n## 2026-01-01T00:01:00Z — bob\nParagraph one.\n\nParagraph two.".to_string();
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
markdown: md.clone(),
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
// Both small messages fit under 10k tokens → one packed chunk.
|
||||
assert_eq!(
|
||||
chunks.len(),
|
||||
1,
|
||||
"small messages should be packed into one chunk; got {chunks:?}"
|
||||
);
|
||||
assert!(
|
||||
chunks[0].content.contains("alice"),
|
||||
"chunk must contain alice's message"
|
||||
);
|
||||
assert!(
|
||||
chunks[0].content.contains("bob"),
|
||||
"chunk must contain bob's message"
|
||||
);
|
||||
assert!(chunks[0].content.contains("Paragraph one."));
|
||||
assert!(chunks[0].content.contains("Paragraph two."));
|
||||
assert!(!chunks[0].partial_message);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_messages_split_at_boundary_when_large() {
|
||||
// Messages that together exceed max_tokens split at message boundaries
|
||||
// into multiple chunks. Each chunk contains whole messages only.
|
||||
// Each message is ~3k tokens at 4 chars/token = 12k chars;
|
||||
// two messages = ~6k tokens > 5k budget → must split.
|
||||
let msg_body = "x".repeat(12_000);
|
||||
let md = format!(
|
||||
"## 2026-01-01T00:00:00Z — alice\n{msg_body}\n\n## 2026-01-01T00:01:00Z — bob\n{msg_body}"
|
||||
);
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
markdown: md,
|
||||
metadata: meta(),
|
||||
};
|
||||
// Use a 5k token budget so two ~3k-token messages don't fit together.
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 });
|
||||
assert_eq!(
|
||||
chunks.len(),
|
||||
2,
|
||||
"two large messages should land in separate chunks; got {chunks:?}"
|
||||
);
|
||||
assert!(chunks[0].content.contains("alice"));
|
||||
assert!(chunks[1].content.contains("bob"));
|
||||
for c in &chunks {
|
||||
assert!(!c.partial_message, "whole messages must not be partial");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_threads_pack_into_one_chunk_when_small() {
|
||||
// Three short emails all fit under default max_tokens → one packed chunk.
|
||||
let md = "---\nFrom: alice@example.com\nSubject: Hello\nDate: 2026-01-01T00:00:00Z\n\nFirst body.\n---\nFrom: bob@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:01:00Z\n\nSecond body.\n---\nFrom: carol@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:02:00Z\n\nThird body.".to_string();
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "gmail:t1".into(),
|
||||
markdown: md,
|
||||
metadata: meta_email(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(
|
||||
chunks.len(),
|
||||
1,
|
||||
"three small emails should pack into one chunk; got {chunks:?}"
|
||||
);
|
||||
assert!(chunks[0].content.contains("First body."));
|
||||
assert!(chunks[0].content.contains("Second body."));
|
||||
assert!(chunks[0].content.contains("Third body."));
|
||||
assert!(!chunks[0].partial_message);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_thread_large_splits_at_email_boundaries() {
|
||||
// Messages totaling >12k tokens split into 2 chunks at email boundaries.
|
||||
// Each email is ~4k tokens (16k chars); 3 emails × 4k = 12k tokens.
|
||||
// With a 5k budget, 2 emails fit per chunk → 2 chunks for 3 emails.
|
||||
let email_body = "y".repeat(16_000); // ~4k tokens
|
||||
let md = format!(
|
||||
"---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{email_body}\n\
|
||||
---\nFrom: b@x.com\nDate: 2026-01-01T00:01:00Z\n\n{email_body}\n\
|
||||
---\nFrom: c@x.com\nDate: 2026-01-01T00:02:00Z\n\n{email_body}"
|
||||
);
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "gmail:t1".into(),
|
||||
markdown: md,
|
||||
metadata: meta_email(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 });
|
||||
assert!(
|
||||
chunks.len() >= 2,
|
||||
"large thread must split into multiple chunks; got {}",
|
||||
chunks.len()
|
||||
);
|
||||
for c in &chunks {
|
||||
assert!(!c.partial_message, "whole-email chunks must not be partial");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversize_single_email_splits_with_partial_flag() {
|
||||
// A single email body > max_tokens must produce partial_message=true pieces.
|
||||
let big_body = "z".repeat(50_000); // ~12.5k tokens at 4 chars/token
|
||||
let md = format!("---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{big_body}");
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "gmail:t1".into(),
|
||||
markdown: md,
|
||||
metadata: meta_email(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 1_000 });
|
||||
assert!(chunks.len() > 1, "oversize email must split");
|
||||
for c in &chunks {
|
||||
assert!(
|
||||
c.partial_message,
|
||||
"all sub-pieces of an oversize email must have partial_message=true"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_units_joined_by_double_newline() {
|
||||
// Two chat messages packed together must be separated by \n\n.
|
||||
let md = "## 2026-01-01T00:00:00Z — alice\nfoo\n\n## 2026-01-01T00:01:00Z — bob\nbar"
|
||||
.to_string();
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "x".into(),
|
||||
markdown: md,
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(chunks.len(), 1);
|
||||
// The two messages must be separated by \n\n in the packed content.
|
||||
assert!(
|
||||
chunks[0].content.contains("\n\n"),
|
||||
"packed units must be joined by \\n\\n; content={:?}",
|
||||
chunks[0].content
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversize_message_falls_back_with_partial_flag() {
|
||||
// Single chat message that is way over max_tokens.
|
||||
let long_body = "x".repeat(8000); // ~2000 tokens at 4 chars/token
|
||||
let md = format!("## 2026-01-01T00:00:00Z — alice\n{long_body}");
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "x".into(),
|
||||
markdown: md,
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 100 });
|
||||
assert!(chunks.len() > 1, "oversize message must split");
|
||||
for c in &chunks {
|
||||
assert!(
|
||||
c.partial_message,
|
||||
"all sub-pieces of an oversize message must have partial_message=true"
|
||||
);
|
||||
}
|
||||
// Reuniting all pieces must reconstruct the message content (minus `## ` line).
|
||||
let rejoined: String = chunks.iter().map(|c| c.content.as_str()).collect();
|
||||
assert!(rejoined.contains(&long_body[..100]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_falls_through_to_paragraph_split() {
|
||||
let para1 = "a".repeat(400); // ~100 tokens
|
||||
let para2 = "b".repeat(400);
|
||||
let para3 = "c".repeat(400);
|
||||
let text = format!("{para1}\n\n{para2}\n\n{para3}");
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "doc1".into(),
|
||||
markdown: text,
|
||||
metadata: meta_doc(),
|
||||
};
|
||||
let chunks = chunk_markdown(
|
||||
&input,
|
||||
&ChunkerOptions {
|
||||
max_tokens: 150, // forces split at paragraph boundary
|
||||
},
|
||||
);
|
||||
assert!(chunks.len() >= 2);
|
||||
for c in &chunks {
|
||||
let first = c.content.chars().next().unwrap();
|
||||
assert!(
|
||||
matches!(first, 'a' | 'b' | 'c'),
|
||||
"document chunk starts with unexpected char: {:?}",
|
||||
c.content.chars().take(10).collect::<String>()
|
||||
);
|
||||
assert!(
|
||||
!c.partial_message,
|
||||
"document chunks must not have partial_message=true"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_line_dropped_in_chat() {
|
||||
// Simulate a blob that has a leading `# Chat transcript` header.
|
||||
let md = "# Chat transcript — slack / #eng\n\n## 2026-01-01T00:00:00Z — alice\nhello"
|
||||
.to_string();
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "x".into(),
|
||||
markdown: md,
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(chunks.len(), 1);
|
||||
// The `# Chat transcript` header must be absent from the chunk content.
|
||||
assert!(
|
||||
!chunks[0].content.contains("# Chat transcript"),
|
||||
"leading `# ` header must be dropped from chunk content"
|
||||
);
|
||||
assert!(chunks[0].content.contains("hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_ids_are_stable_across_runs() {
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
markdown: "## 2026-01-01T00:00:00Z — alice\nhello".into(),
|
||||
metadata: meta(),
|
||||
};
|
||||
let a = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
let b = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(
|
||||
a.iter().map(|c| c.id.clone()).collect::<Vec<_>>(),
|
||||
b.iter().map(|c| c.id.clone()).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sequence_numbers_start_at_zero() {
|
||||
let msgs: String = (0..5)
|
||||
.map(|i| format!("## 2026-01-01T00:0{}:00Z — user{i}\nContent {i}\n\n", i))
|
||||
.collect();
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "x".into(),
|
||||
markdown: msgs,
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
for (idx, c) in chunks.iter().enumerate() {
|
||||
assert_eq!(c.seq_in_source, idx as u32);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn paragraph_boundaries_preferred_for_documents() {
|
||||
// Build something that exceeds token budget so it must split.
|
||||
let para1 = "a".repeat(400); // ~100 tokens
|
||||
let para2 = "b".repeat(400);
|
||||
@@ -225,7 +732,7 @@ mod tests {
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "doc1".into(),
|
||||
markdown: text,
|
||||
metadata: meta(),
|
||||
metadata: meta_doc(),
|
||||
};
|
||||
let chunks = chunk_markdown(
|
||||
&input,
|
||||
@@ -234,8 +741,6 @@ mod tests {
|
||||
},
|
||||
);
|
||||
assert!(chunks.len() >= 2);
|
||||
// Every chunk should be a multiple of complete paragraphs — so starts
|
||||
// with 'a', 'b', or 'c' and doesn't mix them arbitrarily.
|
||||
for c in &chunks {
|
||||
let first = c.content.chars().next().unwrap();
|
||||
assert!(
|
||||
@@ -247,16 +752,16 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_line_split_when_no_paragraphs() {
|
||||
fn falls_back_to_line_split_when_no_paragraphs_document() {
|
||||
let text = (0..30)
|
||||
.map(|i| format!("line-{i}-{}", "x".repeat(40)))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "x".into(),
|
||||
markdown: text,
|
||||
metadata: meta(),
|
||||
metadata: meta_doc(),
|
||||
};
|
||||
let chunks = chunk_markdown(
|
||||
&input,
|
||||
@@ -271,54 +776,14 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_ids_are_stable_across_runs() {
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
markdown: "hello".into(),
|
||||
metadata: meta(),
|
||||
};
|
||||
let a = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
let b = chunk_markdown(&input, &ChunkerOptions::default());
|
||||
assert_eq!(
|
||||
a.iter().map(|c| c.id.clone()).collect::<Vec<_>>(),
|
||||
b.iter().map(|c| c.id.clone()).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sequence_numbers_start_at_zero() {
|
||||
let text = (0..10)
|
||||
.map(|i| format!("p{i}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "d".into(),
|
||||
markdown: text,
|
||||
metadata: meta(),
|
||||
};
|
||||
let chunks = chunk_markdown(
|
||||
&input,
|
||||
&ChunkerOptions {
|
||||
max_tokens: 2, // forces one chunk per paragraph basically
|
||||
},
|
||||
);
|
||||
for (idx, c) in chunks.iter().enumerate() {
|
||||
assert_eq!(c.seq_in_source, idx as u32);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn utf8_boundaries_preserved_on_hard_split() {
|
||||
fn utf8_boundaries_preserved_on_hard_split_document() {
|
||||
// Single long line with no paragraph/line splits → falls to hard cut.
|
||||
// Use a 3-byte-per-char Unicode codepoint (Chinese) and force a cut.
|
||||
let text = "中".repeat(400);
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "d".into(),
|
||||
markdown: text.clone(),
|
||||
metadata: meta(),
|
||||
metadata: meta_doc(),
|
||||
};
|
||||
let chunks = chunk_markdown(
|
||||
&input,
|
||||
@@ -332,12 +797,12 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_token_budget_is_clamped_without_empty_leading_chunk() {
|
||||
fn zero_token_budget_is_clamped_without_empty_leading_chunk_document() {
|
||||
let input = ChunkerInput {
|
||||
source_kind: SourceKind::Document,
|
||||
source_id: "d".into(),
|
||||
markdown: "abcdef".into(),
|
||||
metadata: meta(),
|
||||
metadata: meta_doc(),
|
||||
};
|
||||
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 0 });
|
||||
assert!(!chunks.is_empty());
|
||||
|
||||
@@ -0,0 +1,416 @@
|
||||
//! Atomic content-file writes via tempfile + fsync + rename.
|
||||
//!
|
||||
//! Each chunk body is written to `<parent>/.tmp_<uuid>.md`, then renamed to
|
||||
//! its final path. The rename is atomic on any POSIX filesystem and behaves
|
||||
//! correctly on NTFS (the old file is replaced atomically by the OS).
|
||||
//!
|
||||
//! **Immutability contract**: once a file exists at `abs_path`, it is never
|
||||
//! overwritten. Callers must detect "already exists" and skip the write.
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
|
||||
use super::compose::{compose_summary_md, split_front_matter, SummaryComposeInput};
|
||||
use super::paths::{summary_abs_path, summary_rel_path};
|
||||
|
||||
/// Write `bytes` atomically to `abs_path` if the file does not already exist.
|
||||
///
|
||||
/// Returns `Ok(true)` when the file was newly written, `Ok(false)` when it
|
||||
/// already existed (the existing file is left unchanged).
|
||||
///
|
||||
/// The write uses a sibling tempfile + rename so the final path is never
|
||||
/// visible in a partial state. Parent directories are created automatically.
|
||||
pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result<bool> {
|
||||
// Fast path: file already exists.
|
||||
if abs_path.exists() {
|
||||
log::debug!(
|
||||
"[content_store::atomic] skipping existing file: {}",
|
||||
abs_path.display()
|
||||
);
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| anyhow::anyhow!("create_dir_all {:?}: {e}", parent))?;
|
||||
|
||||
// Write to a temp file in the same directory so rename is atomic.
|
||||
let tmp_name = format!(".tmp_{}.md", uuid_v4_hex());
|
||||
let tmp_path = parent.join(&tmp_name);
|
||||
|
||||
{
|
||||
let mut f = std::fs::File::create(&tmp_path)
|
||||
.map_err(|e| anyhow::anyhow!("create tempfile {:?}: {e}", tmp_path))?;
|
||||
f.write_all(bytes)
|
||||
.map_err(|e| anyhow::anyhow!("write tempfile {:?}: {e}", tmp_path))?;
|
||||
f.sync_all()
|
||||
.map_err(|e| anyhow::anyhow!("fsync tempfile {:?}: {e}", tmp_path))?;
|
||||
}
|
||||
|
||||
// Rename: if the target appeared concurrently (another thread/process beat
|
||||
// us), we lost the race — remove our temp and return false.
|
||||
match std::fs::rename(&tmp_path, abs_path) {
|
||||
Ok(()) => {
|
||||
// fsync the parent directory so the rename (directory entry
|
||||
// update) is durable across a crash or power loss. Without this,
|
||||
// sync_all() on the file alone only durabilises the file data;
|
||||
// the new directory entry can remain in pagecache and be lost if
|
||||
// the system crashes before the OS flushes it. On POSIX (Linux /
|
||||
// macOS) this is required for rename durability. On Windows, NTFS
|
||||
// handles this differently and File::sync_all on a directory
|
||||
// handle is not meaningful, so we restrict the call to Unix.
|
||||
#[cfg(unix)]
|
||||
if let Some(parent) = abs_path.parent() {
|
||||
if let Ok(dir) = std::fs::File::open(parent) {
|
||||
if let Err(e) = dir.sync_all() {
|
||||
// Best-effort: the rename already committed the file;
|
||||
// a dirent fsync failure is logged but not fatal.
|
||||
log::warn!(
|
||||
"[content_store::atomic] parent dir fsync failed for {:?}: {e}",
|
||||
parent
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
log::debug!("[content_store::atomic] wrote {}", abs_path.display());
|
||||
Ok(true)
|
||||
}
|
||||
Err(e) => {
|
||||
// Best-effort cleanup of the temp file on failure.
|
||||
let _ = std::fs::remove_file(&tmp_path);
|
||||
if abs_path.exists() {
|
||||
// Lost the race — another writer created the file first.
|
||||
log::debug!(
|
||||
"[content_store::atomic] lost rename race for {}",
|
||||
abs_path.display()
|
||||
);
|
||||
Ok(false)
|
||||
} else {
|
||||
Err(anyhow::anyhow!(
|
||||
"rename {:?} -> {:?}: {e}",
|
||||
tmp_path,
|
||||
abs_path
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A summary that has been written to disk and is ready for SQLite upsert.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StagedSummary {
|
||||
pub summary_id: String,
|
||||
/// Relative content path (forward-slash, e.g. `"summaries/source/slug/L1/id.md"`).
|
||||
pub content_path: String,
|
||||
/// SHA-256 hex digest over the **body bytes** only (front-matter excluded).
|
||||
pub content_sha256: String,
|
||||
}
|
||||
|
||||
/// Write a summary `.md` file to disk and return a [`StagedSummary`] ready for
|
||||
/// SQLite upsert.
|
||||
///
|
||||
/// The relative path is built from the input metadata and the `tree_kind`. The
|
||||
/// `date_for_global` argument is required when `input.tree_kind ==
|
||||
/// SummaryTreeKind::Global`. The `scope_slug` must already be slugified by the
|
||||
/// caller.
|
||||
///
|
||||
/// If the file already exists with the same body SHA-256 (idempotent re-stage),
|
||||
/// the existing `StagedSummary` is returned without rewriting.
|
||||
pub fn stage_summary(
|
||||
content_root: &Path,
|
||||
input: &SummaryComposeInput<'_>,
|
||||
scope_slug: &str,
|
||||
date_for_global: Option<chrono::DateTime<chrono::Utc>>,
|
||||
) -> anyhow::Result<StagedSummary> {
|
||||
let rel_path = summary_rel_path(
|
||||
input.tree_kind,
|
||||
scope_slug,
|
||||
input.level,
|
||||
input.summary_id,
|
||||
date_for_global,
|
||||
);
|
||||
let abs_path = summary_abs_path(
|
||||
content_root,
|
||||
input.tree_kind,
|
||||
scope_slug,
|
||||
input.level,
|
||||
input.summary_id,
|
||||
date_for_global,
|
||||
);
|
||||
|
||||
let composed = compose_summary_md(input);
|
||||
let body_bytes = composed.body.as_bytes();
|
||||
let sha256 = sha256_hex(body_bytes);
|
||||
|
||||
// Idempotent re-stage: if the file already exists, read and hash its
|
||||
// body bytes. If the on-disk hash matches the new body's hash, return
|
||||
// the StagedSummary unchanged (true idempotency). If the hashes differ
|
||||
// the on-disk file is stale/corrupted — re-write it atomically with the
|
||||
// new content so the db row and disk file are always consistent.
|
||||
//
|
||||
// Not re-writing would leave SQLite storing a content_sha256 that
|
||||
// doesn't match the actual on-disk bytes, breaking integrity checks.
|
||||
if abs_path.exists() {
|
||||
let disk_sha = read_body_sha256(&abs_path).unwrap_or_default();
|
||||
if disk_sha == sha256 {
|
||||
log::debug!(
|
||||
"[content_store::atomic] summary already on disk with matching sha: {}",
|
||||
input.summary_id
|
||||
);
|
||||
return Ok(StagedSummary {
|
||||
summary_id: input.summary_id.to_string(),
|
||||
content_path: rel_path,
|
||||
content_sha256: sha256,
|
||||
});
|
||||
}
|
||||
// Hash mismatch — overwrite atomically.
|
||||
log::debug!(
|
||||
"[content_store::atomic] summary on-disk sha mismatch for {} — re-staging",
|
||||
input.summary_id
|
||||
);
|
||||
// Remove the stale file first; write_if_new's fast-path would skip it.
|
||||
let _ = std::fs::remove_file(&abs_path);
|
||||
}
|
||||
|
||||
let full_bytes = composed.full.as_bytes();
|
||||
write_if_new(&abs_path, full_bytes)?;
|
||||
|
||||
log::debug!(
|
||||
"[content_store::atomic] staged summary {} → {}",
|
||||
input.summary_id,
|
||||
rel_path
|
||||
);
|
||||
|
||||
Ok(StagedSummary {
|
||||
summary_id: input.summary_id.to_string(),
|
||||
content_path: rel_path,
|
||||
content_sha256: sha256,
|
||||
})
|
||||
}
|
||||
|
||||
/// Read a summary/chunk `.md` file from disk, split off the YAML front-matter,
|
||||
/// and return the SHA-256 hex digest of the **body bytes only**. Returns an
|
||||
/// empty string (not an error) if the file cannot be read or parsed, so
|
||||
/// callers can use the result as a cache key without propagating IO errors.
|
||||
fn read_body_sha256(path: &Path) -> anyhow::Result<String> {
|
||||
let raw = std::fs::read(path)?;
|
||||
let content = std::str::from_utf8(&raw)?;
|
||||
let (_fm, body) = split_front_matter(content)
|
||||
.ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", path))?;
|
||||
Ok(sha256_hex(body.as_bytes()))
|
||||
}
|
||||
|
||||
/// Compute the SHA-256 hex digest of `bytes`.
|
||||
pub fn sha256_hex(bytes: &[u8]) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(bytes);
|
||||
hex::encode(hasher.finalize())
|
||||
}
|
||||
|
||||
/// Tiny deterministic-ish hex string for temp file names.
|
||||
fn uuid_v4_hex() -> String {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let t = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
// Use a counter + timestamp for entropy (thread_id::as_u64 is nightly-only).
|
||||
static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
|
||||
let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
format!(
|
||||
"{:08x}{:016x}",
|
||||
t,
|
||||
n.wrapping_mul(0x9e37_79b9_7f4a_7c15).wrapping_add(t as u64)
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store::compose::SummaryComposeInput;
|
||||
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn write_creates_file_and_returns_true() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("sub").join("0.md");
|
||||
let written = write_if_new(&path, b"hello world").unwrap();
|
||||
assert!(written, "first write must return true");
|
||||
assert_eq!(std::fs::read(&path).unwrap(), b"hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_is_idempotent_returns_false_on_second_call() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("0.md");
|
||||
write_if_new(&path, b"first").unwrap();
|
||||
let written = write_if_new(&path, b"second").unwrap();
|
||||
assert!(!written, "second write must return false");
|
||||
assert_eq!(std::fs::read(&path).unwrap(), b"first");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sha256_hex_is_stable() {
|
||||
let a = sha256_hex(b"hello");
|
||||
let b = sha256_hex(b"hello");
|
||||
assert_eq!(a, b);
|
||||
assert_ne!(sha256_hex(b"hello"), sha256_hex(b"world"));
|
||||
assert_eq!(a.len(), 64); // 32 bytes → 64 hex chars
|
||||
}
|
||||
|
||||
fn mk_summary_input<'a>(
|
||||
tree_kind: SummaryTreeKind,
|
||||
scope: &'a str,
|
||||
id: &'a str,
|
||||
body: &'a str,
|
||||
children: &'a [String],
|
||||
) -> SummaryComposeInput<'a> {
|
||||
use chrono::TimeZone;
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
SummaryComposeInput {
|
||||
summary_id: id,
|
||||
tree_kind,
|
||||
tree_id: "tree-001",
|
||||
tree_scope: scope,
|
||||
level: 1,
|
||||
child_ids: children,
|
||||
child_count: children.len(),
|
||||
time_range_start: ts,
|
||||
time_range_end: ts,
|
||||
sealed_at: ts,
|
||||
body,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_summary_writes_file_and_returns_staged() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let children = vec!["c1".to_string()];
|
||||
let input = mk_summary_input(
|
||||
SummaryTreeKind::Source,
|
||||
"gmail:alice@x.com",
|
||||
"summary:L1:test1",
|
||||
"summary body",
|
||||
&children,
|
||||
);
|
||||
let staged = stage_summary(dir.path(), &input, "gmail-alice-x-com", None).unwrap();
|
||||
assert_eq!(staged.summary_id, "summary:L1:test1");
|
||||
assert!(staged.content_path.starts_with("summaries/source/"));
|
||||
assert!(staged.content_path.ends_with(".md"));
|
||||
assert_eq!(staged.content_sha256.len(), 64);
|
||||
|
||||
// File must exist on disk
|
||||
let mut abs = dir.path().to_path_buf();
|
||||
for part in staged.content_path.split('/') {
|
||||
abs.push(part);
|
||||
}
|
||||
assert!(abs.exists(), "staged file must exist");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_summary_is_idempotent() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let children = vec!["c1".to_string()];
|
||||
let input = mk_summary_input(
|
||||
SummaryTreeKind::Topic,
|
||||
"person:alex",
|
||||
"summary:L1:idem",
|
||||
"idempotent body",
|
||||
&children,
|
||||
);
|
||||
let first = stage_summary(dir.path(), &input, "person-alex", None).unwrap();
|
||||
let second = stage_summary(dir.path(), &input, "person-alex", None).unwrap();
|
||||
assert_eq!(first.content_sha256, second.content_sha256);
|
||||
assert_eq!(first.content_path, second.content_path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_summary_global_uses_date_in_path() {
|
||||
use chrono::TimeZone;
|
||||
let dir = TempDir::new().unwrap();
|
||||
let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap();
|
||||
let children = vec![];
|
||||
let input = mk_summary_input(
|
||||
SummaryTreeKind::Global,
|
||||
"global",
|
||||
"summary:L0:daily",
|
||||
"daily recap",
|
||||
&children,
|
||||
);
|
||||
let staged = stage_summary(dir.path(), &input, "global", Some(date)).unwrap();
|
||||
assert!(
|
||||
staged.content_path.contains("2026-04-28"),
|
||||
"global summary path must contain date; got: {}",
|
||||
staged.content_path
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_summary_sha256_is_over_body_only() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let children = vec![];
|
||||
let body = "the body content";
|
||||
let input = mk_summary_input(
|
||||
SummaryTreeKind::Source,
|
||||
"gmail:x@y.com",
|
||||
"summary:L1:sha-test",
|
||||
body,
|
||||
&children,
|
||||
);
|
||||
let staged = stage_summary(dir.path(), &input, "gmail-x-y-com", None).unwrap();
|
||||
let expected = sha256_hex(body.as_bytes());
|
||||
assert_eq!(staged.content_sha256, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_summary_rewrites_stale_on_disk_body() {
|
||||
// Create a tempdir and write a "stale" file at the expected path with
|
||||
// a body that differs from what the new stage_summary call would write.
|
||||
// After stage_summary, the file on disk must match the new body.
|
||||
let dir = TempDir::new().unwrap();
|
||||
let children = vec!["c1".to_string()];
|
||||
let new_body = "fresh body for re-stage test";
|
||||
let input = mk_summary_input(
|
||||
SummaryTreeKind::Source,
|
||||
"gmail:stale@test.com",
|
||||
"summary:L1:stale-test",
|
||||
new_body,
|
||||
&children,
|
||||
);
|
||||
|
||||
// First stage with the real body to get the path.
|
||||
let first = stage_summary(dir.path(), &input, "gmail-stale-test-com", None).unwrap();
|
||||
|
||||
// Corrupt the on-disk file by writing a different body to the path.
|
||||
let mut abs = dir.path().to_path_buf();
|
||||
for part in first.content_path.split('/') {
|
||||
abs.push(part);
|
||||
}
|
||||
// Overwrite with stale content.
|
||||
std::fs::write(&abs, b"---\nstale_key: true\n---\nSTALE BODY CONTENT").unwrap();
|
||||
|
||||
// Now re-stage: must detect sha mismatch and re-write.
|
||||
let second = stage_summary(dir.path(), &input, "gmail-stale-test-com", None).unwrap();
|
||||
|
||||
// The returned sha must match the new body.
|
||||
let expected_sha = sha256_hex(new_body.as_bytes());
|
||||
assert_eq!(
|
||||
second.content_sha256, expected_sha,
|
||||
"re-staged sha must match new body"
|
||||
);
|
||||
|
||||
// The on-disk file must now contain the new body (not the stale one).
|
||||
let disk_bytes = std::fs::read(&abs).unwrap();
|
||||
let disk_str = std::str::from_utf8(&disk_bytes).unwrap();
|
||||
assert!(
|
||||
disk_str.contains(new_body),
|
||||
"on-disk file must contain new body after re-stage"
|
||||
);
|
||||
assert!(
|
||||
!disk_str.contains("STALE BODY CONTENT"),
|
||||
"stale body must be gone after re-stage"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,878 @@
|
||||
//! YAML front-matter + body composition for chunk `.md` files.
|
||||
//!
|
||||
//! Each file written to disk has the form:
|
||||
//! ```text
|
||||
//! ---
|
||||
//! source_kind: chat
|
||||
//! source_id: slack:#eng
|
||||
//! seq: 0
|
||||
//! owner: alice@example.com
|
||||
//! timestamp: 2026-04-28T10:00:00Z
|
||||
//! time_range_start: 2026-04-28T10:00:00Z
|
||||
//! time_range_end: 2026-04-28T10:05:00Z
|
||||
//! source_ref: slack://permalink/…
|
||||
//! tags:
|
||||
//! - person/Alice-Smith
|
||||
//! - project/Phoenix
|
||||
//! ---
|
||||
//! ## 2026-04-28T10:00:00Z — alice
|
||||
//! Message body here.
|
||||
//! ```
|
||||
//!
|
||||
//! For email source_kind, additional fields are emitted:
|
||||
//! ```text
|
||||
//! participants:
|
||||
//! - alice@example.com
|
||||
//! - bob@example.com
|
||||
//! aliases:
|
||||
//! - "alice@example.com <-> bob@example.com: chunk 0"
|
||||
//! ```
|
||||
//! These are parsed from the `source_id` field (format `gmail:{participants}`
|
||||
//! where `participants` is `addr1|addr2|...` pipe-separated) at compose time.
|
||||
//! `sender` and `thread_id` are no longer emitted — they are not meaningful
|
||||
//! with participant-based bucketing.
|
||||
//!
|
||||
//! **SHA-256 is computed over the body bytes only** (everything after `---\n`
|
||||
//! on the second delimiter line). This allows tags to be rewritten atomically
|
||||
//! without invalidating the content hash.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
|
||||
use crate::openhuman::memory::tree::types::{Chunk, SourceKind};
|
||||
|
||||
/// Compose the full file content (front-matter + body) for `chunk`.
|
||||
///
|
||||
/// Returns `(full_file_bytes, body_bytes)`. The caller writes `full_file_bytes`
|
||||
/// to disk; `body_bytes` is what the SHA-256 is computed over.
|
||||
pub fn compose_chunk_file(chunk: &Chunk) -> (Vec<u8>, Vec<u8>) {
|
||||
let front_matter = build_front_matter(chunk);
|
||||
let body = chunk.content.as_bytes().to_vec();
|
||||
|
||||
let mut full = Vec::with_capacity(front_matter.len() + body.len());
|
||||
full.extend_from_slice(&front_matter);
|
||||
full.extend_from_slice(&body);
|
||||
|
||||
(full, body)
|
||||
}
|
||||
|
||||
/// Build the YAML front-matter block (including delimiters) as UTF-8 bytes.
|
||||
fn build_front_matter(chunk: &Chunk) -> Vec<u8> {
|
||||
let meta = &chunk.metadata;
|
||||
let ts = meta.timestamp.to_rfc3339();
|
||||
let ts_start = meta.time_range.0.to_rfc3339();
|
||||
let ts_end = meta.time_range.1.to_rfc3339();
|
||||
|
||||
let mut fm = String::new();
|
||||
fm.push_str("---\n");
|
||||
fm.push_str(&format!("source_kind: {}\n", meta.source_kind.as_str()));
|
||||
// Escape backslashes and quotes in source_id for safety.
|
||||
fm.push_str(&format!("source_id: {}\n", yaml_scalar(&meta.source_id)));
|
||||
fm.push_str(&format!("seq: {}\n", chunk.seq_in_source));
|
||||
fm.push_str(&format!("owner: {}\n", yaml_scalar(&meta.owner)));
|
||||
fm.push_str(&format!("timestamp: {ts}\n"));
|
||||
fm.push_str(&format!("time_range_start: {ts_start}\n"));
|
||||
fm.push_str(&format!("time_range_end: {ts_end}\n"));
|
||||
|
||||
if let Some(ref sr) = meta.source_ref {
|
||||
fm.push_str(&format!("source_ref: {}\n", yaml_scalar(&sr.value)));
|
||||
}
|
||||
|
||||
if meta.tags.is_empty() {
|
||||
fm.push_str("tags: []\n");
|
||||
} else {
|
||||
fm.push_str("tags:\n");
|
||||
for tag in &meta.tags {
|
||||
fm.push_str(&format!(" - {}\n", yaml_scalar(tag)));
|
||||
}
|
||||
}
|
||||
|
||||
// Email-specific fields: participants list + Obsidian alias.
|
||||
// Parsed from source_id which is `gmail:{participants}` for Gmail-ingested
|
||||
// chunks, where participants is `addr1|addr2|...` (sorted, deduped).
|
||||
// If the format doesn't match, these fields are omitted.
|
||||
if meta.source_kind == SourceKind::Email {
|
||||
if let Some(addrs) = parse_gmail_participants_source_id(&meta.source_id) {
|
||||
// participants: YAML list
|
||||
fm.push_str("participants:\n");
|
||||
for addr in &addrs {
|
||||
fm.push_str(&format!(" - {}\n", yaml_scalar(addr)));
|
||||
}
|
||||
// aliases: human-readable conversation label for Obsidian
|
||||
let alias = build_participants_alias(&addrs, chunk.seq_in_source);
|
||||
fm.push_str("aliases:\n");
|
||||
fm.push_str(&format!(" - {}\n", yaml_scalar(&alias)));
|
||||
}
|
||||
}
|
||||
|
||||
fm.push_str("---\n");
|
||||
fm.into_bytes()
|
||||
}
|
||||
|
||||
/// Parse a `gmail:{participants}` source_id into the list of participant addresses.
|
||||
///
|
||||
/// `participants` is `addr1|addr2|...` (sorted, deduped, pipe-separated).
|
||||
/// Returns `Some(Vec<String>)` when the source_id has exactly two
|
||||
/// colon-separated segments (`gmail` prefix + non-empty participants). Returns
|
||||
/// `None` for legacy or malformed source_ids.
|
||||
fn parse_gmail_participants_source_id(source_id: &str) -> Option<Vec<String>> {
|
||||
let (prefix, participants) = source_id.split_once(':')?;
|
||||
if prefix != "gmail" || participants.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let addrs: Vec<String> = participants
|
||||
.split('|')
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
if addrs.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(addrs)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a human-readable alias for an email chunk suitable for Obsidian's
|
||||
/// `aliases:` field.
|
||||
///
|
||||
/// For two participants: `"alice@x.com <-> bob@y.com: chunk 0"`
|
||||
/// For more than two: `"alice@x.com <-> 2 others: chunk 0"`
|
||||
/// (where `alice@x.com` is the first in sorted order)
|
||||
///
|
||||
/// The alias is kept under ~80 characters to avoid YAML rendering issues.
|
||||
fn build_participants_alias(addrs: &[String], seq: u32) -> String {
|
||||
let label = match addrs {
|
||||
[] => "unknown".to_string(),
|
||||
[only] => only.clone(),
|
||||
[first, second] => format!("{} <-> {}", first, second),
|
||||
[first, rest @ ..] => format!("{} <-> {} others", first, rest.len()),
|
||||
};
|
||||
format!("{}: chunk {}", label, seq)
|
||||
}
|
||||
|
||||
/// Rewrite the `tags:` block in an existing file's front-matter, replacing it
|
||||
/// with the new tag list while leaving the body unchanged.
|
||||
///
|
||||
/// Returns the new full file bytes. Errors if the front-matter delimiters
|
||||
/// cannot be found.
|
||||
pub fn rewrite_tags(file_bytes: &[u8], new_tags: &[String]) -> Result<Vec<u8>, String> {
|
||||
let content =
|
||||
std::str::from_utf8(file_bytes).map_err(|e| format!("file is not valid UTF-8: {e}"))?;
|
||||
|
||||
let (front_matter, body) = split_front_matter(content)
|
||||
.ok_or_else(|| "cannot find front-matter delimiters".to_string())?;
|
||||
|
||||
// Rewrite tags: block in the front-matter string.
|
||||
let new_fm = replace_tags_in_front_matter(front_matter, new_tags)?;
|
||||
|
||||
let mut out = Vec::with_capacity(new_fm.len() + body.len() + 4);
|
||||
out.extend_from_slice(new_fm.as_bytes());
|
||||
out.extend_from_slice(body.as_bytes());
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Replace the `tags:` stanza in a front-matter string. Returns the new
|
||||
/// front-matter string (delimiters preserved).
|
||||
fn replace_tags_in_front_matter(fm: &str, new_tags: &[String]) -> Result<String, String> {
|
||||
// Build the replacement block.
|
||||
let replacement = if new_tags.is_empty() {
|
||||
"tags: []".to_string()
|
||||
} else {
|
||||
let mut s = "tags:".to_string();
|
||||
for tag in new_tags {
|
||||
s.push('\n');
|
||||
s.push_str(&format!(" - {}", yaml_scalar(tag)));
|
||||
}
|
||||
s
|
||||
};
|
||||
|
||||
// Locate the `tags:` key and consume through the block.
|
||||
let lines: Vec<&str> = fm.lines().collect();
|
||||
let mut out_lines: Vec<&str> = Vec::new();
|
||||
let mut i = 0;
|
||||
let mut found = false;
|
||||
|
||||
while i < lines.len() {
|
||||
let line = lines[i];
|
||||
if line == "tags: []" || line == "tags:" {
|
||||
found = true;
|
||||
// Skip all subsequent lines that are tag list items (start with ` - `).
|
||||
// The replacement will be inserted wholesale.
|
||||
i += 1;
|
||||
if line == "tags:" {
|
||||
while i < lines.len() && lines[i].starts_with(" - ") {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
// We've consumed the old block; we'll append replacement after the loop.
|
||||
continue;
|
||||
}
|
||||
out_lines.push(line);
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if !found {
|
||||
return Err("tags: key not found in front-matter".to_string());
|
||||
}
|
||||
|
||||
// Rebuild: all non-tag lines + replacement + closing `---`.
|
||||
// Front-matter was: `---\n...\ntags: ...\n---\n`
|
||||
// After loop, out_lines has everything except the tags block.
|
||||
// Insert replacement before the closing `---`.
|
||||
let closing = out_lines
|
||||
.iter()
|
||||
.rposition(|l| *l == "---")
|
||||
.unwrap_or(out_lines.len());
|
||||
|
||||
let mut result_lines: Vec<String> =
|
||||
out_lines[..closing].iter().map(|l| l.to_string()).collect();
|
||||
result_lines.push(replacement);
|
||||
result_lines.push("---".to_string());
|
||||
|
||||
let mut result = result_lines.join("\n");
|
||||
result.push('\n');
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// ── Summary composition ──────────────────────────────────────────────────────
|
||||
|
||||
/// Input data required to compose a summary `.md` file.
|
||||
pub struct SummaryComposeInput<'a> {
|
||||
pub summary_id: &'a str,
|
||||
pub tree_kind: SummaryTreeKind,
|
||||
pub tree_id: &'a str,
|
||||
/// Raw tree scope string, e.g. `"gmail:alice@x.com|bob@y.com"` or `"global"`.
|
||||
pub tree_scope: &'a str,
|
||||
pub level: u32,
|
||||
/// Child ids (chunk_ids at L0 → L1, summary_ids for cascades).
|
||||
pub child_ids: &'a [String],
|
||||
/// Total child count (== child_ids.len() unless truncated).
|
||||
pub child_count: usize,
|
||||
pub time_range_start: DateTime<Utc>,
|
||||
pub time_range_end: DateTime<Utc>,
|
||||
pub sealed_at: DateTime<Utc>,
|
||||
/// Raw summariser output text — the body written to disk.
|
||||
pub body: &'a str,
|
||||
}
|
||||
|
||||
/// The composed front-matter, body, and full file content for a summary.
|
||||
pub struct ComposedSummary {
|
||||
/// The YAML front-matter block (including `---` delimiters), UTF-8 string.
|
||||
pub front_matter: String,
|
||||
/// The body (summariser output), UTF-8 string.
|
||||
pub body: String,
|
||||
/// `front_matter + body` — what gets written to disk.
|
||||
pub full: String,
|
||||
}
|
||||
|
||||
/// Compose the full `.md` content for a summary node.
|
||||
///
|
||||
/// Returns a [`ComposedSummary`] whose `full` field is written to disk.
|
||||
/// SHA-256 is computed over `body` bytes only, not `full`.
|
||||
pub fn compose_summary_md(record: &SummaryComposeInput<'_>) -> ComposedSummary {
|
||||
let fm = build_summary_front_matter(record);
|
||||
let body = record.body.to_string();
|
||||
let full = format!("{}{}", fm, body);
|
||||
ComposedSummary {
|
||||
front_matter: fm,
|
||||
body,
|
||||
full,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the YAML front-matter block for a summary node.
|
||||
fn build_summary_front_matter(r: &SummaryComposeInput<'_>) -> String {
|
||||
let tree_kind_str = match r.tree_kind {
|
||||
SummaryTreeKind::Source => "source",
|
||||
SummaryTreeKind::Global => "global",
|
||||
SummaryTreeKind::Topic => "topic",
|
||||
};
|
||||
|
||||
let trs = r.time_range_start.to_rfc3339();
|
||||
let tre = r.time_range_end.to_rfc3339();
|
||||
let sealed = r.sealed_at.to_rfc3339();
|
||||
|
||||
let mut fm = String::new();
|
||||
fm.push_str("---\n");
|
||||
fm.push_str(&format!("id: {}\n", yaml_scalar(r.summary_id)));
|
||||
fm.push_str("kind: summary\n");
|
||||
fm.push_str(&format!("tree_kind: {tree_kind_str}\n"));
|
||||
fm.push_str(&format!("tree_id: {}\n", yaml_scalar(r.tree_id)));
|
||||
fm.push_str(&format!("tree_scope: {}\n", yaml_scalar(r.tree_scope)));
|
||||
fm.push_str(&format!("level: {}\n", r.level));
|
||||
|
||||
// children: YAML list
|
||||
if r.child_ids.is_empty() {
|
||||
fm.push_str("children: []\n");
|
||||
} else {
|
||||
fm.push_str("children:\n");
|
||||
for id in r.child_ids {
|
||||
fm.push_str(&format!(" - {}\n", yaml_scalar(id)));
|
||||
}
|
||||
}
|
||||
fm.push_str(&format!("child_count: {}\n", r.child_count));
|
||||
fm.push_str(&format!("time_range_start: {trs}\n"));
|
||||
fm.push_str(&format!("time_range_end: {tre}\n"));
|
||||
fm.push_str(&format!("sealed_at: {sealed}\n"));
|
||||
|
||||
// aliases: human-readable title
|
||||
let alias = build_summary_alias(r);
|
||||
fm.push_str("aliases:\n");
|
||||
fm.push_str(&format!(" - {}\n", yaml_scalar(&alias)));
|
||||
|
||||
fm.push_str("tags: []\n");
|
||||
fm.push_str("---\n");
|
||||
fm
|
||||
}
|
||||
|
||||
/// Build a human-readable alias for the summary's `aliases:` front-matter field.
|
||||
fn build_summary_alias(r: &SummaryComposeInput<'_>) -> String {
|
||||
let date_range = format_date_range(r.time_range_start, r.time_range_end);
|
||||
match r.tree_kind {
|
||||
SummaryTreeKind::Source => {
|
||||
let scope_short = scope_short_label(r.tree_scope);
|
||||
format!(
|
||||
"L{} \u{00b7} {} \u{00b7} {} children \u{00b7} {}",
|
||||
r.level, scope_short, r.child_count, date_range
|
||||
)
|
||||
}
|
||||
SummaryTreeKind::Global => {
|
||||
format!(
|
||||
"L{} \u{00b7} global digest \u{00b7} {}",
|
||||
r.level, date_range
|
||||
)
|
||||
}
|
||||
SummaryTreeKind::Topic => {
|
||||
// Strip protocol prefix like "topic:" from scope for readability.
|
||||
let entity = r
|
||||
.tree_scope
|
||||
.split_once(':')
|
||||
.map(|(_, v)| v)
|
||||
.unwrap_or(r.tree_scope);
|
||||
format!(
|
||||
"L{} \u{00b7} topic {} \u{00b7} {} children",
|
||||
r.level, entity, r.child_count
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Format the date range as `"yyyy-mm-dd"` (if start == end date) or
|
||||
/// `"yyyy-mm-dd–yyyy-mm-dd"`.
|
||||
fn format_date_range(start: DateTime<Utc>, end: DateTime<Utc>) -> String {
|
||||
let s = start.format("%Y-%m-%d").to_string();
|
||||
let e = end.format("%Y-%m-%d").to_string();
|
||||
if s == e {
|
||||
s
|
||||
} else {
|
||||
format!("{s}\u{2013}{e}") // en dash
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a short human-readable label for the tree scope used in aliases.
|
||||
///
|
||||
/// For Gmail source scopes like `"gmail:alice@x.com|bob@y.com"`:
|
||||
/// - 2 participants → `"alice@x.com ↔ bob@y.com"`
|
||||
/// - N > 2 → `"alice@x.com + N-1 others"`
|
||||
/// - Otherwise → the raw scope (e.g. `"slack:#eng"`)
|
||||
fn scope_short_label(scope: &str) -> String {
|
||||
if let Some((prefix, participants)) = scope.split_once(':') {
|
||||
if prefix == "gmail" && !participants.is_empty() {
|
||||
let addrs: Vec<&str> = participants.split('|').collect();
|
||||
return match addrs.as_slice() {
|
||||
[] => scope.to_string(),
|
||||
[only] => only.to_string(),
|
||||
[first, second] => format!("{} \u{2194} {}", first, second), // ↔
|
||||
[first, rest @ ..] => format!("{} + {} others", first, rest.len()),
|
||||
};
|
||||
}
|
||||
}
|
||||
scope.to_string()
|
||||
}
|
||||
|
||||
/// Rewrite the `tags:` block in a summary file's front-matter, replacing it
|
||||
/// with the new tag list while leaving the body unchanged.
|
||||
///
|
||||
/// Reuses the generic [`rewrite_tags`] function — the front-matter structure
|
||||
/// is identical for both chunk and summary `.md` files.
|
||||
pub fn rewrite_summary_tags(file_bytes: &[u8], new_tags: &[String]) -> Result<Vec<u8>, String> {
|
||||
rewrite_tags(file_bytes, new_tags)
|
||||
}
|
||||
|
||||
/// Split a file into `(front_matter, body)` at the second `---` delimiter.
|
||||
///
|
||||
/// Returns `None` if the file does not have the expected `---\n...\n---\n` form.
|
||||
pub fn split_front_matter(content: &str) -> Option<(&str, &str)> {
|
||||
// The file must start with `---\n`.
|
||||
if !content.starts_with("---\n") {
|
||||
return None;
|
||||
}
|
||||
// Find the closing `---` line (must be `---` alone on a line after the first line).
|
||||
let rest = &content[4..]; // skip the opening `---\n`
|
||||
let close_idx = rest.find("\n---\n").or_else(|| {
|
||||
// Could be at the very end (no body).
|
||||
rest.strip_suffix("\n---").map(|r| r.len())
|
||||
})?;
|
||||
let fm_end = 4 + close_idx + 5; // include `\n---\n`
|
||||
Some((&content[..fm_end], &content[fm_end..]))
|
||||
}
|
||||
|
||||
/// Format a string as an unquoted YAML scalar when safe, or as a
|
||||
/// double-quoted string when it contains special characters.
|
||||
///
|
||||
/// We conservatively quote strings containing `:`, `#`, `[`, `]`, `{`, `}`,
|
||||
/// `"`, `'`, `\`, leading/trailing whitespace, or that start with special
|
||||
/// YAML indicator characters.
|
||||
fn yaml_scalar(s: &str) -> String {
|
||||
let needs_quoting = s.is_empty()
|
||||
|| s.trim() != s
|
||||
|| s.starts_with(|c: char| {
|
||||
matches!(
|
||||
c,
|
||||
'&' | '*' | '?' | '|' | '-' | '<' | '>' | '=' | '!' | '%' | '@' | '`'
|
||||
)
|
||||
})
|
||||
|| s.contains([':', '#', '[', ']', '{', '}', '"', '\'']);
|
||||
|
||||
if needs_quoting {
|
||||
let escaped = s.replace('\\', "\\\\").replace('"', "\\\"");
|
||||
format!("\"{escaped}\"")
|
||||
} else {
|
||||
s.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
|
||||
use crate::openhuman::memory::tree::types::{Metadata, SourceKind, SourceRef};
|
||||
use chrono::TimeZone;
|
||||
|
||||
fn sample_chunk() -> Chunk {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
Chunk {
|
||||
id: "abc123".into(),
|
||||
content: "## 2026-01-01T00:00:00Z — alice\nhello world".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice@example.com".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec!["person/Alice".into(), "org/Acme".into()],
|
||||
source_ref: Some(SourceRef::new("slack://m1".to_string())),
|
||||
},
|
||||
token_count: 10,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_produces_front_matter_and_body() {
|
||||
let chunk = sample_chunk();
|
||||
let (full, body) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
assert!(full_str.starts_with("---\n"), "must start with ---");
|
||||
assert!(full_str.contains("source_kind: chat"));
|
||||
assert!(full_str.contains("source_id: \"slack:#eng\""));
|
||||
assert!(full_str.contains("seq: 0"));
|
||||
assert!(full_str.contains("tags:"));
|
||||
assert!(full_str.contains(" - person/Alice"));
|
||||
assert!(full_str.ends_with("hello world"));
|
||||
assert_eq!(
|
||||
body,
|
||||
b"## 2026-01-01T00:00:00Z \xe2\x80\x94 alice\nhello world"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_front_matter_round_trips() {
|
||||
let chunk = sample_chunk();
|
||||
let (full, body) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
let (fm, b) = split_front_matter(full_str).expect("split must succeed");
|
||||
assert!(fm.starts_with("---\n"));
|
||||
assert!(fm.ends_with("---\n"));
|
||||
assert_eq!(b.as_bytes(), body.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrite_tags_preserves_body() {
|
||||
let chunk = sample_chunk();
|
||||
let (full, body) = compose_chunk_file(&chunk);
|
||||
let new_tags = vec!["person/Bob".into(), "project/Phoenix".into()];
|
||||
let rewritten = rewrite_tags(&full, &new_tags).unwrap();
|
||||
let rewritten_str = std::str::from_utf8(&rewritten).unwrap();
|
||||
assert!(rewritten_str.contains(" - person/Bob"));
|
||||
assert!(!rewritten_str.contains(" - person/Alice"));
|
||||
// Body must be unchanged.
|
||||
assert!(rewritten_str.ends_with(std::str::from_utf8(&body).unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrite_tags_empty_list() {
|
||||
let chunk = sample_chunk();
|
||||
let (full, _) = compose_chunk_file(&chunk);
|
||||
let rewritten = rewrite_tags(&full, &[]).unwrap();
|
||||
let s = std::str::from_utf8(&rewritten).unwrap();
|
||||
assert!(s.contains("tags: []"));
|
||||
assert!(!s.contains(" - person/"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn yaml_scalar_quotes_special_characters() {
|
||||
assert_eq!(yaml_scalar("slack:#eng"), "\"slack:#eng\"");
|
||||
assert_eq!(yaml_scalar("hello world"), "hello world");
|
||||
assert_eq!(yaml_scalar(""), "\"\"");
|
||||
}
|
||||
|
||||
fn sample_email_chunk() -> Chunk {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
Chunk {
|
||||
id: "emailchunk1".into(),
|
||||
content: "---\nFrom: alice@example.com\nSubject: Hello\n\nHello there.".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "gmail:alice@example.com|bob@example.com".into(),
|
||||
owner: "owner@example.com".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec!["gmail".into()],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 15,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_chunk_has_participants_list_and_alias() {
|
||||
let chunk = sample_email_chunk();
|
||||
let (full, _body) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
// participants block must be a YAML list
|
||||
assert!(
|
||||
full_str.contains("participants:"),
|
||||
"email chunk must have participants field; got:\n{full_str}"
|
||||
);
|
||||
assert!(
|
||||
full_str.contains(" - alice@example.com"),
|
||||
"alice must appear as list item; got:\n{full_str}"
|
||||
);
|
||||
assert!(
|
||||
full_str.contains(" - bob@example.com"),
|
||||
"bob must appear as list item; got:\n{full_str}"
|
||||
);
|
||||
// aliases block must be present
|
||||
assert!(
|
||||
full_str.contains("aliases:"),
|
||||
"email chunk must have aliases field; got:\n{full_str}"
|
||||
);
|
||||
assert!(
|
||||
full_str.contains("alice@example.com <-> bob@example.com: chunk 0"),
|
||||
"alias must encode participants; got:\n{full_str}"
|
||||
);
|
||||
// sender and thread_id must NOT appear
|
||||
assert!(
|
||||
!full_str.contains("sender:"),
|
||||
"email chunk must NOT have sender field; got:\n{full_str}"
|
||||
);
|
||||
assert!(
|
||||
!full_str.contains("thread_id:"),
|
||||
"email chunk must NOT have thread_id field; got:\n{full_str}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_chunk_many_participants_alias_summarises() {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: "em2".into(),
|
||||
content: "body".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "gmail:alice@x.com|bob@y.com|carol@z.com".into(),
|
||||
owner: "owner".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 1,
|
||||
seq_in_source: 3,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let (full, _) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
assert!(
|
||||
full_str.contains("participants:"),
|
||||
"three-party chunk needs participants list; got:\n{full_str}"
|
||||
);
|
||||
// With 3 participants: first + "2 others"
|
||||
assert!(
|
||||
full_str.contains("alice@x.com <-> 2 others: chunk 3"),
|
||||
"alias with 3 participants must summarise; got:\n{full_str}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_chunk_body_bytes_unchanged_by_extra_fields() {
|
||||
// Adding participants/aliases to front-matter must not affect body_bytes
|
||||
// (SHA-256 invariant: the hash is over body only, not front-matter).
|
||||
let chunk = sample_email_chunk();
|
||||
let (full, body) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
// Body must still appear at the end unmodified.
|
||||
assert!(
|
||||
full_str.ends_with(std::str::from_utf8(&body).unwrap()),
|
||||
"body bytes must appear unmodified after front-matter"
|
||||
);
|
||||
// body must equal chunk.content bytes
|
||||
assert_eq!(body, chunk.content.as_bytes());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_chunk_has_no_email_specific_fields() {
|
||||
let chunk = sample_chunk(); // source_kind = Chat
|
||||
let (full, _) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
assert!(
|
||||
!full_str.contains("aliases:"),
|
||||
"chat chunk must not have aliases field"
|
||||
);
|
||||
assert!(
|
||||
!full_str.contains("participants:"),
|
||||
"chat chunk must not have participants field"
|
||||
);
|
||||
assert!(
|
||||
!full_str.contains("sender:"),
|
||||
"chat chunk must not have sender field"
|
||||
);
|
||||
assert!(
|
||||
!full_str.contains("thread_id:"),
|
||||
"chat chunk must not have thread_id field"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_chunk_with_malformed_source_id_omits_extra_fields() {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: "xyz".into(),
|
||||
content: "body".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Email,
|
||||
source_id: "legacysourceid".into(), // no `gmail:` prefix → parse fails
|
||||
owner: "owner".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 1,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let (full, _) = compose_chunk_file(&chunk);
|
||||
let full_str = std::str::from_utf8(&full).unwrap();
|
||||
// Malformed source_id → no email extras, no panic.
|
||||
assert!(!full_str.contains("aliases:"));
|
||||
assert!(!full_str.contains("participants:"));
|
||||
assert!(!full_str.contains("sender:"));
|
||||
}
|
||||
|
||||
// ─── summary compose tests ────────────────────────────────────────────────
|
||||
|
||||
fn sample_summary_input(
|
||||
tree_kind: SummaryTreeKind,
|
||||
scope: &str,
|
||||
level: u32,
|
||||
) -> SummaryComposeInput<'static> {
|
||||
let ts_start = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let ts_end = chrono::Utc.timestamp_millis_opt(1_700_086_400_000).unwrap();
|
||||
let sealed = chrono::Utc.timestamp_millis_opt(1_700_090_000_000).unwrap();
|
||||
// Leak the strings so they have 'static lifetime for this test helper.
|
||||
// Only used in tests, not production code.
|
||||
let scope: &'static str = Box::leak(scope.to_string().into_boxed_str());
|
||||
SummaryComposeInput {
|
||||
summary_id: "summary:L1:abc",
|
||||
tree_kind,
|
||||
tree_id: "tree-id-001",
|
||||
tree_scope: scope,
|
||||
level,
|
||||
child_ids: Box::leak(
|
||||
vec!["child-1".to_string(), "child-2".to_string()].into_boxed_slice(),
|
||||
),
|
||||
child_count: 2,
|
||||
time_range_start: ts_start,
|
||||
time_range_end: ts_end,
|
||||
sealed_at: sealed,
|
||||
body: "This is the summariser output.\n",
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_source_summary_has_required_front_matter() {
|
||||
let input = sample_summary_input(SummaryTreeKind::Source, "gmail:alice@x.com|bob@y.com", 1);
|
||||
let composed = compose_summary_md(&input);
|
||||
let fm = &composed.front_matter;
|
||||
assert!(fm.starts_with("---\n"), "front-matter must start with ---");
|
||||
assert!(fm.ends_with("---\n"), "front-matter must end with ---\\n");
|
||||
assert!(fm.contains("kind: summary"), "must have kind: summary");
|
||||
assert!(
|
||||
fm.contains("tree_kind: source"),
|
||||
"must have tree_kind: source"
|
||||
);
|
||||
assert!(fm.contains("level: 1"), "must have level");
|
||||
assert!(fm.contains("child_count: 2"), "must have child_count");
|
||||
assert!(fm.contains(" - child-1"), "must list child ids");
|
||||
assert!(fm.contains(" - child-2"), "must list child ids");
|
||||
assert!(fm.contains("tags: []"), "must start with empty tags");
|
||||
// aliases must mention the scope
|
||||
assert!(fm.contains("aliases:"), "must have aliases");
|
||||
assert!(
|
||||
composed.body == "This is the summariser output.\n",
|
||||
"body must be the summariser text"
|
||||
);
|
||||
assert!(composed.full.ends_with("This is the summariser output.\n"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_global_summary_alias_format() {
|
||||
let input = sample_summary_input(SummaryTreeKind::Global, "global", 0);
|
||||
let composed = compose_summary_md(&input);
|
||||
assert!(
|
||||
composed.front_matter.contains("tree_kind: global"),
|
||||
"must have tree_kind: global"
|
||||
);
|
||||
assert!(
|
||||
composed.front_matter.contains("global digest"),
|
||||
"alias must mention 'global digest'"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_topic_summary_alias_format() {
|
||||
let input = sample_summary_input(SummaryTreeKind::Topic, "person:alex-johnson", 1);
|
||||
let composed = compose_summary_md(&input);
|
||||
assert!(
|
||||
composed.front_matter.contains("tree_kind: topic"),
|
||||
"must have tree_kind: topic"
|
||||
);
|
||||
assert!(
|
||||
composed.front_matter.contains("topic"),
|
||||
"alias must mention topic entity"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_summary_with_zero_children() {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let input = SummaryComposeInput {
|
||||
summary_id: "summary:L0:empty",
|
||||
tree_kind: SummaryTreeKind::Source,
|
||||
tree_id: "t1",
|
||||
tree_scope: "gmail:alice@x.com",
|
||||
level: 0,
|
||||
child_ids: &[],
|
||||
child_count: 0,
|
||||
time_range_start: ts,
|
||||
time_range_end: ts,
|
||||
sealed_at: ts,
|
||||
body: "empty",
|
||||
};
|
||||
let composed = compose_summary_md(&input);
|
||||
assert!(composed.front_matter.contains("children: []"));
|
||||
assert!(composed.front_matter.contains("child_count: 0"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_summary_same_start_end_date_single_date_alias() {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let input = SummaryComposeInput {
|
||||
summary_id: "summary:L1:sameday",
|
||||
tree_kind: SummaryTreeKind::Global,
|
||||
tree_id: "t1",
|
||||
tree_scope: "global",
|
||||
level: 1,
|
||||
child_ids: &["child-a".to_string()],
|
||||
child_count: 1,
|
||||
time_range_start: ts,
|
||||
time_range_end: ts, // same as start
|
||||
sealed_at: ts,
|
||||
body: "day recap",
|
||||
};
|
||||
let composed = compose_summary_md(&input);
|
||||
// Alias must contain just one date, not "date–date"
|
||||
let alias_line = composed
|
||||
.front_matter
|
||||
.lines()
|
||||
.find(|l| l.contains("L1") && l.contains("global digest"))
|
||||
.expect("alias line must be present");
|
||||
// The date should appear exactly once (no en-dash range)
|
||||
let date_str = ts.format("%Y-%m-%d").to_string();
|
||||
assert!(
|
||||
alias_line.contains(&date_str),
|
||||
"alias must contain the date; got: {alias_line}"
|
||||
);
|
||||
// Must not contain an en-dash (range indicator)
|
||||
assert!(
|
||||
!alias_line.contains('\u{2013}'),
|
||||
"same-day alias must not have en-dash range; got: {alias_line}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scope_short_label_two_participants() {
|
||||
let label = scope_short_label("gmail:alice@x.com|bob@y.com");
|
||||
assert_eq!(label, "alice@x.com \u{2194} bob@y.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scope_short_label_many_participants() {
|
||||
let label = scope_short_label("gmail:alice@x.com|bob@y.com|carol@z.com");
|
||||
assert_eq!(label, "alice@x.com + 2 others");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scope_short_label_non_gmail_returns_raw() {
|
||||
let label = scope_short_label("slack:#general");
|
||||
assert_eq!(label, "slack:#general");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrite_summary_tags_delegates_to_rewrite_tags() {
|
||||
// compose a summary, then rewrite its tags — body must stay unchanged.
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let input = SummaryComposeInput {
|
||||
summary_id: "sum:L1:rwttest",
|
||||
tree_kind: SummaryTreeKind::Source,
|
||||
tree_id: "t1",
|
||||
tree_scope: "gmail:alice@x.com",
|
||||
level: 1,
|
||||
child_ids: &["c1".to_string()],
|
||||
child_count: 1,
|
||||
time_range_start: ts,
|
||||
time_range_end: ts,
|
||||
sealed_at: ts,
|
||||
body: "summary body text",
|
||||
};
|
||||
let composed = compose_summary_md(&input);
|
||||
let file_bytes = composed.full.as_bytes();
|
||||
let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()];
|
||||
let rewritten = rewrite_summary_tags(file_bytes, &new_tags).unwrap();
|
||||
let rewritten_str = std::str::from_utf8(&rewritten).unwrap();
|
||||
assert!(rewritten_str.contains(" - person/Alice-Smith"));
|
||||
assert!(rewritten_str.contains(" - topic/Memory"));
|
||||
assert!(!rewritten_str.contains("tags: []"));
|
||||
// Body must be unchanged
|
||||
assert!(rewritten_str.ends_with("summary body text"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
//! Content store for memory-tree chunk and summary `.md` files (Phase MD-content).
|
||||
//!
|
||||
//! Bodies are stored on disk as `.md` files with YAML front-matter.
|
||||
//! SQLite holds `content_path` (relative, forward-slash) and `content_sha256`
|
||||
//! (over body bytes only) as pointers + integrity tokens.
|
||||
//!
|
||||
//! ## Module layout
|
||||
//!
|
||||
//! - [`paths`] — path generation + `slugify_source_id` + summary path builders
|
||||
//! - [`compose`] — YAML front-matter + body composition; tag rewriting
|
||||
//! - [`atomic`] — tempfile+fsync+rename writes; SHA-256; `stage_summary`
|
||||
//! - [`read`] — read + SHA-256 verification + `split_front_matter`; summary variants
|
||||
//! - [`tags`] — `update_chunk_tags` + `update_summary_tags` + slugifiers
|
||||
|
||||
pub mod atomic;
|
||||
pub mod compose;
|
||||
pub mod paths;
|
||||
pub mod read;
|
||||
pub mod tags;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use crate::openhuman::memory::tree::types::Chunk;
|
||||
|
||||
pub use atomic::StagedSummary;
|
||||
pub use compose::SummaryComposeInput;
|
||||
pub use paths::SummaryTreeKind;
|
||||
|
||||
/// A chunk that has been written to disk and is ready for SQLite upsert.
|
||||
///
|
||||
/// Callers build a `Vec<StagedChunk>` from `stage_chunks`, then pass it to
|
||||
/// `store::upsert_chunks_tx` in the same SQLite transaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StagedChunk {
|
||||
/// The original chunk (metadata + content).
|
||||
pub chunk: Chunk,
|
||||
/// Relative content path (forward-slash, e.g. `"chat/slack-eng/0.md"`).
|
||||
pub content_path: String,
|
||||
/// SHA-256 hex digest over the body bytes only.
|
||||
pub content_sha256: String,
|
||||
}
|
||||
|
||||
/// Update the `tags:` block in a summary's on-disk `.md` file after an
|
||||
/// extraction job runs.
|
||||
///
|
||||
/// Delegates to [`tags::update_summary_tags`].
|
||||
pub fn update_summary_tags(
|
||||
config: &crate::openhuman::config::Config,
|
||||
summary_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
tags::update_summary_tags(config, summary_id)
|
||||
}
|
||||
|
||||
/// Write all chunks in `chunks` to disk and return `StagedChunk` records
|
||||
/// ready for SQLite upsert.
|
||||
///
|
||||
/// Each chunk file is written atomically via a sibling temp-file + rename.
|
||||
/// Already-existing files are skipped (immutable-body contract). Parent
|
||||
/// directories are created on demand.
|
||||
///
|
||||
/// `content_root` — absolute path to the root of the content store.
|
||||
pub fn stage_chunks(content_root: &Path, chunks: &[Chunk]) -> anyhow::Result<Vec<StagedChunk>> {
|
||||
let mut staged = Vec::with_capacity(chunks.len());
|
||||
|
||||
for chunk in chunks {
|
||||
let source_kind = chunk.metadata.source_kind.as_str();
|
||||
let source_id = &chunk.metadata.source_id;
|
||||
|
||||
let rel_path = paths::chunk_rel_path(source_kind, source_id, &chunk.id);
|
||||
let abs_path = paths::chunk_abs_path(content_root, source_kind, source_id, &chunk.id);
|
||||
|
||||
let (full_bytes, body_bytes) = compose::compose_chunk_file(chunk);
|
||||
let sha256 = atomic::sha256_hex(&body_bytes);
|
||||
|
||||
match atomic::write_if_new(&abs_path, &full_bytes) {
|
||||
Ok(written) => {
|
||||
if written {
|
||||
log::debug!("[content_store] wrote chunk {} → {}", chunk.id, rel_path);
|
||||
} else {
|
||||
log::debug!(
|
||||
"[content_store] chunk {} already on disk at {}",
|
||||
chunk.id,
|
||||
rel_path
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[content_store] failed to write chunk {} to {}: {e}",
|
||||
chunk.id,
|
||||
rel_path
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
|
||||
staged.push(StagedChunk {
|
||||
chunk: chunk.clone(),
|
||||
content_path: rel_path,
|
||||
content_sha256: sha256,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(staged)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::types::{Metadata, SourceKind};
|
||||
use chrono::TimeZone;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn sample_chunk(seq: u32) -> Chunk {
|
||||
let ts = chrono::Utc
|
||||
.timestamp_millis_opt(1_700_000_000_000 + seq as i64)
|
||||
.unwrap();
|
||||
Chunk {
|
||||
id: format!("chunk_{seq}"),
|
||||
content: format!("## ts — alice\nMessage {seq}"),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 5,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_chunks_writes_files_and_returns_staged() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunks = vec![sample_chunk(0), sample_chunk(1)];
|
||||
let staged = stage_chunks(dir.path(), &chunks).unwrap();
|
||||
|
||||
assert_eq!(staged.len(), 2);
|
||||
for s in &staged {
|
||||
let abs = paths::chunk_abs_path(
|
||||
dir.path(),
|
||||
s.chunk.metadata.source_kind.as_str(),
|
||||
&s.chunk.metadata.source_id,
|
||||
&s.chunk.id,
|
||||
);
|
||||
assert!(abs.exists(), "file must exist: {}", abs.display());
|
||||
assert!(!s.content_path.is_empty());
|
||||
assert_eq!(s.content_sha256.len(), 64);
|
||||
// Path must be relative with forward slashes.
|
||||
assert!(!s.content_path.starts_with('/'));
|
||||
assert!(s.content_path.contains('/'));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_chunks_is_idempotent() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunks = vec![sample_chunk(0)];
|
||||
let first = stage_chunks(dir.path(), &chunks).unwrap();
|
||||
let second = stage_chunks(dir.path(), &chunks).unwrap();
|
||||
assert_eq!(first[0].content_sha256, second[0].content_sha256);
|
||||
assert_eq!(first[0].content_path, second[0].content_path);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,473 @@
|
||||
//! Content-file path generation.
|
||||
//!
|
||||
//! Each chunk body is stored as a `.md` file under `<content_root>/`. The path
|
||||
//! structure depends on the source kind:
|
||||
//!
|
||||
//! ```text
|
||||
//! Email: <content_root>/email/<participants_slug>/<chunk_id>.md
|
||||
//! Chat: <content_root>/chat/<source_slug>/<chunk_id>.md
|
||||
//! Document: <content_root>/document/<source_slug>/<chunk_id>.md
|
||||
//! ```
|
||||
//!
|
||||
//! Email paths parse `source_id` as `gmail:{participants}` where `participants`
|
||||
//! is `addr1|addr2|...` (sorted, deduped, lowercased bare emails). The
|
||||
//! participants string is slugified as a whole (pipe and `@` both become `-`)
|
||||
//! to produce a single directory level, giving one folder per unique
|
||||
//! conversation set.
|
||||
//!
|
||||
//! Paths are stored in SQLite as **relative** strings with forward slashes so
|
||||
//! they remain valid regardless of where the workspace is mounted.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::openhuman::memory::tree::util::redact::redact;
|
||||
|
||||
/// Which kind of summary tree a summary belongs to. Determines the top-level
|
||||
/// directory under `<content_root>/summaries/`.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum SummaryTreeKind {
|
||||
/// Per-source-tree summary. Layout: `summaries/source/<scope_slug>/L<level>/<id>.md`
|
||||
Source,
|
||||
/// Global digest tree. Layout: `summaries/global/<yyyy-mm-dd>/L<level>/<id>.md`
|
||||
Global,
|
||||
/// Per-topic (entity) tree. Layout: `summaries/topic/<scope_slug>/L<level>/<id>.md`
|
||||
Topic,
|
||||
}
|
||||
|
||||
/// Build the relative content path for a summary, using forward slashes.
|
||||
///
|
||||
/// Path layout depends on tree_kind:
|
||||
/// - Source: `"summaries/source/<scope_slug>/L<level>/<summary_filename>.md"`
|
||||
/// - Global: `"summaries/global/<yyyy-mm-dd>/L<level>/<summary_filename>.md"`
|
||||
/// Falls back to `unknown-date` (with a warn log) if `date_for_global` is
|
||||
/// `None` — preferable to panicking inside a path utility.
|
||||
/// - Topic: `"summaries/topic/<scope_slug>/L<level>/<summary_filename>.md"`
|
||||
///
|
||||
/// `scope_slug` must already be slugified by the caller (use [`slugify_source_id`] or
|
||||
/// a per-kind variant). A trailing `.md` on `summary_id` is stripped if present.
|
||||
///
|
||||
/// The `summary_id` is sanitized into a filesystem-safe filename by replacing
|
||||
/// characters illegal on Windows (`:`, `\`, `*`, `?`, `"`, `<`, `>`, `|`) with `-`.
|
||||
pub fn summary_rel_path(
|
||||
tree_kind: SummaryTreeKind,
|
||||
scope_slug: &str,
|
||||
level: u32,
|
||||
summary_id: &str,
|
||||
date_for_global: Option<DateTime<Utc>>,
|
||||
) -> String {
|
||||
// Strip a trailing `.md` from summary_id if accidentally included.
|
||||
let id = summary_id.strip_suffix(".md").unwrap_or(summary_id);
|
||||
// Sanitize to a cross-platform filename (colons are illegal on Windows NTFS).
|
||||
let filename = sanitize_filename(id);
|
||||
|
||||
match tree_kind {
|
||||
SummaryTreeKind::Source => {
|
||||
format!("summaries/source/{}/L{}/{}.md", scope_slug, level, filename)
|
||||
}
|
||||
SummaryTreeKind::Global => {
|
||||
// Fall back to a sentinel date rather than panic — a path-utility
|
||||
// panic would propagate up through seal/digest/janitor codepaths
|
||||
// and abort otherwise-recoverable work. Callers should always
|
||||
// pass a date for Global; the warn log surfaces the contract
|
||||
// violation without taking the process down.
|
||||
let date_str = match date_for_global {
|
||||
Some(d) => d.format("%Y-%m-%d").to_string(),
|
||||
None => {
|
||||
log::warn!(
|
||||
"[content_store::paths] summary_rel_path called for Global \
|
||||
without date_for_global; using sentinel 'unknown-date'. \
|
||||
Caller bug — please pass a date."
|
||||
);
|
||||
"unknown-date".to_string()
|
||||
}
|
||||
};
|
||||
format!("summaries/global/{}/L{}/{}.md", date_str, level, filename)
|
||||
}
|
||||
SummaryTreeKind::Topic => {
|
||||
format!("summaries/topic/{}/L{}/{}.md", scope_slug, level, filename)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace characters that are illegal in filenames on Windows NTFS with `-`.
|
||||
///
|
||||
/// Illegal characters: `\`, `/`, `:`, `*`, `?`, `"`, `<`, `>`, `|`.
|
||||
/// (Forward slash is not replaced since `summary_id` should not contain path
|
||||
/// separators, but we sanitize it anyway for safety.)
|
||||
fn sanitize_filename(s: &str) -> String {
|
||||
s.chars()
|
||||
.map(|c| match c {
|
||||
'\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '-',
|
||||
c => c,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Build the absolute on-disk path for a summary given the content root.
|
||||
pub fn summary_abs_path(
|
||||
content_root: &Path,
|
||||
tree_kind: SummaryTreeKind,
|
||||
scope_slug: &str,
|
||||
level: u32,
|
||||
summary_id: &str,
|
||||
date_for_global: Option<DateTime<Utc>>,
|
||||
) -> PathBuf {
|
||||
let rel = summary_rel_path(tree_kind, scope_slug, level, summary_id, date_for_global);
|
||||
let mut abs = content_root.to_path_buf();
|
||||
for component in rel.split('/') {
|
||||
abs.push(component);
|
||||
}
|
||||
abs
|
||||
}
|
||||
|
||||
/// Build the relative content path for a chunk, using forward slashes.
|
||||
///
|
||||
/// Path layout depends on source_kind:
|
||||
/// - Email: `"email/<participants_slug>/<chunk_id>.md"`
|
||||
/// Parses `source_id` as `gmail:{participants}` (two colon-separated parts)
|
||||
/// where `participants` is `addr1|addr2|...` (sorted, deduped, lowercased).
|
||||
/// The entire participants string is slugified as a single unit to produce
|
||||
/// one folder level per conversation set (no nested thread subfolder).
|
||||
/// If the source_id lacks a `gmail:` prefix or has no participants segment,
|
||||
/// falls through to the chat/document layout using `slugify_source_id(source_id)`.
|
||||
/// - Chat: `"chat/<source_slug>/<chunk_id>.md"`
|
||||
/// - Document: `"document/<source_slug>/<chunk_id>.md"`
|
||||
///
|
||||
/// `chunk_id` — the deterministic content hash produced by `types::chunk_id`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```text
|
||||
/// chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc")
|
||||
/// → "email/alice-x-com-bob-y-com/abc.md"
|
||||
///
|
||||
/// chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def")
|
||||
/// → "email/notifications-github-com-sanil-x-com/def.md"
|
||||
///
|
||||
/// chunk_rel_path("email", "legacyid", "xyz")
|
||||
/// → "email/legacyid/xyz.md" (malformed — flat fallback)
|
||||
/// ```
|
||||
pub fn chunk_rel_path(source_kind: &str, source_id: &str, chunk_id: &str) -> String {
|
||||
// Sanitize chunk_id into a cross-platform filename. Chunk IDs contain
|
||||
// colons (e.g. `chat:slack:#eng:0`) which are illegal on Windows NTFS;
|
||||
// replace illegal characters with `-` to match summary_rel_path behaviour.
|
||||
let filename = sanitize_filename(chunk_id);
|
||||
match source_kind {
|
||||
"email" => {
|
||||
// Expected format: "gmail:{participants}"
|
||||
// Split on ':' — exactly 2 parts required; part[0] == "gmail".
|
||||
let parts: Vec<&str> = source_id.splitn(2, ':').collect();
|
||||
if parts.len() == 2 && parts[0] == "gmail" && !parts[1].is_empty() {
|
||||
let participants_slug = slugify_source_id(parts[1]);
|
||||
format!("email/{}/{}.md", participants_slug, filename)
|
||||
} else {
|
||||
// Malformed / legacy source_id — fall back to flat layout.
|
||||
// Redact the source_id before logging since it may embed email
|
||||
// addresses.
|
||||
log::debug!(
|
||||
"[content_store::paths] email source_id has unexpected format, falling back to flat layout: source_id_hash={}",
|
||||
redact(source_id)
|
||||
);
|
||||
let slug = slugify_source_id(source_id);
|
||||
format!("email/{}/{}.md", slug, filename)
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// Chat, Document, and any future kinds use a 3-level layout.
|
||||
let slug = slugify_source_id(source_id);
|
||||
format!("{}/{}/{}.md", source_kind, slug, filename)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the absolute on-disk path for a chunk given the content root.
|
||||
pub fn chunk_abs_path(
|
||||
content_root: &Path,
|
||||
source_kind: &str,
|
||||
source_id: &str,
|
||||
chunk_id: &str,
|
||||
) -> PathBuf {
|
||||
let rel = chunk_rel_path(source_kind, source_id, chunk_id);
|
||||
// Convert forward-slash relative path to OS-native path.
|
||||
let mut abs = content_root.to_path_buf();
|
||||
for component in rel.split('/') {
|
||||
abs.push(component);
|
||||
}
|
||||
abs
|
||||
}
|
||||
|
||||
/// Convert a raw `source_id` (e.g. `"slack:#general"`, `"gmail:thread/abc"`)
|
||||
/// into a filesystem-safe slug using only `[a-z0-9_-]` characters.
|
||||
///
|
||||
/// Rules:
|
||||
/// - lowercase the whole string
|
||||
/// - replace any character outside `[a-z0-9_-]` with `-`
|
||||
/// - collapse consecutive `-` to one
|
||||
/// - trim leading/trailing `-`
|
||||
/// - `_` is preserved anywhere in the string (interior underscores are kept)
|
||||
/// - truncate to 120 characters
|
||||
pub fn slugify_source_id(source_id: &str) -> String {
|
||||
let lower = source_id.to_lowercase();
|
||||
let mut out = String::with_capacity(lower.len().min(120));
|
||||
let mut last_dash = true; // avoids leading dash; also suppresses leading underscore runs
|
||||
let mut pending_underscore = false; // deferred `_` to avoid leading underscore
|
||||
|
||||
for ch in lower.chars() {
|
||||
if ch == '_' {
|
||||
// Defer underscores — emit only if we have already emitted a
|
||||
// non-separator character (so `_solo_` becomes `_solo_` once the
|
||||
// `s` is emitted, but a leading `_` is dropped).
|
||||
if !last_dash {
|
||||
// We have real content before this, so emit the underscore now.
|
||||
pending_underscore = true;
|
||||
}
|
||||
// If last_dash is true (nothing emitted yet), silently skip.
|
||||
} else if ch.is_ascii_alphanumeric() {
|
||||
if pending_underscore {
|
||||
out.push('_');
|
||||
pending_underscore = false;
|
||||
}
|
||||
out.push(ch);
|
||||
last_dash = false;
|
||||
} else {
|
||||
// Non-alphanumeric, non-underscore → convert to `-`.
|
||||
pending_underscore = false; // drop any pending underscore before a dash
|
||||
if !last_dash {
|
||||
out.push('-');
|
||||
last_dash = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
// trailing underscore: drop it (trim trailing separators).
|
||||
// trim trailing dash
|
||||
let trimmed = out.trim_end_matches('-');
|
||||
// also trim any trailing underscore
|
||||
let trimmed = trimmed.trim_end_matches('_');
|
||||
let truncated = truncate_at_char(trimmed, 120);
|
||||
if truncated.is_empty() {
|
||||
"unknown".to_string()
|
||||
} else {
|
||||
truncated.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate `s` to at most `max_chars` Unicode code points.
|
||||
fn truncate_at_char(s: &str, max_chars: usize) -> &str {
|
||||
match s.char_indices().nth(max_chars) {
|
||||
Some((idx, _)) => &s[..idx],
|
||||
None => s,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ─── slugify tests ────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn slugify_slack_channel() {
|
||||
assert_eq!(slugify_source_id("slack:#general"), "slack-general");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_gmail_thread() {
|
||||
assert_eq!(
|
||||
slugify_source_id("gmail:thread/abc-123"),
|
||||
"gmail-thread-abc-123"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_collapses_consecutive_separators() {
|
||||
assert_eq!(slugify_source_id("foo::bar"), "foo-bar");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_uppercase_lowercased() {
|
||||
assert_eq!(slugify_source_id("Slack:ABC"), "slack-abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_empty_falls_back_to_unknown() {
|
||||
assert_eq!(slugify_source_id(""), "unknown");
|
||||
assert_eq!(slugify_source_id(":::"), "unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_truncates_at_120_chars() {
|
||||
let long = "a".repeat(200);
|
||||
let slug = slugify_source_id(&long);
|
||||
assert_eq!(slug.len(), 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_preserves_interior_underscore() {
|
||||
// `_solo_` has a leading and trailing underscore; only the interior
|
||||
// `solo` + the part after should survive. When used as a thread key
|
||||
// it arrives as the whole string `_solo_`.
|
||||
// Leading `_` is stripped (it's treated like a leading dash),
|
||||
// trailing `_` is stripped; interior `_` is preserved when sandwiched
|
||||
// between alphanumeric characters.
|
||||
let s = slugify_source_id("_solo_");
|
||||
// "solo" — both outer underscores trimmed, interior underscore has
|
||||
// nothing on the right so it's also trailing and trimmed.
|
||||
assert_eq!(s, "solo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_preserves_interior_underscore_between_chars() {
|
||||
// `foo_bar` — interior underscore stays.
|
||||
assert_eq!(slugify_source_id("foo_bar"), "foo_bar");
|
||||
}
|
||||
|
||||
// ─── chunk_rel_path tests ─────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn email_one_to_one_conversation_path() {
|
||||
// 1:1 conversation between alice and bob.
|
||||
let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc");
|
||||
assert_eq!(p, "email/alice-x-com-bob-y-com/abc.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_group_conversation_path() {
|
||||
// Group conversation with three participants.
|
||||
let p = chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def");
|
||||
assert_eq!(p, "email/notifications-github-com-sanil-x-com/def.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_solo_no_to_path() {
|
||||
// Solo sender (no To), participants = single address.
|
||||
let p = chunk_rel_path("email", "gmail:alice@x.com", "solo123");
|
||||
assert_eq!(p, "email/alice-x-com/solo123.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_malformed_source_id_falls_back_to_flat_layout() {
|
||||
// Malformed: no `gmail:` prefix → flat fallback.
|
||||
let p = chunk_rel_path("email", "legacyid", "xyz");
|
||||
// Falls back to email/<slug>/<chunk_id>.md
|
||||
assert!(p.starts_with("email/"), "must remain under email/");
|
||||
assert!(p.ends_with("/xyz.md"), "chunk_id must be the filename");
|
||||
// Must not panic.
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_three_participant_path() {
|
||||
// Three participants: alice, bob, carol (pipe-separated, sorted).
|
||||
let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com|carol@z.com", "g42");
|
||||
assert_eq!(p, "email/alice-x-com-bob-y-com-carol-z-com/g42.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_path() {
|
||||
let p = chunk_rel_path("chat", "slack:#eng", "xyz789");
|
||||
assert_eq!(p, "chat/slack-eng/xyz789.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_path() {
|
||||
let p = chunk_rel_path("document", "doc:notes.md", "uvw");
|
||||
assert_eq!(p, "document/doc-notes-md/uvw.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_abs_path_uses_os_separator() {
|
||||
use std::path::Path;
|
||||
let root = Path::new("/workspace/content");
|
||||
let abs = chunk_abs_path(root, "email", "gmail:alice@x.com|bob@y.com", "abc");
|
||||
assert!(abs.starts_with(root));
|
||||
assert!(abs.ends_with("abc.md"));
|
||||
}
|
||||
|
||||
// ─── summary_rel_path tests ───────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn summary_rel_path_source() {
|
||||
let p = summary_rel_path(
|
||||
SummaryTreeKind::Source,
|
||||
"gmail-alice-x-com-bob-y-com",
|
||||
1,
|
||||
"summary:L1:abc",
|
||||
None,
|
||||
);
|
||||
// Colons in summary_id are replaced with '-' for cross-platform filenames.
|
||||
assert_eq!(
|
||||
p,
|
||||
"summaries/source/gmail-alice-x-com-bob-y-com/L1/summary-L1-abc.md"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_rel_path_global() {
|
||||
use chrono::TimeZone;
|
||||
let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap();
|
||||
let p = summary_rel_path(
|
||||
SummaryTreeKind::Global,
|
||||
"global",
|
||||
0,
|
||||
"summary:L0:daily",
|
||||
Some(date),
|
||||
);
|
||||
assert_eq!(p, "summaries/global/2026-04-28/L0/summary-L0-daily.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_rel_path_topic() {
|
||||
let p = summary_rel_path(
|
||||
SummaryTreeKind::Topic,
|
||||
"person-alex-johnson",
|
||||
1,
|
||||
"summary:L1:xyz",
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
p,
|
||||
"summaries/topic/person-alex-johnson/L1/summary-L1-xyz.md"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_rel_path_strips_trailing_md_extension() {
|
||||
// If the caller accidentally appends .md to the summary_id, strip it.
|
||||
let p = summary_rel_path(
|
||||
SummaryTreeKind::Topic,
|
||||
"entity-slug",
|
||||
2,
|
||||
"summary:L2:foo.md",
|
||||
None,
|
||||
);
|
||||
assert_eq!(p, "summaries/topic/entity-slug/L2/summary-L2-foo.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_rel_path_global_falls_back_to_sentinel_without_date() {
|
||||
// Caller bug to omit date for Global, but a path utility shouldn't
|
||||
// panic — fall back to a sentinel `unknown-date` segment so the
|
||||
// file lands somewhere predictable rather than aborting the seal.
|
||||
let p = summary_rel_path(SummaryTreeKind::Global, "global", 0, "summary:L0:x", None);
|
||||
assert_eq!(p, "summaries/global/unknown-date/L0/summary-L0-x.md");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_abs_path_rooted_under_content_root() {
|
||||
use chrono::TimeZone;
|
||||
use std::path::Path;
|
||||
let root = Path::new("/workspace/content");
|
||||
let date = chrono::Utc.with_ymd_and_hms(2026, 1, 15, 0, 0, 0).unwrap();
|
||||
let abs = summary_abs_path(
|
||||
root,
|
||||
SummaryTreeKind::Global,
|
||||
"global",
|
||||
0,
|
||||
"daily-123",
|
||||
Some(date),
|
||||
);
|
||||
assert!(abs.starts_with(root));
|
||||
assert!(abs.ends_with("daily-123.md"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
//! Read and verify chunk and summary `.md` files from the content store.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::atomic::sha256_hex;
|
||||
use super::compose::split_front_matter;
|
||||
use crate::openhuman::memory::tree::util::redact::redact;
|
||||
|
||||
/// The result of reading a chunk file from disk.
|
||||
pub struct ChunkFileContents {
|
||||
/// The Markdown body (everything after the closing `---` of the front-matter).
|
||||
pub body: String,
|
||||
/// SHA-256 hex digest over the **body bytes** only.
|
||||
pub sha256: String,
|
||||
}
|
||||
|
||||
/// Read a chunk file and return its body + SHA-256.
|
||||
///
|
||||
/// Returns an error if:
|
||||
/// - the file does not exist
|
||||
/// - the file is not valid UTF-8
|
||||
/// - the front-matter delimiters cannot be found
|
||||
pub fn read_chunk_file(abs_path: &Path) -> anyhow::Result<ChunkFileContents> {
|
||||
let raw = std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?;
|
||||
let content = std::str::from_utf8(&raw)
|
||||
.map_err(|e| anyhow::anyhow!("invalid UTF-8 in {:?}: {e}", abs_path))?;
|
||||
|
||||
let (_fm, body) = split_front_matter(content)
|
||||
.ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", abs_path))?;
|
||||
|
||||
let sha256 = sha256_hex(body.as_bytes());
|
||||
Ok(ChunkFileContents {
|
||||
body: body.to_string(),
|
||||
sha256,
|
||||
})
|
||||
}
|
||||
|
||||
/// Verify that the body of a chunk file matches the expected SHA-256.
|
||||
///
|
||||
/// Returns `Ok(true)` on a match, `Ok(false)` on a mismatch, and an `Err`
|
||||
/// if the file cannot be read or parsed.
|
||||
pub fn verify_chunk_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result<bool> {
|
||||
let contents = read_chunk_file(abs_path)?;
|
||||
let ok = contents.sha256 == expected_sha256;
|
||||
if !ok {
|
||||
// Log the path as a redacted hash — the path may embed email addresses
|
||||
// (participant slugs) after the participant-bucketing change.
|
||||
let path_str = abs_path.to_string_lossy();
|
||||
log::warn!(
|
||||
"[content_store::read] sha256 mismatch for path_hash={}: expected={} actual={}",
|
||||
redact(&path_str),
|
||||
expected_sha256,
|
||||
contents.sha256,
|
||||
);
|
||||
}
|
||||
Ok(ok)
|
||||
}
|
||||
|
||||
// ── Summary reads ────────────────────────────────────────────────────────────
|
||||
|
||||
/// The result of verifying a summary file on disk.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum VerifyResult {
|
||||
/// The on-disk body SHA-256 matches the stored value.
|
||||
Ok,
|
||||
/// The file exists but the body SHA-256 does not match.
|
||||
Mismatch { actual: String },
|
||||
/// The file does not exist at the given path.
|
||||
Missing,
|
||||
}
|
||||
|
||||
/// Read a summary file and return its body + SHA-256.
|
||||
///
|
||||
/// Returns an error if:
|
||||
/// - the file does not exist
|
||||
/// - the file is not valid UTF-8
|
||||
/// - the front-matter delimiters cannot be found
|
||||
pub fn read_summary_file(abs_path: &Path) -> anyhow::Result<ChunkFileContents> {
|
||||
// Reuse the same reader as chunks — the file format is identical.
|
||||
read_chunk_file(abs_path)
|
||||
}
|
||||
|
||||
/// Verify a summary file's body SHA-256 without returning the body itself.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `VerifyResult::Ok` on match
|
||||
/// - `VerifyResult::Mismatch { actual }` on hash mismatch
|
||||
/// - `VerifyResult::Missing` when the file does not exist
|
||||
pub fn verify_summary_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result<VerifyResult> {
|
||||
if !abs_path.exists() {
|
||||
return Ok(VerifyResult::Missing);
|
||||
}
|
||||
let contents = read_summary_file(abs_path)?;
|
||||
if contents.sha256 == expected_sha256 {
|
||||
Ok(VerifyResult::Ok)
|
||||
} else {
|
||||
// Redact the path — it can embed participant slugs (email addresses).
|
||||
let path_str = abs_path.to_string_lossy();
|
||||
log::warn!(
|
||||
"[content_store::read] sha256 mismatch for summary path_hash={}: expected={} actual={}",
|
||||
redact(&path_str),
|
||||
expected_sha256,
|
||||
contents.sha256,
|
||||
);
|
||||
Ok(VerifyResult::Mismatch {
|
||||
actual: contents.sha256,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ── High-level body readers (Config-aware) ───────────────────────────────────
|
||||
//
|
||||
// These helpers resolve the on-disk path from SQLite via
|
||||
// `get_chunk_content_pointers` / `get_summary_content_pointers`, then read the
|
||||
// file body. They are the single authoritative entry-point for every caller
|
||||
// that needs the **full** chunk or summary body (LLM extractor, summariser
|
||||
// inputs, retrieval API, embedder). Preview-only consumers (UI cards, fast
|
||||
// filter scans) continue reading the `content` column directly from SQLite.
|
||||
//
|
||||
// Error policy:
|
||||
// - If `content_path` / `content_sha256` are NULL (legacy rows ingested before
|
||||
// the MD-on-disk migration), return `Err` — callers must handle the
|
||||
// "pre-migration chunk" case explicitly. The job pipeline propagates the
|
||||
// error and retries; retrieval falls back gracefully.
|
||||
// - File-not-found or SHA mismatch → `Err` (propagated to caller for retry /
|
||||
// alerting).
|
||||
|
||||
/// Read the full body of a chunk `.md` file by its chunk id.
|
||||
///
|
||||
/// Looks up `content_path` in SQLite, resolves it to an absolute path under
|
||||
/// `config.memory_tree_content_root()`, reads the file, and returns the body
|
||||
/// string (everything after the YAML front-matter delimiter).
|
||||
///
|
||||
/// Returns `Err` if:
|
||||
/// - The chunk row has no `content_path` recorded (pre-MD-migration row).
|
||||
/// - The file cannot be read or has no valid front-matter.
|
||||
///
|
||||
/// # Preview vs. full body
|
||||
/// The `content` column in `mem_tree_chunks` holds a ≤500-char preview after
|
||||
/// the MD-on-disk migration. Use this function wherever the full body is
|
||||
/// required (LLM extraction, embedding, summariser inputs, retrieval API).
|
||||
pub fn read_chunk_body(
|
||||
config: &crate::openhuman::config::Config,
|
||||
chunk_id: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
use crate::openhuman::memory::tree::store::get_chunk_content_pointers;
|
||||
|
||||
let pointers = get_chunk_content_pointers(config, chunk_id)?.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"[content_store::read] no content_path for chunk_id={} (pre-MD-migration row?)",
|
||||
chunk_id
|
||||
)
|
||||
})?;
|
||||
let (rel_path, expected_sha256) = pointers;
|
||||
|
||||
let content_root = config.memory_tree_content_root();
|
||||
// Reconstruct the absolute path from the stored relative forward-slash path.
|
||||
let abs_path = {
|
||||
let mut p = content_root.clone();
|
||||
for component in rel_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
log::debug!(
|
||||
"[content_store::read] read_chunk_body chunk_id={} path_hash={}",
|
||||
chunk_id,
|
||||
redact(&rel_path),
|
||||
);
|
||||
|
||||
let result = read_chunk_file(&abs_path).with_context(|| {
|
||||
format!(
|
||||
"read_chunk_body: failed to read file for chunk_id={} path_hash={}",
|
||||
chunk_id,
|
||||
redact(&rel_path),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Verify the on-disk body matches the SHA stored at write time. A mismatch
|
||||
// means the file was tampered with, the tx that committed the pointer
|
||||
// raced with a separate writer, or the disk corrupted — all unsafe to
|
||||
// hand back to a consumer. Fail loudly rather than serve stale/corrupt
|
||||
// bytes into the LLM extractor / summariser pipeline.
|
||||
if result.sha256 != expected_sha256 {
|
||||
return Err(anyhow::anyhow!(
|
||||
"[content_store::read] sha256 mismatch for chunk_id={} \
|
||||
expected={} actual={} path_hash={}",
|
||||
chunk_id,
|
||||
expected_sha256,
|
||||
result.sha256,
|
||||
redact(&rel_path),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(result.body)
|
||||
}
|
||||
|
||||
use anyhow::Context as _;
|
||||
|
||||
/// Read the full body of a summary `.md` file by its summary id.
|
||||
///
|
||||
/// Looks up `content_path` in SQLite, resolves it to an absolute path under
|
||||
/// `config.memory_tree_content_root()`, reads the file, and returns the body
|
||||
/// string.
|
||||
///
|
||||
/// Returns `Err` if:
|
||||
/// - The summary row has no `content_path` recorded (pre-MD-migration row).
|
||||
/// - The file cannot be read or has no valid front-matter.
|
||||
///
|
||||
/// # Preview vs. full body
|
||||
/// The `content` column in `mem_tree_summaries` holds a ≤500-char preview after
|
||||
/// the MD-on-disk migration. Use this function wherever the full body is
|
||||
/// required (LLM extraction, embedding, summariser inputs, retrieval API).
|
||||
pub fn read_summary_body(
|
||||
config: &crate::openhuman::config::Config,
|
||||
summary_id: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
use crate::openhuman::memory::tree::store::get_summary_content_pointers;
|
||||
|
||||
let pointers = get_summary_content_pointers(config, summary_id)?.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"[content_store::read] no content_path for summary_id={} (pre-MD-migration row?)",
|
||||
summary_id
|
||||
)
|
||||
})?;
|
||||
let (rel_path, expected_sha256) = pointers;
|
||||
|
||||
let content_root = config.memory_tree_content_root();
|
||||
let abs_path = {
|
||||
let mut p = content_root.clone();
|
||||
for component in rel_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
log::debug!(
|
||||
"[content_store::read] read_summary_body summary_id={} path_hash={}",
|
||||
summary_id,
|
||||
redact(&rel_path),
|
||||
);
|
||||
|
||||
let result = read_summary_file(&abs_path).with_context(|| {
|
||||
format!(
|
||||
"read_summary_body: failed to read file for summary_id={} path_hash={}",
|
||||
summary_id,
|
||||
redact(&rel_path),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Verify the on-disk body matches the SHA stored at seal time. See the
|
||||
// matching guard in `read_chunk_body` for rationale.
|
||||
if result.sha256 != expected_sha256 {
|
||||
return Err(anyhow::anyhow!(
|
||||
"[content_store::read] sha256 mismatch for summary_id={} \
|
||||
expected={} actual={} path_hash={}",
|
||||
summary_id,
|
||||
expected_sha256,
|
||||
result.sha256,
|
||||
redact(&rel_path),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(result.body)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
|
||||
use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file;
|
||||
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind};
|
||||
use chrono::TimeZone;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn sample_chunk() -> Chunk {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
Chunk {
|
||||
id: "read_test".into(),
|
||||
content: "## ts — alice\nhello from read test".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 8,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_returns_body_and_correct_sha256() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunk = sample_chunk();
|
||||
let (full_bytes, body_bytes) = compose_chunk_file(&chunk);
|
||||
let path = dir.path().join("0.md");
|
||||
write_if_new(&path, &full_bytes).unwrap();
|
||||
|
||||
let result = read_chunk_file(&path).unwrap();
|
||||
assert_eq!(result.body, std::str::from_utf8(&body_bytes).unwrap());
|
||||
assert_eq!(result.sha256, sha256_hex(&body_bytes));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_passes_for_correct_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunk = sample_chunk();
|
||||
let (full_bytes, body_bytes) = compose_chunk_file(&chunk);
|
||||
let path = dir.path().join("0.md");
|
||||
write_if_new(&path, &full_bytes).unwrap();
|
||||
|
||||
let expected = sha256_hex(&body_bytes);
|
||||
assert!(verify_chunk_file(&path, &expected).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_fails_for_wrong_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunk = sample_chunk();
|
||||
let (full_bytes, _) = compose_chunk_file(&chunk);
|
||||
let path = dir.path().join("0.md");
|
||||
write_if_new(&path, &full_bytes).unwrap();
|
||||
|
||||
assert!(!verify_chunk_file(&path, "deadbeef").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_missing_file_returns_error() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("nonexistent.md");
|
||||
assert!(read_chunk_file(&path).is_err());
|
||||
}
|
||||
|
||||
// ─── summary read / verify tests ─────────────────────────────────────────
|
||||
|
||||
fn write_summary_file(dir: &TempDir, body: &str) -> (std::path::PathBuf, String) {
|
||||
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
|
||||
use crate::openhuman::memory::tree::content_store::compose::{
|
||||
compose_summary_md, SummaryComposeInput,
|
||||
};
|
||||
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
|
||||
use chrono::TimeZone;
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let input = SummaryComposeInput {
|
||||
summary_id: "sum:L1:readtest",
|
||||
tree_kind: SummaryTreeKind::Source,
|
||||
tree_id: "t1",
|
||||
tree_scope: "gmail:alice@x.com",
|
||||
level: 1,
|
||||
child_ids: &["c1".to_string()],
|
||||
child_count: 1,
|
||||
time_range_start: ts,
|
||||
time_range_end: ts,
|
||||
sealed_at: ts,
|
||||
body,
|
||||
};
|
||||
let composed = compose_summary_md(&input);
|
||||
let path = dir.path().join("sum.md");
|
||||
let sha = sha256_hex(composed.body.as_bytes());
|
||||
write_if_new(&path, composed.full.as_bytes()).unwrap();
|
||||
(path, sha)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_summary_file_returns_body_and_sha() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let body = "summary body content\n";
|
||||
let (path, expected_sha) = write_summary_file(&dir, body);
|
||||
let result = read_summary_file(&path).unwrap();
|
||||
assert_eq!(result.body, body);
|
||||
assert_eq!(result.sha256, expected_sha);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_summary_file_ok_for_correct_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (path, sha) = write_summary_file(&dir, "body text\n");
|
||||
assert_eq!(verify_summary_file(&path, &sha).unwrap(), VerifyResult::Ok);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_summary_file_mismatch_for_wrong_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (path, _) = write_summary_file(&dir, "body text\n");
|
||||
let r = verify_summary_file(&path, "deadbeef").unwrap();
|
||||
assert!(matches!(r, VerifyResult::Mismatch { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_summary_file_missing_for_absent_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("does_not_exist.md");
|
||||
assert_eq!(
|
||||
verify_summary_file(&path, "abc").unwrap(),
|
||||
VerifyResult::Missing
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,420 @@
|
||||
//! Post-extraction tag rewriting for chunk and summary `.md` files.
|
||||
//!
|
||||
//! After the LLM extraction job runs, it produces a list of entities. Each
|
||||
//! entity is converted to an Obsidian-style hierarchical tag (`kind/Value`)
|
||||
//! and written into the `tags:` block in the file's front-matter.
|
||||
//!
|
||||
//! The body bytes (and therefore the SHA-256) are never changed — only the
|
||||
//! front-matter is rewritten.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::compose::{rewrite_summary_tags as compose_rewrite_summary_tags, rewrite_tags};
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::score::store::list_entity_ids_for_node;
|
||||
use crate::openhuman::memory::tree::store::get_summary_content_pointers;
|
||||
|
||||
/// Rewrite the `tags:` block in a chunk's on-disk `.md` file.
|
||||
///
|
||||
/// `abs_path` — absolute path to the chunk file.
|
||||
/// `tags` — new list of tag strings (Obsidian `kind/Value` format).
|
||||
///
|
||||
/// The operation is atomic: the new file is written to a sibling temp path and
|
||||
/// then renamed over the original. If the file does not exist, the call is a
|
||||
/// no-op (returns `Ok(())`).
|
||||
///
|
||||
/// Note: unlike the initial chunk write, tag rewrites MAY overwrite an
|
||||
/// existing file. The immutability contract covers the **body** only; tags are
|
||||
/// explicitly designed to be updated post-extraction.
|
||||
pub fn update_chunk_tags(abs_path: &Path, tags: &[String]) -> anyhow::Result<()> {
|
||||
if !abs_path.exists() {
|
||||
log::debug!(
|
||||
"[content_store::tags] skipping tag update — file not found: {}",
|
||||
abs_path.display()
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let old_bytes =
|
||||
std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?;
|
||||
|
||||
let new_bytes = rewrite_tags(&old_bytes, tags)
|
||||
.map_err(|e| anyhow::anyhow!("rewrite_tags {:?}: {e}", abs_path))?;
|
||||
|
||||
// Write the new content atomically via a sibling temp file.
|
||||
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
|
||||
let tmp_name = format!(".tmp_tags_{}.md", crate_temp_id());
|
||||
let tmp_path = parent.join(&tmp_name);
|
||||
|
||||
{
|
||||
use std::io::Write;
|
||||
let mut f = std::fs::File::create(&tmp_path)
|
||||
.map_err(|e| anyhow::anyhow!("create tag-rewrite tempfile {:?}: {e}", tmp_path))?;
|
||||
f.write_all(&new_bytes)
|
||||
.map_err(|e| anyhow::anyhow!("write tag-rewrite tempfile {:?}: {e}", tmp_path))?;
|
||||
f.sync_all()
|
||||
.map_err(|e| anyhow::anyhow!("fsync tag-rewrite tempfile {:?}: {e}", tmp_path))?;
|
||||
}
|
||||
|
||||
std::fs::rename(&tmp_path, abs_path).map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp_path);
|
||||
anyhow::anyhow!("rename tag-rewrite {:?} -> {:?}: {e}", tmp_path, abs_path)
|
||||
})?;
|
||||
|
||||
log::debug!(
|
||||
"[content_store::tags] updated tags in {}",
|
||||
abs_path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rewrite the `tags:` block in a summary's on-disk `.md` file.
|
||||
///
|
||||
/// Reads entity rows from `mem_tree_entity_index` for `summary_id`, converts
|
||||
/// them to `kind/Value` Obsidian tags, rewrites the YAML `tags:` block
|
||||
/// atomically (tempfile + fsync + rename), and verifies the body SHA-256 is
|
||||
/// unchanged afterwards.
|
||||
///
|
||||
/// Best-effort: tag-rewrite failures should not fail the extraction job. Callers
|
||||
/// should log a warning and continue — the entity index is the authoritative source.
|
||||
pub fn update_summary_tags(config: &Config, summary_id: &str) -> anyhow::Result<()> {
|
||||
// 1. Fetch content_path from SQLite.
|
||||
let pointers = get_summary_content_pointers(config, summary_id)?;
|
||||
let (rel_path, expected_sha) = match pointers {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
log::debug!(
|
||||
"[content_store::tags] update_summary_tags: no content_path for summary {summary_id} — skipping"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let content_root = config.memory_tree_content_root();
|
||||
let abs_path = {
|
||||
let mut p = content_root;
|
||||
for component in rel_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
if !abs_path.exists() {
|
||||
log::debug!(
|
||||
"[content_store::tags] update_summary_tags: file missing for summary {summary_id} \
|
||||
at {} — skipping",
|
||||
abs_path.display()
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 2. Fetch entity_index rows and build the merged tag list.
|
||||
let entity_ids = list_entity_ids_for_node(config, summary_id)?;
|
||||
let tags: Vec<String> = entity_ids
|
||||
.iter()
|
||||
.filter_map(|eid| {
|
||||
// entity_id format: "kind:surface"
|
||||
let (kind, surface) = eid.split_once(':')?;
|
||||
Some(entity_tag(kind, surface))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Sort + dedup for stability.
|
||||
let mut tags = tags;
|
||||
tags.sort();
|
||||
tags.dedup();
|
||||
|
||||
// 3. Read + atomic rewrite of the front-matter `tags:` block.
|
||||
let old_bytes = std::fs::read(&abs_path)
|
||||
.map_err(|e| anyhow::anyhow!("read summary {:?}: {e}", abs_path))?;
|
||||
|
||||
let new_bytes = compose_rewrite_summary_tags(&old_bytes, &tags)
|
||||
.map_err(|e| anyhow::anyhow!("rewrite_summary_tags {:?}: {e}", abs_path))?;
|
||||
|
||||
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
|
||||
let tmp_name = format!(".tmp_sum_tags_{}.md", crate_temp_id());
|
||||
let tmp_path = parent.join(&tmp_name);
|
||||
|
||||
{
|
||||
use std::io::Write;
|
||||
let mut f = std::fs::File::create(&tmp_path).map_err(|e| {
|
||||
anyhow::anyhow!("create summary tag-rewrite tempfile {:?}: {e}", tmp_path)
|
||||
})?;
|
||||
f.write_all(&new_bytes).map_err(|e| {
|
||||
anyhow::anyhow!("write summary tag-rewrite tempfile {:?}: {e}", tmp_path)
|
||||
})?;
|
||||
f.sync_all().map_err(|e| {
|
||||
anyhow::anyhow!("fsync summary tag-rewrite tempfile {:?}: {e}", tmp_path)
|
||||
})?;
|
||||
}
|
||||
|
||||
std::fs::rename(&tmp_path, &abs_path).map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp_path);
|
||||
anyhow::anyhow!(
|
||||
"rename summary tag-rewrite {:?} -> {:?}: {e}",
|
||||
tmp_path,
|
||||
abs_path
|
||||
)
|
||||
})?;
|
||||
|
||||
// 4. Sanity check: body sha must still match after the rewrite.
|
||||
let verify_bytes = std::fs::read(&abs_path)
|
||||
.map_err(|e| anyhow::anyhow!("re-read after tag rewrite {:?}: {e}", abs_path))?;
|
||||
let content = std::str::from_utf8(&verify_bytes)
|
||||
.map_err(|e| anyhow::anyhow!("UTF-8 after tag rewrite {:?}: {e}", abs_path))?;
|
||||
let body_after = super::compose::split_front_matter(content)
|
||||
.ok_or_else(|| anyhow::anyhow!("no front-matter after tag rewrite {:?}", abs_path))?
|
||||
.1;
|
||||
let actual_sha = super::atomic::sha256_hex(body_after.as_bytes());
|
||||
if actual_sha != expected_sha {
|
||||
return Err(anyhow::anyhow!(
|
||||
"[content_store::tags] update_summary_tags body mutated after rewrite \
|
||||
summary_id={summary_id} expected_sha={expected_sha} actual_sha={actual_sha}"
|
||||
));
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[content_store::tags] updated {} tags in summary file summary_id={summary_id} n_tags={}",
|
||||
tags.len(),
|
||||
tags.len()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Slugify an entity kind string for use in an Obsidian hierarchical tag.
|
||||
///
|
||||
/// Output: lowercase, spaces and non-alphanumeric chars replaced with `-`,
|
||||
/// consecutive dashes collapsed, leading/trailing dashes stripped.
|
||||
///
|
||||
/// Example: `"Person"` → `"person"`, `"GitHub Repo"` → `"github-repo"`
|
||||
pub fn slugify_tag_kind(kind: &str) -> String {
|
||||
slugify_tag_component(kind)
|
||||
}
|
||||
|
||||
/// Slugify an entity value string for use in an Obsidian hierarchical tag.
|
||||
///
|
||||
/// Like `slugify_tag_kind`, but capitalises the first letter of each word
|
||||
/// so values are visually distinct from kinds:
|
||||
///
|
||||
/// `"alice johnson"` → `"Alice-Johnson"`,
|
||||
/// `"project Phoenix"` → `"Project-Phoenix"`
|
||||
pub fn slugify_tag_value(value: &str) -> String {
|
||||
// Split on non-alphanumeric boundaries, capitalise first letter of each word.
|
||||
let mut parts: Vec<String> = Vec::new();
|
||||
let mut current = String::new();
|
||||
|
||||
for ch in value.chars() {
|
||||
if ch.is_alphanumeric() || ch == '_' {
|
||||
current.push(ch);
|
||||
} else if !current.is_empty() {
|
||||
parts.push(capitalise(¤t));
|
||||
current.clear();
|
||||
}
|
||||
}
|
||||
if !current.is_empty() {
|
||||
parts.push(capitalise(¤t));
|
||||
}
|
||||
|
||||
let joined = parts.join("-");
|
||||
if joined.is_empty() {
|
||||
"unknown".to_string()
|
||||
} else {
|
||||
joined
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an Obsidian-style `kind/Value` tag string from raw entity kind + surface.
|
||||
pub fn entity_tag(kind: &str, surface: &str) -> String {
|
||||
format!("{}/{}", slugify_tag_kind(kind), slugify_tag_value(surface))
|
||||
}
|
||||
|
||||
fn slugify_tag_component(s: &str) -> String {
|
||||
let lower = s.to_lowercase();
|
||||
let mut out = String::new();
|
||||
let mut last_dash = true;
|
||||
for ch in lower.chars() {
|
||||
if ch.is_ascii_alphanumeric() || ch == '_' {
|
||||
out.push(ch);
|
||||
last_dash = false;
|
||||
} else if !last_dash {
|
||||
out.push('-');
|
||||
last_dash = true;
|
||||
}
|
||||
}
|
||||
let trimmed = out.trim_end_matches('-');
|
||||
if trimmed.is_empty() {
|
||||
"unknown".to_string()
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn capitalise(s: &str) -> String {
|
||||
let mut chars = s.chars();
|
||||
match chars.next() {
|
||||
None => String::new(),
|
||||
Some(first) => {
|
||||
let upper: String = first.to_uppercase().collect();
|
||||
upper + chars.as_str()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn crate_temp_id() -> String {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let ns = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
format!("{ns:08x}")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
|
||||
use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file;
|
||||
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind};
|
||||
use chrono::TimeZone;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn sample_chunk() -> Chunk {
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
Chunk {
|
||||
id: "tags_test".into(),
|
||||
content: "hello from tags test".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec!["old/Tag".into()],
|
||||
source_ref: None,
|
||||
},
|
||||
token_count: 4,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_chunk_tags_replaces_tag_block() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let chunk = sample_chunk();
|
||||
let (full, _) = compose_chunk_file(&chunk);
|
||||
let path = dir.path().join("0.md");
|
||||
write_if_new(&path, &full).unwrap();
|
||||
|
||||
update_chunk_tags(
|
||||
&path,
|
||||
&["person/Alice-Smith".into(), "project/Phoenix".into()],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let updated = std::fs::read_to_string(&path).unwrap();
|
||||
assert!(updated.contains(" - person/Alice-Smith"));
|
||||
assert!(updated.contains(" - project/Phoenix"));
|
||||
assert!(!updated.contains(" - old/Tag"));
|
||||
// Body unchanged.
|
||||
assert!(updated.ends_with("hello from tags test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_chunk_tags_is_noop_for_missing_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("nonexistent.md");
|
||||
assert!(update_chunk_tags(&path, &["p/X".into()]).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_tag_kind_examples() {
|
||||
assert_eq!(slugify_tag_kind("Person"), "person");
|
||||
assert_eq!(slugify_tag_kind("GitHub Repo"), "github-repo");
|
||||
assert_eq!(slugify_tag_kind("EMAIL"), "email");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slugify_tag_value_capitalises_words() {
|
||||
assert_eq!(slugify_tag_value("alice johnson"), "Alice-Johnson");
|
||||
assert_eq!(slugify_tag_value("project Phoenix"), "Project-Phoenix");
|
||||
assert_eq!(slugify_tag_value("OPENAI"), "OPENAI");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entity_tag_builds_obsidian_tag() {
|
||||
assert_eq!(
|
||||
entity_tag("person", "Alice Johnson"),
|
||||
"person/Alice-Johnson"
|
||||
);
|
||||
assert_eq!(entity_tag("ORG", "Tinyhumans AI"), "org/Tinyhumans-AI");
|
||||
}
|
||||
|
||||
// ─── update_summary_tags tests ────────────────────────────────────────────
|
||||
|
||||
/// Write a summary .md file to disk with empty tags and verify rewriting works.
|
||||
#[test]
|
||||
fn rewrite_summary_tags_preserves_body_and_replaces_tags() {
|
||||
use crate::openhuman::memory::tree::content_store::compose::{
|
||||
compose_summary_md, SummaryComposeInput,
|
||||
};
|
||||
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let body = "summary body for tag test\n";
|
||||
let children = vec!["c1".to_string()];
|
||||
let input = SummaryComposeInput {
|
||||
summary_id: "sum:L1:tagtest",
|
||||
tree_kind: SummaryTreeKind::Source,
|
||||
tree_id: "t1",
|
||||
tree_scope: "gmail:alice@x.com",
|
||||
level: 1,
|
||||
child_ids: &children,
|
||||
child_count: 1,
|
||||
time_range_start: ts,
|
||||
time_range_end: ts,
|
||||
sealed_at: ts,
|
||||
body,
|
||||
};
|
||||
let composed = compose_summary_md(&input);
|
||||
let path = dir.path().join("sum.md");
|
||||
write_if_new(&path, composed.full.as_bytes()).unwrap();
|
||||
|
||||
// Original must have `tags: []`
|
||||
let original = std::fs::read_to_string(&path).unwrap();
|
||||
assert!(original.contains("tags: []"));
|
||||
|
||||
// Rewrite the tags block
|
||||
let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()];
|
||||
let file_bytes = std::fs::read(&path).unwrap();
|
||||
let rewritten = super::compose_rewrite_summary_tags(&file_bytes, &new_tags).unwrap();
|
||||
|
||||
// Write rewritten bytes back (simulating atomic rewrite)
|
||||
let tmp = dir.path().join("sum.tmp.md");
|
||||
{
|
||||
use std::io::Write;
|
||||
let mut f = std::fs::File::create(&tmp).unwrap();
|
||||
f.write_all(&rewritten).unwrap();
|
||||
}
|
||||
std::fs::rename(&tmp, &path).unwrap();
|
||||
|
||||
let updated = std::fs::read_to_string(&path).unwrap();
|
||||
assert!(updated.contains(" - person/Alice-Smith"));
|
||||
assert!(updated.contains(" - topic/Memory"));
|
||||
assert!(!updated.contains("tags: []"));
|
||||
// Body unchanged
|
||||
assert!(updated.ends_with(body));
|
||||
|
||||
// Body sha unchanged
|
||||
use crate::openhuman::memory::tree::content_store::compose::split_front_matter;
|
||||
let (_, body_after) = split_front_matter(&updated).unwrap();
|
||||
let sha = sha256_hex(body_after.as_bytes());
|
||||
let expected_sha = sha256_hex(body.as_bytes());
|
||||
assert_eq!(
|
||||
sha, expected_sha,
|
||||
"body sha must be stable after tag rewrite"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -19,11 +19,17 @@
|
||||
//! - Idempotency: if an L0 daily node already exists for the target day,
|
||||
//! return `DigestOutcome::Skipped` rather than emitting a duplicate.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{DateTime, Duration, NaiveDate, TimeZone, Utc};
|
||||
use rusqlite::OptionalExtension;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::{
|
||||
atomic::stage_summary, paths::slugify_source_id, read as content_read, SummaryComposeInput,
|
||||
SummaryTreeKind,
|
||||
};
|
||||
use crate::openhuman::memory::tree::global_tree::registry::get_or_create_global_tree;
|
||||
use crate::openhuman::memory::tree::global_tree::seal::append_daily_and_cascade;
|
||||
use crate::openhuman::memory::tree::global_tree::GLOBAL_TOKEN_BUDGET;
|
||||
@@ -154,6 +160,25 @@ pub async fn end_of_day_digest(
|
||||
.await
|
||||
.context("embed daily summary during end_of_day_digest")?;
|
||||
|
||||
// L0 daily node inherits entities/topics by union of contributing
|
||||
// source-tree summaries. Each input was already labeled at source-tree
|
||||
// seal time, so emergent themes don't need another extractor pass
|
||||
// here — global is a sink; union preserves "days that mentioned X"
|
||||
// retrieval without an extra LLM call. See LabelStrategy in
|
||||
// source_tree::bucket_seal for the full design.
|
||||
let mut entities_set: BTreeSet<String> = BTreeSet::new();
|
||||
let mut topics_set: BTreeSet<String> = BTreeSet::new();
|
||||
for inp in &inputs {
|
||||
for e in &inp.entities {
|
||||
entities_set.insert(e.clone());
|
||||
}
|
||||
for t in &inp.topics {
|
||||
topics_set.insert(t.clone());
|
||||
}
|
||||
}
|
||||
let daily_entities: Vec<String> = entities_set.into_iter().collect();
|
||||
let daily_topics: Vec<String> = topics_set.into_iter().collect();
|
||||
|
||||
let now = Utc::now();
|
||||
let daily_id = new_summary_id(0);
|
||||
let daily = SummaryNode {
|
||||
@@ -165,8 +190,8 @@ pub async fn end_of_day_digest(
|
||||
child_ids: inputs.iter().map(|i| i.id.clone()).collect(),
|
||||
content: output.content,
|
||||
token_count: output.token_count,
|
||||
entities: output.entities,
|
||||
topics: output.topics,
|
||||
entities: daily_entities,
|
||||
topics: daily_topics,
|
||||
time_range_start: day_start,
|
||||
time_range_end: day_end,
|
||||
score,
|
||||
@@ -175,6 +200,44 @@ pub async fn end_of_day_digest(
|
||||
embedding: Some(embedding),
|
||||
};
|
||||
|
||||
// Phase MD-content: stage the L0 daily .md file before the write tx.
|
||||
// `date_for_global` = day_start (the calendar day this digest covers).
|
||||
let daily_compose_input = SummaryComposeInput {
|
||||
summary_id: &daily.id,
|
||||
tree_kind: SummaryTreeKind::Global,
|
||||
tree_id: &daily.tree_id,
|
||||
tree_scope: &global.scope,
|
||||
level: daily.level,
|
||||
child_ids: &daily.child_ids,
|
||||
child_count: daily.child_ids.len(),
|
||||
time_range_start: daily.time_range_start,
|
||||
time_range_end: daily.time_range_end,
|
||||
sealed_at: daily.sealed_at,
|
||||
body: &daily.content,
|
||||
};
|
||||
// Stage the summary .md file — abort the digest on failure so the database
|
||||
// never commits a row with content_path = NULL. The digest job is retried
|
||||
// via the normal job-retry path.
|
||||
let content_root_daily = config.memory_tree_content_root();
|
||||
let global_scope_slug = slugify_source_id(&global.scope);
|
||||
let staged_daily = stage_summary(
|
||||
&content_root_daily,
|
||||
&daily_compose_input,
|
||||
&global_scope_slug,
|
||||
Some(day_start),
|
||||
)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"stage_summary failed for daily {}; digest aborted for retry",
|
||||
daily.id
|
||||
)
|
||||
})?;
|
||||
log::debug!(
|
||||
"[global_tree::digest] staged daily summary {} → {}",
|
||||
daily.id,
|
||||
staged_daily.content_path
|
||||
);
|
||||
|
||||
// Persist the daily node. Note: we do NOT backlink parent_id on the
|
||||
// child summaries here — their parents are their own source trees, not
|
||||
// the global tree. The global-tree child_ids are cross-source
|
||||
@@ -183,7 +246,7 @@ pub async fn end_of_day_digest(
|
||||
let tree_id_clone = global.id.clone();
|
||||
with_connection(config, move |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
store::insert_summary_tx(&tx, &daily_clone)?;
|
||||
store::insert_summary_tx(&tx, &daily_clone, Some(&staged_daily))?;
|
||||
// Index any entities the summariser emitted (no-op under inert).
|
||||
crate::openhuman::memory::tree::score::store::index_summary_entity_ids_tx(
|
||||
&tx,
|
||||
@@ -315,9 +378,23 @@ fn pick_source_contribution(
|
||||
}
|
||||
};
|
||||
|
||||
// Read the full body from disk — `node.content` is a ≤500-char preview
|
||||
// after the MD-on-disk migration. The digest summariser must receive the
|
||||
// complete summary text so the daily recap is not assembled from previews.
|
||||
let body = match content_read::read_summary_body(config, &node.id) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[global_tree::digest] read_summary_body failed for {} — using preview: {e:#}",
|
||||
node.id
|
||||
);
|
||||
// Non-fatal: fall back to preview for pre-MD-migration rows.
|
||||
node.content.clone()
|
||||
}
|
||||
};
|
||||
Ok(Some(SummaryInput {
|
||||
id: node.id,
|
||||
content: format!("[{}]\n{}", source_tree.scope, node.content),
|
||||
content: format!("[{}]\n{}", source_tree.scope, body),
|
||||
token_count: node.token_count,
|
||||
entities: node.entities,
|
||||
topics: node.topics,
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use crate::openhuman::memory::tree::source_tree::types::TreeStatus;
|
||||
@@ -7,6 +10,24 @@ use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Stage a batch of chunks to the content store so that `read_chunk_body`
|
||||
/// can find the on-disk file during seals. Tests that call `upsert_chunks`
|
||||
/// and then trigger a seal MUST also call this helper; otherwise
|
||||
/// `hydrate_leaf_inputs` will fail with "no content_path for chunk_id".
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged =
|
||||
content_store::stage_chunks(&content_root, chunks).expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -39,6 +60,7 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
|
||||
token_count: 6_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let c2 = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, scope, 1, "test-content"),
|
||||
@@ -55,8 +77,10 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
|
||||
token_count: 6_000,
|
||||
seq_in_source: 1,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
stage_test_chunks(cfg, &[c1.clone(), c2.clone()]);
|
||||
|
||||
let leaf1 = LeafRef {
|
||||
chunk_id: c1.id.clone(),
|
||||
@@ -76,8 +100,12 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf(cfg, &tree, &leaf1, &summariser).await.unwrap();
|
||||
append_leaf(cfg, &tree, &leaf2, &summariser).await.unwrap();
|
||||
append_leaf(cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
append_leaf(cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
// 12k tokens > 10k budget → one L1 summary covering `ts`.
|
||||
}
|
||||
|
||||
@@ -216,3 +244,157 @@ async fn seven_days_cascade_to_weekly_seal() {
|
||||
assert_eq!(t.max_level, 1);
|
||||
assert_eq!(t.status, TreeStatus::Active);
|
||||
}
|
||||
|
||||
/// Seed a source tree whose sealed L1 summary carries the given entities
|
||||
/// and topics. Entities are written into `mem_tree_entity_index` (where
|
||||
/// seal-time hydration reads them); topics are stored on chunk metadata
|
||||
/// tags. The seal then unions both into the L1 summary.
|
||||
async fn seed_source_tree_with_labeled_l1(
|
||||
cfg: &Config,
|
||||
scope: &str,
|
||||
ts: DateTime<Utc>,
|
||||
entities: Vec<String>,
|
||||
topics: Vec<String>,
|
||||
) {
|
||||
use crate::openhuman::memory::tree::score::extract::EntityKind;
|
||||
use crate::openhuman::memory::tree::score::resolver::CanonicalEntity;
|
||||
use crate::openhuman::memory::tree::score::store::index_entity;
|
||||
|
||||
let tree = get_or_create_source_tree(cfg, scope).unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
|
||||
let mut chunks: Vec<Chunk> = Vec::new();
|
||||
for seq in 0..2u32 {
|
||||
chunks.push(Chunk {
|
||||
id: chunk_id(SourceKind::Chat, scope, seq, "labeled-test"),
|
||||
content: format!("labeled chunk {seq} in {scope}"),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: scope.into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: topics.clone(),
|
||||
source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))),
|
||||
},
|
||||
token_count: 6_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
});
|
||||
}
|
||||
upsert_chunks(cfg, &chunks).unwrap();
|
||||
stage_test_chunks(cfg, &chunks);
|
||||
|
||||
for chunk in &chunks {
|
||||
for entity_id in &entities {
|
||||
let kind = entity_id
|
||||
.split_once(':')
|
||||
.map_or(EntityKind::Misc, |(k, _)| {
|
||||
EntityKind::parse(k).unwrap_or(EntityKind::Misc)
|
||||
});
|
||||
let surface = entity_id
|
||||
.split_once(':')
|
||||
.map_or(entity_id.as_str(), |(_, v)| v);
|
||||
let e = CanonicalEntity {
|
||||
canonical_id: entity_id.clone(),
|
||||
kind,
|
||||
surface: surface.to_string(),
|
||||
span_start: 0,
|
||||
span_end: surface.len() as u32,
|
||||
score: 1.0,
|
||||
};
|
||||
index_entity(
|
||||
cfg,
|
||||
&e,
|
||||
&chunk.id,
|
||||
"leaf",
|
||||
ts.timestamp_millis(),
|
||||
Some(scope),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// Two 6k-token leaves total 12k → exceeds L0 budget → seal fires on
|
||||
// the second append, producing one L1 summary that unions all leaf
|
||||
// labels (every leaf has the same set, so dedup yields the input set).
|
||||
for chunk in &chunks {
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id.clone(),
|
||||
token_count: 6_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content.clone(),
|
||||
entities: entities.clone(),
|
||||
topics: topics.clone(),
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf(
|
||||
cfg,
|
||||
&tree,
|
||||
&leaf,
|
||||
&summariser,
|
||||
&LabelStrategy::UnionFromChildren,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn daily_digest_unions_labels_from_source_summaries() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let summariser = InertSummariser::new();
|
||||
|
||||
let day = NaiveDate::from_ymd_opt(2025, 5, 1).unwrap();
|
||||
let ts = day.and_hms_opt(10, 0, 0).unwrap().and_utc();
|
||||
|
||||
// Source A's L1 carries (alice, phoenix-migration). Source B's L1
|
||||
// carries (bob, phoenix-migration, qa). The daily L0 should union to
|
||||
// (alice, bob, phoenix-migration) for entities and (phoenix-migration,
|
||||
// qa) for topics — overlap dedup'd.
|
||||
seed_source_tree_with_labeled_l1(
|
||||
&cfg,
|
||||
"slack:#a",
|
||||
ts,
|
||||
vec!["email:alice@example.com".into(), "topic:phoenix".into()],
|
||||
vec!["phoenix-migration".into()],
|
||||
)
|
||||
.await;
|
||||
seed_source_tree_with_labeled_l1(
|
||||
&cfg,
|
||||
"slack:#b",
|
||||
ts,
|
||||
vec!["person:bob".into(), "topic:phoenix".into()],
|
||||
vec!["phoenix-migration".into(), "qa".into()],
|
||||
)
|
||||
.await;
|
||||
|
||||
let outcome = end_of_day_digest(&cfg, day, &summariser).await.unwrap();
|
||||
let daily_id = match outcome {
|
||||
DigestOutcome::Emitted { daily_id, .. } => daily_id,
|
||||
other => panic!("expected Emitted, got {other:?}"),
|
||||
};
|
||||
|
||||
let daily = store::get_summary(&cfg, &daily_id).unwrap().unwrap();
|
||||
let entities: std::collections::BTreeSet<&str> =
|
||||
daily.entities.iter().map(String::as_str).collect();
|
||||
let topics: std::collections::BTreeSet<&str> =
|
||||
daily.topics.iter().map(String::as_str).collect();
|
||||
|
||||
assert!(entities.contains("email:alice@example.com"));
|
||||
assert!(entities.contains("person:bob"));
|
||||
assert!(entities.contains("topic:phoenix"));
|
||||
assert_eq!(
|
||||
entities.len(),
|
||||
3,
|
||||
"expected 3 unique entities (deduped); got {entities:?}"
|
||||
);
|
||||
assert!(topics.contains("phoenix-migration"));
|
||||
assert!(topics.contains("qa"));
|
||||
assert_eq!(
|
||||
topics.len(),
|
||||
2,
|
||||
"expected 2 unique topics (deduped); got {topics:?}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -158,14 +158,31 @@ fn assemble_recap(covering: &[&SummaryNode], level: u32) -> RecapOutput {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, chunks)
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -216,6 +233,7 @@ mod tests {
|
||||
token_count: 6_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let c2 = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, scope, 1, "test-content"),
|
||||
@@ -232,8 +250,10 @@ mod tests {
|
||||
token_count: 6_000,
|
||||
seq_in_source: 1,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
stage_test_chunks(cfg, &[c1.clone(), c2.clone()]);
|
||||
append_leaf(
|
||||
cfg,
|
||||
&tree,
|
||||
@@ -247,6 +267,7 @@ mod tests {
|
||||
score: 0.5,
|
||||
},
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -263,6 +284,7 @@ mod tests {
|
||||
score: 0.5,
|
||||
},
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -11,10 +11,15 @@
|
||||
//! `mem_tree_summaries` on both sides (children and output), since even L0
|
||||
//! is a sealed summary node rather than a raw chunk.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::{
|
||||
atomic::stage_summary, SummaryComposeInput, SummaryTreeKind,
|
||||
};
|
||||
use crate::openhuman::memory::tree::global_tree::{
|
||||
GLOBAL_TOKEN_BUDGET, MONTHLY_SEAL_THRESHOLD, WEEKLY_SEAL_THRESHOLD, YEARLY_SEAL_THRESHOLD,
|
||||
};
|
||||
@@ -186,6 +191,26 @@ async fn seal_one_level(
|
||||
.await
|
||||
.context("summariser failed during global seal")?;
|
||||
|
||||
// Global-tree summaries inherit their entity/topic labels via union
|
||||
// from their already-labeled inputs (source-tree summaries carry
|
||||
// labels from the source-tree seal extractor; global L1+ inputs
|
||||
// carry labels from this same union path one level down). We
|
||||
// deliberately do NOT run an extractor on the daily/weekly/monthly
|
||||
// synthesis: the inputs already cover what the summary represents,
|
||||
// and global is a sink — no second-pass labeling earns its keep.
|
||||
let mut entities_set: BTreeSet<String> = BTreeSet::new();
|
||||
let mut topics_set: BTreeSet<String> = BTreeSet::new();
|
||||
for inp in &inputs {
|
||||
for e in &inp.entities {
|
||||
entities_set.insert(e.clone());
|
||||
}
|
||||
for t in &inp.topics {
|
||||
topics_set.insert(t.clone());
|
||||
}
|
||||
}
|
||||
let node_entities: Vec<String> = entities_set.into_iter().collect();
|
||||
let node_topics: Vec<String> = topics_set.into_iter().collect();
|
||||
|
||||
// Phase 4 (#710): embed BEFORE opening the write tx so an embedder
|
||||
// error aborts the cascade without half-committing the summary.
|
||||
let embedder =
|
||||
@@ -208,8 +233,8 @@ async fn seal_one_level(
|
||||
child_ids: buf.item_ids.clone(),
|
||||
content: output.content,
|
||||
token_count: output.token_count,
|
||||
entities: output.entities,
|
||||
topics: output.topics,
|
||||
entities: node_entities,
|
||||
topics: node_topics,
|
||||
time_range_start,
|
||||
time_range_end,
|
||||
score,
|
||||
@@ -218,6 +243,49 @@ async fn seal_one_level(
|
||||
embedding: Some(embedding),
|
||||
};
|
||||
|
||||
// Phase MD-content: stage the global summary .md file before opening the
|
||||
// write tx. date_for_global = time_range_start date (daily for L0, or
|
||||
// the start of the range for higher levels).
|
||||
let global_date = Some(time_range_start);
|
||||
let compose_input_global = SummaryComposeInput {
|
||||
summary_id: &node.id,
|
||||
tree_kind: SummaryTreeKind::Global,
|
||||
tree_id: &node.tree_id,
|
||||
tree_scope: &tree.scope,
|
||||
level: node.level,
|
||||
child_ids: &node.child_ids,
|
||||
child_count: node.child_ids.len(),
|
||||
time_range_start: node.time_range_start,
|
||||
time_range_end: node.time_range_end,
|
||||
sealed_at: node.sealed_at,
|
||||
body: &node.content,
|
||||
};
|
||||
// Stage the summary .md file — abort the seal on failure so the database
|
||||
// never commits a row with content_path = NULL. The job-retry path will
|
||||
// re-attempt the file write on next execution.
|
||||
let content_root_global = config.memory_tree_content_root();
|
||||
// Global tree scope is typically the literal "global" string.
|
||||
// Use it as-is for the path (slugify passes through short ascii strings unchanged).
|
||||
let global_scope_slug =
|
||||
crate::openhuman::memory::tree::content_store::paths::slugify_source_id(&tree.scope);
|
||||
let staged_global = stage_summary(
|
||||
&content_root_global,
|
||||
&compose_input_global,
|
||||
&global_scope_slug,
|
||||
global_date,
|
||||
)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"stage_summary failed for {}; global-tree seal aborted for retry",
|
||||
node.id
|
||||
)
|
||||
})?;
|
||||
log::debug!(
|
||||
"[global_tree::seal] staged summary {} → {}",
|
||||
node.id,
|
||||
staged_global.content_path
|
||||
);
|
||||
|
||||
// Single write transaction: insert the new summary, clear this level's
|
||||
// buffer, append the new id to the parent buffer, and bump the tree's
|
||||
// max_level/root_id if we just climbed. Re-read `max_level` inside the
|
||||
@@ -238,7 +306,7 @@ async fn seal_one_level(
|
||||
.map(|n| n.max(0) as u32)
|
||||
.context("Failed to read current max_level for global tree")?;
|
||||
|
||||
store::insert_summary_tx(&tx, &node)?;
|
||||
store::insert_summary_tx(&tx, &node, Some(&staged_global))?;
|
||||
// Index any entities the summariser emitted. No-op under
|
||||
// InertSummariser (entities stays empty by design — see
|
||||
// summariser/inert.rs). Becomes active when the Ollama summariser
|
||||
@@ -383,7 +451,7 @@ mod tests {
|
||||
fn insert_daily(cfg: &Config, node: &SummaryNode) {
|
||||
with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
store::insert_summary_tx(&tx, node)?;
|
||||
store::insert_summary_tx(&tx, node, None)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
|
||||
+127
-282
@@ -1,14 +1,12 @@
|
||||
//! Ingest orchestrator (Phase 1 + Phase 2):
|
||||
//! Ingest orchestrator for the async memory-tree pipeline.
|
||||
//!
|
||||
//! canonicalise → chunk → score → admission gate → persist (chunks + scores + entity index)
|
||||
//! The hot path now does:
|
||||
//! `canonicalise -> chunk -> fast score -> persist chunks/score rows -> enqueue extract jobs`
|
||||
//!
|
||||
//! Phase 2 inserts scoring between chunker and persistence. Low-scoring
|
||||
//! chunks are dropped (their rationale is still persisted to
|
||||
//! `mem_tree_score` for diagnostics); surviving chunks get their entities
|
||||
//! indexed so later phases can resolve "which chunks mention Alice?" in
|
||||
//! O(lookup).
|
||||
//! The slower work (full extraction, admission, tree buffering, sealing,
|
||||
//! topic routing, daily digests) runs out of the SQLite-backed jobs queue.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
@@ -19,25 +17,21 @@ use crate::openhuman::memory::tree::canonicalize::{
|
||||
CanonicalisedSource,
|
||||
};
|
||||
use crate::openhuman::memory::tree::chunker::{chunk_markdown, ChunkerInput, ChunkerOptions};
|
||||
use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked};
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::jobs::{self, ExtractChunkPayload, NewJob};
|
||||
use crate::openhuman::memory::tree::score::{self, ScoreResult, ScoringConfig};
|
||||
use crate::openhuman::memory::tree::source_tree::{
|
||||
append_leaf, get_or_create_source_tree, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::store;
|
||||
use crate::openhuman::memory::tree::topic_tree::route_leaf_to_topic_trees;
|
||||
use crate::openhuman::memory::tree::types::Chunk;
|
||||
|
||||
/// Outcome of one ingest call — extended with per-chunk admission info.
|
||||
/// Outcome of one ingest call.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct IngestResult {
|
||||
pub source_id: String,
|
||||
/// Number of chunks that passed the admission gate and were persisted.
|
||||
/// Number of chunks persisted and queued for async extraction.
|
||||
pub chunks_written: usize,
|
||||
/// Number of chunks that failed the admission gate and were NOT persisted
|
||||
/// (their score rationale IS persisted for diagnostics).
|
||||
/// Number of chunks the cheap fast-score path would drop. Final admission
|
||||
/// still happens later in the extract job.
|
||||
pub chunks_dropped: usize,
|
||||
/// IDs of all chunks that were persisted (in source order).
|
||||
/// IDs of all chunks written and queued.
|
||||
pub chunk_ids: Vec<String>,
|
||||
}
|
||||
|
||||
@@ -52,7 +46,6 @@ impl IngestResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Ingest a batch of chat messages scoped to one channel/group.
|
||||
pub async fn ingest_chat(
|
||||
config: &Config,
|
||||
source_id: &str,
|
||||
@@ -60,11 +53,6 @@ pub async fn ingest_chat(
|
||||
tags: Vec<String>,
|
||||
batch: ChatBatch,
|
||||
) -> Result<IngestResult> {
|
||||
log::debug!(
|
||||
"[memory_tree::ingest] chat source_id={} msg_count={}",
|
||||
source_id,
|
||||
batch.messages.len()
|
||||
);
|
||||
let canonical =
|
||||
match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? {
|
||||
Some(c) => c,
|
||||
@@ -73,7 +61,6 @@ pub async fn ingest_chat(
|
||||
persist(config, source_id, canonical).await
|
||||
}
|
||||
|
||||
/// Ingest a single email thread.
|
||||
pub async fn ingest_email(
|
||||
config: &Config,
|
||||
source_id: &str,
|
||||
@@ -81,11 +68,6 @@ pub async fn ingest_email(
|
||||
tags: Vec<String>,
|
||||
thread: EmailThread,
|
||||
) -> Result<IngestResult> {
|
||||
log::debug!(
|
||||
"[memory_tree::ingest] email source_id={} msg_count={}",
|
||||
source_id,
|
||||
thread.messages.len()
|
||||
);
|
||||
let canonical =
|
||||
match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? {
|
||||
Some(c) => c,
|
||||
@@ -94,7 +76,6 @@ pub async fn ingest_email(
|
||||
persist(config, source_id, canonical).await
|
||||
}
|
||||
|
||||
/// Ingest a single standalone document.
|
||||
pub async fn ingest_document(
|
||||
config: &Config,
|
||||
source_id: &str,
|
||||
@@ -102,13 +83,6 @@ pub async fn ingest_document(
|
||||
tags: Vec<String>,
|
||||
doc: DocumentInput,
|
||||
) -> Result<IngestResult> {
|
||||
let title_len = doc.title.chars().count();
|
||||
log::debug!(
|
||||
"[memory_tree::ingest] document source_id={} has_title={} title_len={}",
|
||||
source_id,
|
||||
!doc.title.trim().is_empty(),
|
||||
title_len
|
||||
);
|
||||
let canonical =
|
||||
match document::canonicalise(source_id, owner, &tags, doc).map_err(anyhow::Error::msg)? {
|
||||
Some(c) => c,
|
||||
@@ -122,7 +96,6 @@ async fn persist(
|
||||
source_id: &str,
|
||||
canonical: CanonicalisedSource,
|
||||
) -> Result<IngestResult> {
|
||||
// 1. Chunk
|
||||
let input = ChunkerInput {
|
||||
source_kind: canonical.metadata.source_kind,
|
||||
source_id: source_id.to_string(),
|
||||
@@ -134,15 +107,15 @@ async fn persist(
|
||||
return Ok(IngestResult::empty(source_id));
|
||||
}
|
||||
|
||||
// 2. Score (async; extractor chain driven by config.memory_tree —
|
||||
// regex-only unless llm_extractor_endpoint + model are set, in
|
||||
// which case LlmEntityExtractor runs as second-pass on borderline
|
||||
// chunks with soft-fallback to regex on LLM failure).
|
||||
let scoring_cfg = ScoringConfig::from_config(config);
|
||||
let scores = score::score_chunks(&chunks, &scoring_cfg).await?;
|
||||
// Phase MD-content: write chunk bodies to disk before the SQLite upsert.
|
||||
// stage_chunks is sync I/O; run it here (still on the tokio thread) before
|
||||
// spawn_blocking so errors surface before the DB transaction opens.
|
||||
let content_root = config.memory_tree_content_root();
|
||||
let staged = content_store::stage_chunks(&content_root, &chunks)
|
||||
.map_err(|e| anyhow::anyhow!("[memory_tree::ingest] stage_chunks failed: {e}"))?;
|
||||
|
||||
// Fail fast on scorer length mismatch — silently truncating via zip would
|
||||
// drop chunks (or their score rationale) without trace.
|
||||
let scoring_cfg = ScoringConfig::from_config(config);
|
||||
let scores = score::score_chunks_fast(&chunks, &scoring_cfg).await?;
|
||||
if scores.len() != chunks.len() {
|
||||
anyhow::bail!(
|
||||
"[memory_tree::ingest] scorer length mismatch: chunks={} scores={}",
|
||||
@@ -151,78 +124,79 @@ async fn persist(
|
||||
);
|
||||
}
|
||||
|
||||
// 3. Partition kept vs dropped
|
||||
let mut kept_chunks: Vec<Chunk> = Vec::new();
|
||||
let mut all_results: Vec<(ScoreResult, i64)> = Vec::new();
|
||||
for (chunk, result) in chunks.iter().zip(scores.into_iter()) {
|
||||
let ts_ms = chunk.metadata.timestamp.timestamp_millis();
|
||||
if result.kept {
|
||||
kept_chunks.push(chunk.clone());
|
||||
}
|
||||
all_results.push((result, ts_ms));
|
||||
}
|
||||
|
||||
let all_results: Vec<(ScoreResult, i64)> = chunks
|
||||
.iter()
|
||||
.zip(scores.into_iter())
|
||||
.map(|(chunk, result)| (result, chunk.metadata.timestamp.timestamp_millis()))
|
||||
.collect();
|
||||
let dropped = all_results.iter().filter(|(r, _)| !r.kept).count();
|
||||
log::debug!(
|
||||
"[memory_tree::ingest] scoring source_id={} kept={} dropped={}",
|
||||
source_id,
|
||||
kept_chunks.len(),
|
||||
dropped
|
||||
);
|
||||
|
||||
// 3.5. Phase 4 (#710) — embed every kept chunk BEFORE writing. A failed
|
||||
// embed aborts the whole batch so a retry stays idempotent on
|
||||
// `chunk_id` (upsert_chunks_tx ON CONFLICT REPLACE). Serial per
|
||||
// chunk for now; parallelism is an explicit follow-up. We pack
|
||||
// each vector into its SQLite BLOB representation up front so the
|
||||
// write tx below doesn't have to juggle floats.
|
||||
let embedder = build_embedder_from_config(config).context("build embedder during ingest")?;
|
||||
log::debug!(
|
||||
"[memory_tree::ingest] embedding source_id={} provider={} kept={}",
|
||||
source_id,
|
||||
embedder.name(),
|
||||
kept_chunks.len()
|
||||
);
|
||||
let mut chunk_embedding_blobs: Vec<(String, Vec<u8>)> = Vec::with_capacity(kept_chunks.len());
|
||||
for chunk in &kept_chunks {
|
||||
let vector = embedder
|
||||
.embed(&chunk.content)
|
||||
.await
|
||||
.with_context(|| format!("embed chunk_id={} during ingest", chunk.id))?;
|
||||
let packed = pack_checked(&vector)
|
||||
.with_context(|| format!("pack embedding for chunk_id={} during ingest", chunk.id))?;
|
||||
chunk_embedding_blobs.push((chunk.id.clone(), packed));
|
||||
}
|
||||
|
||||
// 4. Persist (blocking SQLite — isolate on a dedicated thread).
|
||||
// Chunks + scores + embeddings all commit in one tx so retries are
|
||||
// idempotent. The chunk upsert preserves any prior embedding on
|
||||
// conflict (see store::upsert_chunks_tx), so we follow with a
|
||||
// deliberate UPDATE that writes the fresh blob.
|
||||
let config_owned = config.clone();
|
||||
let kept_for_store = kept_chunks.clone();
|
||||
let staged_for_store = staged.clone();
|
||||
let results_for_store = all_results.clone();
|
||||
let embeddings_for_store = chunk_embedding_blobs.clone();
|
||||
let written = tokio::task::spawn_blocking(move || -> Result<usize> {
|
||||
use std::collections::{HashMap, HashSet};
|
||||
store::with_connection(&config_owned, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
let n = store::upsert_chunks_tx(&tx, &kept_for_store)?;
|
||||
for (chunk_id, blob) in &embeddings_for_store {
|
||||
let changed = tx.execute(
|
||||
"UPDATE mem_tree_chunks SET embedding = ?1 WHERE id = ?2",
|
||||
rusqlite::params![blob, chunk_id],
|
||||
)?;
|
||||
if changed == 0 {
|
||||
log::warn!(
|
||||
"[memory_tree::ingest] embedding update affected 0 rows chunk_id={chunk_id} — \
|
||||
upsert missed the row?"
|
||||
);
|
||||
|
||||
// Read each chunk's CURRENT lifecycle BEFORE the upsert. This
|
||||
// is the "did this chunk exist before this batch" snapshot,
|
||||
// because `upsert_staged_chunks_tx` will either preserve the
|
||||
// existing row's lifecycle (UPDATE doesn't touch the column) or
|
||||
// insert a new row that picks up the column DEFAULT — so reading
|
||||
// post-upsert can't distinguish "brand new" from
|
||||
// "already-admitted-from-prior-ingest".
|
||||
let mut prior: HashMap<String, Option<String>> = HashMap::new();
|
||||
for s in &staged_for_store {
|
||||
let status = store::get_chunk_lifecycle_status_tx(&tx, &s.chunk.id)?;
|
||||
prior.insert(s.chunk.id.clone(), status);
|
||||
}
|
||||
|
||||
let n = store::upsert_staged_chunks_tx(&tx, &staged_for_store)?;
|
||||
|
||||
// Re-ingest of identical content (same chunk_id) must NOT
|
||||
// downgrade chunks that have already progressed through the
|
||||
// async pipeline. Without this guard, a re-ingest would reset
|
||||
// every chunk to 'pending_extraction' and enqueue a fresh
|
||||
// `extract_chunk` job — sending already-buffered/sealed
|
||||
// chunks back through extract → admit → append, ultimately
|
||||
// duplicating them into a second summary in the same tree.
|
||||
//
|
||||
// Schedule a chunk for processing when its PRE-upsert state
|
||||
// was either absent (genuinely new) or already
|
||||
// `pending_extraction` (a prior ingest crashed before extract
|
||||
// ran). Anything else — `admitted`, `buffered`, `sealed`,
|
||||
// `dropped` — is past the point of accepting new work, so
|
||||
// leave the lifecycle alone and skip the extract enqueue.
|
||||
let mut to_schedule: HashSet<String> = HashSet::new();
|
||||
for s in &staged_for_store {
|
||||
let pre = prior.get(&s.chunk.id).cloned().flatten();
|
||||
let needs_processing = matches!(
|
||||
pre.as_deref(),
|
||||
None | Some(store::CHUNK_STATUS_PENDING_EXTRACTION),
|
||||
);
|
||||
if needs_processing {
|
||||
store::set_chunk_lifecycle_status_tx(
|
||||
&tx,
|
||||
&s.chunk.id,
|
||||
store::CHUNK_STATUS_PENDING_EXTRACTION,
|
||||
)?;
|
||||
to_schedule.insert(s.chunk.id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
for (result, ts_ms) in &results_for_store {
|
||||
// Persist rationale for EVERY chunk (kept or dropped).
|
||||
// Index entities only for kept chunks (handled inside persist_score_tx).
|
||||
if !to_schedule.contains(&result.chunk_id) {
|
||||
// Chunk has already progressed past pending_extraction
|
||||
// on a prior ingest — skip score re-persist and don't
|
||||
// enqueue a duplicate extract job.
|
||||
continue;
|
||||
}
|
||||
score::persist_score_tx(&tx, result, *ts_ms, None)?;
|
||||
let extract = NewJob::extract_chunk(&ExtractChunkPayload {
|
||||
chunk_id: result.chunk_id.clone(),
|
||||
})?;
|
||||
let _ = jobs::enqueue_tx(&tx, &extract)?;
|
||||
}
|
||||
tx.commit()?;
|
||||
Ok(n)
|
||||
@@ -231,98 +205,26 @@ async fn persist(
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("persist join error: {e}"))??;
|
||||
|
||||
// 5. Source-tree append (Phase 3a #709). Each kept leaf pushes into
|
||||
// the tree's L0 buffer and cascades upward when token_sum crosses
|
||||
// the budget. Entities/topics from the scorer are threaded in so
|
||||
// sealed summaries inherit the child signal set. Failures here
|
||||
// log at warn level but don't fail the ingest — leaves are already
|
||||
// persisted, and a later flush/retry can still rebuild the tree.
|
||||
if let Err(e) = append_leaves_to_tree(config, source_id, &kept_chunks, &all_results).await {
|
||||
log::warn!(
|
||||
"[memory_tree::ingest] source_tree append failed source_id={} err={:#}",
|
||||
source_id,
|
||||
e
|
||||
);
|
||||
}
|
||||
jobs::wake_workers();
|
||||
|
||||
Ok(IngestResult {
|
||||
source_id: source_id.to_string(),
|
||||
chunks_written: written,
|
||||
chunks_dropped: dropped,
|
||||
chunk_ids: kept_chunks.iter().map(|c| c.id.clone()).collect(),
|
||||
chunk_ids: staged.iter().map(|s| s.chunk.id.clone()).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Push every kept chunk into its source tree. Scoped to Phase 3a — all
|
||||
/// chunks from one ingest batch share the same `source_id`, so they share
|
||||
/// one tree lookup. The summariser comes from `build_summariser(config)`:
|
||||
/// Ollama-backed [`LlmSummariser`] when `memory_tree.llm_summariser_*` is
|
||||
/// set, else deterministic [`InertSummariser`]. Both share the same trait,
|
||||
/// so the seal cascade doesn't notice the difference.
|
||||
async fn append_leaves_to_tree(
|
||||
config: &Config,
|
||||
source_id: &str,
|
||||
kept_chunks: &[Chunk],
|
||||
all_results: &[(ScoreResult, i64)],
|
||||
) -> Result<()> {
|
||||
if kept_chunks.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let tree = get_or_create_source_tree(config, source_id)?;
|
||||
let summariser = crate::openhuman::memory::tree::source_tree::build_summariser(config);
|
||||
|
||||
// Build a chunk_id → (score, entities, topics) map for quick lookup.
|
||||
use std::collections::HashMap;
|
||||
let mut score_by_id: HashMap<String, &ScoreResult> = HashMap::new();
|
||||
for (r, _) in all_results {
|
||||
score_by_id.insert(r.chunk_id.clone(), r);
|
||||
}
|
||||
|
||||
for chunk in kept_chunks {
|
||||
let (score_value, entities, topics) = match score_by_id.get(&chunk.id) {
|
||||
Some(r) => (
|
||||
r.total,
|
||||
r.canonical_entities
|
||||
.iter()
|
||||
.map(|e| e.canonical_id.clone())
|
||||
.collect::<Vec<_>>(),
|
||||
chunk.metadata.tags.clone(),
|
||||
),
|
||||
None => (0.0, Vec::new(), chunk.metadata.tags.clone()),
|
||||
};
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id.clone(),
|
||||
token_count: chunk.token_count,
|
||||
timestamp: chunk.metadata.timestamp,
|
||||
content: chunk.content.clone(),
|
||||
entities: entities.clone(),
|
||||
topics,
|
||||
score: score_value,
|
||||
};
|
||||
append_leaf(config, &tree, &leaf, summariser.as_ref()).await?;
|
||||
|
||||
// Phase 3c (#709): route the leaf to every matching topic tree
|
||||
// and tick the curator for each entity. Non-fatal on error —
|
||||
// the source-tree append has already succeeded above.
|
||||
if let Err(e) =
|
||||
route_leaf_to_topic_trees(config, &leaf, &entities, summariser.as_ref()).await
|
||||
{
|
||||
log::warn!(
|
||||
"[memory_tree::ingest] topic_tree routing failed chunk_id={} err={:#}",
|
||||
chunk.id,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::canonicalize::chat::ChatMessage;
|
||||
use crate::openhuman::memory::tree::jobs::drain_until_idle;
|
||||
use crate::openhuman::memory::tree::score::store::{count_scores, lookup_entity};
|
||||
use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery};
|
||||
use crate::openhuman::memory::tree::store::{
|
||||
count_chunks, count_chunks_by_lifecycle_status, get_chunk_embedding, list_chunks,
|
||||
ListChunksQuery, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED,
|
||||
};
|
||||
use crate::openhuman::memory::tree::types::SourceKind;
|
||||
use chrono::{TimeZone, Utc};
|
||||
use tempfile::TempDir;
|
||||
@@ -331,16 +233,12 @@ mod tests {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
// Phase 4 (#710): disable Ollama-backed embeddings in tests.
|
||||
// Ingest/seal call `build_embedder_from_config`; falling back to
|
||||
// inert (zero vectors) keeps tests deterministic and network-free.
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
/// Build a substantive batch that reliably passes the admission gate.
|
||||
fn substantive_batch() -> ChatBatch {
|
||||
ChatBatch {
|
||||
platform: "slack".into(),
|
||||
@@ -349,17 +247,14 @@ mod tests {
|
||||
ChatMessage {
|
||||
author: "alice".into(),
|
||||
timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
|
||||
text: "We are planning to ship the Phoenix migration on Friday \
|
||||
after reviewing the runbook and staging results. Please \
|
||||
confirm availability by replying here. alice@example.com"
|
||||
text: "We are planning to ship the Phoenix migration on Friday after reviewing the runbook and staging results. alice@example.com"
|
||||
.into(),
|
||||
source_ref: Some("slack://m1".into()),
|
||||
},
|
||||
ChatMessage {
|
||||
author: "bob".into(),
|
||||
timestamp: Utc.timestamp_millis_opt(1_700_000_010_000).unwrap(),
|
||||
text: "Confirmed — I'll handle the coordination and cut a release \
|
||||
candidate tonight. #launch-q2 will be tracked in Notion."
|
||||
text: "Confirmed, I will handle the coordination and launch tracking tonight."
|
||||
.into(),
|
||||
source_ref: None,
|
||||
},
|
||||
@@ -368,25 +263,40 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ingest_chat_writes_substantive_chunks() {
|
||||
async fn ingest_chat_writes_and_queue_drains_to_admitted_chunk() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch())
|
||||
.await
|
||||
.unwrap();
|
||||
// Greedy packing: both small messages fit under 10k token budget
|
||||
// and are packed into a single chunk.
|
||||
assert_eq!(out.chunks_written, 1);
|
||||
assert_eq!(out.chunks_dropped, 0);
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 1);
|
||||
// Score row persisted for the kept chunk
|
||||
assert_eq!(count_scores(&cfg).unwrap(), 1);
|
||||
// Entity index populated from regex extraction (alice@example.com + hashtag)
|
||||
let alice_hits = lookup_entity(&cfg, "email:alice@example.com", None).unwrap();
|
||||
assert_eq!(alice_hits.len(), 1);
|
||||
|
||||
drain_until_idle(&cfg).await.unwrap();
|
||||
|
||||
// Final lifecycle is `buffered`: extract → admitted → append_buffer → buffered.
|
||||
// The single packed chunk does not cross TOKEN_BUDGET so no seal fires.
|
||||
assert_eq!(
|
||||
count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_BUFFERED).unwrap(),
|
||||
1
|
||||
);
|
||||
assert!(count_scores(&cfg).unwrap() >= 1);
|
||||
assert_eq!(
|
||||
lookup_entity(&cfg, "email:alice@example.com", None)
|
||||
.unwrap()
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
|
||||
assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat);
|
||||
assert!(get_chunk_embedding(&cfg, &out.chunk_ids[0])
|
||||
.unwrap()
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn low_signal_chunks_are_dropped_but_score_persists() {
|
||||
async fn low_signal_chunks_end_up_dropped_after_queue_processing() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let batch = ChatBatch {
|
||||
platform: "slack".into(),
|
||||
@@ -394,18 +304,22 @@ mod tests {
|
||||
messages: vec![ChatMessage {
|
||||
author: "alice".into(),
|
||||
timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
|
||||
text: "+1".into(), // extremely low-signal
|
||||
text: "+1".into(),
|
||||
source_ref: None,
|
||||
}],
|
||||
};
|
||||
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(out.chunks_written, 0);
|
||||
assert_eq!(out.chunks_dropped, 1);
|
||||
// Chunk NOT in chunks table
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 0);
|
||||
// Score row IS persisted for diagnostics
|
||||
assert_eq!(out.chunks_written, 1);
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 1);
|
||||
|
||||
drain_until_idle(&cfg).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_DROPPED).unwrap(),
|
||||
1
|
||||
);
|
||||
assert_eq!(count_scores(&cfg).unwrap(), 1);
|
||||
}
|
||||
|
||||
@@ -421,7 +335,6 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(out.chunks_written, 0);
|
||||
assert_eq!(out.chunks_dropped, 0);
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 0);
|
||||
assert_eq!(count_scores(&cfg).unwrap(), 0);
|
||||
}
|
||||
@@ -432,9 +345,7 @@ mod tests {
|
||||
let doc = DocumentInput {
|
||||
provider: "notion".into(),
|
||||
title: "Launch plan".into(),
|
||||
body: "We are planning to ship Phoenix on Friday after review. \
|
||||
Coordination is via email and the launch thread tracks \
|
||||
the relevant decisions. alice@example.com owns this."
|
||||
body: "We are planning to ship Phoenix on Friday after review. alice@example.com owns this."
|
||||
.into(),
|
||||
modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
|
||||
source_ref: Some("notion://page/abc".into()),
|
||||
@@ -445,74 +356,8 @@ mod tests {
|
||||
ingest_document(&cfg, "notion:abc", "alice", vec![], doc)
|
||||
.await
|
||||
.unwrap();
|
||||
drain_until_idle(&cfg).await.unwrap();
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 1);
|
||||
assert_eq!(count_scores(&cfg).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunks_preserve_source_ref_when_kept() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let doc = DocumentInput {
|
||||
provider: "notion".into(),
|
||||
title: "t".into(),
|
||||
body: "Phoenix launch plan with enough substance to pass the admission \
|
||||
gate: we are reviewing the migration runbook alice@example.com \
|
||||
on Friday evening."
|
||||
.into(),
|
||||
modified_at: Utc::now(),
|
||||
source_ref: Some("notion://x".into()),
|
||||
};
|
||||
ingest_document(&cfg, "notion:x", "alice", vec![], doc)
|
||||
.await
|
||||
.unwrap();
|
||||
let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(
|
||||
rows[0].metadata.source_ref.as_ref().unwrap().value,
|
||||
"notion://x"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Phase 4 (#710) ──────────────────────────────────────────────
|
||||
|
||||
/// Ingesting with an inert embedder must still write embeddings for
|
||||
/// every kept chunk. Use a trimmed test_config so we go through the
|
||||
/// full embed → persist pipeline.
|
||||
#[tokio::test]
|
||||
async fn ingest_writes_chunk_embeddings() {
|
||||
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
|
||||
use crate::openhuman::memory::tree::store::get_chunk_embedding;
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.chunks_written >= 1);
|
||||
for id in &out.chunk_ids {
|
||||
let v = get_chunk_embedding(&cfg, id)
|
||||
.unwrap()
|
||||
.unwrap_or_else(|| panic!("missing embedding for {id}"));
|
||||
assert_eq!(v.len(), EMBEDDING_DIM);
|
||||
}
|
||||
}
|
||||
|
||||
/// When the embedder errors (here: Ollama required but unavailable),
|
||||
/// ingest must fail and persist nothing. Point embedding at a dead
|
||||
/// port so the HTTP call refuses, and flip `strict=true` so the
|
||||
/// factory returns Ollama (not the inert fallback).
|
||||
#[tokio::test]
|
||||
async fn ingest_fails_when_embedder_fails_and_persists_nothing() {
|
||||
let (_tmp, mut cfg) = test_config();
|
||||
cfg.memory_tree.embedding_endpoint = Some("http://127.0.0.1:1".into());
|
||||
cfg.memory_tree.embedding_model = Some("nomic-embed-text".into());
|
||||
cfg.memory_tree.embedding_timeout_ms = Some(500);
|
||||
cfg.memory_tree.embedding_strict = true;
|
||||
|
||||
let result = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()).await;
|
||||
assert!(result.is_err(), "expected ingest to bail when embed fails");
|
||||
|
||||
// No chunks, no scores persisted — retry stays clean.
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 0);
|
||||
assert_eq!(count_scores(&cfg).unwrap(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,861 @@
|
||||
use anyhow::{Context, Result};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::{
|
||||
self as content_store, read as content_read, tags as content_tags,
|
||||
};
|
||||
use crate::openhuman::memory::tree::global_tree::digest::{self, DigestOutcome};
|
||||
use crate::openhuman::memory::tree::jobs::store;
|
||||
use crate::openhuman::memory::tree::jobs::types::{
|
||||
AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload,
|
||||
Job, JobKind, NewJob, NodeRef, SealPayload, TopicRoutePayload,
|
||||
};
|
||||
use crate::openhuman::memory::tree::score;
|
||||
use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked};
|
||||
use crate::openhuman::memory::tree::score::extract::build_summary_extractor;
|
||||
use crate::openhuman::memory::tree::score::store as score_store;
|
||||
use crate::openhuman::memory::tree::source_tree::{
|
||||
build_summariser, get_or_create_source_tree, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::store as chunk_store;
|
||||
use crate::openhuman::memory::tree::topic_tree::curator;
|
||||
|
||||
pub async fn handle_job(config: &Config, job: &Job) -> Result<()> {
|
||||
match job.kind {
|
||||
JobKind::ExtractChunk => handle_extract(config, job).await,
|
||||
JobKind::AppendBuffer => handle_append_buffer(config, job).await,
|
||||
JobKind::Seal => handle_seal(config, job).await,
|
||||
JobKind::TopicRoute => handle_topic_route(config, job).await,
|
||||
JobKind::DigestDaily => handle_digest_daily(config, job).await,
|
||||
JobKind::FlushStale => handle_flush_stale(config, job).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_extract(config: &Config, job: &Job) -> Result<()> {
|
||||
let payload: ExtractChunkPayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse ExtractChunk payload")?;
|
||||
let Some(chunk) = chunk_store::get_chunk(config, &payload.chunk_id)? else {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] extract chunk missing chunk_id={}",
|
||||
payload.chunk_id
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
// Read the full body from disk (the `content` column in SQLite holds a
|
||||
// ≤500-char preview after the MD-on-disk migration). Both the scorer and
|
||||
// the embedder need the complete text so extraction and semantic indexing
|
||||
// operate over the full chunk body, not a truncated preview.
|
||||
let body = content_read::read_chunk_body(config, &chunk.id)
|
||||
.with_context(|| format!("read full body for extract chunk_id={}", chunk.id))?;
|
||||
// Score a clone of the chunk with the full body swapped in.
|
||||
let chunk_with_body = {
|
||||
let mut c = chunk.clone();
|
||||
c.content = body.clone();
|
||||
c
|
||||
};
|
||||
|
||||
let scoring_cfg = score::ScoringConfig::from_config(config);
|
||||
let result = score::score_chunk(&chunk_with_body, &scoring_cfg).await?;
|
||||
let packed_embedding = if result.kept {
|
||||
let embedder =
|
||||
build_embedder_from_config(config).context("build embedder in extract handler")?;
|
||||
// Reuse the body already read — avoid a second disk read.
|
||||
let vector = embedder
|
||||
.embed(&body)
|
||||
.await
|
||||
.with_context(|| format!("embed chunk_id={} in extract handler", chunk.id))?;
|
||||
Some(
|
||||
pack_checked(&vector)
|
||||
.with_context(|| format!("pack embedding for chunk_id={}", chunk.id))?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Build follow-up job payloads before opening the tx — construction is
|
||||
// cheap and doesn't require a database connection. The two jobs are
|
||||
// enqueued inside the SAME transaction that commits the lifecycle update,
|
||||
// so a crash anywhere rolls everything back together and prevents the
|
||||
// "lifecycle committed but job lost" crash window.
|
||||
let source_job = if result.kept {
|
||||
Some(NewJob::append_buffer(&AppendBufferPayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: chunk.id.clone(),
|
||||
},
|
||||
target: AppendTarget::Source {
|
||||
source_id: chunk.metadata.source_id.clone(),
|
||||
},
|
||||
})?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let route_job = if result.kept {
|
||||
Some(NewJob::topic_route(&TopicRoutePayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: chunk.id.clone(),
|
||||
},
|
||||
})?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let (did_enqueue_source, did_enqueue_route) = chunk_store::with_connection(config, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
score::persist_score_tx(
|
||||
&tx,
|
||||
&result,
|
||||
chunk.metadata.timestamp.timestamp_millis(),
|
||||
None,
|
||||
)?;
|
||||
|
||||
if result.kept {
|
||||
tx.execute(
|
||||
"UPDATE mem_tree_chunks
|
||||
SET embedding = ?1,
|
||||
lifecycle_status = ?2
|
||||
WHERE id = ?3",
|
||||
rusqlite::params![
|
||||
packed_embedding,
|
||||
chunk_store::CHUNK_STATUS_ADMITTED,
|
||||
chunk.id,
|
||||
],
|
||||
)?;
|
||||
} else {
|
||||
tx.execute(
|
||||
"UPDATE mem_tree_chunks
|
||||
SET lifecycle_status = ?1
|
||||
WHERE id = ?2",
|
||||
rusqlite::params![chunk_store::CHUNK_STATUS_DROPPED, chunk.id],
|
||||
)?;
|
||||
}
|
||||
|
||||
// Enqueue follow-up jobs inside the SAME transaction so they are
|
||||
// atomically visible with the lifecycle update.
|
||||
let mut eq_src = false;
|
||||
let mut eq_route = false;
|
||||
if let Some(ref j) = source_job {
|
||||
eq_src = store::enqueue_tx(&tx, j)?.is_some();
|
||||
}
|
||||
if let Some(ref j) = route_job {
|
||||
eq_route = store::enqueue_tx(&tx, j)?.is_some();
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok((eq_src, eq_route))
|
||||
})?;
|
||||
|
||||
// Phase MD-content: rewrite the `tags:` block in the on-disk chunk file
|
||||
// with Obsidian-style hierarchical tags derived from the extracted entities.
|
||||
// This runs after the tx commits so the entity index is visible to readers.
|
||||
// It is a filesystem op and therefore lives outside the SQL tx — best-effort.
|
||||
if result.kept {
|
||||
if let Some(content_path) = chunk_store::get_chunk_content_path(config, &chunk.id)? {
|
||||
let content_root = config.memory_tree_content_root();
|
||||
let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?;
|
||||
let obsidian_tags: Vec<String> = entity_ids
|
||||
.iter()
|
||||
.filter_map(|eid| {
|
||||
// entity_id format: "kind:surface"
|
||||
let (kind, surface) = eid.split_once(':')?;
|
||||
Some(content_tags::entity_tag(kind, surface))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Build the absolute path from the stored relative path.
|
||||
let abs_path = {
|
||||
let mut p = content_root.clone();
|
||||
for component in content_path.split('/') {
|
||||
p.push(component);
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
if let Err(e) = content_tags::update_chunk_tags(&abs_path, &obsidian_tags) {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] failed to update tags in chunk file chunk_id={} path_hash={}: {e}",
|
||||
chunk.id,
|
||||
crate::openhuman::memory::tree::util::redact::redact(&content_path),
|
||||
);
|
||||
// Non-fatal: tag rewrite failure does not block the pipeline.
|
||||
} else {
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] updated {} obsidian tags in chunk file chunk_id={}",
|
||||
obsidian_tags.len(),
|
||||
chunk.id,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Signal workers after the tx commits (no atomicity requirement on signaling).
|
||||
if did_enqueue_source {
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
if did_enqueue_route {
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_append_buffer(config: &Config, job: &Job) -> Result<()> {
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::should_seal;
|
||||
use crate::openhuman::memory::tree::source_tree::store as src_store;
|
||||
|
||||
let payload: AppendBufferPayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse AppendBuffer payload")?;
|
||||
|
||||
// Hydrate the leaf-shaped record from either a chunk row or a summary
|
||||
// row. The downstream buffer-push doesn't care which kind produced
|
||||
// the LeafRef.
|
||||
let (leaf, chunk_id_for_lifecycle): (LeafRef, Option<String>) = match &payload.node {
|
||||
NodeRef::Leaf { chunk_id } => {
|
||||
let Some(chunk) = chunk_store::get_chunk(config, chunk_id)? else {
|
||||
log::warn!("[memory_tree::jobs] append_buffer chunk missing chunk_id={chunk_id}");
|
||||
return Ok(());
|
||||
};
|
||||
let score_row = score_store::get_score(config, &chunk.id)?
|
||||
.ok_or_else(|| anyhow::anyhow!("missing score row for chunk {}", chunk.id))?;
|
||||
let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?;
|
||||
// Read the full body from disk — the `content` column in SQLite
|
||||
// is a ≤500-char preview after the MD-on-disk migration. The
|
||||
// summariser receives this LeafRef and must see the complete text.
|
||||
let body = content_read::read_chunk_body(config, chunk_id)
|
||||
.with_context(|| format!("read chunk body in append_buffer chunk_id={chunk_id}"))?;
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id.clone(),
|
||||
token_count: chunk.token_count,
|
||||
timestamp: chunk.metadata.timestamp,
|
||||
content: body,
|
||||
entities: entity_ids,
|
||||
topics: chunk.metadata.tags.clone(),
|
||||
score: score_row.total,
|
||||
};
|
||||
(leaf, Some(chunk.id))
|
||||
}
|
||||
NodeRef::Summary { summary_id } => {
|
||||
let Some(summary) = src_store::get_summary(config, summary_id)? else {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] append_buffer summary missing summary_id={summary_id}"
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
// Read the full body from disk — `summary.content` is a ≤500-char
|
||||
// preview after the MD-on-disk migration. The summariser receives
|
||||
// this LeafRef when sealing higher-level nodes and must see the
|
||||
// complete summary text.
|
||||
let body = content_read::read_summary_body(config, summary_id).with_context(|| {
|
||||
format!("read summary body in append_buffer summary_id={summary_id}")
|
||||
})?;
|
||||
// Build a LeafRef from the summary's already-populated fields.
|
||||
// `chunk_id` carries the source-node id (any string); buffer
|
||||
// accounting uses it as the item id only.
|
||||
let leaf = LeafRef {
|
||||
chunk_id: summary.id.clone(),
|
||||
token_count: summary.token_count,
|
||||
timestamp: summary.time_range_start,
|
||||
content: body,
|
||||
entities: summary.entities.clone(),
|
||||
topics: summary.topics.clone(),
|
||||
score: summary.score,
|
||||
};
|
||||
(leaf, None) // summaries have no chunk lifecycle to update
|
||||
}
|
||||
};
|
||||
|
||||
// Resolve target tree (no tx open yet — this can create a row).
|
||||
let tree = match &payload.target {
|
||||
AppendTarget::Source { source_id } => Some(get_or_create_source_tree(config, source_id)?),
|
||||
AppendTarget::Topic { tree_id } => src_store::get_tree(config, tree_id)?,
|
||||
};
|
||||
let Some(tree) = tree else {
|
||||
// Target topic tree doesn't exist (e.g. archived between
|
||||
// topic_route and this append). Drop on the floor — the
|
||||
// topic_route was advisory and the source-tree path already
|
||||
// ran for this leaf.
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let is_source_target = matches!(payload.target, AppendTarget::Source { .. });
|
||||
let leaf_for_tx = leaf.clone();
|
||||
let tree_for_tx = tree.clone();
|
||||
let lifecycle_chunk_id = chunk_id_for_lifecycle.clone();
|
||||
|
||||
// ATOMIC: buffer push + seal enqueue (if gate met) + lifecycle update
|
||||
// happen in a single SQLite transaction. Eliminates the crash window
|
||||
// where the buffer commits but the seal job is lost — which can
|
||||
// duplicate the leaf into two summaries on retry-after-seal-cleared.
|
||||
let did_enqueue_seal = chunk_store::with_connection(config, move |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
|
||||
// 1. Push leaf into L0 buffer (idempotent on (tree, level, item_id)).
|
||||
let mut buf = src_store::get_buffer_conn(&tx, &tree_for_tx.id, 0)?;
|
||||
if !buf.item_ids.iter().any(|x| x == &leaf_for_tx.chunk_id) {
|
||||
buf.item_ids.push(leaf_for_tx.chunk_id.clone());
|
||||
buf.token_sum = buf.token_sum.saturating_add(leaf_for_tx.token_count as i64);
|
||||
buf.oldest_at = match buf.oldest_at {
|
||||
Some(existing) => Some(existing.min(leaf_for_tx.timestamp)),
|
||||
None => Some(leaf_for_tx.timestamp),
|
||||
};
|
||||
src_store::upsert_buffer_tx(&tx, &buf)?;
|
||||
}
|
||||
|
||||
// 2. If the gate is met, enqueue a seal job atomically.
|
||||
let did_enqueue = if should_seal(&buf) {
|
||||
let seal = SealPayload {
|
||||
tree_id: tree_for_tx.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
store::enqueue_tx(&tx, &NewJob::seal(&seal)?)?.is_some()
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
// 3. Lifecycle transition (Source target with a leaf chunk).
|
||||
// Last step in the tx — its presence is the "this handler
|
||||
// finished" marker. Same tx as the push + seal-enqueue, so a
|
||||
// crash anywhere rolls everything back together.
|
||||
if is_source_target {
|
||||
if let Some(chunk_id) = lifecycle_chunk_id.as_deref() {
|
||||
chunk_store::set_chunk_lifecycle_status_tx(
|
||||
&tx,
|
||||
chunk_id,
|
||||
chunk_store::CHUNK_STATUS_BUFFERED,
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(did_enqueue)
|
||||
})?;
|
||||
|
||||
if did_enqueue_seal {
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_seal(config: &Config, job: &Job) -> Result<()> {
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{seal_one_level, should_seal};
|
||||
use crate::openhuman::memory::tree::source_tree::store as src_store;
|
||||
use crate::openhuman::memory::tree::source_tree::types::TreeKind;
|
||||
|
||||
let payload: SealPayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse Seal payload")?;
|
||||
let Some(tree) = src_store::get_tree(config, &payload.tree_id)? else {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] seal tree missing tree_id={}",
|
||||
payload.tree_id
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
// Seal exactly one level. Parents only get sealed via a follow-up job
|
||||
// so each level is its own crash-recovery checkpoint and each LLM
|
||||
// summariser call competes for a fresh slot from the global semaphore.
|
||||
let buf = src_store::get_buffer(config, &tree.id, payload.level)?;
|
||||
let forced = payload.force_now_ms.is_some();
|
||||
if buf.is_empty() {
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] seal skipped — empty buffer tree_id={} level={}",
|
||||
tree.id,
|
||||
payload.level
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
if !forced && !should_seal(&buf) {
|
||||
// Another job sealed this level out from under us (or the buffer
|
||||
// hasn't crossed the gate yet); idempotent no-op.
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] seal gate not met tree_id={} level={} token_sum={}",
|
||||
tree.id,
|
||||
payload.level,
|
||||
buf.token_sum
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Pick the labeling strategy for this tree kind. Source trees mint
|
||||
// emergent themes via the seal-time extractor; topic trees stay empty
|
||||
// by design (scope already pins the canonical id). Global trees never
|
||||
// reach here — `digest_daily` handles them — but Empty is a safe
|
||||
// defensive default.
|
||||
let strategy = match tree.kind {
|
||||
TreeKind::Source => LabelStrategy::ExtractFromContent(build_summary_extractor(config)),
|
||||
TreeKind::Topic => LabelStrategy::Empty,
|
||||
TreeKind::Global => LabelStrategy::Empty,
|
||||
};
|
||||
|
||||
let summariser = build_summariser(config);
|
||||
// `seal_one_level` with `enqueue_follow_ups: true` atomically inserts
|
||||
// the parent-cascade seal (if the parent buffer now meets its gate)
|
||||
// and the summary-side `topic_route` (for source trees) inside the
|
||||
// same SQLite transaction that commits the seal. This eliminates the
|
||||
// crash window where the seal succeeds but the follow-up enqueues
|
||||
// are silently lost.
|
||||
let summary_id =
|
||||
seal_one_level(config, &tree, &buf, summariser.as_ref(), &strategy, true).await?;
|
||||
|
||||
// Phase MD-content: rewrite the `tags:` block in the sealed summary's
|
||||
// on-disk .md file. Entity index rows were committed inside
|
||||
// `seal_one_level` (via `index_summary_entity_ids_tx`), so they are
|
||||
// visible here. Best-effort: failure does not abort the seal.
|
||||
if let Err(e) = content_store::update_summary_tags(config, &summary_id) {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] update_summary_tags failed for summary_id={summary_id}: {e:#}"
|
||||
);
|
||||
}
|
||||
|
||||
super::worker::wake_workers();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_topic_route(config: &Config, job: &Job) -> Result<()> {
|
||||
let payload: TopicRoutePayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse TopicRoute payload")?;
|
||||
|
||||
// Resolve the source node id and verify it exists. `mem_tree_entity_index`
|
||||
// already indexes both chunks and summaries via `node_kind`, so the
|
||||
// canonical-id loop below is identical for either case.
|
||||
let node_id: String = match &payload.node {
|
||||
NodeRef::Leaf { chunk_id } => {
|
||||
if chunk_store::get_chunk(config, chunk_id)?.is_none() {
|
||||
log::warn!("[memory_tree::jobs] topic_route chunk missing chunk_id={chunk_id}");
|
||||
return Ok(());
|
||||
}
|
||||
chunk_id.clone()
|
||||
}
|
||||
NodeRef::Summary { summary_id } => {
|
||||
if crate::openhuman::memory::tree::source_tree::store::get_summary(config, summary_id)?
|
||||
.is_none()
|
||||
{
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] topic_route summary missing summary_id={summary_id}"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
summary_id.clone()
|
||||
}
|
||||
};
|
||||
|
||||
let entity_ids = score_store::list_entity_ids_for_node(config, &node_id)?;
|
||||
if entity_ids.is_empty() {
|
||||
log::debug!("[memory_tree::jobs] topic_route no entities for node_id={node_id} — skipping");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let summariser = build_summariser(config);
|
||||
for entity_id in entity_ids {
|
||||
let _ = curator::maybe_spawn_topic_tree(config, &entity_id, summariser.as_ref()).await?;
|
||||
if let Some(tree) = crate::openhuman::memory::tree::source_tree::store::get_tree_by_scope(
|
||||
config,
|
||||
crate::openhuman::memory::tree::source_tree::types::TreeKind::Topic,
|
||||
&entity_id,
|
||||
)? {
|
||||
let job = NewJob::append_buffer(&AppendBufferPayload {
|
||||
node: payload.node.clone(),
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: tree.id.clone(),
|
||||
},
|
||||
})?;
|
||||
if store::enqueue(config, &job)?.is_some() {
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_digest_daily(config: &Config, job: &Job) -> Result<()> {
|
||||
let payload: DigestDailyPayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse DigestDaily payload")?;
|
||||
let day = chrono::NaiveDate::parse_from_str(&payload.date_iso, "%Y-%m-%d")
|
||||
.with_context(|| format!("invalid digest date {}", payload.date_iso))?;
|
||||
let summariser = build_summariser(config);
|
||||
match digest::end_of_day_digest(config, day, summariser.as_ref()).await? {
|
||||
DigestOutcome::Emitted { daily_id, .. } => {
|
||||
log::info!("[memory_tree::jobs] emitted digest daily_id={daily_id}");
|
||||
}
|
||||
DigestOutcome::EmptyDay => {}
|
||||
DigestOutcome::Skipped { existing_id } => {
|
||||
log::debug!("[memory_tree::jobs] digest skipped existing_id={existing_id}");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_flush_stale(config: &Config, job: &Job) -> Result<()> {
|
||||
let payload: FlushStalePayload =
|
||||
serde_json::from_str(&job.payload_json).context("parse FlushStale payload")?;
|
||||
let age_secs = payload
|
||||
.max_age_secs
|
||||
.unwrap_or(crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS);
|
||||
let cutoff = chrono::Utc::now() - chrono::Duration::seconds(age_secs);
|
||||
let buffers =
|
||||
crate::openhuman::memory::tree::source_tree::store::list_stale_buffers(config, cutoff)?;
|
||||
for buf in buffers {
|
||||
let seal = SealPayload {
|
||||
tree_id: buf.tree_id.clone(),
|
||||
level: buf.level,
|
||||
force_now_ms: Some(chrono::Utc::now().timestamp_millis()),
|
||||
};
|
||||
if store::enqueue(config, &NewJob::seal(&seal)?)?.is_some() {
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::jobs::store::{count_by_status, count_total};
|
||||
use crate::openhuman::memory::tree::jobs::types::JobStatus;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf_deferred, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::store as src_store;
|
||||
use crate::openhuman::memory::tree::store::with_connection;
|
||||
use chrono::TimeZone;
|
||||
use rusqlite::params;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
/// Build a minimal `Job` row for direct handler invocation. Mirrors
|
||||
/// what `claim_next` would produce for a freshly-claimed row.
|
||||
fn mk_running_job(kind: JobKind, payload_json: String) -> Job {
|
||||
let now_ms = chrono::Utc::now().timestamp_millis();
|
||||
Job {
|
||||
id: "test-job-id".into(),
|
||||
kind,
|
||||
payload_json,
|
||||
dedupe_key: None,
|
||||
status: JobStatus::Running,
|
||||
attempts: 1,
|
||||
max_attempts: 5,
|
||||
available_at_ms: now_ms,
|
||||
locked_until_ms: Some(now_ms + 60_000),
|
||||
last_error: None,
|
||||
created_at_ms: now_ms,
|
||||
started_at_ms: Some(now_ms),
|
||||
completed_at_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Count rows in `mem_tree_jobs` matching a specific kind.
|
||||
fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 {
|
||||
with_connection(cfg, |conn| {
|
||||
let n: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1",
|
||||
params![kind],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Seed a source tree and push enough labeled leaves into its L0 buffer
|
||||
/// to cross `TOKEN_BUDGET`, returning the tree. The caller can then
|
||||
/// fire `handle_seal` and inspect the result.
|
||||
async fn seed_source_tree_ready_to_seal(
|
||||
cfg: &Config,
|
||||
) -> crate::openhuman::memory::tree::source_tree::types::Tree {
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"),
|
||||
content: "alice@example.com leading the rollout".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
// Bust budget so the L0 buffer is "ready" for seal.
|
||||
token_count: 10_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body via
|
||||
// `read_chunk_body` when `handle_seal` fires and calls `seal_one_level`.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 10_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
// append_leaf_deferred only buffers; doesn't seal. handle_seal will.
|
||||
let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap();
|
||||
tree
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn source_tree_seal_handler_enqueues_summary_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = seed_source_tree_ready_to_seal(&cfg).await;
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
// Pre-condition: queue has no topic_route jobs.
|
||||
assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0);
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Post-condition: source-tree seal must enqueue exactly one
|
||||
// topic_route job carrying NodeRef::Summary { summary_id: <new> }.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
1,
|
||||
"source-tree seal must enqueue summary-side topic_route"
|
||||
);
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
|
||||
|
||||
// Inspect the enqueued payload to confirm it's a Summary variant.
|
||||
let payload_json: String = with_connection(&cfg, |conn| {
|
||||
let s: String = conn
|
||||
.query_row(
|
||||
"SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
Ok(s)
|
||||
})
|
||||
.unwrap();
|
||||
let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap();
|
||||
match p.node {
|
||||
NodeRef::Summary { summary_id } => {
|
||||
assert!(summary_id.starts_with("summary:L1:"));
|
||||
}
|
||||
other => panic!("expected NodeRef::Summary, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topic_tree_seal_handler_does_not_enqueue_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
// Spawn a topic tree directly via the registry (skipping curator's
|
||||
// hotness gate — we just need a TreeKind::Topic with leaves).
|
||||
let topic_tree =
|
||||
crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"topic:phoenix-migration",
|
||||
)
|
||||
.unwrap();
|
||||
// Push a single 10k-token leaf so L0 is gate-ready.
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"),
|
||||
content: "topic content".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 10_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// when `handle_seal` fires.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 10_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap();
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Topic-tree seals are sinks: must not enqueue any topic_route.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
0,
|
||||
"topic-tree seal must NOT enqueue topic_route (trees are sinks)"
|
||||
);
|
||||
// The seal itself should still have produced a summary node.
|
||||
assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
|
||||
// 1. Create a target topic tree with a clean L0 buffer.
|
||||
let topic_tree =
|
||||
crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"email:alice@example.com",
|
||||
)
|
||||
.unwrap();
|
||||
let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert!(l0_before.is_empty());
|
||||
|
||||
// 2. Manually insert a summary node we can route. The simplest way
|
||||
// is to create a separate source tree, push two 6k leaves into
|
||||
// it, and let the seal produce a summary we can address.
|
||||
let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::seal_one_level;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
for seq in 0..2 {
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"),
|
||||
content: format!("source content {seq}"),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 6_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// during `seal_one_level`.
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 6_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap();
|
||||
}
|
||||
// Force-seal the source tree's L0 to mint the summary.
|
||||
let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap();
|
||||
let summariser = build_summariser(&cfg);
|
||||
let summary_id = seal_one_level(
|
||||
&cfg,
|
||||
&source_tree,
|
||||
&buf,
|
||||
summariser.as_ref(),
|
||||
&crate::openhuman::memory::tree::source_tree::bucket_seal::LabelStrategy::Empty,
|
||||
// No follow-up enqueues — the test scopes assertions to the
|
||||
// append_buffer handler, not seal-side fan-out.
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// 3. Build an append_buffer payload routing the summary into the
|
||||
// topic tree.
|
||||
let payload = AppendBufferPayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: summary_id.clone(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
},
|
||||
};
|
||||
let job = mk_running_job(
|
||||
JobKind::AppendBuffer,
|
||||
serde_json::to_string(&payload).unwrap(),
|
||||
);
|
||||
|
||||
// Clear out any pending append_buffer jobs minted upstream so the
|
||||
// post-condition assertion below is unambiguous.
|
||||
let pre = count_total(&cfg).unwrap();
|
||||
|
||||
super::handle_append_buffer(&cfg, &job).await.unwrap();
|
||||
|
||||
// 4. Topic tree's L0 buffer should now hold the summary id.
|
||||
let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert_eq!(l0_after.item_ids, vec![summary_id]);
|
||||
assert!(l0_after.token_sum > 0);
|
||||
|
||||
// No new jobs should have been enqueued (buffer didn't cross gate).
|
||||
assert_eq!(count_total(&cfg).unwrap(), pre);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
//! Async job pipeline for memory-tree work.
|
||||
//!
|
||||
//! Replaces the previous synchronous `append_leaf → cascade_seal → LLM
|
||||
//! summarise` chain on the ingest hot path with a SQLite-backed job queue
|
||||
//! and a worker pool. The shape is:
|
||||
//!
|
||||
//! ```text
|
||||
//! ingest::persist
|
||||
//! └── writes chunk row (lifecycle = pending_extraction)
|
||||
//! enqueues `extract_chunk`
|
||||
//!
|
||||
//! worker pool (3 tasks) ──► claims jobs by kind:
|
||||
//! extract_chunk → LLM extraction → admission decision → enqueue append_buffer
|
||||
//! append_buffer → push to L0 → enqueue seal if gate met → enqueue topic_route
|
||||
//! seal → seal one level → enqueue parent seal if cascading
|
||||
//! topic_route → match topics → enqueue per-topic append_buffer
|
||||
//! digest_daily → call global_tree::digest::end_of_day_digest
|
||||
//! flush_stale → enqueue seals for time-stale buffers
|
||||
//!
|
||||
//! scheduler (1 task) ──► daily wall-clock tick:
|
||||
//! enqueues digest_daily(yesterday) + flush_stale(today)
|
||||
//! ```
|
||||
//!
|
||||
//! All persistence lives in the same `chunks.db` as `mem_tree_chunks` so a
|
||||
//! producer can insert its side-effect and its follow-up job in one tx.
|
||||
//! See [`store::enqueue_tx`] for the in-tx producer entry point.
|
||||
|
||||
mod handlers;
|
||||
pub mod scheduler;
|
||||
pub mod store;
|
||||
pub mod testing;
|
||||
pub mod types;
|
||||
mod worker;
|
||||
|
||||
pub use scheduler::{backfill_missing_digests, trigger_digest};
|
||||
pub use store::{
|
||||
claim_next, count_by_status, count_total, enqueue, enqueue_tx, get_job, mark_done, mark_failed,
|
||||
recover_stale_locks, DEFAULT_LOCK_DURATION_MS,
|
||||
};
|
||||
pub use testing::drain_until_idle;
|
||||
pub use types::{
|
||||
AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload,
|
||||
Job, JobKind, JobStatus, NewJob, NodeRef, SealPayload, TopicRoutePayload,
|
||||
};
|
||||
pub use worker::{start, wake_workers};
|
||||
@@ -0,0 +1,216 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use chrono::{Datelike, Duration as ChronoDuration, NaiveDate, TimeZone, Utc};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::jobs::store;
|
||||
use crate::openhuman::memory::tree::jobs::types::{DigestDailyPayload, FlushStalePayload, NewJob};
|
||||
|
||||
static STARTED: std::sync::Once = std::sync::Once::new();
|
||||
|
||||
/// Start the daily wall-clock scheduler. Takes the full `Config` so the
|
||||
/// digest enqueues match the same workspace + LLM settings the workers
|
||||
/// see — not `Config::default()`.
|
||||
pub fn start(config: Config) {
|
||||
STARTED.call_once(|| {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if let Err(err) = enqueue_daily_jobs(&config) {
|
||||
log::warn!("[memory_tree::jobs] scheduler enqueue failed: {err:#}");
|
||||
}
|
||||
tokio::time::sleep(next_sleep_duration()).await;
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
fn enqueue_daily_jobs(config: &Config) -> anyhow::Result<()> {
|
||||
let now = Utc::now();
|
||||
let yesterday = now.date_naive() - ChronoDuration::days(1);
|
||||
let date_iso = yesterday.format("%Y-%m-%d").to_string();
|
||||
|
||||
if store::enqueue(
|
||||
config,
|
||||
&NewJob::digest_daily(&DigestDailyPayload {
|
||||
date_iso: date_iso.clone(),
|
||||
})?,
|
||||
)?
|
||||
.is_some()
|
||||
{
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
|
||||
let today_iso = now.date_naive().format("%Y-%m-%d").to_string();
|
||||
if store::enqueue(
|
||||
config,
|
||||
&NewJob::flush_stale(&FlushStalePayload::default(), &today_iso)?,
|
||||
)?
|
||||
.is_some()
|
||||
{
|
||||
super::worker::wake_workers();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Manually enqueue a `digest_daily` job for `date`. Idempotent — if a
|
||||
/// digest already ran for that day, the handler's `find_existing_daily`
|
||||
/// check will return `Skipped` without doing any work; if a job for the
|
||||
/// same date is already queued or running, the partial unique index on
|
||||
/// `dedupe_key` suppresses the duplicate.
|
||||
///
|
||||
/// Useful for catch-up after the process was down across midnight, or
|
||||
/// to force a re-run for testing / debugging.
|
||||
pub fn trigger_digest(config: &Config, date: NaiveDate) -> Result<Option<String>> {
|
||||
let payload = DigestDailyPayload {
|
||||
date_iso: date.format("%Y-%m-%d").to_string(),
|
||||
};
|
||||
let job_id = store::enqueue(config, &NewJob::digest_daily(&payload)?)?;
|
||||
if job_id.is_some() {
|
||||
log::info!(
|
||||
"[memory_tree::jobs] manual digest trigger enqueued date={} id={:?}",
|
||||
payload.date_iso,
|
||||
job_id.as_deref()
|
||||
);
|
||||
super::worker::wake_workers();
|
||||
} else {
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] manual digest trigger dedupe-suppressed date={} \
|
||||
(an active job for this date already exists)",
|
||||
payload.date_iso
|
||||
);
|
||||
}
|
||||
Ok(job_id)
|
||||
}
|
||||
|
||||
/// Enqueue `digest_daily` jobs for the last `days_back` calendar days
|
||||
/// (excluding today). Catch-up helper for cases where the scheduler
|
||||
/// missed days because the process was down.
|
||||
///
|
||||
/// Returns the number of jobs newly enqueued. Days that already have a
|
||||
/// completed digest are still re-enqueued — the handler is idempotent
|
||||
/// and skips them — so this is safe to call repeatedly.
|
||||
pub fn backfill_missing_digests(config: &Config, days_back: i64) -> Result<usize> {
|
||||
if days_back <= 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let today = Utc::now().date_naive();
|
||||
let mut enqueued = 0usize;
|
||||
for offset in 1..=days_back {
|
||||
let date = today - ChronoDuration::days(offset);
|
||||
if trigger_digest(config, date)?.is_some() {
|
||||
enqueued += 1;
|
||||
}
|
||||
}
|
||||
log::info!(
|
||||
"[memory_tree::jobs] backfill_missing_digests window={}d enqueued={}",
|
||||
days_back,
|
||||
enqueued
|
||||
);
|
||||
Ok(enqueued)
|
||||
}
|
||||
|
||||
fn next_sleep_duration() -> Duration {
|
||||
let now = Utc::now();
|
||||
let tomorrow = now.date_naive() + ChronoDuration::days(1);
|
||||
let next = Utc
|
||||
.with_ymd_and_hms(tomorrow.year(), tomorrow.month(), tomorrow.day(), 0, 5, 0)
|
||||
// UTC has no DST gaps/overlaps, so `single()` always returns
|
||||
// `Some` for any valid (Y, M, D, h, m, s). Fallback retained
|
||||
// only as a defensive belt-and-braces against future API churn.
|
||||
.single()
|
||||
.unwrap_or_else(|| now + ChronoDuration::hours(24));
|
||||
(next - now)
|
||||
.to_std()
|
||||
.unwrap_or_else(|_| Duration::from_secs(60))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::jobs::store::{
|
||||
claim_next, count_by_status, count_total, mark_done, DEFAULT_LOCK_DURATION_MS,
|
||||
};
|
||||
use crate::openhuman::memory::tree::jobs::types::JobStatus;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trigger_digest_enqueues_a_job() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
|
||||
let id = trigger_digest(&cfg, date).unwrap();
|
||||
assert!(id.is_some(), "first trigger must enqueue");
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trigger_digest_dedupes_active_jobs() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
|
||||
let first = trigger_digest(&cfg, date).unwrap();
|
||||
let second = trigger_digest(&cfg, date).unwrap();
|
||||
assert!(first.is_some());
|
||||
assert!(
|
||||
second.is_none(),
|
||||
"duplicate trigger must be dedupe-suppressed while active"
|
||||
);
|
||||
assert_eq!(count_total(&cfg).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trigger_digest_after_done_creates_fresh_row() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
|
||||
let id1 = trigger_digest(&cfg, date).unwrap().unwrap();
|
||||
// Simulate a worker finishing the job — claim it first so we have a
|
||||
// Job snapshot for the claim-token-gated mark_done.
|
||||
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(claimed.id, id1);
|
||||
mark_done(&cfg, &claimed).unwrap();
|
||||
|
||||
let id2 = trigger_digest(&cfg, date).unwrap();
|
||||
assert!(
|
||||
id2.is_some(),
|
||||
"after the prior job completes, a fresh trigger must enqueue"
|
||||
);
|
||||
assert_ne!(id2.unwrap(), id1);
|
||||
assert_eq!(count_total(&cfg).unwrap(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backfill_missing_digests_enqueues_one_per_day() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let n = backfill_missing_digests(&cfg, 5).unwrap();
|
||||
assert_eq!(n, 5, "expected one job per day in the 5-day window");
|
||||
assert_eq!(count_total(&cfg).unwrap(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backfill_missing_digests_zero_window_is_noop() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let n = backfill_missing_digests(&cfg, 0).unwrap();
|
||||
assert_eq!(n, 0);
|
||||
assert_eq!(count_total(&cfg).unwrap(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backfill_missing_digests_is_idempotent_while_active() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let n1 = backfill_missing_digests(&cfg, 3).unwrap();
|
||||
let n2 = backfill_missing_digests(&cfg, 3).unwrap();
|
||||
assert_eq!(n1, 3);
|
||||
assert_eq!(n2, 0, "second call must be fully dedupe-suppressed");
|
||||
assert_eq!(count_total(&cfg).unwrap(), 3);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,618 @@
|
||||
//! SQLite persistence for the memory-tree job queue.
|
||||
//!
|
||||
//! Producers call [`enqueue`] inside their own writes (or with a fresh tx)
|
||||
//! to atomically commit the side-effect plus its follow-up job. The worker
|
||||
//! pool calls [`claim_next`] to lease a job, [`mark_done`] / [`mark_failed`]
|
||||
//! to settle it, and [`recover_stale_locks`] on startup to flip rows whose
|
||||
//! `locked_until_ms` expired without a settle.
|
||||
//!
|
||||
//! Concurrency:
|
||||
//! - The dedupe key is enforced by a partial `UNIQUE` index that only
|
||||
//! covers `status IN ('ready', 'running')`. Producers use `INSERT OR
|
||||
//! IGNORE` so a duplicate enqueue while a job is in flight or queued is
|
||||
//! a silent no-op; a duplicate enqueue after the first completes is
|
||||
//! accepted and creates a fresh row.
|
||||
//! - `claim_next` is one statement: `UPDATE … WHERE id = (SELECT … LIMIT 1)
|
||||
//! RETURNING …`. SQLite serialises writes, so no two workers can claim
|
||||
//! the same row.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::Utc;
|
||||
use rusqlite::{params, Connection, OptionalExtension, Transaction};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::jobs::types::{Job, JobKind, JobStatus, NewJob};
|
||||
use crate::openhuman::memory::tree::store::with_connection;
|
||||
|
||||
/// Default visibility lock — a worker that crashes mid-job will have its
|
||||
/// row recovered after this window. 5 min is comfortably larger than any
|
||||
/// expected single-job runtime (LLM extract or summarise) without leaving
|
||||
/// real failures stuck for hours.
|
||||
pub const DEFAULT_LOCK_DURATION_MS: i64 = 5 * 60 * 1_000;
|
||||
|
||||
/// Backoff math for retry. Returns `now + min(base * 2^attempts, cap)`.
|
||||
const RETRY_BASE_MS: i64 = 60 * 1_000;
|
||||
const RETRY_CAP_MS: i64 = 60 * 60 * 1_000;
|
||||
const DEFAULT_MAX_ATTEMPTS: u32 = 5;
|
||||
|
||||
/// Enqueue one job. Idempotent on `dedupe_key` while another active row
|
||||
/// (status `ready`/`running`) shares it. Returns `Some(id)` if the row
|
||||
/// was inserted, `None` if a duplicate was suppressed.
|
||||
pub fn enqueue(config: &Config, job: &NewJob) -> Result<Option<String>> {
|
||||
with_connection(config, |conn| enqueue_conn(conn, job))
|
||||
}
|
||||
|
||||
/// Enqueue inside a caller-owned transaction. Use this when the producer
|
||||
/// is already mid-tx (e.g. `ingest::persist` writing chunks + jobs in one
|
||||
/// commit) so the queue insert lands atomically with the side-effect.
|
||||
/// `Transaction` derefs to `Connection`, so callers just pass `&tx`.
|
||||
pub fn enqueue_tx(tx: &Transaction<'_>, job: &NewJob) -> Result<Option<String>> {
|
||||
enqueue_conn(tx, job)
|
||||
}
|
||||
|
||||
pub(crate) fn enqueue_conn(conn: &Connection, job: &NewJob) -> Result<Option<String>> {
|
||||
let id = format!("job:{}", Uuid::new_v4());
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
let available_at = job.available_at_ms.unwrap_or(now_ms);
|
||||
let max_attempts = job.max_attempts.unwrap_or(DEFAULT_MAX_ATTEMPTS) as i64;
|
||||
|
||||
let inserted = conn.execute(
|
||||
"INSERT OR IGNORE INTO mem_tree_jobs (
|
||||
id, kind, payload_json, dedupe_key, status, attempts, max_attempts,
|
||||
available_at_ms, locked_until_ms, last_error,
|
||||
created_at_ms, started_at_ms, completed_at_ms
|
||||
) VALUES (?1, ?2, ?3, ?4, 'ready', 0, ?5, ?6, NULL, NULL, ?7, NULL, NULL)",
|
||||
params![
|
||||
id,
|
||||
job.kind.as_str(),
|
||||
job.payload_json,
|
||||
job.dedupe_key,
|
||||
max_attempts,
|
||||
available_at,
|
||||
now_ms,
|
||||
],
|
||||
)?;
|
||||
|
||||
if inserted == 0 {
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] enqueue suppressed by dedupe kind={} key={:?}",
|
||||
job.kind.as_str(),
|
||||
job.dedupe_key
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] enqueued id={} kind={} avail_at_ms={} dedupe={:?}",
|
||||
id,
|
||||
job.kind.as_str(),
|
||||
available_at,
|
||||
job.dedupe_key
|
||||
);
|
||||
Ok(Some(id))
|
||||
}
|
||||
|
||||
/// Atomically claim the next ready job whose `available_at_ms` has come
|
||||
/// due. Sets `status=running`, bumps `attempts`, stamps `started_at_ms`
|
||||
/// and `locked_until_ms`. Returns `None` when the queue is empty / not
|
||||
/// yet due.
|
||||
pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result<Option<Job>> {
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
let lock_until = now_ms.saturating_add(lock_duration_ms);
|
||||
|
||||
let row = conn
|
||||
.query_row(
|
||||
"UPDATE mem_tree_jobs
|
||||
SET status = 'running',
|
||||
attempts = attempts + 1,
|
||||
started_at_ms = ?1,
|
||||
locked_until_ms = ?2,
|
||||
last_error = NULL
|
||||
WHERE id = (
|
||||
SELECT id FROM mem_tree_jobs
|
||||
WHERE status = 'ready'
|
||||
AND available_at_ms <= ?1
|
||||
ORDER BY available_at_ms ASC
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING id, kind, payload_json, dedupe_key, status, attempts,
|
||||
max_attempts, available_at_ms, locked_until_ms, last_error,
|
||||
created_at_ms, started_at_ms, completed_at_ms",
|
||||
params![now_ms, lock_until],
|
||||
row_to_job,
|
||||
)
|
||||
.optional()
|
||||
.context("Failed to claim next mem_tree_jobs row")?;
|
||||
if let Some(j) = &row {
|
||||
log::debug!(
|
||||
"[memory_tree::jobs] claimed id={} kind={} attempt={}/{}",
|
||||
j.id,
|
||||
j.kind.as_str(),
|
||||
j.attempts,
|
||||
j.max_attempts
|
||||
);
|
||||
}
|
||||
Ok(row)
|
||||
})
|
||||
}
|
||||
|
||||
/// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`.
|
||||
///
|
||||
/// The UPDATE is gated on `attempts` and `started_at_ms` matching the values
|
||||
/// in `job` (the snapshot returned by [`claim_next`]). If the lease expired
|
||||
/// and another worker re-claimed the row, `rows_affected` will be 0 — the
|
||||
/// stale worker's settlement is a silent no-op rather than clobbering the new
|
||||
/// lessee's state.
|
||||
pub fn mark_done(config: &Config, job: &Job) -> Result<()> {
|
||||
let job_id = &job.id;
|
||||
let claim_attempts = job.attempts as i64;
|
||||
let claim_started_at = job.started_at_ms;
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
let n = conn.execute(
|
||||
"UPDATE mem_tree_jobs
|
||||
SET status = 'done',
|
||||
completed_at_ms = ?1,
|
||||
locked_until_ms = NULL,
|
||||
last_error = NULL
|
||||
WHERE id = ?2
|
||||
AND attempts = ?3
|
||||
AND started_at_ms IS ?4",
|
||||
params![now_ms, job_id, claim_attempts, claim_started_at],
|
||||
)?;
|
||||
if n == 0 {
|
||||
// Either the job row was deleted (shouldn't happen) or the lease
|
||||
// expired and a second worker re-claimed the row. Log and move on —
|
||||
// this is a known race outcome, not a bug in the current worker.
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] mark_done id={job_id} was a no-op \
|
||||
(stale lease: attempts={claim_attempts} started_at_ms={claim_started_at:?})"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Settle a failed job. If `attempts < max_attempts`, the row goes back
|
||||
/// to `ready` with an exponential-backoff `available_at_ms`. Otherwise
|
||||
/// it terminates as `failed`. Either way `last_error` is recorded.
|
||||
///
|
||||
/// Like [`mark_done`], the UPDATE is gated on the claim-token
|
||||
/// (`attempts` + `started_at_ms`) so a stale worker's failure settlement
|
||||
/// cannot clobber an active lessee's row — rows_affected == 0 is a silent
|
||||
/// no-op.
|
||||
pub fn mark_failed(config: &Config, job: &Job, error: &str) -> Result<()> {
|
||||
let job_id = &job.id;
|
||||
let attempts = job.attempts as i64;
|
||||
let max_attempts = job.max_attempts as i64;
|
||||
let claim_started_at = job.started_at_ms;
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
|
||||
if attempts >= max_attempts {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] terminal failure id={job_id} \
|
||||
attempts={attempts}/{max_attempts} err={error}"
|
||||
);
|
||||
let n = conn.execute(
|
||||
"UPDATE mem_tree_jobs
|
||||
SET status = 'failed',
|
||||
completed_at_ms = ?1,
|
||||
locked_until_ms = NULL,
|
||||
last_error = ?2
|
||||
WHERE id = ?3
|
||||
AND attempts = ?4
|
||||
AND started_at_ms IS ?5",
|
||||
params![now_ms, error, job_id, attempts, claim_started_at],
|
||||
)?;
|
||||
if n == 0 {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] mark_failed(terminal) id={job_id} was a no-op \
|
||||
(stale lease: attempts={attempts} started_at_ms={claim_started_at:?})"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
let backoff = backoff_ms(attempts as u32);
|
||||
let next_at = now_ms.saturating_add(backoff);
|
||||
log::info!(
|
||||
"[memory_tree::jobs] retry id={job_id} attempt={attempts}/{max_attempts} \
|
||||
next_at_ms={next_at} err={error}"
|
||||
);
|
||||
let n = conn.execute(
|
||||
"UPDATE mem_tree_jobs
|
||||
SET status = 'ready',
|
||||
available_at_ms = ?1,
|
||||
locked_until_ms = NULL,
|
||||
last_error = ?2
|
||||
WHERE id = ?3
|
||||
AND attempts = ?4
|
||||
AND started_at_ms IS ?5",
|
||||
params![next_at, error, job_id, attempts, claim_started_at],
|
||||
)?;
|
||||
if n == 0 {
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] mark_failed(retry) id={job_id} was a no-op \
|
||||
(stale lease: attempts={attempts} started_at_ms={claim_started_at:?})"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Flip any `running` row whose `locked_until_ms` has expired back to
|
||||
/// `ready`. Called once at worker startup so a process crash mid-job
|
||||
/// doesn't leave work stranded. Returns the number of rows recovered.
|
||||
pub fn recover_stale_locks(config: &Config) -> Result<usize> {
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
let n = conn.execute(
|
||||
"UPDATE mem_tree_jobs
|
||||
SET status = 'ready',
|
||||
last_error = COALESCE(last_error, 'recovered_from_stale_lock')
|
||||
WHERE status = 'running'
|
||||
AND locked_until_ms IS NOT NULL
|
||||
AND locked_until_ms < ?1",
|
||||
params![now_ms],
|
||||
)?;
|
||||
if n > 0 {
|
||||
log::warn!("[memory_tree::jobs] recovered {n} stale-locked job(s) at startup");
|
||||
}
|
||||
Ok(n)
|
||||
})
|
||||
}
|
||||
|
||||
/// Quick count helper for tests / diagnostics.
|
||||
pub fn count_by_status(config: &Config, status: JobStatus) -> Result<u64> {
|
||||
with_connection(config, |conn| {
|
||||
let n: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_jobs WHERE status = ?1",
|
||||
params![status.as_str()],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
}
|
||||
|
||||
/// Total count regardless of status — handy for assertions.
|
||||
pub fn count_total(config: &Config) -> Result<u64> {
|
||||
with_connection(config, |conn| {
|
||||
let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0))?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch one job by id (test/diagnostic helper).
|
||||
pub fn get_job(config: &Config, id: &str) -> Result<Option<Job>> {
|
||||
with_connection(config, |conn| {
|
||||
let job = conn
|
||||
.query_row(
|
||||
"SELECT id, kind, payload_json, dedupe_key, status, attempts, max_attempts,
|
||||
available_at_ms, locked_until_ms, last_error,
|
||||
created_at_ms, started_at_ms, completed_at_ms
|
||||
FROM mem_tree_jobs WHERE id = ?1",
|
||||
params![id],
|
||||
row_to_job,
|
||||
)
|
||||
.optional()?;
|
||||
Ok(job)
|
||||
})
|
||||
}
|
||||
|
||||
fn row_to_job(row: &rusqlite::Row<'_>) -> rusqlite::Result<Job> {
|
||||
let id: String = row.get(0)?;
|
||||
let kind_s: String = row.get(1)?;
|
||||
let payload_json: String = row.get(2)?;
|
||||
let dedupe_key: Option<String> = row.get(3)?;
|
||||
let status_s: String = row.get(4)?;
|
||||
let attempts: i64 = row.get(5)?;
|
||||
let max_attempts: i64 = row.get(6)?;
|
||||
let available_at_ms: i64 = row.get(7)?;
|
||||
let locked_until_ms: Option<i64> = row.get(8)?;
|
||||
let last_error: Option<String> = row.get(9)?;
|
||||
let created_at_ms: i64 = row.get(10)?;
|
||||
let started_at_ms: Option<i64> = row.get(11)?;
|
||||
let completed_at_ms: Option<i64> = row.get(12)?;
|
||||
|
||||
let kind = JobKind::parse(&kind_s).map_err(|e| {
|
||||
rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into())
|
||||
})?;
|
||||
let status = JobStatus::parse(&status_s).map_err(|e| {
|
||||
rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, e.into())
|
||||
})?;
|
||||
|
||||
Ok(Job {
|
||||
id,
|
||||
kind,
|
||||
payload_json,
|
||||
dedupe_key,
|
||||
status,
|
||||
attempts: attempts.max(0) as u32,
|
||||
max_attempts: max_attempts.max(0) as u32,
|
||||
available_at_ms,
|
||||
locked_until_ms,
|
||||
last_error,
|
||||
created_at_ms,
|
||||
started_at_ms,
|
||||
completed_at_ms,
|
||||
})
|
||||
}
|
||||
|
||||
/// Exponential backoff: attempt 1 → 60s, 2 → 120s, 3 → 240s, capped at 1h.
|
||||
fn backoff_ms(attempts_so_far: u32) -> i64 {
|
||||
// attempts_so_far is the count BEFORE the next retry's attempt — so the
|
||||
// first retry uses attempts_so_far=1, giving base*2^0 = 60s.
|
||||
let exp = attempts_so_far.saturating_sub(1).min(20); // cap shift
|
||||
let mult = 1i64 << exp; // 1, 2, 4, …
|
||||
let raw = RETRY_BASE_MS.saturating_mul(mult);
|
||||
raw.min(RETRY_CAP_MS)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::jobs::types::{
|
||||
AppendBufferPayload, AppendTarget, ExtractChunkPayload, NodeRef,
|
||||
};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enqueue_and_claim_roundtrip() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c1".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
|
||||
|
||||
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(claimed.id, id);
|
||||
assert_eq!(claimed.status, JobStatus::Running);
|
||||
assert_eq!(claimed.attempts, 1);
|
||||
assert!(claimed.locked_until_ms.is_some());
|
||||
|
||||
// Second claim should see no eligible row (the only one is now running).
|
||||
let again = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap();
|
||||
assert!(again.is_none());
|
||||
|
||||
mark_done(&cfg, &claimed).unwrap();
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(row.status, JobStatus::Done);
|
||||
assert!(row.completed_at_ms.is_some());
|
||||
assert!(row.locked_until_ms.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enqueue_dedupes_active_jobs() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c1".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id1 = enqueue(&cfg, &nj).unwrap();
|
||||
let id2 = enqueue(&cfg, &nj).unwrap();
|
||||
assert!(id1.is_some());
|
||||
assert!(id2.is_none(), "duplicate should be suppressed while ready");
|
||||
assert_eq!(count_total(&cfg).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enqueue_after_done_creates_fresh_row() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c1".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id1 = enqueue(&cfg, &nj).unwrap().unwrap();
|
||||
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(claimed.id, id1);
|
||||
mark_done(&cfg, &claimed).unwrap();
|
||||
|
||||
// Now the dedupe key is free (partial index excludes 'done').
|
||||
let id2 = enqueue(&cfg, &nj).unwrap();
|
||||
assert!(id2.is_some());
|
||||
assert_ne!(id2.unwrap(), id1);
|
||||
assert_eq!(count_total(&cfg).unwrap(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mark_failed_retries_then_terminates() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = AppendBufferPayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: "c1".into(),
|
||||
},
|
||||
target: AppendTarget::Source {
|
||||
source_id: "slack:#x".into(),
|
||||
},
|
||||
};
|
||||
let mut nj = NewJob::append_buffer(&payload).unwrap();
|
||||
nj.max_attempts = Some(2);
|
||||
let id = enqueue(&cfg, &nj).unwrap().unwrap();
|
||||
|
||||
// Fail #1 — should bounce back to 'ready' with future available_at.
|
||||
let attempt1 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
mark_failed(&cfg, &attempt1, "boom").unwrap();
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(row.status, JobStatus::Ready);
|
||||
assert!(row.available_at_ms > Utc::now().timestamp_millis());
|
||||
assert_eq!(row.last_error.as_deref(), Some("boom"));
|
||||
|
||||
// Force the row available again so the test doesn't hinge on sleep.
|
||||
with_connection(&cfg, |c| {
|
||||
c.execute(
|
||||
"UPDATE mem_tree_jobs SET available_at_ms = 0 WHERE id = ?1",
|
||||
params![id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Fail #2 — exceeds max_attempts → terminal 'failed'.
|
||||
let attempt2 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
mark_failed(&cfg, &attempt2, "fatal").unwrap();
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(row.status, JobStatus::Failed);
|
||||
assert_eq!(row.last_error.as_deref(), Some("fatal"));
|
||||
assert!(row.completed_at_ms.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recover_stale_locks_resets_running_rows() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c1".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id = enqueue(&cfg, &nj).unwrap().unwrap();
|
||||
|
||||
// Claim with a lock window that's already in the past so recovery
|
||||
// sees it as expired.
|
||||
let _ = claim_next(&cfg, -1).unwrap().unwrap();
|
||||
|
||||
let recovered = recover_stale_locks(&cfg).unwrap();
|
||||
assert_eq!(recovered, 1);
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(row.status, JobStatus::Ready);
|
||||
}
|
||||
|
||||
/// Happy path: a non-stale settlement still succeeds after the claim-token
|
||||
/// check is applied. Regression guard so the common case isn't broken.
|
||||
#[test]
|
||||
fn mark_done_succeeds_for_current_lessee() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c-happy".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
|
||||
|
||||
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(claimed.id, id);
|
||||
|
||||
// Current lessee should settle successfully.
|
||||
mark_done(&cfg, &claimed).unwrap();
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(row.status, JobStatus::Done);
|
||||
assert!(row.completed_at_ms.is_some());
|
||||
assert!(row.locked_until_ms.is_none());
|
||||
}
|
||||
|
||||
/// Stale-worker settlement is a no-op: after a lock expires and a second
|
||||
/// worker re-claims the job, the first worker's `mark_done` must not
|
||||
/// clobber the new lessee's row.
|
||||
#[test]
|
||||
fn stale_worker_settlement_is_noop() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c-stale".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
|
||||
|
||||
// Worker A claims with a lock that's already expired (negative window).
|
||||
let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap();
|
||||
assert_eq!(worker_a_job.id, id);
|
||||
assert_eq!(worker_a_job.attempts, 1);
|
||||
|
||||
// Simulate lease expiry: recover_stale_locks resets the row to 'ready'.
|
||||
let recovered = recover_stale_locks(&cfg).unwrap();
|
||||
assert_eq!(recovered, 1);
|
||||
|
||||
// Worker B claims the reset row — different lease token (attempts=2).
|
||||
let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(worker_b_job.id, id);
|
||||
assert_eq!(worker_b_job.attempts, 2);
|
||||
|
||||
// Worker A (stale) tries to mark done using its old claim snapshot.
|
||||
mark_done(&cfg, &worker_a_job).unwrap(); // must NOT return Err
|
||||
|
||||
// Worker B's row must be untouched — still 'running' with attempts=2.
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(
|
||||
row.status,
|
||||
JobStatus::Running,
|
||||
"stale settlement must not clobber Worker B's running row"
|
||||
);
|
||||
assert_eq!(
|
||||
row.attempts, 2,
|
||||
"attempts must still reflect Worker B's claim"
|
||||
);
|
||||
}
|
||||
|
||||
/// Same contract as stale_worker_settlement_is_noop but for mark_failed.
|
||||
#[test]
|
||||
fn stale_worker_mark_failed_is_noop() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let payload = ExtractChunkPayload {
|
||||
chunk_id: "c-stale-fail".into(),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&payload).unwrap();
|
||||
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
|
||||
|
||||
// Worker A claims with an already-expired lock.
|
||||
let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap();
|
||||
assert_eq!(worker_a_job.attempts, 1);
|
||||
|
||||
// Lease expires, recovered, Worker B re-claims.
|
||||
let recovered = recover_stale_locks(&cfg).unwrap();
|
||||
assert_eq!(recovered, 1);
|
||||
let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
assert_eq!(worker_b_job.attempts, 2);
|
||||
|
||||
// Worker A (stale) tries to record a failure — must be a no-op.
|
||||
mark_failed(&cfg, &worker_a_job, "stale error").unwrap();
|
||||
|
||||
// Worker B's row must be untouched.
|
||||
let row = get_job(&cfg, &id).unwrap().unwrap();
|
||||
assert_eq!(
|
||||
row.status,
|
||||
JobStatus::Running,
|
||||
"stale mark_failed must not clobber Worker B's running row"
|
||||
);
|
||||
assert_ne!(
|
||||
row.last_error.as_deref(),
|
||||
Some("stale error"),
|
||||
"stale error must not be written to the row"
|
||||
);
|
||||
assert_eq!(row.attempts, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_grows_then_caps() {
|
||||
assert_eq!(backoff_ms(1), 60_000);
|
||||
assert_eq!(backoff_ms(2), 120_000);
|
||||
assert_eq!(backoff_ms(3), 240_000);
|
||||
// Eventually clamps at the cap.
|
||||
assert_eq!(backoff_ms(20), RETRY_CAP_MS);
|
||||
assert_eq!(backoff_ms(99), RETRY_CAP_MS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_by_status_reports_each_state() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
for i in 0..3 {
|
||||
let p = ExtractChunkPayload {
|
||||
chunk_id: format!("c{i}"),
|
||||
};
|
||||
let nj = NewJob::extract_chunk(&p).unwrap();
|
||||
enqueue(&cfg, &nj).unwrap();
|
||||
}
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 3);
|
||||
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
|
||||
mark_done(&cfg, &claimed).unwrap();
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Done).unwrap(), 1);
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
use anyhow::Result;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
|
||||
/// Deterministically run queued memory-tree jobs until no immediately
|
||||
/// claimable work remains. Intended for tests that need the async pipeline
|
||||
/// to settle without spawning background tasks.
|
||||
pub async fn drain_until_idle(config: &Config) -> Result<()> {
|
||||
loop {
|
||||
if !super::worker::run_once(config).await? {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,439 @@
|
||||
//! Job types for the async memory-tree pipeline.
|
||||
//!
|
||||
//! Each `Job` row in `mem_tree_jobs` stores its discriminator as a string
|
||||
//! `kind` plus a JSON-encoded `payload`. The strongly-typed payload structs
|
||||
//! below own (de)serialisation; handlers parse the payload by branching on
|
||||
//! [`JobKind`] and calling the matching `from_payload_json`.
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Discriminator persisted in `mem_tree_jobs.kind`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum JobKind {
|
||||
/// Run LLM entity extraction over a single chunk and decide admission.
|
||||
ExtractChunk,
|
||||
/// Push an admitted chunk into a tree's L0 buffer.
|
||||
AppendBuffer,
|
||||
/// Seal exactly one buffer level; cascades enqueue a follow-up.
|
||||
Seal,
|
||||
/// Match a chunk's entities against active topic trees and enqueue
|
||||
/// per-topic `AppendBuffer` jobs.
|
||||
TopicRoute,
|
||||
/// Build the global tree's daily digest for a given UTC date.
|
||||
DigestDaily,
|
||||
/// Walk stale buffers and enqueue `Seal` jobs for any over the age cap.
|
||||
FlushStale,
|
||||
}
|
||||
|
||||
impl JobKind {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
JobKind::ExtractChunk => "extract_chunk",
|
||||
JobKind::AppendBuffer => "append_buffer",
|
||||
JobKind::Seal => "seal",
|
||||
JobKind::TopicRoute => "topic_route",
|
||||
JobKind::DigestDaily => "digest_daily",
|
||||
JobKind::FlushStale => "flush_stale",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse(s: &str) -> Result<Self> {
|
||||
Ok(match s {
|
||||
"extract_chunk" => JobKind::ExtractChunk,
|
||||
"append_buffer" => JobKind::AppendBuffer,
|
||||
"seal" => JobKind::Seal,
|
||||
"topic_route" => JobKind::TopicRoute,
|
||||
"digest_daily" => JobKind::DigestDaily,
|
||||
"flush_stale" => JobKind::FlushStale,
|
||||
other => return Err(anyhow!("unknown JobKind '{other}'")),
|
||||
})
|
||||
}
|
||||
|
||||
/// True when handling this kind should hold a slot from the global
|
||||
/// LLM concurrency semaphore. `TopicRoute` is bound because
|
||||
/// `maybe_spawn_topic_tree → backfill_topic_tree` can transitively
|
||||
/// trigger summariser LLM calls when an entity first crosses the
|
||||
/// hotness threshold.
|
||||
pub fn is_llm_bound(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
JobKind::ExtractChunk | JobKind::Seal | JobKind::DigestDaily | JobKind::TopicRoute
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Lifecycle states persisted on `mem_tree_jobs.status`. Workers transition
|
||||
/// `ready → running → done|failed`. `Cancelled` is reserved for explicit
|
||||
/// admin actions (none surfaced yet).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum JobStatus {
|
||||
Ready,
|
||||
Running,
|
||||
Done,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl JobStatus {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
JobStatus::Ready => "ready",
|
||||
JobStatus::Running => "running",
|
||||
JobStatus::Done => "done",
|
||||
JobStatus::Failed => "failed",
|
||||
JobStatus::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse(s: &str) -> Result<Self> {
|
||||
Ok(match s {
|
||||
"ready" => JobStatus::Ready,
|
||||
"running" => JobStatus::Running,
|
||||
"done" => JobStatus::Done,
|
||||
"failed" => JobStatus::Failed,
|
||||
"cancelled" => JobStatus::Cancelled,
|
||||
other => return Err(anyhow!("unknown JobStatus '{other}'")),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
JobStatus::Done | JobStatus::Failed | JobStatus::Cancelled
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Payloads ───────────────────────────────────────────────────────────────
|
||||
|
||||
/// Reference to either a leaf chunk or a sealed summary node. Used by
|
||||
/// payloads that route content through the pipeline regardless of which
|
||||
/// kind of source produced it.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum NodeRef {
|
||||
Leaf { chunk_id: String },
|
||||
Summary { summary_id: String },
|
||||
}
|
||||
|
||||
impl NodeRef {
|
||||
/// Stringified id with kind prefix, suitable for dedupe-key composition.
|
||||
pub fn dedupe_fragment(&self) -> String {
|
||||
match self {
|
||||
NodeRef::Leaf { chunk_id } => format!("leaf:{chunk_id}"),
|
||||
NodeRef::Summary { summary_id } => format!("summary:{summary_id}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct ExtractChunkPayload {
|
||||
pub chunk_id: String,
|
||||
}
|
||||
|
||||
impl ExtractChunkPayload {
|
||||
pub fn dedupe_key(&self) -> String {
|
||||
format!("extract:{}", self.chunk_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Where an `AppendBuffer` job should land its node. Source-tree appends
|
||||
/// are keyed by `source_id`; topic-tree appends are keyed by `tree_id`
|
||||
/// because there can be many topic trees per node.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum AppendTarget {
|
||||
Source { source_id: String },
|
||||
Topic { tree_id: String },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct AppendBufferPayload {
|
||||
pub node: NodeRef,
|
||||
pub target: AppendTarget,
|
||||
}
|
||||
|
||||
impl AppendBufferPayload {
|
||||
pub fn dedupe_key(&self) -> String {
|
||||
let node_part = self.node.dedupe_fragment();
|
||||
match &self.target {
|
||||
AppendTarget::Source { source_id } => {
|
||||
format!("append:source:{source_id}:{node_part}")
|
||||
}
|
||||
AppendTarget::Topic { tree_id } => {
|
||||
format!("append:topic:{tree_id}:{node_part}")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct SealPayload {
|
||||
pub tree_id: String,
|
||||
pub level: u32,
|
||||
/// When `Some`, the seal handler bypasses the buffer-budget check and
|
||||
/// force-seals — used by the time-based flush path. The wall-clock is
|
||||
/// passed through so the seal stamps a deterministic `sealed_at`.
|
||||
pub force_now_ms: Option<i64>,
|
||||
}
|
||||
|
||||
impl SealPayload {
|
||||
pub fn dedupe_key(&self) -> String {
|
||||
// Active seal-job uniqueness is enforced per (tree, level): a seal
|
||||
// already in flight suppresses duplicate enqueues. Once the job
|
||||
// completes the partial index releases the key, so the next time
|
||||
// the buffer crosses its gate a fresh seal can be enqueued.
|
||||
format!("seal:{}:{}", self.tree_id, self.level)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct TopicRoutePayload {
|
||||
pub node: NodeRef,
|
||||
}
|
||||
|
||||
impl TopicRoutePayload {
|
||||
pub fn dedupe_key(&self) -> String {
|
||||
format!("topic_route:{}", self.node.dedupe_fragment())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct DigestDailyPayload {
|
||||
/// UTC calendar date in `YYYY-MM-DD` form. Stored as a string so the
|
||||
/// dedupe key doesn't need to know about chrono.
|
||||
pub date_iso: String,
|
||||
}
|
||||
|
||||
impl DigestDailyPayload {
|
||||
pub fn dedupe_key(&self) -> String {
|
||||
format!("digest_daily:{}", self.date_iso)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
|
||||
pub struct FlushStalePayload {
|
||||
/// Override the configured `DEFAULT_FLUSH_AGE_SECS`. Optional so the
|
||||
/// scheduler can enqueue with `None` and let the handler use the
|
||||
/// configured default.
|
||||
pub max_age_secs: Option<i64>,
|
||||
}
|
||||
|
||||
impl FlushStalePayload {
|
||||
pub fn dedupe_key(&self, date_iso: &str) -> String {
|
||||
format!("flush_stale:{date_iso}")
|
||||
}
|
||||
}
|
||||
|
||||
/// One row in `mem_tree_jobs`. `payload_json` is left as a raw string so
|
||||
/// callers parse it lazily based on `kind`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Job {
|
||||
pub id: String,
|
||||
pub kind: JobKind,
|
||||
pub payload_json: String,
|
||||
pub dedupe_key: Option<String>,
|
||||
pub status: JobStatus,
|
||||
pub attempts: u32,
|
||||
pub max_attempts: u32,
|
||||
pub available_at_ms: i64,
|
||||
pub locked_until_ms: Option<i64>,
|
||||
pub last_error: Option<String>,
|
||||
pub created_at_ms: i64,
|
||||
pub started_at_ms: Option<i64>,
|
||||
pub completed_at_ms: Option<i64>,
|
||||
}
|
||||
|
||||
/// Caller-side bundle for `enqueue` — `Job` minus the persistence-only
|
||||
/// columns. Keeps producers from having to mint timestamps and ids by hand.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct NewJob {
|
||||
pub kind: JobKind,
|
||||
pub payload_json: String,
|
||||
pub dedupe_key: Option<String>,
|
||||
/// `None` means "available immediately." Set this for delayed jobs
|
||||
/// (retries, scheduled work).
|
||||
pub available_at_ms: Option<i64>,
|
||||
pub max_attempts: Option<u32>,
|
||||
}
|
||||
|
||||
impl NewJob {
|
||||
pub fn extract_chunk(p: &ExtractChunkPayload) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::ExtractChunk,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key()),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn append_buffer(p: &AppendBufferPayload) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::AppendBuffer,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key()),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn seal(p: &SealPayload) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::Seal,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key()),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn topic_route(p: &TopicRoutePayload) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::TopicRoute,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key()),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn digest_daily(p: &DigestDailyPayload) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::DigestDaily,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key()),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn flush_stale(p: &FlushStalePayload, date_iso: &str) -> Result<Self> {
|
||||
Ok(Self {
|
||||
kind: JobKind::FlushStale,
|
||||
payload_json: serde_json::to_string(p)?,
|
||||
dedupe_key: Some(p.dedupe_key(date_iso)),
|
||||
available_at_ms: None,
|
||||
max_attempts: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn job_kind_roundtrip() {
|
||||
for k in [
|
||||
JobKind::ExtractChunk,
|
||||
JobKind::AppendBuffer,
|
||||
JobKind::Seal,
|
||||
JobKind::TopicRoute,
|
||||
JobKind::DigestDaily,
|
||||
JobKind::FlushStale,
|
||||
] {
|
||||
assert_eq!(JobKind::parse(k.as_str()).unwrap(), k);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_status_terminality() {
|
||||
assert!(!JobStatus::Ready.is_terminal());
|
||||
assert!(!JobStatus::Running.is_terminal());
|
||||
assert!(JobStatus::Done.is_terminal());
|
||||
assert!(JobStatus::Failed.is_terminal());
|
||||
assert!(JobStatus::Cancelled.is_terminal());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupe_keys_distinguish_targets() {
|
||||
let p_src = AppendBufferPayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: "c1".into(),
|
||||
},
|
||||
target: AppendTarget::Source {
|
||||
source_id: "slack:#eng".into(),
|
||||
},
|
||||
};
|
||||
let p_topic = AppendBufferPayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: "c1".into(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: "topic:abc".into(),
|
||||
},
|
||||
};
|
||||
assert_ne!(p_src.dedupe_key(), p_topic.dedupe_key());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupe_keys_distinguish_node_kinds() {
|
||||
let p_leaf = AppendBufferPayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: "x".into(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: "t".into(),
|
||||
},
|
||||
};
|
||||
let p_summary = AppendBufferPayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: "x".into(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: "t".into(),
|
||||
},
|
||||
};
|
||||
assert_ne!(p_leaf.dedupe_key(), p_summary.dedupe_key());
|
||||
|
||||
let r_leaf = TopicRoutePayload {
|
||||
node: NodeRef::Leaf {
|
||||
chunk_id: "x".into(),
|
||||
},
|
||||
};
|
||||
let r_summary = TopicRoutePayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: "x".into(),
|
||||
},
|
||||
};
|
||||
assert_ne!(r_leaf.dedupe_key(), r_summary.dedupe_key());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_bound_kinds() {
|
||||
assert!(JobKind::ExtractChunk.is_llm_bound());
|
||||
assert!(JobKind::Seal.is_llm_bound());
|
||||
assert!(JobKind::DigestDaily.is_llm_bound());
|
||||
assert!(JobKind::TopicRoute.is_llm_bound());
|
||||
assert!(!JobKind::AppendBuffer.is_llm_bound());
|
||||
assert!(!JobKind::FlushStale.is_llm_bound());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn node_ref_serializes_with_kind_tag() {
|
||||
let leaf = NodeRef::Leaf {
|
||||
chunk_id: "x".into(),
|
||||
};
|
||||
let s = serde_json::to_string(&leaf).unwrap();
|
||||
assert!(s.contains("\"kind\":\"leaf\""));
|
||||
let back: NodeRef = serde_json::from_str(&s).unwrap();
|
||||
assert_eq!(back, leaf);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_target_serializes_with_kind_tag() {
|
||||
let p = AppendTarget::Source {
|
||||
source_id: "x".into(),
|
||||
};
|
||||
let s = serde_json::to_string(&p).unwrap();
|
||||
assert!(s.contains("\"kind\":\"source\""));
|
||||
assert!(s.contains("\"source_id\":\"x\""));
|
||||
let back: AppendTarget = serde_json::from_str(&s).unwrap();
|
||||
match back {
|
||||
AppendTarget::Source { source_id } => assert_eq!(source_id, "x"),
|
||||
_ => panic!("wrong variant"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use tokio::sync::{Notify, Semaphore};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::jobs::handlers;
|
||||
use crate::openhuman::memory::tree::jobs::store::{
|
||||
claim_next, mark_done, mark_failed, recover_stale_locks, DEFAULT_LOCK_DURATION_MS,
|
||||
};
|
||||
|
||||
const WORKER_COUNT: usize = 3;
|
||||
const POLL_INTERVAL: Duration = Duration::from_secs(5);
|
||||
|
||||
static WORKER_NOTIFY: OnceLock<Arc<Notify>> = OnceLock::new();
|
||||
static STARTED: std::sync::Once = std::sync::Once::new();
|
||||
|
||||
pub fn wake_workers() {
|
||||
if let Some(notify) = WORKER_NOTIFY.get() {
|
||||
notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the worker pool + daily scheduler. Takes the full `Config` so
|
||||
/// each spawned task sees the user's actual settings (LLM endpoints,
|
||||
/// embedder model, timeouts) — not `Config::default()`. Without this,
|
||||
/// workers fall back to inert/regex-only behavior regardless of what's
|
||||
/// in `config.toml`, defeating the entire async pipeline.
|
||||
///
|
||||
/// Idempotent (`Once`-guarded) so repeat calls during bootstrap are
|
||||
/// safe no-ops after the first.
|
||||
pub fn start(config: Config) {
|
||||
STARTED.call_once(|| {
|
||||
let notify = WORKER_NOTIFY
|
||||
.get_or_init(|| Arc::new(Notify::new()))
|
||||
.clone();
|
||||
let llm_slots = Arc::new(Semaphore::new(3));
|
||||
if let Err(err) = recover_stale_locks(&config) {
|
||||
log::warn!("[memory_tree::jobs] recover_stale_locks failed at startup: {err:#}");
|
||||
}
|
||||
|
||||
for idx in 0..WORKER_COUNT {
|
||||
let notify = notify.clone();
|
||||
let llm_slots = llm_slots.clone();
|
||||
let cfg = config.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match run_once_with_semaphore(&cfg, llm_slots.clone()).await {
|
||||
Ok(true) => continue,
|
||||
Ok(false) => {
|
||||
tokio::select! {
|
||||
_ = notify.notified() => {}
|
||||
_ = tokio::time::sleep(POLL_INTERVAL) => {}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
log::error!("[memory_tree::jobs] worker={} loop error: {:#}", idx, err);
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
super::scheduler::start(config);
|
||||
});
|
||||
}
|
||||
|
||||
pub async fn run_once(config: &Config) -> Result<bool> {
|
||||
let llm_slots = Arc::new(Semaphore::new(1));
|
||||
run_once_with_semaphore(config, llm_slots).await
|
||||
}
|
||||
|
||||
async fn run_once_with_semaphore(config: &Config, llm_slots: Arc<Semaphore>) -> Result<bool> {
|
||||
let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
let permit = if job.kind.is_llm_bound() {
|
||||
Some(llm_slots.acquire().await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let result = handlers::handle_job(config, &job).await;
|
||||
drop(permit);
|
||||
|
||||
match result {
|
||||
Ok(()) => {
|
||||
mark_done(config, &job)?;
|
||||
}
|
||||
Err(err) => {
|
||||
// Preserve the full anyhow cause chain in the persisted
|
||||
// last_error so a reader of mem_tree_jobs can see the root
|
||||
// cause, not just the top-level message. Mirrors the {:#}
|
||||
// log format used right above.
|
||||
let message = format!("{err:#}");
|
||||
log::warn!(
|
||||
"[memory_tree::jobs] job failed id={} kind={} err={:#}",
|
||||
job.id,
|
||||
job.kind.as_str(),
|
||||
err
|
||||
);
|
||||
mark_failed(config, &job, &message)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
@@ -24,8 +24,10 @@
|
||||
|
||||
pub mod canonicalize;
|
||||
pub mod chunker;
|
||||
pub mod content_store;
|
||||
pub mod global_tree;
|
||||
pub mod ingest;
|
||||
pub mod jobs;
|
||||
pub mod retrieval;
|
||||
pub mod rpc;
|
||||
pub mod schemas;
|
||||
@@ -34,6 +36,7 @@ pub mod source_tree;
|
||||
pub mod store;
|
||||
pub mod topic_tree;
|
||||
pub mod types;
|
||||
pub mod util;
|
||||
|
||||
pub use retrieval::{all_retrieval_controller_schemas, all_retrieval_registered_controllers};
|
||||
pub use schemas::{
|
||||
|
||||
@@ -23,6 +23,7 @@ use std::collections::VecDeque;
|
||||
use anyhow::Result;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
use crate::openhuman::memory::tree::retrieval::types::{
|
||||
hit_from_chunk, hit_from_summary, RetrievalHit,
|
||||
};
|
||||
@@ -179,15 +180,27 @@ fn walk_with_embeddings(
|
||||
continue;
|
||||
}
|
||||
// Is it a summary?
|
||||
if let Some(summary) = store::get_summary(config, &id)? {
|
||||
if let Some(mut summary) = store::get_summary(config, &id)? {
|
||||
let scope = store::get_tree(config, &summary.tree_id)?
|
||||
.map(|t| t.scope)
|
||||
.unwrap_or_else(|| root_tree_scope.clone());
|
||||
// Hydrate the full body from disk — `summary.content` is a
|
||||
// ≤500-char preview after the MD-on-disk migration.
|
||||
// Non-fatal fallback for pre-MD-migration rows.
|
||||
match content_read::read_summary_body(config, &id) {
|
||||
Ok(body) => summary.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::drill_down] read_summary_body failed — serving preview: {e:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Summary embeddings live on the struct directly (Phase 4 amend).
|
||||
embeddings.push(summary.embedding.clone());
|
||||
let child_ids = summary.child_ids.clone();
|
||||
out.push(hit_from_summary(&summary, &scope));
|
||||
if depth < max_depth {
|
||||
for next in summary.child_ids {
|
||||
for next in child_ids {
|
||||
frontier.push_back((next, depth + 1));
|
||||
}
|
||||
}
|
||||
@@ -195,11 +208,21 @@ fn walk_with_embeddings(
|
||||
}
|
||||
// Else try as a chunk (leaf). Chunk embeddings live in a separate
|
||||
// blob column — fetch via the existing accessor.
|
||||
if let Some(chunk) = get_chunk(config, &id)? {
|
||||
if let Some(mut chunk) = get_chunk(config, &id)? {
|
||||
// Propagate DB errors rather than silently treating them as
|
||||
// "no embedding" — the caller should know if the store is broken.
|
||||
let emb = get_chunk_embedding(config, &chunk.id)?;
|
||||
embeddings.push(emb);
|
||||
// Hydrate the full body from disk — `chunk.content` is a
|
||||
// ≤500-char preview after the MD-on-disk migration.
|
||||
match content_read::read_chunk_body(config, &id) {
|
||||
Ok(body) => chunk.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::drill_down] read_chunk_body failed — serving preview: {e:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Score unknown here; 0.0 neutral placeholder.
|
||||
out.push(hit_from_chunk(&chunk, "", &chunk.metadata.source_id, 0.0));
|
||||
continue;
|
||||
@@ -215,7 +238,10 @@ fn walk_with_embeddings(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use crate::openhuman::memory::tree::source_tree::types::TreeKind;
|
||||
@@ -240,6 +266,8 @@ mod tests {
|
||||
let ts = Utc::now();
|
||||
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let mut leaf_ids: Vec<String> = Vec::new();
|
||||
for seq in 0..2u32 {
|
||||
let c = Chunk {
|
||||
@@ -254,18 +282,32 @@ mod tests {
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 6_000,
|
||||
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET * 6
|
||||
/ 10,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// via `read_chunk_body` during the seal triggered by `append_leaf`.
|
||||
let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap();
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
leaf_ids.push(c.id.clone());
|
||||
append_leaf(
|
||||
cfg,
|
||||
&tree,
|
||||
&LeafRef {
|
||||
chunk_id: c.id.clone(),
|
||||
token_count: 6_000,
|
||||
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET
|
||||
* 6
|
||||
/ 10,
|
||||
timestamp: ts,
|
||||
content: c.content.clone(),
|
||||
entities: vec![],
|
||||
@@ -273,6 +315,7 @@ mod tests {
|
||||
score: 0.5,
|
||||
},
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -401,6 +444,7 @@ mod tests {
|
||||
token_count: 10,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let leaf_a_2 = Chunk {
|
||||
id: "chat:slack:#eng:1".into(),
|
||||
@@ -423,15 +467,24 @@ mod tests {
|
||||
seq_in_source: 3,
|
||||
..leaf_a_1.clone()
|
||||
};
|
||||
upsert_chunks(
|
||||
cfg,
|
||||
&[
|
||||
leaf_a_1.clone(),
|
||||
leaf_a_2.clone(),
|
||||
leaf_b_1.clone(),
|
||||
leaf_b_2.clone(),
|
||||
],
|
||||
)
|
||||
let all_leaves = [
|
||||
leaf_a_1.clone(),
|
||||
leaf_a_2.clone(),
|
||||
leaf_b_1.clone(),
|
||||
leaf_b_2.clone(),
|
||||
];
|
||||
upsert_chunks(cfg, &all_leaves).unwrap();
|
||||
// Stage to disk so `walk_with_embeddings` can read full bodies via
|
||||
// `read_chunk_body` for leaf hits returned by the drill-down.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &all_leaves).unwrap();
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let l1_a = SummaryNode {
|
||||
@@ -471,9 +524,9 @@ mod tests {
|
||||
with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
tree_store::insert_tree_conn(&tx, &tree)?;
|
||||
tree_store::insert_summary_tx(&tx, &l1_a)?;
|
||||
tree_store::insert_summary_tx(&tx, &l1_b)?;
|
||||
tree_store::insert_summary_tx(&tx, &root)?;
|
||||
tree_store::insert_summary_tx(&tx, &l1_a, None)?;
|
||||
tree_store::insert_summary_tx(&tx, &l1_b, None)?;
|
||||
tree_store::insert_summary_tx(&tx, &root, None)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
use anyhow::Result;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
use crate::openhuman::memory::tree::retrieval::types::{hit_from_chunk, RetrievalHit};
|
||||
use crate::openhuman::memory::tree::score::store::get_score;
|
||||
use crate::openhuman::memory::tree::store::get_chunk;
|
||||
@@ -66,7 +67,22 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result<Vec<R
|
||||
// chunk row. `scope` falls back to the chunk's own source_id so
|
||||
// consumers still see provenance (e.g. "slack:#eng").
|
||||
let scope = chunk.metadata.source_id.clone();
|
||||
out.push(hit_from_chunk(&chunk, "", &scope, score));
|
||||
// Hydrate the full body from disk before building the hit.
|
||||
// The `content` column in SQLite holds a ≤500-char preview after
|
||||
// the MD-on-disk migration; the retrieval API must return the
|
||||
// complete chunk text so the LLM sees untruncated content.
|
||||
let mut chunk_with_body = chunk;
|
||||
match content_read::read_chunk_body(&config_owned, id) {
|
||||
Ok(body) => chunk_with_body.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::fetch] read_chunk_body failed for chunk — serving preview: {e:#}"
|
||||
);
|
||||
// Non-fatal: fall back to the preview already in the struct.
|
||||
// This handles pre-MD-migration rows gracefully.
|
||||
}
|
||||
}
|
||||
out.push(hit_from_chunk(&chunk_with_body, "", &scope, score));
|
||||
}
|
||||
Ok(out)
|
||||
})
|
||||
@@ -80,11 +96,26 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result<Vec<R
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
|
||||
use chrono::{TimeZone, Utc};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, chunks)
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -113,6 +144,7 @@ mod tests {
|
||||
token_count: 20,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -129,6 +161,7 @@ mod tests {
|
||||
let c1 = sample_chunk("slack:#eng", 0);
|
||||
let c2 = sample_chunk("slack:#eng", 1);
|
||||
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
|
||||
let out = fetch_leaves(&cfg, &[c1.id.clone(), c2.id.clone()])
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -142,6 +175,7 @@ mod tests {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let c1 = sample_chunk("slack:#eng", 0);
|
||||
upsert_chunks(&cfg, &[c1.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c1.clone()]);
|
||||
let out = fetch_leaves(
|
||||
&cfg,
|
||||
&[c1.id.clone(), "ghost:nonexistent".into(), c1.id.clone()],
|
||||
@@ -159,6 +193,7 @@ mod tests {
|
||||
for i in 0..(MAX_BATCH + 5) as u32 {
|
||||
let c = sample_chunk("slack:#eng", i);
|
||||
upsert_chunks(&cfg, &[c.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c.clone()]);
|
||||
ids.push(c.id);
|
||||
}
|
||||
let out = fetch_leaves(&cfg, &ids).await.unwrap();
|
||||
@@ -170,6 +205,7 @@ mod tests {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let c = sample_chunk("slack:#eng", 0);
|
||||
upsert_chunks(&cfg, &[c.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c.clone()]);
|
||||
let out = fetch_leaves(&cfg, &[c.id.clone()]).await.unwrap();
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].source_ref.as_deref(), Some("slack://slack:#eng/0"));
|
||||
|
||||
@@ -87,8 +87,11 @@ fn recap_to_hits(recap: RecapOutput, tree_id: &str, tree_scope: &str) -> Vec<Ret
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
@@ -96,6 +99,20 @@ mod tests {
|
||||
use chrono::{DateTime, Utc};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, chunks)
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -134,8 +151,10 @@ mod tests {
|
||||
token_count: 6_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c.clone()]).unwrap();
|
||||
stage_test_chunks(cfg, &[c.clone()]);
|
||||
append_leaf(
|
||||
cfg,
|
||||
&tree,
|
||||
@@ -149,6 +168,7 @@ mod tests {
|
||||
score: 0.5,
|
||||
},
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -160,11 +160,11 @@ async fn topic_entity_surfaces_after_ingest() {
|
||||
// ── Phase 4 (#710): embedding + semantic rerank tests ───────────────────
|
||||
|
||||
/// Ingest with an inert embedder must populate every kept chunk's
|
||||
/// `embedding` column. This guards against regressions where the embed
|
||||
/// step is silently skipped (e.g. future refactors threading embeddings
|
||||
/// through a different code path).
|
||||
/// `embedding` column. Embeddings are written by the async `extract_chunk`
|
||||
/// handler, so the test drains the queue before inspecting.
|
||||
#[tokio::test]
|
||||
async fn ingest_populates_chunk_embeddings() {
|
||||
use crate::openhuman::memory::tree::jobs::drain_until_idle;
|
||||
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
|
||||
use crate::openhuman::memory::tree::store::get_chunk_embedding;
|
||||
|
||||
@@ -172,7 +172,11 @@ async fn ingest_populates_chunk_embeddings() {
|
||||
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], chat_about_phoenix(0))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.chunks_written >= 1, "expected at least one kept chunk");
|
||||
assert!(
|
||||
out.chunks_written >= 1,
|
||||
"expected at least one persisted chunk"
|
||||
);
|
||||
drain_until_idle(&cfg).await.unwrap();
|
||||
for id in &out.chunk_ids {
|
||||
let emb = get_chunk_embedding(&cfg, id).unwrap();
|
||||
let v = emb.unwrap_or_else(|| panic!("embedding missing for chunk_id={id}"));
|
||||
@@ -188,8 +192,11 @@ async fn ingest_populates_chunk_embeddings() {
|
||||
/// the seal from firing on short batches.
|
||||
#[tokio::test]
|
||||
async fn seal_populates_summary_embedding() {
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::store as src_store;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
@@ -216,10 +223,24 @@ async fn seal_populates_summary_embedding() {
|
||||
token_count: tokens,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let c1 = mk_chunk(0, 6_000);
|
||||
let c2 = mk_chunk(1, 6_000);
|
||||
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
{
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, &[c1.clone(), c2.clone()])
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
let leaf_of = |c: &Chunk| LeafRef {
|
||||
chunk_id: c.id.clone(),
|
||||
@@ -230,12 +251,24 @@ async fn seal_populates_summary_embedding() {
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf(&cfg, &tree, &leaf_of(&c1), &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf_of(&c2), &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
append_leaf(
|
||||
&cfg,
|
||||
&tree,
|
||||
&leaf_of(&c1),
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let sealed = append_leaf(
|
||||
&cfg,
|
||||
&tree,
|
||||
&leaf_of(&c2),
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sealed.len(), 1, "expected one seal at the budget crossing");
|
||||
|
||||
let summary = src_store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
|
||||
|
||||
@@ -282,11 +282,26 @@ mod tests {
|
||||
//! initialises the schema idempotently on first access, so read-only
|
||||
//! calls return empty responses rather than erroring.
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceRef};
|
||||
use chrono::{TimeZone, Utc};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, chunks)
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -316,6 +331,7 @@ mod tests {
|
||||
token_count: 20,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -539,6 +555,7 @@ mod tests {
|
||||
let c1 = sample_chunk("slack:#eng", 0);
|
||||
let c2 = sample_chunk("slack:#eng", 1);
|
||||
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
|
||||
let req = FetchLeavesRequest {
|
||||
chunk_ids: vec![c1.id.clone(), c2.id.clone()],
|
||||
};
|
||||
@@ -552,6 +569,7 @@ mod tests {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let c1 = sample_chunk("slack:#eng", 0);
|
||||
upsert_chunks(&cfg, &[c1.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c1.clone()]);
|
||||
let req = FetchLeavesRequest {
|
||||
chunk_ids: vec![c1.id.clone(), "ghost:nonexistent".into()],
|
||||
};
|
||||
|
||||
@@ -21,6 +21,7 @@ use anyhow::Result;
|
||||
use chrono::{Duration, Utc};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
use crate::openhuman::memory::tree::retrieval::types::{
|
||||
hit_from_summary, QueryResponse, RetrievalHit,
|
||||
};
|
||||
@@ -124,7 +125,19 @@ fn collect_hits_and_nodes(
|
||||
// finest-grained summary layer above raw leaves.
|
||||
for level in 1..=tree.max_level {
|
||||
let level_nodes = store::list_summaries_at_level(config, &tree.id, level)?;
|
||||
for node in level_nodes {
|
||||
for mut node in level_nodes {
|
||||
// Hydrate the full body from disk — `node.content` is a
|
||||
// ≤500-char preview after the MD-on-disk migration. Callers
|
||||
// (including the LLM) must receive the complete summary text.
|
||||
// Non-fatal fallback for pre-MD-migration rows.
|
||||
match content_read::read_summary_body(config, &node.id) {
|
||||
Ok(body) => node.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::source] read_summary_body failed — serving preview: {e:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
hits.push(hit_from_summary(&node, &tree.scope));
|
||||
nodes.push((node, tree.scope.clone()));
|
||||
}
|
||||
@@ -293,7 +306,10 @@ fn filter_by_window(hits: Vec<RetrievalHit>, window_days: u32) -> Vec<RetrievalH
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
@@ -315,6 +331,8 @@ mod tests {
|
||||
async fn seed_source(cfg: &Config, scope: &str, ts: DateTime<Utc>) {
|
||||
let tree = get_or_create_source_tree(cfg, scope).unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
for seq in 0..2u32 {
|
||||
let c = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, scope, seq, "test-content"),
|
||||
@@ -328,17 +346,32 @@ mod tests {
|
||||
tags: vec!["eng".into()],
|
||||
source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))),
|
||||
},
|
||||
token_count: 6_000,
|
||||
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET * 6
|
||||
/ 10,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// via `read_chunk_body` during the seal triggered by `append_leaf`,
|
||||
// and `collect_hits_and_nodes` can read summary bodies for the API.
|
||||
let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap();
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
append_leaf(
|
||||
cfg,
|
||||
&tree,
|
||||
&LeafRef {
|
||||
chunk_id: c.id.clone(),
|
||||
token_count: 6_000,
|
||||
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET
|
||||
* 6
|
||||
/ 10,
|
||||
timestamp: ts,
|
||||
content: c.content.clone(),
|
||||
entities: vec![],
|
||||
@@ -346,6 +379,7 @@ mod tests {
|
||||
score: 0.5,
|
||||
},
|
||||
&summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -17,6 +17,7 @@ use anyhow::Result;
|
||||
use chrono::{Duration, TimeZone, Utc};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
use crate::openhuman::memory::tree::retrieval::types::{
|
||||
hit_from_summary, QueryResponse, RetrievalHit,
|
||||
};
|
||||
@@ -229,7 +230,7 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result<Opt
|
||||
Some(id) => id.clone(),
|
||||
None => return Ok(None),
|
||||
};
|
||||
let summary = match store::get_summary(config, &root_id)? {
|
||||
let mut summary = match store::get_summary(config, &root_id)? {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
log::warn!(
|
||||
@@ -238,6 +239,17 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result<Opt
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
// Hydrate the full body from disk — `summary.content` is a ≤500-char
|
||||
// preview after the MD-on-disk migration. Non-fatal fallback for
|
||||
// pre-MD-migration rows.
|
||||
match content_read::read_summary_body(config, &root_id) {
|
||||
Ok(body) => summary.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::topic] read_summary_body failed for topic root — serving preview: {e:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(Some(hit_from_summary(&summary, &tree.scope)))
|
||||
}
|
||||
|
||||
@@ -258,13 +270,23 @@ async fn entity_hit_to_retrieval_hit(
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<RetrievalHit>> {
|
||||
if node_kind == "summary" {
|
||||
let summary = match store::get_summary(&config_owned, &node_id)? {
|
||||
let mut summary = match store::get_summary(&config_owned, &node_id)? {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
log::warn!("[retrieval::topic] entity index points at missing summary row");
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
// Hydrate the full body from disk — `summary.content` is a
|
||||
// ≤500-char preview after the MD-on-disk migration.
|
||||
match content_read::read_summary_body(&config_owned, &node_id) {
|
||||
Ok(body) => summary.content = body,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[retrieval::topic] read_summary_body failed — serving preview: {e:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Prefer tree scope from the summary's parent tree if resolvable.
|
||||
let scope = if let Some(tid) = &tree_id_opt {
|
||||
store::get_tree(&config_owned, tid)?
|
||||
@@ -283,13 +305,21 @@ async fn entity_hit_to_retrieval_hit(
|
||||
// Leaf: fetch chunk and hydrate.
|
||||
use crate::openhuman::memory::tree::retrieval::types::hit_from_chunk;
|
||||
use crate::openhuman::memory::tree::store::get_chunk;
|
||||
let chunk = match get_chunk(&config_owned, &node_id)? {
|
||||
let mut chunk = match get_chunk(&config_owned, &node_id)? {
|
||||
Some(c) => c,
|
||||
None => {
|
||||
log::warn!("[retrieval::topic] entity index points at missing chunk row");
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
// Hydrate the full body from disk — `chunk.content` is a ≤500-char
|
||||
// preview after the MD-on-disk migration.
|
||||
match content_read::read_chunk_body(&config_owned, &node_id) {
|
||||
Ok(body) => chunk.content = body,
|
||||
Err(e) => {
|
||||
log::warn!("[retrieval::topic] read_chunk_body failed — serving preview: {e:#}");
|
||||
}
|
||||
}
|
||||
let scope = if let Some(tid) = &tree_id_opt {
|
||||
store::get_tree(&config_owned, tid)?
|
||||
.map(|t: Tree| t.scope)
|
||||
@@ -545,7 +575,7 @@ mod tests {
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
tree_store::insert_tree_conn(&tx, &tree)?;
|
||||
tree_store::insert_summary_tx(&tx, &summary)?;
|
||||
tree_store::insert_summary_tx(&tx, &summary, None)?;
|
||||
score_store::index_entities_tx(
|
||||
&tx,
|
||||
&[entity],
|
||||
|
||||
@@ -178,3 +178,143 @@ pub async fn get_chunk_rpc(
|
||||
format!("memory_tree: get_chunk id={}", req.id),
|
||||
))
|
||||
}
|
||||
|
||||
/// Manual-trigger surface for the global tree's daily digest. Default
|
||||
/// behavior (no `date_iso`) targets yesterday in UTC, matching the
|
||||
/// scheduler's autonomous behavior. Pass an explicit `YYYY-MM-DD` to
|
||||
/// re-run a specific date (idempotent — the handler skips if a daily
|
||||
/// node already exists for that day).
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
pub struct TriggerDigestRequest {
|
||||
/// UTC calendar date in `YYYY-MM-DD` form. When omitted, defaults to
|
||||
/// `yesterday` (today minus one day, UTC).
|
||||
#[serde(default)]
|
||||
pub date_iso: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct TriggerDigestResponse {
|
||||
/// True when the job was newly enqueued; false when an active job for
|
||||
/// the same date was suppressed by the dedupe partial unique index.
|
||||
pub enqueued: bool,
|
||||
/// ID of the freshly-inserted job row (None when dedupe-suppressed).
|
||||
pub job_id: Option<String>,
|
||||
/// The actual date the digest will run for, echoed back as
|
||||
/// `YYYY-MM-DD`. Useful when the caller didn't pass `date_iso` and
|
||||
/// wants to know what default got chosen.
|
||||
pub date_iso: String,
|
||||
}
|
||||
|
||||
pub async fn trigger_digest_rpc(
|
||||
config: &Config,
|
||||
req: TriggerDigestRequest,
|
||||
) -> Result<RpcOutcome<TriggerDigestResponse>, String> {
|
||||
use crate::openhuman::memory::tree::jobs;
|
||||
use chrono::{Duration as ChronoDuration, NaiveDate, Utc};
|
||||
|
||||
let date = match req
|
||||
.date_iso
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
Some(s) => NaiveDate::parse_from_str(s, "%Y-%m-%d")
|
||||
.map_err(|e| format!("invalid date_iso (expected YYYY-MM-DD): {e}"))?,
|
||||
None => Utc::now().date_naive() - ChronoDuration::days(1),
|
||||
};
|
||||
let date_iso = date.format("%Y-%m-%d").to_string();
|
||||
|
||||
// Run the synchronous enqueue on a blocking thread — `trigger_digest`
|
||||
// touches SQLite and we don't want to block the async runtime even
|
||||
// for the few-microsecond INSERT.
|
||||
let cfg_clone = config.clone();
|
||||
let date_for_blocking = date;
|
||||
let job_id =
|
||||
tokio::task::spawn_blocking(move || jobs::trigger_digest(&cfg_clone, date_for_blocking))
|
||||
.await
|
||||
.map_err(|e| format!("trigger_digest join error: {e}"))?
|
||||
.map_err(|e| format!("trigger_digest: {e}"))?;
|
||||
|
||||
let enqueued = job_id.is_some();
|
||||
Ok(RpcOutcome::single_log(
|
||||
TriggerDigestResponse {
|
||||
enqueued,
|
||||
job_id,
|
||||
date_iso: date_iso.clone(),
|
||||
},
|
||||
format!("memory_tree: trigger_digest date={date_iso} enqueued={enqueued}"),
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::jobs::store::count_total;
|
||||
use chrono::{Duration as ChronoDuration, Utc};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trigger_digest_with_explicit_date_enqueues() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let req = TriggerDigestRequest {
|
||||
date_iso: Some("2026-04-27".into()),
|
||||
};
|
||||
let outcome = trigger_digest_rpc(&cfg, req).await.unwrap();
|
||||
let resp = outcome.value;
|
||||
assert!(resp.enqueued);
|
||||
assert!(resp.job_id.is_some());
|
||||
assert_eq!(resp.date_iso, "2026-04-27");
|
||||
assert_eq!(count_total(&cfg).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trigger_digest_with_no_date_defaults_to_yesterday() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let req = TriggerDigestRequest::default();
|
||||
let outcome = trigger_digest_rpc(&cfg, req).await.unwrap();
|
||||
let resp = outcome.value;
|
||||
assert!(resp.enqueued);
|
||||
let expected = (Utc::now().date_naive() - ChronoDuration::days(1))
|
||||
.format("%Y-%m-%d")
|
||||
.to_string();
|
||||
assert_eq!(resp.date_iso, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trigger_digest_rejects_malformed_date() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let req = TriggerDigestRequest {
|
||||
date_iso: Some("not-a-date".into()),
|
||||
};
|
||||
let err = trigger_digest_rpc(&cfg, req).await.unwrap_err();
|
||||
assert!(
|
||||
err.contains("invalid date_iso"),
|
||||
"expected schema-shaped error message, got: {err}"
|
||||
);
|
||||
assert_eq!(count_total(&cfg).unwrap(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trigger_digest_dedupes_active_jobs() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let req = TriggerDigestRequest {
|
||||
date_iso: Some("2026-04-27".into()),
|
||||
};
|
||||
let first = trigger_digest_rpc(&cfg, req.clone()).await.unwrap().value;
|
||||
let second = trigger_digest_rpc(&cfg, req).await.unwrap().value;
|
||||
assert!(first.enqueued);
|
||||
assert!(!second.enqueued, "duplicate must be dedupe-suppressed");
|
||||
assert!(second.job_id.is_none());
|
||||
assert_eq!(count_total(&cfg).unwrap(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,6 +23,7 @@ pub fn all_controller_schemas() -> Vec<ControllerSchema> {
|
||||
schemas("ingest"),
|
||||
schemas("list_chunks"),
|
||||
schemas("get_chunk"),
|
||||
schemas("trigger_digest"),
|
||||
]
|
||||
}
|
||||
|
||||
@@ -40,6 +41,10 @@ pub fn all_registered_controllers() -> Vec<RegisteredController> {
|
||||
schema: schemas("get_chunk"),
|
||||
handler: handle_get_chunk,
|
||||
},
|
||||
RegisteredController {
|
||||
schema: schemas("trigger_digest"),
|
||||
handler: handle_trigger_digest,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
@@ -184,6 +189,44 @@ pub fn schemas(function: &str) -> ControllerSchema {
|
||||
required: false,
|
||||
}],
|
||||
},
|
||||
"trigger_digest" => ControllerSchema {
|
||||
namespace: NAMESPACE,
|
||||
function: "trigger_digest",
|
||||
description: "Manually enqueue a daily-digest job for the global \
|
||||
tree. Idempotent — re-running for a day that already has a \
|
||||
digest is a no-op (the handler skips). When no date is \
|
||||
supplied, defaults to yesterday in UTC, matching the \
|
||||
scheduler's autonomous behavior.",
|
||||
inputs: vec![FieldSchema {
|
||||
name: "date_iso",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "UTC calendar date as `YYYY-MM-DD`. Optional; \
|
||||
defaults to yesterday when omitted.",
|
||||
required: false,
|
||||
}],
|
||||
outputs: vec![
|
||||
FieldSchema {
|
||||
name: "enqueued",
|
||||
ty: TypeSchema::Bool,
|
||||
comment: "True when a fresh job row was inserted; false \
|
||||
when an active job for the same date suppressed it.",
|
||||
required: true,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "job_id",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "ID of the newly enqueued job row, when enqueued.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "date_iso",
|
||||
ty: TypeSchema::String,
|
||||
comment: "The date the digest will cover, echoed back \
|
||||
as `YYYY-MM-DD`.",
|
||||
required: true,
|
||||
},
|
||||
],
|
||||
},
|
||||
_ => ControllerSchema {
|
||||
namespace: NAMESPACE,
|
||||
function: "unknown",
|
||||
@@ -228,6 +271,14 @@ fn handle_get_chunk(params: Map<String, Value>) -> ControllerFuture {
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_trigger_digest(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let config = config_rpc::load_config_with_timeout().await?;
|
||||
let req = parse_value::<tree_rpc::TriggerDigestRequest>(Value::Object(params))?;
|
||||
to_json(tree_rpc::trigger_digest_rpc(&config, req).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_value<T: DeserializeOwned>(v: Value) -> Result<T, String> {
|
||||
serde_json::from_value(v).map_err(|e| format!("invalid params: {e}"))
|
||||
}
|
||||
|
||||
@@ -36,7 +36,7 @@ use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::types::{EntityKind, ExtractedEntities, ExtractedEntity};
|
||||
use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic};
|
||||
use super::EntityExtractor;
|
||||
|
||||
// ── Configuration ────────────────────────────────────────────────────────
|
||||
@@ -59,6 +59,15 @@ pub struct LlmExtractorConfig {
|
||||
/// If true, drop entities whose declared kind isn't in `allowed_kinds`
|
||||
/// instead of falling back to [`EntityKind::Misc`].
|
||||
pub strict_kinds: bool,
|
||||
/// If true, the system prompt asks the model to also emit a
|
||||
/// `topics` array (free-form theme labels), and the response parser
|
||||
/// populates [`ExtractedEntities::topics`]. Default `false` — the
|
||||
/// extractor's primary job is named-entity extraction; topics are
|
||||
/// an opt-in side-channel for callers that need a thematic
|
||||
/// summary in the same call (e.g. running over a sealed summary's
|
||||
/// content). Adds prompt tokens and gives the model one more
|
||||
/// schema field to keep track of, so leave off unless needed.
|
||||
pub emit_topics: bool,
|
||||
}
|
||||
|
||||
impl Default for LlmExtractorConfig {
|
||||
@@ -73,8 +82,13 @@ impl Default for LlmExtractorConfig {
|
||||
EntityKind::Location,
|
||||
EntityKind::Event,
|
||||
EntityKind::Product,
|
||||
EntityKind::Datetime,
|
||||
EntityKind::Technology,
|
||||
EntityKind::Artifact,
|
||||
EntityKind::Quantity,
|
||||
],
|
||||
strict_kinds: false,
|
||||
emit_topics: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -103,7 +117,7 @@ impl LlmEntityExtractor {
|
||||
messages: vec![
|
||||
OllamaMessage {
|
||||
role: "system".to_string(),
|
||||
content: SYSTEM_PROMPT.to_string(),
|
||||
content: build_system_prompt(self.cfg.emit_topics),
|
||||
},
|
||||
OllamaMessage {
|
||||
role: "user".to_string(),
|
||||
@@ -128,19 +142,54 @@ impl EntityExtractor for LlmEntityExtractor {
|
||||
// JSON parse) is logged as a warn and returns an empty
|
||||
// `ExtractedEntities` rather than `Err`. This makes the extractor
|
||||
// safe to call from any context, not just `score_chunk` (which
|
||||
// separately catches errors from its own extractor chain). A caller
|
||||
// distinguishes "LLM had nothing to say" from "LLM ran and returned
|
||||
// zero entities" by inspecting `llm_importance` — `None` means the
|
||||
// call didn't complete successfully.
|
||||
Ok(self.extract_or_empty(text).await)
|
||||
// separately catches errors from its own extractor chain).
|
||||
//
|
||||
// Transport failures get bounded retry-with-backoff before falling
|
||||
// back to empty — see [`Self::try_extract`]. Non-transport failures
|
||||
// (HTTP non-success, malformed JSON) fall back immediately because
|
||||
// retrying the same input would yield the same bad response.
|
||||
const MAX_ATTEMPTS: u32 = 3;
|
||||
const BASE_BACKOFF_MS: u64 = 250;
|
||||
|
||||
for attempt in 0..MAX_ATTEMPTS {
|
||||
match self.try_extract(text).await {
|
||||
Some(extracted) => return Ok(extracted),
|
||||
None => {
|
||||
// Transport failure. Retry with exponential backoff
|
||||
// unless we've exhausted attempts.
|
||||
if attempt + 1 < MAX_ATTEMPTS {
|
||||
let delay_ms = BASE_BACKOFF_MS * 2u64.pow(attempt);
|
||||
log::warn!(
|
||||
"[memory_tree::extract::llm] transport failure, retrying in \
|
||||
{delay_ms}ms (attempt {}/{})",
|
||||
attempt + 2,
|
||||
MAX_ATTEMPTS
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
log::warn!(
|
||||
"[memory_tree::extract::llm] transport failed after {} attempts — \
|
||||
returning empty extraction",
|
||||
MAX_ATTEMPTS
|
||||
);
|
||||
Ok(ExtractedEntities::default())
|
||||
}
|
||||
}
|
||||
|
||||
impl LlmEntityExtractor {
|
||||
/// Internal: wraps the actual HTTP call and returns `ExtractedEntities`
|
||||
/// for every failure mode via soft-fallback. Split out of `extract` so
|
||||
/// the error branches can share logging without `?`-propagation.
|
||||
async fn extract_or_empty(&self, text: &str) -> ExtractedEntities {
|
||||
/// Internal: one attempt at calling Ollama.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `Some(extracted)` — call completed (HTTP returned). Includes the
|
||||
/// "HTTP non-success" and "malformed JSON" cases, which return
|
||||
/// `Some(empty)` because retrying the same input won't help.
|
||||
/// - `None` — transport-level failure (DNS, connect refused, timeout
|
||||
/// before any HTTP response). Caller may retry.
|
||||
async fn try_extract(&self, text: &str) -> Option<ExtractedEntities> {
|
||||
let url = format!("{}/api/chat", self.cfg.endpoint.trim_end_matches('/'));
|
||||
let body = self.build_request(text);
|
||||
log::debug!(
|
||||
@@ -152,11 +201,8 @@ impl LlmEntityExtractor {
|
||||
let resp = match self.http.post(&url).json(&body).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[memory_tree::extract::llm] transport failure to {url}: {e} — \
|
||||
returning empty extraction"
|
||||
);
|
||||
return ExtractedEntities::default();
|
||||
log::warn!("[memory_tree::extract::llm] transport failure to {url}: {e}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -168,7 +214,7 @@ impl LlmEntityExtractor {
|
||||
returning empty extraction",
|
||||
truncate_for_log(&body, 200)
|
||||
);
|
||||
return ExtractedEntities::default();
|
||||
return Some(ExtractedEntities::default());
|
||||
}
|
||||
|
||||
let envelope: OllamaChatResponse = match resp.json().await {
|
||||
@@ -178,7 +224,7 @@ impl LlmEntityExtractor {
|
||||
"[memory_tree::extract::llm] response body not Ollama-shaped JSON: {e} — \
|
||||
returning empty extraction"
|
||||
);
|
||||
return ExtractedEntities::default();
|
||||
return Some(ExtractedEntities::default());
|
||||
}
|
||||
};
|
||||
log::debug!(
|
||||
@@ -194,38 +240,97 @@ impl LlmEntityExtractor {
|
||||
response: {e}; content was: {} — returning empty extraction",
|
||||
truncate_for_log(&envelope.message.content, 400)
|
||||
);
|
||||
return ExtractedEntities::default();
|
||||
return Some(ExtractedEntities::default());
|
||||
}
|
||||
};
|
||||
|
||||
parsed.into_extracted_entities(text, &self.cfg)
|
||||
Some(parsed.into_extracted_entities(text, &self.cfg))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Prompt ───────────────────────────────────────────────────────────────
|
||||
|
||||
const SYSTEM_PROMPT: &str = "\
|
||||
You are a named-entity extractor and importance rater. Return JSON only — \
|
||||
/// Build the system prompt for the extractor. When `emit_topics` is true
|
||||
/// the schema, required-fields list, and example outputs include a
|
||||
/// `topics` array (free-form theme labels). When false the prompt
|
||||
/// matches the pre-flag behaviour exactly — no mention of topics
|
||||
/// anywhere — so the small model isn't asked to produce a field the
|
||||
/// caller doesn't want.
|
||||
fn build_system_prompt(emit_topics: bool) -> String {
|
||||
let topics_schema_line = if emit_topics {
|
||||
" \"topics\": [\"<short theme label>\"],\n"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
let topics_required = if emit_topics { "topics, " } else { "" };
|
||||
let fields_count = if emit_topics { "four" } else { "three" };
|
||||
let topics_guide = if emit_topics {
|
||||
"Topics are short free-form theme labels for what the text is ABOUT \
|
||||
(e.g. \"rate limiting\", \"memory tree\", \"auth flow\"). They are \
|
||||
distinct from entities — entities are specific named things mentioned \
|
||||
in the text; topics are the abstract themes those things relate to.\n"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
let example1_topics = if emit_topics {
|
||||
",\"topics\":[\"shipping\",\"auth\"]"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
let example2_topics = if emit_topics {
|
||||
",\"topics\":[\"product launch\",\"revenue\"]"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
format!(
|
||||
"You are a named-entity extractor and importance rater. Return JSON only — \
|
||||
no prose, no markdown, no commentary. Do not summarize. Extract every named \
|
||||
entity mention you find, including duplicates, and rate the chunk's overall \
|
||||
importance as a float in [0.0, 1.0].
|
||||
|
||||
Schema:
|
||||
{
|
||||
{{
|
||||
\"entities\": [
|
||||
{ \"kind\": \"person|organization|location|event|product\",
|
||||
\"text\": \"<exact surface form as it appears in the text>\" }
|
||||
{{ \"kind\": \"person|organization|location|event|product|datetime|technology|artifact|quantity\",
|
||||
\"text\": \"<exact surface form as it appears in the text>\" }}
|
||||
],
|
||||
\"importance\": 0.0,
|
||||
{topics_schema_line} \"importance\": 0.0,
|
||||
\"importance_reason\": \"<one short sentence explaining the rating>\"
|
||||
}
|
||||
}}
|
||||
|
||||
Kinds guide:
|
||||
person named human (\"Alice\", \"Steven Enamakel\")
|
||||
organization company / team / project (\"Anthropic\", \"TinyHumans\")
|
||||
location place (\"SF office\", \"London\")
|
||||
event scheduled occurrence (\"Q2 launch\", \"design review\")
|
||||
product commercial offering (\"Claude Code\", \"OpenHuman\")
|
||||
datetime temporal expression (\"Friday\", \"Q2 2026\", \"EOD tomorrow\")
|
||||
technology tool / framework / language / service (\"Rust\", \"OAuth\", \"Slack API\")
|
||||
artifact code / ticket / doc reference (\"PR #934\", \"src/foo.rs\", \"OH-42\")
|
||||
quantity amount / metric / money (\"$5K\", \"20/min\", \"10k tokens\")
|
||||
|
||||
{topics_guide}
|
||||
If a mention doesn't clearly fit a kind above, omit it rather than guessing.
|
||||
Always emit ALL {fields_count} top-level fields (entities, {topics_required}importance, importance_reason),
|
||||
even when entities is empty.
|
||||
|
||||
Examples:
|
||||
|
||||
Input: alice and bob shipped the auth migration friday. PR #42 ships OAuth refactor in src/auth/.
|
||||
Output: {{\"entities\":[{{\"kind\":\"person\",\"text\":\"alice\"}},{{\"kind\":\"person\",\"text\":\"bob\"}},{{\"kind\":\"event\",\"text\":\"auth migration\"}},{{\"kind\":\"datetime\",\"text\":\"friday\"}},{{\"kind\":\"artifact\",\"text\":\"PR #42\"}},{{\"kind\":\"technology\",\"text\":\"OAuth\"}},{{\"kind\":\"artifact\",\"text\":\"src/auth/\"}}]{example1_topics},\"importance\":0.9,\"importance_reason\":\"explicit shipping commitment\"}}
|
||||
|
||||
Input: Anthropic shipped Claude Code in SF — $20M ARR target by Q2.
|
||||
Output: {{\"entities\":[{{\"kind\":\"organization\",\"text\":\"Anthropic\"}},{{\"kind\":\"product\",\"text\":\"Claude Code\"}},{{\"kind\":\"location\",\"text\":\"SF\"}},{{\"kind\":\"quantity\",\"text\":\"$20M ARR\"}},{{\"kind\":\"datetime\",\"text\":\"Q2\"}}]{example2_topics},\"importance\":0.85,\"importance_reason\":\"factual content with key business metric\"}}
|
||||
|
||||
Importance guide:
|
||||
0.9+ actionable decisions, key information, explicit commitments
|
||||
0.6+ substantive discussion, factual content, named entities
|
||||
0.3+ ambient context, low-density prose
|
||||
<0.3 reactions, acknowledgments, bots, trivial exchanges
|
||||
";
|
||||
"
|
||||
)
|
||||
}
|
||||
|
||||
// ── Wire types (Ollama API) ──────────────────────────────────────────────
|
||||
|
||||
@@ -265,6 +370,11 @@ struct OllamaResponseMessage {
|
||||
struct LlmExtractionOutput {
|
||||
#[serde(default)]
|
||||
entities: Vec<LlmEntity>,
|
||||
/// Free-form theme labels — populated only when the extractor is
|
||||
/// configured with `emit_topics = true`. Always tolerant of absence
|
||||
/// so models that ignore the field don't fail parsing.
|
||||
#[serde(default)]
|
||||
topics: Vec<String>,
|
||||
#[serde(default)]
|
||||
importance: Option<f32>,
|
||||
#[serde(default)]
|
||||
@@ -355,9 +465,26 @@ impl LlmExtractionOutput {
|
||||
|
||||
let llm_importance = self.importance.map(|v| v.clamp(0.0, 1.0));
|
||||
|
||||
// Topics: only populated when the caller enabled `emit_topics`
|
||||
// (the prompt asked for them). Otherwise this is empty by
|
||||
// default — the model didn't know to emit topics, so any value
|
||||
// here would be hallucination.
|
||||
let topics = self
|
||||
.topics
|
||||
.into_iter()
|
||||
.filter_map(|raw| {
|
||||
let label = raw.trim().to_string();
|
||||
if label.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ExtractedTopic { label, score: 0.85 })
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
ExtractedEntities {
|
||||
entities,
|
||||
topics: Vec::new(),
|
||||
topics,
|
||||
llm_importance,
|
||||
llm_importance_reason: self.importance_reason,
|
||||
}
|
||||
@@ -373,6 +500,14 @@ fn parse_kind(s: &str) -> Option<EntityKind> {
|
||||
"location" | "place" | "loc" => Some(EntityKind::Location),
|
||||
"event" => Some(EntityKind::Event),
|
||||
"product" => Some(EntityKind::Product),
|
||||
"datetime" | "date" | "time" | "timestamp" => Some(EntityKind::Datetime),
|
||||
"technology" | "tech" | "tool" | "framework" | "library" | "language" | "service" => {
|
||||
Some(EntityKind::Technology)
|
||||
}
|
||||
"artifact" | "reference" | "ref" | "pr" | "ticket" | "file" | "commit" => {
|
||||
Some(EntityKind::Artifact)
|
||||
}
|
||||
"quantity" | "amount" | "metric" | "number" | "money" => Some(EntityKind::Quantity),
|
||||
"misc" | "miscellaneous" | "other" => Some(EntityKind::Misc),
|
||||
_ => None,
|
||||
}
|
||||
|
||||
@@ -1,5 +1,39 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn build_system_prompt_default_omits_topics() {
|
||||
let p = build_system_prompt(false);
|
||||
assert!(!p.contains("\"topics\""));
|
||||
assert!(!p.contains("Topics are"));
|
||||
assert!(p.contains("ALL three top-level fields"));
|
||||
assert!(p.contains("entities, importance"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_prompt_with_flag_includes_topics() {
|
||||
let p = build_system_prompt(true);
|
||||
assert!(p.contains("\"topics\""));
|
||||
assert!(p.contains("Topics are short free-form theme labels"));
|
||||
assert!(p.contains("ALL four top-level fields"));
|
||||
assert!(p.contains("entities, topics, importance"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extraction_output_parses_topics_when_present() {
|
||||
let json = r#"{"entities":[],"topics":["rate limiting","memory tree"],"importance":0.6,"importance_reason":"r"}"#;
|
||||
let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(parsed.topics, vec!["rate limiting", "memory tree"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extraction_output_tolerates_missing_topics() {
|
||||
// Default extractor (emit_topics=false) — model won't emit topics
|
||||
// and parsing must still succeed.
|
||||
let json = r#"{"entities":[],"importance":0.6,"importance_reason":"r"}"#;
|
||||
let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap();
|
||||
assert!(parsed.topics.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_kind_normalisation() {
|
||||
assert_eq!(parse_kind("Person"), Some(EntityKind::Person));
|
||||
@@ -8,6 +42,42 @@ fn parse_kind_normalisation() {
|
||||
assert!(parse_kind("Spaceship").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_kind_accepts_new_semantic_kinds_and_synonyms() {
|
||||
// Datetime
|
||||
for s in ["datetime", "date", "time", "timestamp", " DateTime "] {
|
||||
assert_eq!(parse_kind(s), Some(EntityKind::Datetime), "input={s:?}");
|
||||
}
|
||||
// Technology
|
||||
for s in [
|
||||
"technology",
|
||||
"tech",
|
||||
"tool",
|
||||
"framework",
|
||||
"library",
|
||||
"language",
|
||||
"service",
|
||||
] {
|
||||
assert_eq!(parse_kind(s), Some(EntityKind::Technology), "input={s:?}");
|
||||
}
|
||||
// Artifact
|
||||
for s in [
|
||||
"artifact",
|
||||
"reference",
|
||||
"ref",
|
||||
"pr",
|
||||
"ticket",
|
||||
"file",
|
||||
"commit",
|
||||
] {
|
||||
assert_eq!(parse_kind(s), Some(EntityKind::Artifact), "input={s:?}");
|
||||
}
|
||||
// Quantity
|
||||
for s in ["quantity", "amount", "metric", "number", "money"] {
|
||||
assert_eq!(parse_kind(s), Some(EntityKind::Quantity), "input={s:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_char_span_handles_unicode() {
|
||||
let text = "中 Alice met Bob";
|
||||
@@ -71,6 +141,7 @@ fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() {
|
||||
text: "Alice".into(),
|
||||
},
|
||||
],
|
||||
topics: vec![],
|
||||
importance: None,
|
||||
importance_reason: None,
|
||||
};
|
||||
@@ -100,6 +171,7 @@ fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() {
|
||||
text: "Alice".into(),
|
||||
},
|
||||
],
|
||||
topics: vec![],
|
||||
importance: None,
|
||||
importance_reason: None,
|
||||
};
|
||||
@@ -138,6 +210,7 @@ fn into_extracted_entities_drops_hallucinations() {
|
||||
text: "ImaginaryPerson".into(),
|
||||
},
|
||||
],
|
||||
topics: vec![],
|
||||
importance: Some(0.7),
|
||||
importance_reason: Some("substantive".into()),
|
||||
};
|
||||
@@ -154,6 +227,7 @@ fn into_extracted_entities_drops_hallucinations() {
|
||||
fn into_extracted_entities_clamps_importance() {
|
||||
let out = LlmExtractionOutput {
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
importance: Some(1.5),
|
||||
importance_reason: None,
|
||||
};
|
||||
@@ -169,6 +243,7 @@ fn into_extracted_entities_strict_drops_unknown_kinds() {
|
||||
kind: "spaceship".into(),
|
||||
text: "Enterprise".into(),
|
||||
}],
|
||||
topics: vec![],
|
||||
importance: None,
|
||||
importance_reason: None,
|
||||
};
|
||||
@@ -187,6 +262,7 @@ fn into_extracted_entities_lenient_falls_back_to_misc() {
|
||||
kind: "spaceship".into(),
|
||||
text: "Enterprise".into(),
|
||||
}],
|
||||
topics: vec![],
|
||||
importance: None,
|
||||
importance_reason: None,
|
||||
};
|
||||
@@ -204,6 +280,7 @@ fn into_extracted_entities_disallowed_known_kind_falls_back_to_misc() {
|
||||
kind: "person".into(),
|
||||
text: "Alice".into(),
|
||||
}],
|
||||
topics: vec![],
|
||||
importance: None,
|
||||
importance_reason: None,
|
||||
};
|
||||
|
||||
@@ -10,6 +10,83 @@ pub mod llm;
|
||||
pub mod regex;
|
||||
pub mod types;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::util::redact::redact_endpoint;
|
||||
|
||||
pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor};
|
||||
pub use llm::{LlmEntityExtractor, LlmExtractorConfig};
|
||||
pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic};
|
||||
|
||||
/// Build the extractor used by seal handlers to label new summary nodes.
|
||||
///
|
||||
/// Composition:
|
||||
/// - regex extractor — always on, mechanical, near-zero cost
|
||||
/// - LLM extractor with `emit_topics: true` — added when
|
||||
/// `memory_tree.llm_extractor_endpoint` and `..._model` are both set
|
||||
///
|
||||
/// Differs from [`super::ScoringConfig::from_config`] (the chunk-admission
|
||||
/// builder) in two ways: returns *just* an extractor (no thresholds /
|
||||
/// weights / drop logic — none of which apply at seal time), and flips
|
||||
/// `emit_topics` on so summaries surface thematic labels alongside
|
||||
/// entities. Leaf-side scoring is unchanged.
|
||||
pub fn build_summary_extractor(config: &Config) -> Arc<dyn EntityExtractor> {
|
||||
let endpoint = config
|
||||
.memory_tree
|
||||
.llm_extractor_endpoint
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty());
|
||||
let model = config
|
||||
.memory_tree
|
||||
.llm_extractor_model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty());
|
||||
|
||||
let (Some(endpoint), Some(model)) = (endpoint, model) else {
|
||||
log::debug!(
|
||||
"[memory_tree::extract] summary extractor: LLM not configured — using regex-only"
|
||||
);
|
||||
return Arc::new(CompositeExtractor::regex_only());
|
||||
};
|
||||
|
||||
let timeout_ms = config
|
||||
.memory_tree
|
||||
.llm_extractor_timeout_ms
|
||||
.unwrap_or(15_000);
|
||||
|
||||
let cfg = LlmExtractorConfig {
|
||||
endpoint: endpoint.to_string(),
|
||||
model: model.to_string(),
|
||||
timeout: std::time::Duration::from_millis(timeout_ms),
|
||||
emit_topics: true,
|
||||
..LlmExtractorConfig::default()
|
||||
};
|
||||
|
||||
match LlmEntityExtractor::new(cfg) {
|
||||
Ok(llm) => {
|
||||
// Drop to debug (diagnostic, not always-on) and redact the endpoint
|
||||
// so embedded credentials (e.g. api keys in URL) don't leak.
|
||||
log::debug!(
|
||||
"[memory_tree::extract] summary extractor: regex + LLM endpoint={} model={} \
|
||||
timeout_ms={} emit_topics=true",
|
||||
redact_endpoint(endpoint),
|
||||
model,
|
||||
timeout_ms
|
||||
);
|
||||
Arc::new(CompositeExtractor::new(vec![
|
||||
Box::new(RegexEntityExtractor),
|
||||
Box::new(llm),
|
||||
]))
|
||||
}
|
||||
Err(err) => {
|
||||
log::warn!(
|
||||
"[memory_tree::extract] summary extractor: LlmEntityExtractor construction \
|
||||
failed: {err:#} — falling back to regex-only"
|
||||
);
|
||||
Arc::new(CompositeExtractor::regex_only())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,12 +26,22 @@ pub enum EntityKind {
|
||||
Url,
|
||||
Handle,
|
||||
Hashtag,
|
||||
// Semantic (reserved — not emitted in Phase 2)
|
||||
// Semantic — emitted by the LLM extractor.
|
||||
Person,
|
||||
Organization,
|
||||
Location,
|
||||
Event,
|
||||
Product,
|
||||
/// Temporal expressions: "Friday", "Q2 2026", "EOD tomorrow", "next sprint".
|
||||
Datetime,
|
||||
/// Tools / frameworks / programming languages / services:
|
||||
/// "Rust", "OAuth", "Slack API", "nomic-embed".
|
||||
Technology,
|
||||
/// Code / ticket / doc references that point at something addressable:
|
||||
/// "PR #934", "src/openhuman/...", "OH-42", "ab7da2e2".
|
||||
Artifact,
|
||||
/// Amounts / metrics / money: "$5K", "20/min", "10k tokens", "52 chunks".
|
||||
Quantity,
|
||||
Misc,
|
||||
// Thematic — scorer-surfaced topics (hashtag-like short phrases or
|
||||
// LLM-extracted themes). Promoted into the canonical entity stream
|
||||
@@ -54,6 +64,10 @@ impl EntityKind {
|
||||
Self::Location => "location",
|
||||
Self::Event => "event",
|
||||
Self::Product => "product",
|
||||
Self::Datetime => "datetime",
|
||||
Self::Technology => "technology",
|
||||
Self::Artifact => "artifact",
|
||||
Self::Quantity => "quantity",
|
||||
Self::Misc => "misc",
|
||||
Self::Topic => "topic",
|
||||
}
|
||||
@@ -70,6 +84,10 @@ impl EntityKind {
|
||||
"location" => Ok(Self::Location),
|
||||
"event" => Ok(Self::Event),
|
||||
"product" => Ok(Self::Product),
|
||||
"datetime" => Ok(Self::Datetime),
|
||||
"technology" => Ok(Self::Technology),
|
||||
"artifact" => Ok(Self::Artifact),
|
||||
"quantity" => Ok(Self::Quantity),
|
||||
"misc" => Ok(Self::Misc),
|
||||
"topic" => Ok(Self::Topic),
|
||||
other => Err(format!("unknown entity kind: {other}")),
|
||||
@@ -207,7 +225,12 @@ mod tests {
|
||||
EntityKind::Location,
|
||||
EntityKind::Event,
|
||||
EntityKind::Product,
|
||||
EntityKind::Datetime,
|
||||
EntityKind::Technology,
|
||||
EntityKind::Artifact,
|
||||
EntityKind::Quantity,
|
||||
EntityKind::Misc,
|
||||
EntityKind::Topic,
|
||||
] {
|
||||
assert_eq!(EntityKind::parse(k.as_str()).unwrap(), k);
|
||||
}
|
||||
|
||||
@@ -336,6 +336,22 @@ pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result<Vec<S
|
||||
try_join_all(chunks.iter().map(|chunk| score_chunk(chunk, cfg))).await
|
||||
}
|
||||
|
||||
/// Cheap-only batch scoring path used by the async queue ingest pipeline.
|
||||
///
|
||||
/// This preserves the same thresholds and admission gate as [`score_chunks`]
|
||||
/// but guarantees no LLM extractor is consulted on the ingest hot path.
|
||||
pub async fn score_chunks_fast(chunks: &[Chunk], cfg: &ScoringConfig) -> Result<Vec<ScoreResult>> {
|
||||
let fast_cfg = ScoringConfig {
|
||||
extractor: cfg.extractor.clone(),
|
||||
weights: cfg.weights.clone(),
|
||||
drop_threshold: cfg.drop_threshold,
|
||||
llm_extractor: None,
|
||||
definite_keep_threshold: cfg.definite_keep_threshold,
|
||||
definite_drop_threshold: cfg.definite_drop_threshold,
|
||||
};
|
||||
score_chunks(chunks, &fast_cfg).await
|
||||
}
|
||||
|
||||
// ── Persistence helpers used by the ingest orchestrator ─────────────────
|
||||
|
||||
/// Persist the score row + entity-index rows for one kept chunk.
|
||||
|
||||
@@ -11,6 +11,7 @@ fn test_chunk(content: &str) -> Chunk {
|
||||
metadata: meta,
|
||||
seq_in_source: 0,
|
||||
created_at: Utc::now(),
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -371,6 +371,21 @@ pub fn lookup_entity(
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_entity_ids_for_node(config: &Config, node_id: &str) -> Result<Vec<String>> {
|
||||
with_connection(config, |conn| {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT DISTINCT entity_id
|
||||
FROM mem_tree_entity_index
|
||||
WHERE node_id = ?1
|
||||
ORDER BY score DESC, timestamp_ms DESC, entity_id ASC",
|
||||
)?;
|
||||
let rows = stmt
|
||||
.query_map(params![node_id], |row| row.get::<_, String>(0))?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(rows)
|
||||
})
|
||||
}
|
||||
|
||||
/// Count rows in the entity index (for tests / diagnostics).
|
||||
pub fn count_entity_index(config: &Config) -> Result<u64> {
|
||||
with_connection(config, |conn| {
|
||||
|
||||
@@ -26,12 +26,20 @@
|
||||
//! summariser does no real I/O; when a networked summariser lands, wrap
|
||||
//! DB calls in `tokio::task::spawn_blocking` to keep the runtime healthy.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{DateTime, Utc};
|
||||
use rusqlite::Transaction;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::{
|
||||
atomic::stage_summary, paths::slugify_source_id, SummaryComposeInput, SummaryTreeKind,
|
||||
};
|
||||
use crate::openhuman::memory::tree::score::embed::build_embedder_from_config;
|
||||
use crate::openhuman::memory::tree::score::extract::EntityExtractor;
|
||||
use crate::openhuman::memory::tree::score::resolver::canonicalise;
|
||||
use crate::openhuman::memory::tree::source_tree::registry::new_summary_id;
|
||||
use crate::openhuman::memory::tree::source_tree::store;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::{
|
||||
@@ -47,6 +55,88 @@ use crate::openhuman::memory::tree::store::with_connection;
|
||||
/// realistic source.
|
||||
const MAX_CASCADE_DEPTH: u32 = 32;
|
||||
|
||||
/// How a sealed summary node's `entities` and `topics` fields get populated.
|
||||
///
|
||||
/// Each tree kind has different correct semantics:
|
||||
/// - **Source** trees use [`LabelStrategy::ExtractFromContent`] so the
|
||||
/// summariser's freshly-synthesised text gets its own pass through an
|
||||
/// extractor. Captures emergent themes that no individual leaf expressed.
|
||||
/// - **Global** trees use [`LabelStrategy::UnionFromChildren`] — their
|
||||
/// inputs are already-labeled source-tree summaries; union preserves
|
||||
/// labels for time-based retrieval ("days that mentioned Alice")
|
||||
/// without an LLM call.
|
||||
/// - **Topic** trees use [`LabelStrategy::Empty`] — their scope already
|
||||
/// pins the dominant theme; inheriting auxiliary entities would
|
||||
/// cross-pollinate unrelated topic trees and noise the entity index.
|
||||
#[derive(Clone)]
|
||||
pub enum LabelStrategy {
|
||||
/// Run the extractor on the new summary's content; canonicalise the
|
||||
/// result into `entities` (canonical_ids) and `topics` (labels).
|
||||
ExtractFromContent(Arc<dyn EntityExtractor>),
|
||||
/// Dedup-merge each input's `entities` and `topics` into the parent.
|
||||
UnionFromChildren,
|
||||
/// Leave both fields empty regardless of inputs.
|
||||
Empty,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for LabelStrategy {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ExtractFromContent(ex) => write!(f, "ExtractFromContent({})", ex.name()),
|
||||
Self::UnionFromChildren => f.write_str("UnionFromChildren"),
|
||||
Self::Empty => f.write_str("Empty"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve `entities` and `topics` for a freshly-summarised node according
|
||||
/// to the chosen strategy. Errors propagate from the extractor (when used).
|
||||
async fn resolve_labels(
|
||||
strategy: &LabelStrategy,
|
||||
inputs: &[SummaryInput],
|
||||
summary_content: &str,
|
||||
) -> Result<(Vec<String>, Vec<String>)> {
|
||||
match strategy {
|
||||
LabelStrategy::ExtractFromContent(extractor) => {
|
||||
let extracted = extractor
|
||||
.extract(summary_content)
|
||||
.await
|
||||
.context("seal-time extractor failed")?;
|
||||
let canonical = canonicalise(&extracted);
|
||||
let mut entities: Vec<String> = canonical
|
||||
.into_iter()
|
||||
.map(|c| c.canonical_id)
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
entities.sort();
|
||||
let mut topics: Vec<String> = extracted
|
||||
.topics
|
||||
.into_iter()
|
||||
.map(|t| t.label)
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
topics.sort();
|
||||
Ok((entities, topics))
|
||||
}
|
||||
LabelStrategy::UnionFromChildren => {
|
||||
let mut entities: BTreeSet<String> = BTreeSet::new();
|
||||
let mut topics: BTreeSet<String> = BTreeSet::new();
|
||||
for inp in inputs {
|
||||
for e in &inp.entities {
|
||||
entities.insert(e.clone());
|
||||
}
|
||||
for t in &inp.topics {
|
||||
topics.insert(t.clone());
|
||||
}
|
||||
}
|
||||
Ok((entities.into_iter().collect(), topics.into_iter().collect()))
|
||||
}
|
||||
LabelStrategy::Empty => Ok((Vec::new(), Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// A single leaf being appended to an L0 buffer.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LeafRef {
|
||||
@@ -61,17 +151,22 @@ pub struct LeafRef {
|
||||
|
||||
/// Append a leaf to the source tree for `tree`, sealing buffers as they
|
||||
/// fill. Returns the ids of any summaries that sealed during this call.
|
||||
///
|
||||
/// `strategy` controls how each sealed summary's `entities` and `topics`
|
||||
/// are populated — see [`LabelStrategy`].
|
||||
pub async fn append_leaf(
|
||||
config: &Config,
|
||||
tree: &Tree,
|
||||
leaf: &LeafRef,
|
||||
summariser: &dyn Summariser,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<Vec<String>> {
|
||||
log::debug!(
|
||||
"[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={}",
|
||||
"[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={} strategy={:?}",
|
||||
tree.id,
|
||||
leaf.chunk_id,
|
||||
leaf.token_count
|
||||
leaf.token_count,
|
||||
strategy
|
||||
);
|
||||
|
||||
// 1. Push leaf into L0 buffer (transactional).
|
||||
@@ -85,7 +180,24 @@ pub async fn append_leaf(
|
||||
)?;
|
||||
|
||||
// 2. Cascade seals upward until a level stays under budget.
|
||||
cascade_seals(config, tree, summariser).await
|
||||
cascade_seals(config, tree, summariser, strategy).await
|
||||
}
|
||||
|
||||
/// Queue-oriented variant of [`append_leaf`].
|
||||
///
|
||||
/// This only appends the leaf to the L0 buffer and returns whether the
|
||||
/// caller should enqueue a follow-up seal job for level 0.
|
||||
pub fn append_leaf_deferred(config: &Config, tree: &Tree, leaf: &LeafRef) -> Result<bool> {
|
||||
append_to_buffer(
|
||||
config,
|
||||
&tree.id,
|
||||
0,
|
||||
&leaf.chunk_id,
|
||||
leaf.token_count as i64,
|
||||
leaf.timestamp,
|
||||
)?;
|
||||
let buf = store::get_buffer(config, &tree.id, 0)?;
|
||||
Ok(should_seal(&buf))
|
||||
}
|
||||
|
||||
/// Transactionally append a single item to `(tree_id, level)`'s buffer.
|
||||
@@ -127,20 +239,25 @@ async fn cascade_seals(
|
||||
config: &Config,
|
||||
tree: &Tree,
|
||||
summariser: &dyn Summariser,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<Vec<String>> {
|
||||
cascade_all_from(config, tree, 0, summariser, None).await
|
||||
cascade_all_from(config, tree, 0, summariser, None, strategy).await
|
||||
}
|
||||
|
||||
/// Seal buffers starting at `start_level` and cascade upward. When
|
||||
/// `force_now` is `Some`, the buffer at `start_level` is sealed regardless
|
||||
/// of token budget (used by time-based flush). Upper levels are sealed
|
||||
/// only when they cross the budget.
|
||||
///
|
||||
/// `strategy` is forwarded to every sealed level — same semantics as
|
||||
/// [`append_leaf`].
|
||||
pub async fn cascade_all_from(
|
||||
config: &Config,
|
||||
tree: &Tree,
|
||||
start_level: u32,
|
||||
summariser: &dyn Summariser,
|
||||
force_now: Option<DateTime<Utc>>,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<Vec<String>> {
|
||||
let mut sealed_ids: Vec<String> = Vec::new();
|
||||
let mut level: u32 = start_level;
|
||||
@@ -169,7 +286,9 @@ pub async fn cascade_all_from(
|
||||
break;
|
||||
}
|
||||
|
||||
let summary_id = seal_one_level(config, tree, &buf, summariser).await?;
|
||||
// Sync cascade — drives the level walk itself; doesn't need the
|
||||
// queue follow-ups (we'll hit `seal_one_level` again next iter).
|
||||
let summary_id = seal_one_level(config, tree, &buf, summariser, strategy, false).await?;
|
||||
sealed_ids.push(summary_id);
|
||||
level += 1;
|
||||
}
|
||||
@@ -185,7 +304,7 @@ pub async fn cascade_all_from(
|
||||
/// summarisers that emit at the full token budget (e.g. the inert
|
||||
/// fallback) collapse the cascade into a 1:1:1 chain instead of a real
|
||||
/// tree.
|
||||
fn should_seal(buf: &Buffer) -> bool {
|
||||
pub(crate) fn should_seal(buf: &Buffer) -> bool {
|
||||
if buf.is_empty() {
|
||||
return false;
|
||||
}
|
||||
@@ -198,11 +317,31 @@ fn should_seal(buf: &Buffer) -> bool {
|
||||
|
||||
/// Seal `buf` at `level` into one summary at `level + 1`. Returns the new
|
||||
/// summary id.
|
||||
async fn seal_one_level(
|
||||
///
|
||||
/// `strategy` decides how `entities` and `topics` get populated on the new
|
||||
/// summary node — see [`LabelStrategy`].
|
||||
///
|
||||
/// When `enqueue_follow_ups` is `true`, the function additionally inserts
|
||||
/// follow-up job rows **inside the same transaction** that commits the
|
||||
/// seal:
|
||||
/// - `seal { tree_id, level: parent_level }` if the parent buffer's gate
|
||||
/// is now met (parent-cascade enqueue)
|
||||
/// - `topic_route { NodeRef::Summary { summary_id } }` for source trees
|
||||
/// (so summary-level entities feed the topic-tree spawn pipeline)
|
||||
///
|
||||
/// Atomic enqueue eliminates the crash window where a seal commits but
|
||||
/// the post-commit follow-up enqueues are silently lost on a worker
|
||||
/// crash. The async-pipeline handler (`handle_seal`) passes `true`. The
|
||||
/// synchronous in-process cascade caller ([`cascade_all_from`]) passes
|
||||
/// `false` because it drives the cascade itself and topic_route isn't
|
||||
/// part of the test/flush sync path.
|
||||
pub(crate) async fn seal_one_level(
|
||||
config: &Config,
|
||||
tree: &Tree,
|
||||
buf: &Buffer,
|
||||
summariser: &dyn Summariser,
|
||||
strategy: &LabelStrategy,
|
||||
enqueue_follow_ups: bool,
|
||||
) -> Result<String> {
|
||||
let level = buf.level;
|
||||
let target_level = level + 1;
|
||||
@@ -237,7 +376,7 @@ async fn seal_one_level(
|
||||
// Run summariser — async, OUTSIDE any DB transaction.
|
||||
let ctx = SummaryContext {
|
||||
tree_id: &tree.id,
|
||||
tree_kind: TreeKind::Source,
|
||||
tree_kind: tree.kind,
|
||||
target_level,
|
||||
token_budget: TOKEN_BUDGET,
|
||||
};
|
||||
@@ -246,6 +385,12 @@ async fn seal_one_level(
|
||||
.await
|
||||
.context("summariser failed during seal")?;
|
||||
|
||||
// Resolve labels (entities/topics) for the new summary node according
|
||||
// to the chosen strategy. Done before the write tx so an extractor
|
||||
// failure aborts the seal cleanly — same shape as the embedder guard
|
||||
// below.
|
||||
let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?;
|
||||
|
||||
// Phase 4 (#710): embed the summary BEFORE opening the write tx so an
|
||||
// embedder failure aborts the seal cleanly — nothing is persisted,
|
||||
// the buffer stays intact, and a retry re-embeds from scratch. The
|
||||
@@ -291,14 +436,18 @@ async fn seal_one_level(
|
||||
let node = SummaryNode {
|
||||
id: summary_id.clone(),
|
||||
tree_id: tree.id.clone(),
|
||||
tree_kind: TreeKind::Source,
|
||||
// `seal_one_level` runs for source AND topic trees (handle_seal,
|
||||
// cascade_all_from, flush). Hardcoding Source here would write
|
||||
// topic-tree summaries with tree_kind='source' in
|
||||
// mem_tree_summaries, breaking any query filtering on tree_kind.
|
||||
tree_kind: tree.kind,
|
||||
level: target_level,
|
||||
parent_id: None,
|
||||
child_ids: buf.item_ids.clone(),
|
||||
content: output.content,
|
||||
token_count: output.token_count,
|
||||
entities: output.entities,
|
||||
topics: output.topics,
|
||||
entities: node_entities,
|
||||
topics: node_topics,
|
||||
time_range_start,
|
||||
time_range_end,
|
||||
score,
|
||||
@@ -307,13 +456,88 @@ async fn seal_one_level(
|
||||
embedding: Some(embedding),
|
||||
};
|
||||
|
||||
// Phase MD-content: stage the summary .md file BEFORE opening the write
|
||||
// tx. A staging failure aborts the seal cleanly — nothing is persisted
|
||||
// and the buffer stays intact for retry.
|
||||
//
|
||||
// `bucket_seal.rs` handles both Source and Topic tree seals (Topic trees
|
||||
// use the same cascade machinery via `handle_seal` in the job handler).
|
||||
// Map TreeKind to SummaryTreeKind accordingly.
|
||||
let summary_tree_kind = match tree.kind {
|
||||
TreeKind::Topic => SummaryTreeKind::Topic,
|
||||
_ => SummaryTreeKind::Source,
|
||||
};
|
||||
let scope_slug = {
|
||||
// Path slug semantics per source kind:
|
||||
//
|
||||
// - Gmail source trees: scope is `"gmail:<participants>"` where
|
||||
// participants is `addr1|addr2|...`. Strip the `gmail:` prefix so the
|
||||
// path is `summaries/source/<participants_slug>/...` and mirrors the
|
||||
// chunk layout under `email/<participants_slug>/`.
|
||||
//
|
||||
// - Topic trees: scope is the canonical entity_id (e.g.
|
||||
// `"email:alice@example.com"`). Slugify the FULL string so topic-tree
|
||||
// summaries and source-tree summaries don't share a path prefix.
|
||||
//
|
||||
// - All other source kinds (slack:, discord:, document:, …): slugify the
|
||||
// FULL scope string. Stripping the prefix for non-Gmail sources was a
|
||||
// bug — `"slack:#eng"` and `"discord:#eng"` would both produce slug
|
||||
// `"eng"` and collide in `summaries/source/eng/`.
|
||||
let s = &tree.scope;
|
||||
match tree.kind {
|
||||
TreeKind::Topic => slugify_source_id(s),
|
||||
_ => {
|
||||
if s.starts_with("gmail:") {
|
||||
// Strip "gmail:" prefix; slugify the participants portion.
|
||||
slugify_source_id(&s["gmail:".len()..])
|
||||
} else {
|
||||
// All other source kinds: slugify the full scope string.
|
||||
slugify_source_id(s)
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
let compose_input = SummaryComposeInput {
|
||||
summary_id: &node.id,
|
||||
tree_kind: summary_tree_kind,
|
||||
tree_id: &node.tree_id,
|
||||
tree_scope: &tree.scope,
|
||||
level: node.level,
|
||||
child_ids: &node.child_ids,
|
||||
child_count: node.child_ids.len(),
|
||||
time_range_start: node.time_range_start,
|
||||
time_range_end: node.time_range_end,
|
||||
sealed_at: node.sealed_at,
|
||||
body: &node.content,
|
||||
};
|
||||
// Stage the summary .md file and propagate any error — a staging failure
|
||||
// aborts the seal entirely so the database never commits a row with
|
||||
// content_path = NULL. The buffer stays unsealed and the job-retry path
|
||||
// will re-attempt the file write on next execution.
|
||||
let content_root = config.memory_tree_content_root();
|
||||
let staged =
|
||||
stage_summary(&content_root, &compose_input, &scope_slug, None).with_context(|| {
|
||||
format!(
|
||||
"stage_summary failed for {}; seal aborted, buffer stays unsealed for retry",
|
||||
node.id
|
||||
)
|
||||
})?;
|
||||
log::debug!(
|
||||
"[source_tree::bucket_seal] staged summary {} → {}",
|
||||
node.id,
|
||||
staged.content_path
|
||||
);
|
||||
|
||||
// Single write transaction: insert summary, clear this buffer, append
|
||||
// summary id to parent buffer, bump tree max_level/root if needed.
|
||||
// summary id to parent buffer, bump tree max_level/root if needed,
|
||||
// and (when `enqueue_follow_ups`) atomically enqueue parent-seal +
|
||||
// topic_route follow-ups so they can never desync from the commit.
|
||||
// Re-read `max_level` from inside the tx so cascading seals within
|
||||
// one call see the updated value from earlier levels.
|
||||
let summary_id_for_closure = summary_id.clone();
|
||||
let target_level_for_closure = target_level;
|
||||
let tree_id = tree.id.clone();
|
||||
let tree_kind = tree.kind;
|
||||
with_connection(config, move |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
|
||||
@@ -326,7 +550,7 @@ async fn seal_one_level(
|
||||
.map(|n| n.max(0) as u32)
|
||||
.context("Failed to read current max_level for tree")?;
|
||||
|
||||
store::insert_summary_tx(&tx, &node)?;
|
||||
store::insert_summary_tx(&tx, &node, Some(&staged))?;
|
||||
// Forward-compat: index any entities the summariser emitted into
|
||||
// `mem_tree_entity_index` so Phase 4 retrieval can resolve
|
||||
// "summaries mentioning Alice" via the same inverted index as
|
||||
@@ -375,6 +599,43 @@ async fn seal_one_level(
|
||||
};
|
||||
store::upsert_buffer_tx(&tx, &parent)?;
|
||||
|
||||
// Atomic follow-up enqueues. Done INSIDE this tx — if the commit
|
||||
// rolls back, the queue rows go with it; if it succeeds, the
|
||||
// rows are durably visible to the worker pool. Eliminates the
|
||||
// crash window where the seal commits but post-commit enqueues
|
||||
// are lost.
|
||||
if enqueue_follow_ups {
|
||||
// Parent-cascade: if the new summary made the parent buffer
|
||||
// cross its gate, enqueue the next level's seal. Dedupe key
|
||||
// `seal:{tree_id}:{parent_level}` prevents duplicates if a
|
||||
// parallel path already queued it.
|
||||
if should_seal(&parent) {
|
||||
use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx;
|
||||
use crate::openhuman::memory::tree::jobs::types::{NewJob, SealPayload};
|
||||
let parent_seal = SealPayload {
|
||||
tree_id: tree_id.clone(),
|
||||
level: target_level_for_closure,
|
||||
force_now_ms: None,
|
||||
};
|
||||
enqueue_job_tx(&tx, &NewJob::seal(&parent_seal)?)?;
|
||||
}
|
||||
// Source-tree summary routing: feed the new summary's
|
||||
// entities back into the topic-tree spawn pipeline. Topic
|
||||
// and global trees are sinks — no fan-out from their seals.
|
||||
if matches!(tree_kind, TreeKind::Source) {
|
||||
use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx;
|
||||
use crate::openhuman::memory::tree::jobs::types::{
|
||||
NewJob, NodeRef, TopicRoutePayload,
|
||||
};
|
||||
let route = TopicRoutePayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: summary_id_for_closure.clone(),
|
||||
},
|
||||
};
|
||||
enqueue_job_tx(&tx, &NewJob::topic_route(&route)?)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Update tree root / max_level if we just climbed.
|
||||
if target_level_for_closure > current_max {
|
||||
store::update_tree_after_seal_tx(
|
||||
@@ -447,7 +708,8 @@ fn hydrate_inputs(config: &Config, level: u32, item_ids: &[String]) -> Result<Ve
|
||||
}
|
||||
|
||||
fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<SummaryInput>> {
|
||||
use crate::openhuman::memory::tree::score::store::get_score;
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
use crate::openhuman::memory::tree::score::store::{get_score, list_entity_ids_for_node};
|
||||
use crate::openhuman::memory::tree::store::get_chunk;
|
||||
|
||||
let mut out: Vec<SummaryInput> = Vec::with_capacity(chunk_ids.len());
|
||||
@@ -461,17 +723,30 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<Summ
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let score = get_score(config, id)?;
|
||||
let (score_value, entities, topics) = match &score {
|
||||
Some(row) => (row.total, Vec::new(), chunk.metadata.tags.clone()),
|
||||
None => (0.0, Vec::new(), chunk.metadata.tags.clone()),
|
||||
};
|
||||
let score_value = get_score(config, id)?.map(|row| row.total).unwrap_or(0.0);
|
||||
// Pull canonical entity ids from the inverted index — that's the
|
||||
// authoritative source for "what entities are attached to this
|
||||
// chunk." Topics live on the chunk's metadata tags.
|
||||
// [`LabelStrategy::UnionFromChildren`] reads these fields off
|
||||
// each `SummaryInput` to roll labels up the tree.
|
||||
let entities = list_entity_ids_for_node(config, id).unwrap_or_default();
|
||||
// Read the full body from disk — the `content` column in SQLite holds
|
||||
// a ≤500-char preview after the MD-on-disk migration. The summariser
|
||||
// must receive the complete chunk text so the seal output is not a
|
||||
// summary of previews.
|
||||
//
|
||||
// For pre-MD-migration chunks (no content_path recorded) this call
|
||||
// returns Err; callers that want to handle legacy rows should check
|
||||
// content_path presence before calling hydrate_inputs.
|
||||
let body = content_read::read_chunk_body(config, id).with_context(|| {
|
||||
format!("[source_tree::bucket_seal] hydrate_leaf_inputs: read body for chunk {id}")
|
||||
})?;
|
||||
out.push(SummaryInput {
|
||||
id: chunk.id.clone(),
|
||||
content: chunk.content.clone(),
|
||||
content: body,
|
||||
token_count: chunk.token_count,
|
||||
entities,
|
||||
topics,
|
||||
topics: chunk.metadata.tags.clone(),
|
||||
time_range_start: chunk.metadata.time_range.0,
|
||||
time_range_end: chunk.metadata.time_range.1,
|
||||
score: score_value,
|
||||
@@ -481,6 +756,8 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<Summ
|
||||
}
|
||||
|
||||
fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result<Vec<SummaryInput>> {
|
||||
use crate::openhuman::memory::tree::content_store::read as content_read;
|
||||
|
||||
let mut out: Vec<SummaryInput> = Vec::with_capacity(summary_ids.len());
|
||||
for id in summary_ids {
|
||||
let node = match store::get_summary(config, id)? {
|
||||
@@ -492,9 +769,15 @@ fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result<Vec
|
||||
continue;
|
||||
}
|
||||
};
|
||||
// Read the full body from disk — `node.content` is a ≤500-char preview
|
||||
// after the MD-on-disk migration. Higher-level seals (L2+) summarise
|
||||
// over L1 summary content and need the full text, not a preview.
|
||||
let body = content_read::read_summary_body(config, id).with_context(|| {
|
||||
format!("[source_tree::bucket_seal] hydrate_summary_inputs: read body for summary {id}")
|
||||
})?;
|
||||
out.push(SummaryInput {
|
||||
id: node.id.clone(),
|
||||
content: node.content.clone(),
|
||||
content: body,
|
||||
token_count: node.token_count,
|
||||
entities: node.entities.clone(),
|
||||
topics: node.topics.clone(),
|
||||
|
||||
@@ -1,8 +1,29 @@
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Stage a batch of chunks to the content store so that `read_chunk_body`
|
||||
/// can find the on-disk file during seals. Tests that call `upsert_chunks`
|
||||
/// and then trigger a seal MUST also call this helper; otherwise
|
||||
/// `hydrate_leaf_inputs` will fail with "no content_path for chunk_id".
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[crate::openhuman::memory::tree::types::Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged =
|
||||
content_store::stage_chunks(&content_root, chunks).expect("stage_chunks for test chunks");
|
||||
// Record the content_path + content_sha256 pointers in SQLite so the
|
||||
// store's `get_chunk_content_pointers` can resolve them later.
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -36,7 +57,9 @@ async fn append_below_budget_does_not_seal() {
|
||||
// Chunks don't exist in DB — we're only exercising the buffer
|
||||
// accounting, which doesn't require leaf rows until a seal fires.
|
||||
let leaf = mk_leaf("leaf-1", 100, 1_700_000_000_000);
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(sealed.is_empty());
|
||||
|
||||
let buf = store::get_buffer(&cfg, &tree.id, 0).unwrap();
|
||||
@@ -72,15 +95,21 @@ async fn crossing_budget_triggers_seal() {
|
||||
token_count: tokens,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let c1 = mk_chunk(0, 6_000);
|
||||
let c2 = mk_chunk(1, 6_000);
|
||||
// Budget-relative sizes so the test stays correct as TOKEN_BUDGET shifts:
|
||||
// each leaf is 60% of budget, so the second append crosses the threshold.
|
||||
let per_leaf = TOKEN_BUDGET * 6 / 10;
|
||||
let c1 = mk_chunk(0, per_leaf);
|
||||
let c2 = mk_chunk(1, per_leaf);
|
||||
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
|
||||
// Stage both chunks to disk so the seal's hydrator can read full bodies.
|
||||
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
|
||||
|
||||
// Two leaves whose combined token_sum (12k) exceeds the 10k budget.
|
||||
let leaf1 = LeafRef {
|
||||
chunk_id: c1.id.clone(),
|
||||
token_count: 6_000,
|
||||
token_count: per_leaf,
|
||||
timestamp: ts,
|
||||
content: c1.content.clone(),
|
||||
entities: vec![],
|
||||
@@ -89,7 +118,7 @@ async fn crossing_budget_triggers_seal() {
|
||||
};
|
||||
let leaf2 = LeafRef {
|
||||
chunk_id: c2.id.clone(),
|
||||
token_count: 6_000,
|
||||
token_count: per_leaf,
|
||||
timestamp: ts,
|
||||
content: c2.content.clone(),
|
||||
entities: vec![],
|
||||
@@ -97,10 +126,14 @@ async fn crossing_budget_triggers_seal() {
|
||||
score: 0.5,
|
||||
};
|
||||
|
||||
let first = append_leaf(&cfg, &tree, &leaf1, &summariser).await.unwrap();
|
||||
let first = append_leaf(&cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(first.is_empty(), "first append below budget — no seal");
|
||||
|
||||
let second = append_leaf(&cfg, &tree, &leaf2, &summariser).await.unwrap();
|
||||
let second = append_leaf(&cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(second.len(), 1, "second append crosses budget — one seal");
|
||||
|
||||
let summary_id = &second[0];
|
||||
@@ -169,6 +202,7 @@ async fn fanout_at_l1_triggers_l2_seal() {
|
||||
token_count: 10_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
};
|
||||
|
||||
@@ -177,6 +211,8 @@ async fn fanout_at_l1_triggers_l2_seal() {
|
||||
for seq in 0..fanout {
|
||||
let chunk = mk_chunk(seq);
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so the seal hydrator can read the full body.
|
||||
stage_test_chunks(&cfg, &[chunk.clone()]);
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id.clone(),
|
||||
token_count: chunk.token_count,
|
||||
@@ -186,7 +222,9 @@ async fn fanout_at_l1_triggers_l2_seal() {
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
all_sealed.extend(sealed);
|
||||
}
|
||||
|
||||
@@ -253,8 +291,11 @@ async fn upper_level_does_not_seal_below_fanout() {
|
||||
token_count: 10_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so the seal hydrator can read the full body.
|
||||
stage_test_chunks(&cfg, &[chunk.clone()]);
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: chunk.token_count,
|
||||
@@ -264,7 +305,9 @@ async fn upper_level_does_not_seal_below_fanout() {
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
let _ = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
|
||||
let _ = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap();
|
||||
@@ -281,3 +324,334 @@ async fn upper_level_does_not_seal_below_fanout() {
|
||||
stop_before as u64
|
||||
);
|
||||
}
|
||||
|
||||
// ── LabelStrategy tests (#TBD) ────────────────────────────────────────────
|
||||
//
|
||||
// These exercise the three labeling modes seal_one_level supports. We use
|
||||
// a short token budget so the seal fires on a single leaf — keeps the
|
||||
// arithmetic of "what entities/topics end up on the parent" obvious.
|
||||
|
||||
/// Helper: persist a substantive chunk and return a `LeafRef` referencing
|
||||
/// it, with caller-supplied entity/topic labels (used by Union/Empty tests).
|
||||
///
|
||||
/// To match production, entity labels are written into `mem_tree_entity_index`
|
||||
/// (where seal-time hydration reads them from) and topic labels are stored
|
||||
/// on `chunk.metadata.tags` (the production source of leaf-level topics).
|
||||
fn seed_leaf(
|
||||
cfg: &Config,
|
||||
seq: u32,
|
||||
content: &str,
|
||||
entities: Vec<String>,
|
||||
topics: Vec<String>,
|
||||
) -> LeafRef {
|
||||
use crate::openhuman::memory::tree::score::extract::EntityKind;
|
||||
use crate::openhuman::memory::tree::score::resolver::CanonicalEntity;
|
||||
use crate::openhuman::memory::tree::score::store::index_entity;
|
||||
use crate::openhuman::memory::tree::store::upsert_chunks;
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
|
||||
use chrono::TimeZone;
|
||||
let ts = Utc
|
||||
.timestamp_millis_opt(1_700_000_000_000 + seq as i64)
|
||||
.unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, content),
|
||||
content: content.to_string(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: topics.clone(),
|
||||
source_ref: Some(SourceRef::new(format!("slack://x{seq}"))),
|
||||
},
|
||||
// Bust TOKEN_BUDGET in one leaf so the seal fires immediately.
|
||||
token_count: 10_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage the chunk to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// via `read_chunk_body` during a seal triggered by `append_leaf`.
|
||||
stage_test_chunks(cfg, &[chunk.clone()]);
|
||||
// Mirror production indexing: entities go into mem_tree_entity_index
|
||||
// so the seal hydrator can pull them via list_entity_ids_for_node.
|
||||
for entity_id in &entities {
|
||||
let kind = entity_id
|
||||
.split_once(':')
|
||||
.map_or(EntityKind::Misc, |(k, _)| {
|
||||
EntityKind::parse(k).unwrap_or(EntityKind::Misc)
|
||||
});
|
||||
let surface = entity_id
|
||||
.split_once(':')
|
||||
.map_or(entity_id.as_str(), |(_, v)| v);
|
||||
let e = CanonicalEntity {
|
||||
canonical_id: entity_id.clone(),
|
||||
kind,
|
||||
surface: surface.to_string(),
|
||||
span_start: 0,
|
||||
span_end: surface.len() as u32,
|
||||
score: 1.0,
|
||||
};
|
||||
index_entity(cfg, &e, &chunk.id, "leaf", ts.timestamp_millis(), None).unwrap();
|
||||
}
|
||||
LeafRef {
|
||||
chunk_id: chunk.id.clone(),
|
||||
token_count: chunk.token_count,
|
||||
timestamp: ts,
|
||||
content: chunk.content.clone(),
|
||||
entities,
|
||||
topics,
|
||||
score: 0.5,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seal_with_extract_strategy_populates_entities_and_topics() {
|
||||
use crate::openhuman::memory::tree::score::extract::{CompositeExtractor, EntityExtractor};
|
||||
use std::sync::Arc;
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
|
||||
// Content the regex extractor can find: an email and a hashtag. The
|
||||
// inert summariser concatenates leaf content into the L1 summary, so
|
||||
// these tokens survive into the summary text and the extractor finds
|
||||
// them when run on the summary content.
|
||||
let leaf = seed_leaf(
|
||||
&cfg,
|
||||
0,
|
||||
"alice@example.com is leading the #launch sprint this week.",
|
||||
vec![],
|
||||
vec![],
|
||||
);
|
||||
|
||||
let extractor: Arc<dyn EntityExtractor> = Arc::new(CompositeExtractor::regex_only());
|
||||
let strategy = LabelStrategy::ExtractFromContent(extractor);
|
||||
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &strategy)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sealed.len(), 1, "single 10k-token leaf should seal L0→L1");
|
||||
|
||||
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
|
||||
assert!(
|
||||
summary
|
||||
.entities
|
||||
.iter()
|
||||
.any(|e| e == "email:alice@example.com"),
|
||||
"ExtractFromContent should surface the email entity from summary text; got entities={:?}",
|
||||
summary.entities
|
||||
);
|
||||
assert!(
|
||||
summary.topics.iter().any(|t| t == "launch"),
|
||||
"ExtractFromContent should surface the hashtag-derived topic; got topics={:?}",
|
||||
summary.topics
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seal_with_union_strategy_inherits_labels_from_children() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
|
||||
// Two leaves with overlapping + distinct labels. Union should
|
||||
// dedup-merge them into the parent.
|
||||
let leaf1 = seed_leaf(
|
||||
&cfg,
|
||||
0,
|
||||
"first leaf body",
|
||||
vec!["email:alice@example.com".into(), "topic:phoenix".into()],
|
||||
vec!["phoenix".into(), "launch".into()],
|
||||
);
|
||||
let leaf2 = seed_leaf(
|
||||
&cfg,
|
||||
1,
|
||||
"second leaf body",
|
||||
vec!["email:alice@example.com".into(), "person:bob".into()],
|
||||
vec!["launch".into(), "qa".into()],
|
||||
);
|
||||
|
||||
// L0 seals when the budget is crossed. With each leaf at 10k tokens,
|
||||
// the first append triggers a seal containing only leaf1; we want a
|
||||
// seal containing both, so use UnionFromChildren and a single seal of
|
||||
// both leaves at once. The simplest way is to lower budget by sealing
|
||||
// two leaves into one buffer — the second append crosses budget, so
|
||||
// the seal contains [leaf1, leaf2].
|
||||
//
|
||||
// Adjust by using smaller token counts so both fit in L0 first, then
|
||||
// a third append triggers a seal containing both. Reuse the helper
|
||||
// and override the leaf's token_count for this test.
|
||||
// Each leaf at half the budget so two together hit threshold exactly.
|
||||
let per_leaf = TOKEN_BUDGET / 2;
|
||||
let leaf1 = LeafRef {
|
||||
token_count: per_leaf,
|
||||
..leaf1
|
||||
};
|
||||
let leaf2 = LeafRef {
|
||||
token_count: per_leaf,
|
||||
..leaf2
|
||||
};
|
||||
|
||||
// First leaf: under budget, no seal.
|
||||
let sealed_1 = append_leaf(
|
||||
&cfg,
|
||||
&tree,
|
||||
&leaf1,
|
||||
&summariser,
|
||||
&LabelStrategy::UnionFromChildren,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(sealed_1.is_empty());
|
||||
// Second leaf: crosses budget → one seal covering both leaves.
|
||||
let sealed_2 = append_leaf(
|
||||
&cfg,
|
||||
&tree,
|
||||
&leaf2,
|
||||
&summariser,
|
||||
&LabelStrategy::UnionFromChildren,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sealed_2.len(), 1);
|
||||
|
||||
let summary = store::get_summary(&cfg, &sealed_2[0]).unwrap().unwrap();
|
||||
let entities: std::collections::BTreeSet<&str> =
|
||||
summary.entities.iter().map(String::as_str).collect();
|
||||
let topics: std::collections::BTreeSet<&str> =
|
||||
summary.topics.iter().map(String::as_str).collect();
|
||||
assert!(entities.contains("email:alice@example.com"));
|
||||
assert!(entities.contains("topic:phoenix"));
|
||||
assert!(entities.contains("person:bob"));
|
||||
assert_eq!(
|
||||
entities.len(),
|
||||
3,
|
||||
"expected 3 unique entities; got {entities:?}"
|
||||
);
|
||||
assert!(topics.contains("phoenix"));
|
||||
assert!(topics.contains("launch"));
|
||||
assert!(topics.contains("qa"));
|
||||
assert_eq!(topics.len(), 3, "expected 3 unique topics; got {topics:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seal_with_empty_strategy_leaves_labels_empty() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
|
||||
// Leaf carries labels — Empty strategy should ignore them.
|
||||
let leaf = seed_leaf(
|
||||
&cfg,
|
||||
0,
|
||||
"alice@example.com discussing #launch",
|
||||
vec!["email:alice@example.com".into(), "topic:launch".into()],
|
||||
vec!["launch".into()],
|
||||
);
|
||||
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sealed.len(), 1);
|
||||
|
||||
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
|
||||
assert!(
|
||||
summary.entities.is_empty(),
|
||||
"Empty strategy must leave entities empty; got {:?}",
|
||||
summary.entities
|
||||
);
|
||||
assert!(
|
||||
summary.topics.is_empty(),
|
||||
"Empty strategy must leave topics empty; got {:?}",
|
||||
summary.topics
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topic_tree_seal_persists_topic_kind_not_source() {
|
||||
use crate::openhuman::memory::tree::source_tree::types::TreeStatus;
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
// Build a topic tree directly — `seal_one_level` runs for both
|
||||
// source and topic trees, and previously hardcoded Source on the
|
||||
// resulting summary regardless of the parent tree's kind.
|
||||
let tree = Tree {
|
||||
id: "topic-tree-test-id".to_string(),
|
||||
kind: TreeKind::Topic,
|
||||
scope: "topic:launch".to_string(),
|
||||
root_id: None,
|
||||
max_level: 0,
|
||||
status: TreeStatus::Active,
|
||||
created_at: Utc::now(),
|
||||
last_sealed_at: None,
|
||||
};
|
||||
store::insert_tree(&cfg, &tree).unwrap();
|
||||
|
||||
let summariser = InertSummariser::new();
|
||||
let leaf = seed_leaf(&cfg, 0, "topic content", vec![], vec![]);
|
||||
|
||||
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sealed.len(), 1);
|
||||
|
||||
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
|
||||
assert_eq!(
|
||||
summary.tree_kind,
|
||||
TreeKind::Topic,
|
||||
"topic-tree summary must persist tree_kind=Topic, not Source"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scope_slug_non_gmail_uses_full_scope() {
|
||||
// slack:#eng and discord:#eng must NOT produce the same scope slug.
|
||||
// Previously, stripping everything before ':' made both → "eng".
|
||||
// After Fix K, only gmail: strips the prefix — others use the full string.
|
||||
use crate::openhuman::memory::tree::content_store::paths::slugify_source_id;
|
||||
|
||||
// Verify that the slug logic produces distinct values for different platforms.
|
||||
let slack_slug = slugify_source_id("slack:#eng");
|
||||
let discord_slug = slugify_source_id("discord:#eng");
|
||||
assert_ne!(
|
||||
slack_slug, discord_slug,
|
||||
"slack:#eng and discord:#eng must produce distinct slugs; got slack={slack_slug:?} discord={discord_slug:?}"
|
||||
);
|
||||
// Both must include their platform prefix in the slug.
|
||||
assert!(
|
||||
slack_slug.contains("slack"),
|
||||
"slack slug must include 'slack'; got {slack_slug:?}"
|
||||
);
|
||||
assert!(
|
||||
discord_slug.contains("discord"),
|
||||
"discord slug must include 'discord'; got {discord_slug:?}"
|
||||
);
|
||||
|
||||
// Confirm gmail: correctly strips the "gmail:" prefix so the participants
|
||||
// portion (used as the bucket key) matches the chunk path layout.
|
||||
// scope_slug for a gmail source tree is built by stripping "gmail:" and
|
||||
// slugifying the remainder; the result must equal slugify of just the
|
||||
// participants string.
|
||||
let participants = "alice@x.com|bob@y.com";
|
||||
let participants_slug = slugify_source_id(participants);
|
||||
let gmail_scope = format!("gmail:{participants}");
|
||||
// Strip "gmail:" prefix as bucket_seal.rs does.
|
||||
let gmail_slug = slugify_source_id(&gmail_scope["gmail:".len()..]);
|
||||
assert_eq!(
|
||||
participants_slug, gmail_slug,
|
||||
"gmail scope_slug must equal slugify of participants portion; \
|
||||
participants_slug={participants_slug:?} gmail_slug={gmail_slug:?}"
|
||||
);
|
||||
|
||||
// Also assert the full-scope slug for gmail is DIFFERENT (shows the bug
|
||||
// would still exist if we used the full string for gmail).
|
||||
let gmail_full_slug = slugify_source_id(&gmail_scope);
|
||||
assert_ne!(
|
||||
gmail_full_slug, participants_slug,
|
||||
"slugifying the full 'gmail:...' scope must differ from the participants-only slug"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -13,7 +13,7 @@ use anyhow::Result;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::cascade_all_from;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{cascade_all_from, LabelStrategy};
|
||||
use crate::openhuman::memory::tree::source_tree::store;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
|
||||
use crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS;
|
||||
@@ -25,6 +25,7 @@ pub async fn flush_stale_buffers(
|
||||
config: &Config,
|
||||
max_age: Duration,
|
||||
summariser: &dyn Summariser,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<usize> {
|
||||
let now = Utc::now();
|
||||
let cutoff = now - max_age;
|
||||
@@ -48,7 +49,8 @@ pub async fn flush_stale_buffers(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let sealed = cascade_all_from(config, &tree, buf.level, summariser, Some(now)).await?;
|
||||
let sealed =
|
||||
cascade_all_from(config, &tree, buf.level, summariser, Some(now), strategy).await?;
|
||||
seals += sealed.len();
|
||||
}
|
||||
Ok(seals)
|
||||
@@ -58,11 +60,13 @@ pub async fn flush_stale_buffers(
|
||||
pub async fn flush_stale_buffers_default(
|
||||
config: &Config,
|
||||
summariser: &dyn Summariser,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<usize> {
|
||||
flush_stale_buffers(
|
||||
config,
|
||||
Duration::seconds(DEFAULT_FLUSH_AGE_SECS),
|
||||
summariser,
|
||||
strategy,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -74,15 +78,17 @@ pub async fn force_flush_tree(
|
||||
tree_id: &str,
|
||||
summariser: &dyn Summariser,
|
||||
now: Option<DateTime<Utc>>,
|
||||
strategy: &LabelStrategy,
|
||||
) -> Result<Vec<String>> {
|
||||
let tree = store::get_tree(config, tree_id)?
|
||||
.ok_or_else(|| anyhow::anyhow!("no tree with id {tree_id}"))?;
|
||||
cascade_all_from(config, &tree, 0, summariser, now).await
|
||||
cascade_all_from(config, &tree, 0, summariser, now, strategy).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::content_store;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
|
||||
@@ -90,6 +96,20 @@ mod tests {
|
||||
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).expect("create content_root for test");
|
||||
let staged = content_store::stage_chunks(&content_root, chunks)
|
||||
.expect("stage_chunks for test chunks");
|
||||
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.expect("persist staged chunk pointers");
|
||||
}
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
@@ -124,8 +144,10 @@ mod tests {
|
||||
token_count: 100,
|
||||
seq_in_source: 0,
|
||||
created_at: old_ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[c.clone()]).unwrap();
|
||||
stage_test_chunks(&cfg, &[c.clone()]);
|
||||
|
||||
let leaf = LeafRef {
|
||||
chunk_id: c.id.clone(),
|
||||
@@ -136,12 +158,15 @@ mod tests {
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
|
||||
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
|
||||
|
||||
let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser)
|
||||
append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
|
||||
|
||||
let seals =
|
||||
flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(seals, 1);
|
||||
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 1);
|
||||
|
||||
@@ -172,6 +197,7 @@ mod tests {
|
||||
token_count: 50,
|
||||
seq_in_source: 0,
|
||||
created_at: now,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[c.clone()]).unwrap();
|
||||
let leaf = LeafRef {
|
||||
@@ -183,11 +209,14 @@ mod tests {
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
|
||||
|
||||
let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser)
|
||||
append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let seals =
|
||||
flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(seals, 0);
|
||||
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
|
||||
}
|
||||
|
||||
@@ -22,7 +22,7 @@ pub mod store;
|
||||
pub mod summariser;
|
||||
pub mod types;
|
||||
|
||||
pub use bucket_seal::{append_leaf, LeafRef};
|
||||
pub use bucket_seal::{append_leaf, append_leaf_deferred, LabelStrategy, LeafRef};
|
||||
pub use registry::get_or_create_source_tree;
|
||||
pub use store::{get_summary_embedding, set_summary_embedding};
|
||||
pub use summariser::{build_summariser, inert::InertSummariser, llm::LlmSummariser, Summariser};
|
||||
|
||||
@@ -21,6 +21,7 @@ use chrono::{DateTime, TimeZone, Utc};
|
||||
use rusqlite::{params, Connection, OptionalExtension, Transaction};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::StagedSummary;
|
||||
use crate::openhuman::memory::tree::score::embed::{decode_optional_blob, pack_checked};
|
||||
use crate::openhuman::memory::tree::source_tree::types::{
|
||||
Buffer, SummaryNode, Tree, TreeKind, TreeStatus,
|
||||
@@ -180,7 +181,16 @@ fn row_to_tree(row: &rusqlite::Row<'_>) -> rusqlite::Result<Tree> {
|
||||
/// Phase 4 (#710): if `node.embedding` is `Some`, the packed vector is
|
||||
/// written to the `embedding` blob column; `None` writes NULL so legacy
|
||||
/// rows from Phases 1-3 (no embed) read back identically.
|
||||
pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Result<()> {
|
||||
///
|
||||
/// Phase MD-content: if `staged` is `Some`, writes `content_path` and
|
||||
/// `content_sha256` and truncates `content` to a ≤500-char preview. Callers
|
||||
/// that have not yet staged the file pass `None`, in which case the full
|
||||
/// `node.content` is stored (legacy behaviour).
|
||||
pub(crate) fn insert_summary_tx(
|
||||
tx: &Transaction<'_>,
|
||||
node: &SummaryNode,
|
||||
staged: Option<&StagedSummary>,
|
||||
) -> Result<()> {
|
||||
let embedding_blob: Option<Vec<u8>> = match node.embedding.as_deref() {
|
||||
Some(v) => Some(
|
||||
pack_checked(v)
|
||||
@@ -188,14 +198,30 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
|
||||
// Phase MD-content: when a staged file exists, truncate `content` to a
|
||||
// ≤500-char plain-text preview (char boundary safe via chars().take(500)).
|
||||
let (content_preview, content_path, content_sha256) = match staged {
|
||||
Some(s) => {
|
||||
let preview: String = node.content.chars().take(500).collect();
|
||||
(
|
||||
preview,
|
||||
Some(s.content_path.clone()),
|
||||
Some(s.content_sha256.clone()),
|
||||
)
|
||||
}
|
||||
None => (node.content.clone(), None, None),
|
||||
};
|
||||
|
||||
tx.execute(
|
||||
"INSERT OR IGNORE INTO mem_tree_summaries (
|
||||
id, tree_id, tree_kind, level, parent_id,
|
||||
child_ids_json, content, token_count,
|
||||
entities_json, topics_json,
|
||||
time_range_start_ms, time_range_end_ms,
|
||||
score, sealed_at_ms, deleted, embedding
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
|
||||
score, sealed_at_ms, deleted, embedding,
|
||||
content_path, content_sha256
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
|
||||
params![
|
||||
node.id,
|
||||
node.tree_id,
|
||||
@@ -203,7 +229,7 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
|
||||
node.level,
|
||||
node.parent_id,
|
||||
serde_json::to_string(&node.child_ids)?,
|
||||
node.content,
|
||||
content_preview,
|
||||
node.token_count,
|
||||
serde_json::to_string(&node.entities)?,
|
||||
serde_json::to_string(&node.topics)?,
|
||||
@@ -213,6 +239,8 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
|
||||
node.sealed_at.timestamp_millis(),
|
||||
node.deleted as i64,
|
||||
embedding_blob,
|
||||
content_path,
|
||||
content_sha256,
|
||||
],
|
||||
)
|
||||
.with_context(|| format!("Failed to insert summary id={}", node.id))?;
|
||||
|
||||
@@ -71,7 +71,7 @@ fn summary_insert_and_fetch() {
|
||||
let node = sample_summary("sum-1", "tree-1", 1);
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
insert_summary_tx(&tx, &node)?;
|
||||
insert_summary_tx(&tx, &node, None)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
@@ -90,8 +90,8 @@ fn summary_insert_is_idempotent_on_id() {
|
||||
let node = sample_summary("sum-1", "tree-1", 1);
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
insert_summary_tx(&tx, &node)?;
|
||||
insert_summary_tx(&tx, &node)?;
|
||||
insert_summary_tx(&tx, &node, None)?;
|
||||
insert_summary_tx(&tx, &node, None)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
|
||||
@@ -6,10 +6,12 @@
|
||||
//! When the source / topic / global tree's bucket-seal cascade decides to
|
||||
//! fold N contributions (raw leaves at L0→L1, or lower-level summaries at
|
||||
//! L_n→L_{n+1}), this summariser is asked to produce the parent node's
|
||||
//! `content` + derived `entities` + `topics`. The seal machinery itself
|
||||
//! (bucket budgeting, level promotion, `mem_tree_summaries` persistence)
|
||||
//! is unchanged — only the text inside the summary row differs from
|
||||
//! [`super::inert::InertSummariser`].
|
||||
//! `content`. The seal machinery itself (bucket budgeting, level
|
||||
//! promotion, `mem_tree_summaries` persistence) is unchanged — only the
|
||||
//! text inside the summary row differs from [`super::inert::InertSummariser`].
|
||||
//! Entities and topics on `SummaryOutput` are always emitted empty by
|
||||
//! this summariser; canonical entity ids are populated separately by the
|
||||
//! entity extractor.
|
||||
//!
|
||||
//! ## Soft-fallback contract
|
||||
//!
|
||||
@@ -18,14 +20,12 @@
|
||||
//! parent row. We therefore promise **never** to return `Err`: every
|
||||
//! failure (transport, HTTP status, JSON shape) falls back to the same
|
||||
//! deterministic concat-and-truncate behaviour as `InertSummariser` and
|
||||
//! logs a warn. Callers distinguish "LLM ran fine" from "we fell back"
|
||||
//! only by observing whether the returned `entities`/`topics` are
|
||||
//! populated — the inert branch emits empty vecs.
|
||||
//! logs a warn.
|
||||
//!
|
||||
//! ## Prompt shape
|
||||
//!
|
||||
//! The system prompt commits the model to returning JSON with the shape
|
||||
//! `{ summary, entities, topics }`. We use Ollama's `format: "json"` +
|
||||
//! `{ summary }`. We use Ollama's `format: "json"` +
|
||||
//! `temperature: 0.0` to maximise determinism — same knobs the entity
|
||||
//! extractor already uses with success.
|
||||
|
||||
@@ -40,15 +40,24 @@ use super::inert::InertSummariser;
|
||||
use super::{Summariser, SummaryContext, SummaryInput, SummaryOutput};
|
||||
use crate::openhuman::memory::tree::types::approx_token_count;
|
||||
|
||||
/// Hard cap on summary OUTPUT tokens, regardless of the seal's
|
||||
/// `ctx.token_budget`. Driven by the embedder's context window —
|
||||
/// `nomic-embed-text-v1.5` accepts up to 8192 tokens and Phase 4
|
||||
/// (`source_tree::bucket_seal`) embeds the summary right after we
|
||||
/// produce it. If our summary overshoots, the embedder returns 500
|
||||
/// and the whole seal transaction rolls back → no summary persists.
|
||||
/// 6000 leaves a safety margin for tokenizer differences and the
|
||||
/// JSON wrapper.
|
||||
const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 6_000;
|
||||
/// Hard cap on summariser output length (in approximate tokens).
|
||||
///
|
||||
/// Two constraints set this:
|
||||
///
|
||||
/// 1. The downstream embedder (`nomic-embed-text-v1.5`) accepts up to
|
||||
/// 8192 tokens, and Phase 4 (`source_tree::bucket_seal`) embeds the
|
||||
/// summary right after we produce it. An overshoot returns HTTP 500
|
||||
/// and rolls back the whole seal transaction.
|
||||
/// 2. Empirically, small instruction-tuned models running locally
|
||||
/// degrade quickly past ~3500 tokens — they drift, hallucinate, or
|
||||
/// produce repetitive boilerplate as they extend toward longer
|
||||
/// targets. Keeping the cap below that breakeven keeps output
|
||||
/// quality stable on local Ollama deployments.
|
||||
///
|
||||
/// 3500 sits comfortably under the embedder ceiling AND below the local
|
||||
/// LLM quality cliff. The post-generation [`clamp_to_budget`] enforces
|
||||
/// this regardless of what the model produces.
|
||||
const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 3_500;
|
||||
|
||||
/// Context window we ask Ollama for. Must match the value below in
|
||||
/// [`OllamaOptions::num_ctx`] so the per-input clamp computed in
|
||||
@@ -261,31 +270,18 @@ impl Summariser for LlmSummariser {
|
||||
|
||||
let (content, token_count) = clamp_to_budget(&parsed.summary, effective_budget);
|
||||
log::debug!(
|
||||
"[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={} \
|
||||
surface_entities_dropped={} topics={}",
|
||||
"[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={}",
|
||||
ctx.tree_id,
|
||||
ctx.target_level,
|
||||
inputs.len(),
|
||||
token_count,
|
||||
parsed.entities.len(),
|
||||
parsed.topics.len()
|
||||
token_count
|
||||
);
|
||||
|
||||
// Drop LLM-emitted entities. The model returns surface forms
|
||||
// ("Alice", "she"), but `SummaryNode.entities` is indexed via
|
||||
// `index_summary_entity_ids_tx` as canonical ids. Surface forms
|
||||
// would silently corrupt that index — searches by canonical id
|
||||
// would not find these summaries. Canonicalisation is the
|
||||
// entity extractor's job, not the summariser's. Topics stay
|
||||
// because they're free-form labels, not indexed as canonical
|
||||
// ids. The `entities` field stays in the prompt to nudge the
|
||||
// model toward entity-aware summarisation; we just don't
|
||||
// persist its output.
|
||||
Ok(SummaryOutput {
|
||||
content,
|
||||
token_count,
|
||||
entities: Vec::new(),
|
||||
topics: dedupe_sorted(parsed.topics),
|
||||
topics: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -314,20 +310,21 @@ fn build_user_prompt(inputs: &[SummaryInput], per_input_cap_tokens: u32) -> Stri
|
||||
out
|
||||
}
|
||||
|
||||
/// System prompt. Token budget is templated in so the model aims under it.
|
||||
fn system_prompt(budget: u32) -> String {
|
||||
format!(
|
||||
"You are a precise summariser. Summarise the user-provided contributions into a \
|
||||
single cohesive passage that preserves concrete facts, decisions, named entities, \
|
||||
and temporal ordering. Do not invent facts. Stay well under {budget} tokens.\n\
|
||||
\n\
|
||||
Return JSON only — no prose, no markdown, no commentary. Schema:\n\
|
||||
{{\n\
|
||||
\x20 \"summary\": \"<summary body>\",\n\
|
||||
\x20 \"entities\": [\"<named entity surface forms mentioned in the summary>\"],\n\
|
||||
\x20 \"topics\": [\"<short topic labels covered by the summary>\"]\n\
|
||||
}}"
|
||||
)
|
||||
/// System prompt. Length isn't templated in — empirically, telling small
|
||||
/// instruction-tuned models "stay under N tokens" makes them produce
|
||||
/// curt, generic output even when the input has plenty of substance.
|
||||
/// Output is clamped post-generation by [`clamp_to_budget`] in the
|
||||
/// caller, so we don't need the model to self-police length.
|
||||
fn system_prompt(_budget: u32) -> String {
|
||||
"You are a precise summariser. Summarise the user-provided contributions into a \
|
||||
single cohesive passage that preserves concrete facts, decisions, \
|
||||
and temporal ordering. Do not invent facts.\n\
|
||||
\n\
|
||||
Return JSON only — no prose, no markdown, no commentary. Schema:\n\
|
||||
{\n\
|
||||
\x20 \"summary\": \"<summary body>\"\n\
|
||||
}"
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Truncate to the caller's token budget using the same ~4 chars/token
|
||||
@@ -343,16 +340,6 @@ fn clamp_to_budget(text: &str, budget: u32) -> (String, u32) {
|
||||
(truncated, tokens)
|
||||
}
|
||||
|
||||
fn dedupe_sorted(mut items: Vec<String>) -> Vec<String> {
|
||||
for item in items.iter_mut() {
|
||||
*item = item.trim().to_string();
|
||||
}
|
||||
items.retain(|s| !s.is_empty());
|
||||
items.sort();
|
||||
items.dedup();
|
||||
items
|
||||
}
|
||||
|
||||
fn truncate_for_log(s: &str, max_chars: usize) -> String {
|
||||
if s.chars().count() <= max_chars {
|
||||
return s.to_string();
|
||||
@@ -405,10 +392,6 @@ struct OllamaResponseMessage {
|
||||
struct LlmSummaryOutput {
|
||||
#[serde(default)]
|
||||
summary: String,
|
||||
#[serde(default)]
|
||||
entities: Vec<String>,
|
||||
#[serde(default)]
|
||||
topics: Vec<String>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -493,12 +476,16 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn system_prompt_templates_budget() {
|
||||
fn system_prompt_describes_schema() {
|
||||
// Budget is no longer templated into the prompt — small models
|
||||
// produced overly curt output when told to "stay under N tokens".
|
||||
// The clamp in `clamp_to_budget` handles enforcement instead.
|
||||
let p = system_prompt(4096);
|
||||
assert!(p.contains("4096"));
|
||||
assert!(!p.contains("4096"));
|
||||
assert!(!p.contains("Stay well under"));
|
||||
assert!(p.contains("\"summary\""));
|
||||
assert!(p.contains("\"entities\""));
|
||||
assert!(p.contains("\"topics\""));
|
||||
assert!(!p.contains("\"entities\""));
|
||||
assert!(!p.contains("\"topics\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -516,19 +503,6 @@ mod tests {
|
||||
assert!(t <= 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupe_sorted_trims_and_dedupes() {
|
||||
let out = dedupe_sorted(vec![
|
||||
"Bob".into(),
|
||||
" Alice ".into(),
|
||||
"Bob".into(),
|
||||
"".into(),
|
||||
" ".into(),
|
||||
"Alice".into(),
|
||||
]);
|
||||
assert_eq!(out, vec!["Alice", "Bob"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_for_log_short_input_unchanged() {
|
||||
assert_eq!(truncate_for_log("hi", 10), "hi");
|
||||
@@ -592,16 +566,24 @@ mod tests {
|
||||
assert!(!req.stream);
|
||||
assert_eq!(req.options.temperature, 0.0);
|
||||
assert_eq!(req.messages[0].role, "system");
|
||||
assert!(req.messages[0].content.contains("2048"));
|
||||
assert!(req.messages[0].content.contains("\"summary\""));
|
||||
assert_eq!(req.messages[1].role, "user");
|
||||
assert_eq!(req.messages[1].content, "body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_output_deserialises_with_missing_fields() {
|
||||
fn llm_output_deserialises_with_only_summary() {
|
||||
let v: LlmSummaryOutput = serde_json::from_str(r#"{"summary":"hi"}"#).unwrap();
|
||||
assert_eq!(v.summary, "hi");
|
||||
assert!(v.entities.is_empty());
|
||||
assert!(v.topics.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_output_ignores_extraneous_fields() {
|
||||
// Prompt no longer asks for entities/topics, but if the model
|
||||
// emits them anyway we should still parse `summary` cleanly.
|
||||
let v: LlmSummaryOutput =
|
||||
serde_json::from_str(r#"{"summary":"hi","entities":["Alice"],"topics":["x"]}"#)
|
||||
.unwrap();
|
||||
assert_eq!(v.summary, "hi");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -172,13 +172,17 @@ impl Buffer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Token ceiling for one summariser invocation — aligned with the Phase 1
|
||||
/// chunker ceiling so a single leaf never busts a seal on its own.
|
||||
/// Token ceiling for one summariser invocation.
|
||||
///
|
||||
/// Sized for the local 1B summariser (`gemma3:1b-it-qat`), which produces
|
||||
/// noticeably better summaries with ≤4-5k input than at higher caps. The
|
||||
/// chunker's `DEFAULT_CHUNK_MAX_TOKENS` (3_000) sits below this so each
|
||||
/// L0 buffer accumulates roughly 1-3 chunks before sealing.
|
||||
///
|
||||
/// Gates only the L0 → L1 seal: leaves are fan-in by raw token volume so
|
||||
/// the summariser input stays bounded. Summaries above L0 use
|
||||
/// [`SUMMARY_FANOUT`] instead — see `bucket_seal::should_seal`.
|
||||
pub const TOKEN_BUDGET: u32 = 10_000;
|
||||
pub const TOKEN_BUDGET: u32 = 4_500;
|
||||
|
||||
/// Sibling count that triggers a seal at level ≥ 1 (summaries → next level).
|
||||
///
|
||||
|
||||
@@ -13,6 +13,7 @@ use rusqlite::{params, Connection, OptionalExtension, Transaction};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::content_store::StagedChunk;
|
||||
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind, SourceRef};
|
||||
|
||||
const DB_DIR: &str = "memory_tree";
|
||||
@@ -21,6 +22,12 @@ const DEFAULT_LIST_LIMIT: usize = 100;
|
||||
const MAX_LIST_LIMIT: usize = 10_000;
|
||||
const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction";
|
||||
pub const CHUNK_STATUS_ADMITTED: &str = "admitted";
|
||||
pub const CHUNK_STATUS_BUFFERED: &str = "buffered";
|
||||
pub const CHUNK_STATUS_SEALED: &str = "sealed";
|
||||
pub const CHUNK_STATUS_DROPPED: &str = "dropped";
|
||||
|
||||
const SCHEMA: &str = "
|
||||
PRAGMA foreign_keys = ON;
|
||||
|
||||
@@ -174,6 +181,35 @@ CREATE TABLE IF NOT EXISTS mem_tree_entity_hotness (
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_hotness_score
|
||||
ON mem_tree_entity_hotness(last_hotness);
|
||||
|
||||
-- Async job queue for memory-tree work (extract → admit → buffer → seal →
|
||||
-- topic-route → daily digest). Producers (ingest, schedulers, handlers)
|
||||
-- enqueue rows transactionally; the worker pool claims them via the
|
||||
-- `(status, available_at_ms)` index. `dedupe_key` is enforced as unique
|
||||
-- only for ready/running rows so a completed job's key can be re-used.
|
||||
CREATE TABLE IF NOT EXISTS mem_tree_jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
kind TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL,
|
||||
dedupe_key TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'ready',
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
max_attempts INTEGER NOT NULL DEFAULT 5,
|
||||
available_at_ms INTEGER NOT NULL,
|
||||
locked_until_ms INTEGER,
|
||||
last_error TEXT,
|
||||
created_at_ms INTEGER NOT NULL,
|
||||
started_at_ms INTEGER,
|
||||
completed_at_ms INTEGER
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_ready
|
||||
ON mem_tree_jobs(status, available_at_ms);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_kind
|
||||
ON mem_tree_jobs(kind);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS idx_mem_tree_jobs_dedupe_active
|
||||
ON mem_tree_jobs(dedupe_key)
|
||||
WHERE dedupe_key IS NOT NULL AND status IN ('ready', 'running');
|
||||
";
|
||||
|
||||
/// Upsert a batch of chunks atomically.
|
||||
@@ -249,6 +285,66 @@ pub(crate) fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result
|
||||
Ok(chunks.len())
|
||||
}
|
||||
|
||||
/// Upsert staged chunks (with content_path + content_sha256) using an existing transaction.
|
||||
///
|
||||
/// Identical to `upsert_chunks_tx` but also writes the Phase MD-content pointer columns.
|
||||
/// `content` column receives a ≤500-char plain-text preview of the body (the full body
|
||||
/// lives on disk at `content_path`).
|
||||
pub(crate) fn upsert_staged_chunks_tx(
|
||||
tx: &Transaction<'_>,
|
||||
staged: &[StagedChunk],
|
||||
) -> Result<usize> {
|
||||
if staged.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let mut stmt = tx.prepare(
|
||||
"INSERT INTO mem_tree_chunks (
|
||||
id, source_kind, source_id, source_ref, owner,
|
||||
timestamp_ms, time_range_start_ms, time_range_end_ms,
|
||||
tags_json, content, token_count, seq_in_source, created_at_ms,
|
||||
content_path, content_sha256
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
source_kind = excluded.source_kind,
|
||||
source_id = excluded.source_id,
|
||||
source_ref = excluded.source_ref,
|
||||
owner = excluded.owner,
|
||||
timestamp_ms = excluded.timestamp_ms,
|
||||
time_range_start_ms = excluded.time_range_start_ms,
|
||||
time_range_end_ms = excluded.time_range_end_ms,
|
||||
tags_json = excluded.tags_json,
|
||||
content = excluded.content,
|
||||
token_count = excluded.token_count,
|
||||
seq_in_source = excluded.seq_in_source,
|
||||
created_at_ms = excluded.created_at_ms,
|
||||
content_path = excluded.content_path,
|
||||
content_sha256 = excluded.content_sha256",
|
||||
)?;
|
||||
for s in staged {
|
||||
let chunk = &s.chunk;
|
||||
// Store a ≤500-char preview in the `content` column; full body is on disk.
|
||||
let preview: String = chunk.content.chars().take(500).collect();
|
||||
stmt.execute(params![
|
||||
chunk.id,
|
||||
chunk.metadata.source_kind.as_str(),
|
||||
chunk.metadata.source_id,
|
||||
chunk.metadata.source_ref.as_ref().map(|r| r.value.as_str()),
|
||||
chunk.metadata.owner,
|
||||
chunk.metadata.timestamp.timestamp_millis(),
|
||||
chunk.metadata.time_range.0.timestamp_millis(),
|
||||
chunk.metadata.time_range.1.timestamp_millis(),
|
||||
serde_json::to_string(&chunk.metadata.tags)?,
|
||||
preview,
|
||||
chunk.token_count,
|
||||
chunk.seq_in_source,
|
||||
chunk.created_at.timestamp_millis(),
|
||||
s.content_path,
|
||||
s.content_sha256,
|
||||
])?;
|
||||
}
|
||||
Ok(staged.len())
|
||||
}
|
||||
|
||||
fn upsert_chunks_with_statement(
|
||||
stmt: &mut rusqlite::Statement<'_>,
|
||||
chunks: &[Chunk],
|
||||
@@ -361,6 +457,70 @@ pub fn count_chunks(config: &Config) -> Result<u64> {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_chunk_lifecycle_status(config: &Config, chunk_id: &str, status: &str) -> Result<()> {
|
||||
with_connection(config, |conn| {
|
||||
set_chunk_lifecycle_status_conn(conn, chunk_id, status)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn set_chunk_lifecycle_status_tx(
|
||||
tx: &Transaction<'_>,
|
||||
chunk_id: &str,
|
||||
status: &str,
|
||||
) -> Result<()> {
|
||||
set_chunk_lifecycle_status_conn(tx, chunk_id, status)
|
||||
}
|
||||
|
||||
pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result<Option<String>> {
|
||||
with_connection(config, |conn| {
|
||||
get_chunk_lifecycle_status_conn(conn, chunk_id)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn get_chunk_lifecycle_status_tx(
|
||||
tx: &Transaction<'_>,
|
||||
chunk_id: &str,
|
||||
) -> Result<Option<String>> {
|
||||
get_chunk_lifecycle_status_conn(tx, chunk_id)
|
||||
}
|
||||
|
||||
fn get_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str) -> Result<Option<String>> {
|
||||
let row = conn
|
||||
.query_row(
|
||||
"SELECT lifecycle_status FROM mem_tree_chunks WHERE id = ?1",
|
||||
params![chunk_id],
|
||||
|r| r.get::<_, String>(0),
|
||||
)
|
||||
.optional()?;
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
pub fn count_chunks_by_lifecycle_status(config: &Config, status: &str) -> Result<u64> {
|
||||
with_connection(config, |conn| {
|
||||
let n: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_chunks WHERE lifecycle_status = ?1",
|
||||
params![status],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
}
|
||||
|
||||
fn set_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str, status: &str) -> Result<()> {
|
||||
let changed = conn.execute(
|
||||
"UPDATE mem_tree_chunks SET lifecycle_status = ?1 WHERE id = ?2",
|
||||
params![status, chunk_id],
|
||||
)?;
|
||||
if changed == 0 {
|
||||
log::warn!(
|
||||
"[memory_tree::store] lifecycle update affected 0 rows chunk_id={} status={}",
|
||||
chunk_id,
|
||||
status
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result<Chunk> {
|
||||
let id: String = row.get(0)?;
|
||||
let source_kind_s: String = row.get(1)?;
|
||||
@@ -401,6 +561,10 @@ fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result<Chunk> {
|
||||
token_count: token_count.max(0) as u32,
|
||||
seq_in_source: seq.max(0) as u32,
|
||||
created_at,
|
||||
// partial_message is not stored in SQLite — it's a transient chunker
|
||||
// signal. Chunks read back from DB always get false (the column doesn't
|
||||
// exist; callers that need this flag hold the Chunk in memory).
|
||||
partial_message: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -451,9 +615,124 @@ pub(crate) fn with_connection<T>(
|
||||
// legacy summaries from Phases 1-3 read back as None; retrieval
|
||||
// tolerates NULL by dropping the row to the bottom of a rerank.
|
||||
add_column_if_missing(&conn, "mem_tree_summaries", "embedding", "BLOB")?;
|
||||
// Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested
|
||||
// before the queue migration stay queryable. New writes start at
|
||||
// 'pending_extraction'; the extract handler advances them to 'admitted'
|
||||
// (then 'buffered' / 'sealed') or 'dropped'.
|
||||
add_column_if_missing(
|
||||
&conn,
|
||||
"mem_tree_chunks",
|
||||
"lifecycle_status",
|
||||
"TEXT NOT NULL DEFAULT 'admitted'",
|
||||
)?;
|
||||
conn.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \
|
||||
ON mem_tree_chunks(lifecycle_status);",
|
||||
)
|
||||
.context("Failed to create mem_tree_chunks lifecycle index")?;
|
||||
// Phase MD-content (#TBD): pointer + integrity hash. Body lives at
|
||||
// <content_root>/<content_path> as a .md file. Both nullable so chunks
|
||||
// ingested before this migration read back with NULL (body still in
|
||||
// `content`). New writes populate both columns. The `content` column
|
||||
// stores a 500-char plain-text preview instead of the full body.
|
||||
add_column_if_missing(&conn, "mem_tree_chunks", "content_path", "TEXT")?;
|
||||
add_column_if_missing(&conn, "mem_tree_chunks", "content_sha256", "TEXT")?;
|
||||
// Phase MD-content (summaries): same pointer pattern for summary nodes.
|
||||
// `content_path` is the relative path to the .md file under
|
||||
// `<content_root>/summaries/...`. `content_sha256` is the SHA-256 hex
|
||||
// of the body bytes only (front-matter excluded). Both nullable so
|
||||
// legacy rows (from before this migration) read back with NULL — callers
|
||||
// fall back to the `content` column for those rows.
|
||||
add_column_if_missing(&conn, "mem_tree_summaries", "content_path", "TEXT")?;
|
||||
add_column_if_missing(&conn, "mem_tree_summaries", "content_sha256", "TEXT")?;
|
||||
f(&conn)
|
||||
}
|
||||
|
||||
/// Return both `content_path` and `content_sha256` stored in SQLite for `chunk_id`.
|
||||
///
|
||||
/// Returns `Ok(None)` if the chunk does not exist or has no content_path recorded yet.
|
||||
pub fn get_chunk_content_pointers(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
) -> Result<Option<(String, String)>> {
|
||||
with_connection(config, |conn| {
|
||||
let row = conn
|
||||
.query_row(
|
||||
"SELECT content_path, content_sha256 FROM mem_tree_chunks WHERE id = ?1",
|
||||
params![chunk_id],
|
||||
|r| {
|
||||
let path: Option<String> = r.get(0)?;
|
||||
let sha: Option<String> = r.get(1)?;
|
||||
Ok((path, sha))
|
||||
},
|
||||
)
|
||||
.optional()?;
|
||||
Ok(row.and_then(|(p, s)| p.zip(s)))
|
||||
})
|
||||
}
|
||||
|
||||
/// Return the `content_path` stored in SQLite for `chunk_id`, if any.
|
||||
pub fn get_chunk_content_path(config: &Config, chunk_id: &str) -> Result<Option<String>> {
|
||||
with_connection(config, |conn| {
|
||||
let row = conn
|
||||
.query_row(
|
||||
"SELECT content_path FROM mem_tree_chunks WHERE id = ?1",
|
||||
params![chunk_id],
|
||||
|r| r.get::<_, Option<String>>(0),
|
||||
)
|
||||
.optional()?
|
||||
.flatten();
|
||||
Ok(row)
|
||||
})
|
||||
}
|
||||
|
||||
/// Return both `content_path` and `content_sha256` stored in SQLite for `summary_id`.
|
||||
///
|
||||
/// Returns `Ok(None)` if the summary does not exist or has no content_path recorded yet
|
||||
/// (legacy rows pre-MD-content migration).
|
||||
pub fn get_summary_content_pointers(
|
||||
config: &Config,
|
||||
summary_id: &str,
|
||||
) -> Result<Option<(String, String)>> {
|
||||
with_connection(config, |conn| {
|
||||
let row = conn
|
||||
.query_row(
|
||||
"SELECT content_path, content_sha256 FROM mem_tree_summaries WHERE id = ?1",
|
||||
params![summary_id],
|
||||
|r| {
|
||||
let path: Option<String> = r.get(0)?;
|
||||
let sha: Option<String> = r.get(1)?;
|
||||
Ok((path, sha))
|
||||
},
|
||||
)
|
||||
.optional()?;
|
||||
Ok(row.and_then(|(p, s)| p.zip(s)))
|
||||
})
|
||||
}
|
||||
|
||||
/// List all summary rows that have a non-NULL `content_path`. Used by the
|
||||
/// bin integrity checker.
|
||||
pub fn list_summaries_with_content_path(config: &Config) -> Result<Vec<(String, String, String)>> {
|
||||
with_connection(config, |conn| {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT id, content_path, content_sha256
|
||||
FROM mem_tree_summaries
|
||||
WHERE content_path IS NOT NULL AND content_sha256 IS NOT NULL
|
||||
AND deleted = 0",
|
||||
)?;
|
||||
let rows = stmt
|
||||
.query_map([], |r| {
|
||||
let id: String = r.get(0)?;
|
||||
let path: String = r.get(1)?;
|
||||
let sha: String = r.get(2)?;
|
||||
Ok((id, path, sha))
|
||||
})?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()
|
||||
.context("Failed to list summaries with content_path")?;
|
||||
Ok(rows)
|
||||
})
|
||||
}
|
||||
|
||||
fn normalized_limit(requested: Option<usize>) -> i64 {
|
||||
let clamped = requested
|
||||
.unwrap_or(DEFAULT_LIST_LIMIT)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree::types::chunk_id;
|
||||
use chrono::TimeZone;
|
||||
use rusqlite::params;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
@@ -27,6 +28,7 @@ fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk {
|
||||
token_count: 12,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -164,3 +166,37 @@ fn empty_batch_is_noop() {
|
||||
assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0);
|
||||
assert_eq!(count_chunks(&cfg).unwrap(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schema_has_content_path_and_content_sha256_columns() {
|
||||
// Phase MD-content: verify that with_connection applies the additive
|
||||
// migrations for the new pointer + hash columns on a fresh DB.
|
||||
let (_tmp, cfg) = test_config();
|
||||
with_connection(&cfg, |conn| {
|
||||
let mut has_content_path = false;
|
||||
let mut has_content_sha256 = false;
|
||||
let mut stmt = conn.prepare("PRAGMA table_info(mem_tree_chunks)")?;
|
||||
let names: Vec<String> = stmt
|
||||
.query_map(params![], |row| row.get::<_, String>(1))?
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
for name in &names {
|
||||
if name == "content_path" {
|
||||
has_content_path = true;
|
||||
}
|
||||
if name == "content_sha256" {
|
||||
has_content_sha256 = true;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
has_content_path,
|
||||
"mem_tree_chunks must have content_path column after migration; found: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
has_content_sha256,
|
||||
"mem_tree_chunks must have content_sha256 column after migration; found: {names:?}"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
@@ -1,30 +1,52 @@
|
||||
//! Topic-tree backfill — hydrate a freshly-materialised topic tree with
|
||||
//! every historical leaf mentioning the entity (#709 Phase 3c).
|
||||
//! recent leaves mentioning the entity (#709 Phase 3c).
|
||||
//!
|
||||
//! When the curator decides an entity has crossed the hotness threshold
|
||||
//! for the first time, we create a fresh topic tree AND walk the
|
||||
//! `mem_tree_entity_index` inverted index to append every prior leaf into
|
||||
//! `mem_tree_entity_index` inverted index to append matching leaves into
|
||||
//! its L0 buffer. Reusing `bucket_seal::append_leaf` means the cascade
|
||||
//! fires automatically — a well-established entity may seal several
|
||||
//! levels as soon as the tree is spawned.
|
||||
//! fires automatically.
|
||||
//!
|
||||
//! ## Why bounded by hotness window
|
||||
//!
|
||||
//! Hotness uses a 30-day recency decay (see `topic_tree::hotness`). Leaves
|
||||
//! older than 30 days contribute zero to current hotness, so by definition
|
||||
//! they cannot be the reason a tree is spawning *now*. Including them
|
||||
//! bloats the spawn latency, wastes summariser LLM calls, and amplifies
|
||||
//! ancient signal that has already decayed away. We cap the backfill
|
||||
//! window at [`BACKFILL_WINDOW_DAYS`] to align with the hotness math.
|
||||
//!
|
||||
//! Older content is still queryable through source-tree retrieval and the
|
||||
//! entity index — it just doesn't get its own slot in the topic tree.
|
||||
//!
|
||||
//! Backfill is intentionally best-effort: missing chunks are skipped with
|
||||
//! a warn log rather than failing the whole spawn, because Phase 3c is
|
||||
//! additive — a partial topic tree is still useful.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::Utc;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::score::store::lookup_entity;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
|
||||
use crate::openhuman::memory::tree::source_tree::types::Tree;
|
||||
use crate::openhuman::memory::tree::store::get_chunk;
|
||||
use crate::openhuman::memory::tree::util::redact::redact;
|
||||
|
||||
/// Max leaves to pull from the entity index during backfill. A hard cap
|
||||
/// keeps initial spawn latency bounded even for very active entities.
|
||||
const BACKFILL_LIMIT: usize = 500;
|
||||
|
||||
/// Backfill window in days — matches `topic_tree::hotness::recency_decay`'s
|
||||
/// hard cliff. Leaves older than this contribute zero to current hotness
|
||||
/// so they cannot have driven the spawn decision.
|
||||
pub const BACKFILL_WINDOW_DAYS: i64 = 30;
|
||||
|
||||
const DAY_MS: i64 = 24 * 60 * 60 * 1_000;
|
||||
|
||||
/// Walk the entity index for `entity_id` and append every discovered leaf
|
||||
/// to `tree`. Returns the number of leaves appended (NOT the number of
|
||||
/// summaries sealed). Idempotent: `append_leaf` itself is a no-op when a
|
||||
@@ -35,19 +57,65 @@ pub async fn backfill_topic_tree(
|
||||
entity_id: &str,
|
||||
summariser: &dyn Summariser,
|
||||
) -> Result<usize> {
|
||||
log::info!(
|
||||
"[topic_tree::backfill] start entity_id={} tree_id={}",
|
||||
backfill_topic_tree_at(
|
||||
config,
|
||||
tree,
|
||||
entity_id,
|
||||
tree.id
|
||||
summariser,
|
||||
Utc::now().timestamp_millis(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Deterministic variant — backfill against a caller-supplied `now_ms`
|
||||
/// for the recency window. Used by tests so the 30-day cutoff doesn't
|
||||
/// depend on the wall clock.
|
||||
pub async fn backfill_topic_tree_at(
|
||||
config: &Config,
|
||||
tree: &Tree,
|
||||
entity_id: &str,
|
||||
summariser: &dyn Summariser,
|
||||
now_ms: i64,
|
||||
) -> Result<usize> {
|
||||
let cutoff_ms = now_ms.saturating_sub(BACKFILL_WINDOW_DAYS.saturating_mul(DAY_MS));
|
||||
log::info!(
|
||||
"[topic_tree::backfill] start entity_id_hash={} tree_id={} window_days={} cutoff_ms={}",
|
||||
redact(entity_id),
|
||||
tree.id,
|
||||
BACKFILL_WINDOW_DAYS,
|
||||
cutoff_ms
|
||||
);
|
||||
|
||||
let hits = lookup_entity(config, entity_id, Some(BACKFILL_LIMIT))
|
||||
.with_context(|| format!("failed to lookup entity {entity_id}"))?;
|
||||
.with_context(|| format!("failed to lookup entity {}", redact(entity_id)))?;
|
||||
|
||||
if hits.is_empty() {
|
||||
log::debug!(
|
||||
"[topic_tree::backfill] no entity-index hits for entity_id={} — empty backfill",
|
||||
entity_id
|
||||
"[topic_tree::backfill] no entity-index hits for entity_id_hash={} — empty backfill",
|
||||
redact(entity_id)
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Drop hits older than the hotness recency window — see module docs.
|
||||
let total_hits = hits.len();
|
||||
let mut hits: Vec<_> = hits
|
||||
.into_iter()
|
||||
.filter(|h| h.timestamp_ms >= cutoff_ms)
|
||||
.collect();
|
||||
let dropped = total_hits - hits.len();
|
||||
if dropped > 0 {
|
||||
log::debug!(
|
||||
"[topic_tree::backfill] dropped {dropped} hits older than {BACKFILL_WINDOW_DAYS}d \
|
||||
for entity_id_hash={}",
|
||||
redact(entity_id)
|
||||
);
|
||||
}
|
||||
if hits.is_empty() {
|
||||
log::debug!(
|
||||
"[topic_tree::backfill] all entity-index hits fell outside the {BACKFILL_WINDOW_DAYS}d \
|
||||
window for entity_id_hash={} — empty backfill",
|
||||
redact(entity_id)
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
@@ -55,7 +123,6 @@ pub async fn backfill_topic_tree(
|
||||
// Sort by timestamp ASC so the buffer's `oldest_at` and the sealed
|
||||
// summary's `time_range_start` reflect the true historical order, not
|
||||
// the DESC ordering `lookup_entity` returns.
|
||||
let mut hits = hits;
|
||||
hits.sort_by_key(|h| h.timestamp_ms);
|
||||
|
||||
let mut appended = 0usize;
|
||||
@@ -77,9 +144,9 @@ pub async fn backfill_topic_tree(
|
||||
Some(c) => c,
|
||||
None => {
|
||||
log::warn!(
|
||||
"[topic_tree::backfill] missing chunk {} for entity {} — skipping",
|
||||
"[topic_tree::backfill] missing chunk {} for entity_id_hash={} — skipping",
|
||||
hit.node_id,
|
||||
entity_id
|
||||
redact(entity_id)
|
||||
);
|
||||
continue;
|
||||
}
|
||||
@@ -95,20 +162,25 @@ pub async fn backfill_topic_tree(
|
||||
score: hit.score,
|
||||
};
|
||||
|
||||
append_leaf(config, tree, &leaf, summariser)
|
||||
// Topic-tree backfill: empty labels for sealed summaries — the
|
||||
// tree's scope already pins the canonical id, so cross-pollinating
|
||||
// descendants' entities would noise the index. See LabelStrategy.
|
||||
append_leaf(config, tree, &leaf, summariser, &LabelStrategy::Empty)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"backfill append_leaf failed tree_id={} chunk_id={}",
|
||||
tree.id, chunk.id
|
||||
"backfill append_leaf failed entity_id_hash={} tree_id={} chunk_id={}",
|
||||
redact(entity_id),
|
||||
tree.id,
|
||||
chunk.id
|
||||
)
|
||||
})?;
|
||||
appended += 1;
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"[topic_tree::backfill] done entity_id={} tree_id={} appended={}",
|
||||
entity_id,
|
||||
"[topic_tree::backfill] done entity_id_hash={} tree_id={} appended={}",
|
||||
redact(entity_id),
|
||||
tree.id,
|
||||
appended
|
||||
);
|
||||
@@ -158,6 +230,7 @@ mod tests {
|
||||
token_count: tokens,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -172,6 +245,11 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Deterministic "now" used by the windowed-backfill tests: 1 hour
|
||||
/// after the latest seeded leaf so all three sit inside the 30-day
|
||||
/// cutoff. Lets us keep the legacy 2023-era timestamps unchanged.
|
||||
const TEST_NOW_MS: i64 = 1_700_000_020_000 + 3_600_000;
|
||||
|
||||
#[tokio::test]
|
||||
async fn backfill_appends_all_entity_leaves() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
@@ -212,9 +290,15 @@ mod tests {
|
||||
|
||||
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
let n = backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(n, 3);
|
||||
|
||||
// L0 buffer should hold all three leaves (combined tokens well
|
||||
@@ -226,6 +310,38 @@ mod tests {
|
||||
assert_eq!(buf.oldest_at.unwrap().timestamp_millis(), 1_700_000_000_000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backfill_drops_leaves_older_than_window() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
// c_old is 60d before TEST_NOW_MS — outside the 30d cutoff.
|
||||
// c_new is 5d before TEST_NOW_MS — inside the window.
|
||||
let old_ts = TEST_NOW_MS - 60 * DAY_MS;
|
||||
let new_ts = TEST_NOW_MS - 5 * DAY_MS;
|
||||
let c_old = mk_chunk("slack:#eng", 0, old_ts, 100);
|
||||
let c_new = mk_chunk("slack:#eng", 1, new_ts, 100);
|
||||
upsert_chunks(&cfg, &[c_old.clone(), c_new.clone()]).unwrap();
|
||||
|
||||
let e = sample_entity("email:alice@example.com", "alice@example.com");
|
||||
index_entity(&cfg, &e, &c_old.id, "leaf", old_ts, Some("source:slack")).unwrap();
|
||||
index_entity(&cfg, &e, &c_new.id, "leaf", new_ts, Some("source:slack")).unwrap();
|
||||
|
||||
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let n = backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(n, 1, "only the in-window leaf should be appended");
|
||||
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
|
||||
assert_eq!(buf.item_ids.len(), 1);
|
||||
assert_eq!(buf.item_ids[0], c_new.id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backfill_skips_missing_chunks_without_failing() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
@@ -247,9 +363,15 @@ mod tests {
|
||||
|
||||
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
let n = backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(n, 1, "only the existing chunk should be appended");
|
||||
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
|
||||
assert_eq!(buf.item_ids.len(), 1);
|
||||
@@ -273,12 +395,24 @@ mod tests {
|
||||
|
||||
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
// append_leaf is idempotent so the buffer still has exactly one row.
|
||||
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
|
||||
assert_eq!(buf.item_ids.len(), 1);
|
||||
@@ -300,9 +434,15 @@ mod tests {
|
||||
.unwrap();
|
||||
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
|
||||
let summariser = InertSummariser::new();
|
||||
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
|
||||
.await
|
||||
.unwrap();
|
||||
let n = backfill_topic_tree_at(
|
||||
&cfg,
|
||||
&tree,
|
||||
"email:alice@example.com",
|
||||
&summariser,
|
||||
TEST_NOW_MS,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(n, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -186,7 +186,11 @@ mod tests {
|
||||
}
|
||||
|
||||
fn seed_leaf_for_entity(cfg: &Config, entity_id: &str, source_tree: &str, seq: u32) {
|
||||
let ts_ms = 1_700_000_000_000 + (seq as i64) * 1_000;
|
||||
// Use a "now-anchored" timestamp so backfill's 30-day window
|
||||
// (see topic_tree::backfill::BACKFILL_WINDOW_DAYS) always
|
||||
// includes these seeded leaves. Spread by seq to keep ordering
|
||||
// deterministic.
|
||||
let ts_ms = Utc::now().timestamp_millis() - (seq as i64) * 1_000;
|
||||
let ts = Utc.timestamp_millis_opt(ts_ms).unwrap();
|
||||
let c = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, source_tree, seq, "test-content"),
|
||||
@@ -203,6 +207,7 @@ mod tests {
|
||||
token_count: 50,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[c.clone()]).unwrap();
|
||||
let e = CanonicalEntity {
|
||||
|
||||
@@ -21,7 +21,9 @@
|
||||
use anyhow::Result;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
|
||||
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
|
||||
append_leaf, LabelStrategy, LeafRef,
|
||||
};
|
||||
use crate::openhuman::memory::tree::source_tree::store as src_store;
|
||||
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
|
||||
use crate::openhuman::memory::tree::source_tree::types::{TreeKind, TreeStatus};
|
||||
@@ -90,7 +92,16 @@ async fn route_one_entity(
|
||||
entities: vec![entity_id.to_string()],
|
||||
..leaf.clone()
|
||||
};
|
||||
append_leaf(config, &tree, &topic_leaf, summariser).await?;
|
||||
// Topic-tree seals leave entities/topics empty: the tree's
|
||||
// scope already pins the canonical id this tree represents.
|
||||
append_leaf(
|
||||
config,
|
||||
&tree,
|
||||
&topic_leaf,
|
||||
summariser,
|
||||
&LabelStrategy::Empty,
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
log::debug!(
|
||||
"[topic_tree::routing] skip archived topic tree id={} entity={}",
|
||||
@@ -162,6 +173,7 @@ mod tests {
|
||||
token_count: tokens,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
let id = c.id.clone();
|
||||
upsert_chunks(cfg, &[c]).unwrap();
|
||||
@@ -295,9 +307,15 @@ mod tests {
|
||||
crate::openhuman::memory::tree::topic_tree::types::TOPIC_RECHECK_EVERY - 1;
|
||||
crate::openhuman::memory::tree::topic_tree::store::upsert(&cfg, &counters).unwrap();
|
||||
|
||||
// Seed a leaf in slack and gmail referencing Alice.
|
||||
let c1 = persist_chunk(&cfg, "slack:#eng", 0, 1_700_000_000_000, 100);
|
||||
let c2 = persist_chunk(&cfg, "gmail:alice", 0, 1_700_000_010_000, 100);
|
||||
// Seed leaves in slack and gmail referencing Alice. Anchor the
|
||||
// timestamps to "now" so the 30-day backfill window
|
||||
// (topic_tree::backfill::BACKFILL_WINDOW_DAYS) covers them.
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
let ts_c1 = now_ms - 20_000;
|
||||
let ts_c2 = now_ms - 10_000;
|
||||
let ts_c3 = now_ms;
|
||||
let c1 = persist_chunk(&cfg, "slack:#eng", 0, ts_c1, 100);
|
||||
let c2 = persist_chunk(&cfg, "gmail:alice", 0, ts_c2, 100);
|
||||
let e = CanonicalEntity {
|
||||
canonical_id: entity_id.into(),
|
||||
kind: EntityKind::Email,
|
||||
@@ -306,24 +324,16 @@ mod tests {
|
||||
span_end: entity_id.len() as u32,
|
||||
score: 1.0,
|
||||
};
|
||||
index_entity(&cfg, &e, &c1, "leaf", 1_700_000_000_000, Some("slack:#eng")).unwrap();
|
||||
index_entity(
|
||||
&cfg,
|
||||
&e,
|
||||
&c2,
|
||||
"leaf",
|
||||
1_700_000_010_000,
|
||||
Some("gmail:alice"),
|
||||
)
|
||||
.unwrap();
|
||||
index_entity(&cfg, &e, &c1, "leaf", ts_c1, Some("slack:#eng")).unwrap();
|
||||
index_entity(&cfg, &e, &c2, "leaf", ts_c2, Some("gmail:alice")).unwrap();
|
||||
|
||||
// A third leaf arrives — should both fan out to (future) topic tree
|
||||
// and push the curator over the recheck cadence, materialising it.
|
||||
let c3 = persist_chunk(&cfg, "slack:#eng", 1, 1_700_000_020_000, 100);
|
||||
let c3 = persist_chunk(&cfg, "slack:#eng", 1, ts_c3, 100);
|
||||
let leaf = LeafRef {
|
||||
chunk_id: c3.clone(),
|
||||
token_count: 100,
|
||||
timestamp: Utc.timestamp_millis_opt(1_700_000_020_000).unwrap(),
|
||||
timestamp: Utc.timestamp_millis_opt(ts_c3).unwrap(),
|
||||
content: "new mention".into(),
|
||||
entities: vec![entity_id.into()],
|
||||
topics: vec![],
|
||||
|
||||
@@ -232,6 +232,12 @@ pub struct Chunk {
|
||||
/// When this chunk was persisted to the local store.
|
||||
#[serde(with = "chrono::serde::ts_milliseconds")]
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// True when this chunk is a sub-split of a single logical unit (e.g. a
|
||||
/// chat message or email body that exceeded `max_tokens`). The full logical
|
||||
/// unit was split into multiple pieces; each piece carries this flag so
|
||||
/// downstream scorers can lower its weight relative to whole-unit chunks.
|
||||
#[serde(default)]
|
||||
pub partial_message: bool,
|
||||
}
|
||||
|
||||
/// Deterministic chunk id.
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
//! Shared utility helpers for the memory-tree subsystem.
|
||||
|
||||
pub mod redact;
|
||||
@@ -0,0 +1,136 @@
|
||||
//! PII redaction helpers for log output.
|
||||
//!
|
||||
//! Per project rule (CLAUDE.md): "Never log secrets or full PII."
|
||||
//! After the participant-bucketing change introduced in the MD-content PR,
|
||||
//! source_ids and content_paths can embed full email addresses, so any log
|
||||
//! line that prints them needs to redact.
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// Redact a string by hashing it to 8 hex chars. Stable across runs for the
|
||||
/// same input — safe to grep for in logs when debugging with the raw value
|
||||
/// available externally.
|
||||
///
|
||||
/// Use for source_ids, entity_ids, content_paths and similar PII-bearing
|
||||
/// strings in log output.
|
||||
pub fn redact(s: &str) -> String {
|
||||
let mut h = Sha256::new();
|
||||
h.update(s.as_bytes());
|
||||
let d = h.finalize();
|
||||
format!("{:08x}", u32::from_be_bytes([d[0], d[1], d[2], d[3]]))
|
||||
}
|
||||
|
||||
/// Redact a URL/endpoint by stripping path, query, fragment and credentials,
|
||||
/// keeping only the host (and port if present).
|
||||
///
|
||||
/// Examples:
|
||||
/// - `"http://localhost:11434/api/chat"` → `"localhost:11434"`
|
||||
/// - `"https://user:pass@example.com/foo?q=1"` → `"example.com"`
|
||||
/// - `"ollama://host:1234"` → `"host:1234"`
|
||||
///
|
||||
/// Does not pull in a URL-parsing crate; uses cheap string splitting which is
|
||||
/// sufficient for the endpoint-config strings this codebase passes around.
|
||||
pub fn redact_endpoint(url: &str) -> String {
|
||||
// Strip scheme (everything before "://").
|
||||
let after_scheme = url.split_once("://").map(|(_, r)| r).unwrap_or(url);
|
||||
// Take only the authority (everything up to the first '/', '?', or '#') so
|
||||
// any '@' in the path / query (e.g. `?email=foo@bar`) doesn't get treated
|
||||
// as a userinfo separator.
|
||||
let authority = after_scheme
|
||||
.split(['/', '?', '#'])
|
||||
.next()
|
||||
.unwrap_or(after_scheme);
|
||||
// Within the authority, the LAST '@' separates userinfo from host:port.
|
||||
// (RFC 3986: userinfo may itself contain '@' — split-on-first would
|
||||
// truncate the host. Use rsplit so `user:p@ss@example.com` extracts
|
||||
// `example.com` correctly.)
|
||||
let host_port = authority
|
||||
.rsplit_once('@')
|
||||
.map(|(_, r)| r)
|
||||
.unwrap_or(authority);
|
||||
host_port.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── redact ───────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn redact_returns_eight_hex_chars() {
|
||||
let r = redact("alice@example.com");
|
||||
assert_eq!(r.len(), 8, "must be 8 hex chars; got {r:?}");
|
||||
assert!(r.chars().all(|c| c.is_ascii_hexdigit()), "must be hex");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_is_stable_across_calls() {
|
||||
assert_eq!(redact("alice@example.com"), redact("alice@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_is_different_for_different_inputs() {
|
||||
assert_ne!(redact("alice@example.com"), redact("bob@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_empty_string_does_not_panic() {
|
||||
let r = redact("");
|
||||
assert_eq!(r.len(), 8);
|
||||
}
|
||||
|
||||
// ── redact_endpoint ─────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_strips_path_and_query() {
|
||||
assert_eq!(
|
||||
redact_endpoint("http://localhost:11434/api/chat"),
|
||||
"localhost:11434"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_strips_credentials() {
|
||||
assert_eq!(
|
||||
redact_endpoint("https://user:pass@example.com/foo"),
|
||||
"example.com"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_no_scheme_passthrough() {
|
||||
// No "://" present — treat the whole string as host/path; still strip path.
|
||||
assert_eq!(redact_endpoint("localhost:11434/api"), "localhost:11434");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_just_host() {
|
||||
assert_eq!(redact_endpoint("https://example.com"), "example.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_strips_fragment() {
|
||||
assert_eq!(redact_endpoint("http://host:9090/path#frag"), "host:9090");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_strips_query() {
|
||||
assert_eq!(redact_endpoint("http://host/path?q=1"), "host");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_empty_does_not_panic() {
|
||||
let r = redact_endpoint("");
|
||||
// Empty input: no scheme, no host — returns empty string.
|
||||
assert_eq!(r, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_ollama_style() {
|
||||
assert_eq!(
|
||||
redact_endpoint("http://127.0.0.1:11434/v1/chat/completions"),
|
||||
"127.0.0.1:11434"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -825,6 +825,7 @@ async fn json_rpc_memory_tree_end_to_end() {
|
||||
"openhuman.memory_tree_ingest".to_string(),
|
||||
"openhuman.memory_tree_list_chunks".to_string(),
|
||||
"openhuman.memory_tree_get_chunk".to_string(),
|
||||
"openhuman.memory_tree_trigger_digest".to_string(),
|
||||
];
|
||||
assert_eq!(controllers.len(), expected_methods.len());
|
||||
for method in &expected_methods {
|
||||
|
||||
Reference in New Issue
Block a user