From 5fff5568d3699769ab0454d96f8c2a6a7424e5ce Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Wed, 22 Apr 2026 04:47:27 +0530 Subject: [PATCH] feat(memory): Phase 2 memory tree - preprocessing, scoring, admission gate (#708) (#733) * feat(memory): phase 1 memory tree - multi-source ingestion & canonical chunks (#707) Adds an isolated memory tree layer under src/openhuman/memory/tree/ implementing Phase 1 of the new memory architecture (umbrella #711). Zero edits to existing memory/*.rs files - the new layer coexists with the legacy TinyHumans-backed client. - Source adapters: chat / email / document -> canonical Markdown - Token-bounded chunker with deterministic SHA-256 chunk IDs - SQLite persistence at /memory_tree/chunks.db with full provenance metadata (source_kind, source_id, owner, timestamps, tags, time_range) and back-pointer to raw source - Unified JSON-RPC ingest (dispatches on source_kind + JSON payload): openhuman.memory_tree_ingest, _list_chunks, _get_chunk - DataSource enum covering the 8 providers from m.excalidraw step 1 (Discord/Telegram/Whatsapp/Gmail/OtherEmail/Notion/MeetingNotes/DriveDocs) - ~40 unit tests (chunk ID stability, UTF-8-safe splitting, canonicalisation idempotence, store round-trip, filter behavior) Additive only: new tables in a new DB file, new JSON-RPC namespace, no existing behavior changes. Feeds #708 (scoring), #709 (summary trees), #710 (query tools). Closes #707. Parent: #711. Co-Authored-By: Claude Opus 4.7 (1M context) * feat(memory): phase 2 memory tree - preprocessing, scoring, admission gate (#708) Adds the scoring / admission layer between Phase 1's chunker and store. Stacked on feat/707-memory-ingestion (PR #732) - depends on Phase 1's chunk substrate. - Pluggable EntityExtractor trait + CompositeExtractor chain - RegexEntityExtractor: mechanical entities (emails, URLs, @handles, #hashtags) - Always on, deterministic, zero deps, UTF-8-safe char spans - Five weighted signals: token count, unique-word ratio, metadata weight, source weight (per-DataSource), interaction (reply/sent/mention/dm tags), entity density - Exact-match entity canonicalisation (email lowercased, @ and # stripped) - Admission gate drops chunks below configurable threshold (default 0.3) - Score rationale persists for EVERY chunk (kept or dropped) for debugging - Entities indexed for KEPT chunks only - Two new SQLite tables added to the memory_tree DB: - mem_tree_score: per-chunk score rationale with all signal values - mem_tree_entity_index: inverted index entity_id -> node_id - Idempotent ALTER TABLE migration adds embedding BLOB column to mem_tree_chunks (used in Phase 3 retrieval, wired but not populated here) - Ingest pipeline converted to async to accommodate the extractor trait; blocking SQLite work isolated on spawn_blocking; JSON-RPC surface unchanged (same memory_tree_ingest / list / get methods) - Phase 2 deliberately ships without GLiNER/semantic NER - per-chunk semantic entities land later behind a cargo feature flag; the composite extractor interface keeps that drop-in trivial Additive only: new tables, new columns, new module. Existing Phase 1 behavior unchanged except that low-signal chunks are now dropped before reaching mem_tree_chunks. Raise score_drop_threshold to 0 to disable the gate and restore Phase-1-identical behavior. Closes #708. Parent: #711. Depends on: #707 (#732). Co-Authored-By: Claude Opus 4.7 (1M context) * Fix memory tree scoring persistence issues * Fix memory tree scoring robustness issues from PR review - ingest: fail fast if scorer returns fewer/more results than chunks (silent zip truncation would drop chunks or their score rationale) - score::persist_score{,_tx}: clear stale entity-index rows before re-indexing a re-scored chunk, since INSERT OR REPLACE never deletes rows whose entity_id is no longer in the new extraction - score::store::lookup_entity: clamp limit to i64::MAX before casting to prevent a large usize wrapping into a negative LIMIT Adds clear_entity_index_drops_stale_rows regression test. --------- Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Steven Enamakel --- Cargo.lock | 2 +- src/openhuman/memory/tree/ingest.rs | 247 ++++++--- src/openhuman/memory/tree/mod.rs | 1 + src/openhuman/memory/tree/rpc.rs | 51 +- src/openhuman/memory/tree/schemas.rs | 10 +- .../memory/tree/score/extract/extractor.rs | 106 ++++ .../memory/tree/score/extract/mod.rs | 13 + .../memory/tree/score/extract/regex.rs | 209 ++++++++ .../memory/tree/score/extract/types.rs | 259 ++++++++++ src/openhuman/memory/tree/score/mod.rs | 305 +++++++++++ src/openhuman/memory/tree/score/resolver.rs | 136 +++++ .../memory/tree/score/signals/interaction.rs | 101 ++++ .../tree/score/signals/metadata_weight.rs | 49 ++ .../memory/tree/score/signals/mod.rs | 20 + .../memory/tree/score/signals/ops.rs | 142 ++++++ .../tree/score/signals/source_weight.rs | 110 ++++ .../memory/tree/score/signals/token_count.rs | 79 +++ .../memory/tree/score/signals/types.rs | 37 ++ .../memory/tree/score/signals/unique_words.rs | 85 ++++ src/openhuman/memory/tree/score/store.rs | 473 ++++++++++++++++++ src/openhuman/memory/tree/store.rs | 222 +++++++- tests/json_rpc_e2e.rs | 3 +- 22 files changed, 2539 insertions(+), 121 deletions(-) create mode 100644 src/openhuman/memory/tree/score/extract/extractor.rs create mode 100644 src/openhuman/memory/tree/score/extract/mod.rs create mode 100644 src/openhuman/memory/tree/score/extract/regex.rs create mode 100644 src/openhuman/memory/tree/score/extract/types.rs create mode 100644 src/openhuman/memory/tree/score/mod.rs create mode 100644 src/openhuman/memory/tree/score/resolver.rs create mode 100644 src/openhuman/memory/tree/score/signals/interaction.rs create mode 100644 src/openhuman/memory/tree/score/signals/metadata_weight.rs create mode 100644 src/openhuman/memory/tree/score/signals/mod.rs create mode 100644 src/openhuman/memory/tree/score/signals/ops.rs create mode 100644 src/openhuman/memory/tree/score/signals/source_weight.rs create mode 100644 src/openhuman/memory/tree/score/signals/token_count.rs create mode 100644 src/openhuman/memory/tree/score/signals/types.rs create mode 100644 src/openhuman/memory/tree/score/signals/unique_words.rs create mode 100644 src/openhuman/memory/tree/score/store.rs diff --git a/Cargo.lock b/Cargo.lock index 3590980a2..a2379a3dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4420,7 +4420,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "openhuman" -version = "0.52.26" +version = "0.52.27" dependencies = [ "aes-gcm", "anyhow", diff --git a/src/openhuman/memory/tree/ingest.rs b/src/openhuman/memory/tree/ingest.rs index 9b68862fd..4b3cbeeb3 100644 --- a/src/openhuman/memory/tree/ingest.rs +++ b/src/openhuman/memory/tree/ingest.rs @@ -1,7 +1,12 @@ -//! Ingest orchestrator: canonicalise → chunk → persist (Phase 1 / #707). +//! Ingest orchestrator (Phase 1 + Phase 2): //! -//! Consumers call one `ingest_*` function per source kind. Each returns an -//! [`IngestResult`] so the RPC layer can report how many chunks landed. +//! canonicalise → chunk → score → admission gate → persist (chunks + scores + entity index) +//! +//! 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). use anyhow::Result; use serde::{Deserialize, Serialize}; @@ -11,31 +16,39 @@ use crate::openhuman::memory::tree::canonicalize::{ chat::{self, ChatBatch}, document::{self, DocumentInput}, email::{self, EmailThread}, + CanonicalisedSource, }; use crate::openhuman::memory::tree::chunker::{chunk_markdown, ChunkerInput, ChunkerOptions}; +use crate::openhuman::memory::tree::score::{self, ScoreResult, ScoringConfig}; use crate::openhuman::memory::tree::store; use crate::openhuman::memory::tree::types::Chunk; -/// Outcome of one ingest call. +/// Outcome of one ingest call — extended with per-chunk admission info. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct IngestResult { pub source_id: String, + /// Number of chunks that passed the admission gate and were persisted. pub chunks_written: usize, + /// Number of chunks that failed the admission gate and were NOT persisted + /// (their score rationale IS persisted for diagnostics). + pub chunks_dropped: usize, + /// IDs of all chunks that were persisted (in source order). pub chunk_ids: Vec, } impl IngestResult { - fn from_chunks(source_id: String, chunks: &[Chunk], chunks_written: usize) -> Self { + fn empty(source_id: &str) -> Self { Self { - source_id, - chunks_written, - chunk_ids: chunks.iter().map(|c| c.id.clone()).collect(), + source_id: source_id.to_string(), + chunks_written: 0, + chunks_dropped: 0, + chunk_ids: Vec::new(), } } } /// Ingest a batch of chat messages scoped to one channel/group. -pub fn ingest_chat( +pub async fn ingest_chat( config: &Config, source_id: &str, owner: &str, @@ -50,19 +63,13 @@ pub fn ingest_chat( let canonical = match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? { Some(c) => c, - None => { - return Ok(IngestResult { - source_id: source_id.to_string(), - chunks_written: 0, - chunk_ids: Vec::new(), - }); - } + None => return Ok(IngestResult::empty(source_id)), }; - persist(config, source_id, canonical) + persist(config, source_id, canonical).await } /// Ingest a single email thread. -pub fn ingest_email( +pub async fn ingest_email( config: &Config, source_id: &str, owner: &str, @@ -77,19 +84,13 @@ pub fn ingest_email( let canonical = match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? { Some(c) => c, - None => { - return Ok(IngestResult { - source_id: source_id.to_string(), - chunks_written: 0, - chunk_ids: Vec::new(), - }); - } + None => return Ok(IngestResult::empty(source_id)), }; - persist(config, source_id, canonical) + persist(config, source_id, canonical).await } /// Ingest a single standalone document. -pub fn ingest_document( +pub async fn ingest_document( config: &Config, source_id: &str, owner: &str, @@ -106,22 +107,17 @@ pub fn ingest_document( let canonical = match document::canonicalise(source_id, owner, &tags, doc).map_err(anyhow::Error::msg)? { Some(c) => c, - None => { - return Ok(IngestResult { - source_id: source_id.to_string(), - chunks_written: 0, - chunk_ids: Vec::new(), - }); - } + None => return Ok(IngestResult::empty(source_id)), }; - persist(config, source_id, canonical) + persist(config, source_id, canonical).await } -fn persist( +async fn persist( config: &Config, source_id: &str, - canonical: crate::openhuman::memory::tree::canonicalize::CanonicalisedSource, + canonical: CanonicalisedSource, ) -> Result { + // 1. Chunk let input = ChunkerInput { source_kind: canonical.metadata.source_kind, source_id: source_id.to_string(), @@ -129,23 +125,76 @@ fn persist( metadata: canonical.metadata, }; let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - let written = store::upsert_chunks(config, &chunks)?; + if chunks.is_empty() { + return Ok(IngestResult::empty(source_id)); + } + + // 2. Score (async; uses configured extractor) + let scoring_cfg = ScoringConfig::default_regex_only(); + let scores = score::score_chunks(&chunks, &scoring_cfg).await?; + + // Fail fast on scorer length mismatch — silently truncating via zip would + // drop chunks (or their score rationale) without trace. + if scores.len() != chunks.len() { + anyhow::bail!( + "[memory_tree::ingest] scorer length mismatch: chunks={} scores={}", + chunks.len(), + scores.len() + ); + } + + // 3. Partition kept vs dropped + let mut kept_chunks: Vec = 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 dropped = all_results.iter().filter(|(r, _)| !r.kept).count(); log::debug!( - "[memory_tree::ingest] persisted source_id={} chunks={}", + "[memory_tree::ingest] scoring source_id={} kept={} dropped={}", source_id, - written + kept_chunks.len(), + dropped ); - Ok(IngestResult::from_chunks( - source_id.to_string(), - &chunks, - written, - )) + + // 4. Persist (blocking SQLite — isolate on a dedicated thread) + let config_owned = config.clone(); + let kept_for_store = kept_chunks.clone(); + let results_for_store = all_results.clone(); + let written = tokio::task::spawn_blocking(move || -> Result { + store::with_connection(&config_owned, |conn| { + let tx = conn.unchecked_transaction()?; + let n = store::upsert_chunks_tx(&tx, &kept_for_store)?; + 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). + score::persist_score_tx(&tx, result, *ts_ms, None)?; + } + tx.commit()?; + Ok(n) + }) + }) + .await + .map_err(|e| anyhow::anyhow!("persist join error: {e}"))??; + + 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(), + }) } #[cfg(test)] mod tests { use super::*; use crate::openhuman::memory::tree::canonicalize::chat::ChatMessage; + 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::types::SourceKind; use chrono::{TimeZone, Utc}; @@ -158,75 +207,133 @@ mod tests { (tmp, cfg) } - #[test] - fn ingest_chat_writes_chunks() { - let (_tmp, cfg) = test_config(); - let batch = ChatBatch { + /// Build a substantive batch that reliably passes the admission gate. + fn substantive_batch() -> ChatBatch { + ChatBatch { platform: "slack".into(), channel_label: "#eng".into(), messages: vec![ ChatMessage { author: "alice".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "hello".into(), + 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" + .into(), source_ref: Some("slack://m1".into()), }, ChatMessage { author: "bob".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_010_000).unwrap(), - text: "world".into(), + text: "Confirmed — I'll handle the coordination and cut a release \ + candidate tonight. #launch-q2 will be tracked in Notion." + .into(), source_ref: None, }, ], - }; - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch).unwrap(); - assert_eq!(out.chunks_written, 1); - assert_eq!(count_chunks(&cfg).unwrap(), 1); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat); - assert_eq!(rows[0].metadata.source_id, "slack:#eng"); + } } - #[test] - fn ingest_chat_empty_batch_is_noop() { + #[tokio::test] + async fn ingest_chat_writes_substantive_chunks() { + let (_tmp, cfg) = test_config(); + let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()) + .await + .unwrap(); + 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); + let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); + assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat); + } + + #[tokio::test] + async fn low_signal_chunks_are_dropped_but_score_persists() { + let (_tmp, cfg) = test_config(); + let batch = ChatBatch { + platform: "slack".into(), + channel_label: "#eng".into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + text: "+1".into(), // extremely low-signal + 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!(count_scores(&cfg).unwrap(), 1); + } + + #[tokio::test] + async fn ingest_chat_empty_batch_is_noop() { let (_tmp, cfg) = test_config(); let batch = ChatBatch { platform: "slack".into(), channel_label: "#eng".into(), messages: vec![], }; - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch).unwrap(); + let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch) + .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); } - #[test] - fn re_ingest_is_idempotent() { + #[tokio::test] + async fn re_ingest_is_idempotent_on_chunks_and_scores() { let (_tmp, cfg) = test_config(); let doc = DocumentInput { provider: "notion".into(), title: "Launch plan".into(), - body: "content here".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." + .into(), modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), source_ref: Some("notion://page/abc".into()), }; - ingest_document(&cfg, "notion:abc", "alice", vec![], doc.clone()).unwrap(); - ingest_document(&cfg, "notion:abc", "alice", vec![], doc).unwrap(); + ingest_document(&cfg, "notion:abc", "alice", vec![], doc.clone()) + .await + .unwrap(); + ingest_document(&cfg, "notion:abc", "alice", vec![], doc) + .await + .unwrap(); assert_eq!(count_chunks(&cfg).unwrap(), 1); + assert_eq!(count_scores(&cfg).unwrap(), 1); } - #[test] - fn chunks_preserve_source_ref() { + #[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: "b".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).unwrap(); + 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" diff --git a/src/openhuman/memory/tree/mod.rs b/src/openhuman/memory/tree/mod.rs index 7c2ac7bcd..d8b157141 100644 --- a/src/openhuman/memory/tree/mod.rs +++ b/src/openhuman/memory/tree/mod.rs @@ -27,6 +27,7 @@ pub mod chunker; pub mod ingest; pub mod rpc; pub mod schemas; +pub mod score; pub mod store; pub mod types; diff --git a/src/openhuman/memory/tree/rpc.rs b/src/openhuman/memory/tree/rpc.rs index d420efe04..2c9132d4f 100644 --- a/src/openhuman/memory/tree/rpc.rs +++ b/src/openhuman/memory/tree/rpc.rs @@ -63,33 +63,32 @@ pub async fn ingest_rpc( source_id ); - let result = tokio::task::spawn_blocking({ - let config = config.clone(); - let source_id = source_id.clone(); - let owner = owner.clone(); - move || -> anyhow::Result { - match source_kind { - SourceKind::Chat => { - let batch: ChatBatch = serde_json::from_value(payload) - .map_err(|e| anyhow::anyhow!("invalid chat payload: {e}"))?; - do_ingest_chat(&config, &source_id, &owner, tags, batch) - } - SourceKind::Email => { - let thread: EmailThread = serde_json::from_value(payload) - .map_err(|e| anyhow::anyhow!("invalid email payload: {e}"))?; - do_ingest_email(&config, &source_id, &owner, tags, thread) - } - SourceKind::Document => { - let doc: DocumentInput = serde_json::from_value(payload) - .map_err(|e| anyhow::anyhow!("invalid document payload: {e}"))?; - do_ingest_document(&config, &source_id, &owner, tags, doc) - } - } + // Phase 2: ingest functions are async. Their scoring stage awaits the + // extractor (cheap for regex, not-cheap for future GLiNER/LLM impls) + // and the DB work is isolated on `spawn_blocking` inside `persist`. + let result = match source_kind { + SourceKind::Chat => { + let batch: ChatBatch = serde_json::from_value(payload) + .map_err(|e| format!("invalid chat payload: {e}"))?; + do_ingest_chat(config, &source_id, &owner, tags, batch) + .await + .map_err(|e| format!("ingest: {e}"))? } - }) - .await - .map_err(|e| format!("ingest join error: {e}"))? - .map_err(|e| format!("ingest: {e}"))?; + SourceKind::Email => { + let thread: EmailThread = serde_json::from_value(payload) + .map_err(|e| format!("invalid email payload: {e}"))?; + do_ingest_email(config, &source_id, &owner, tags, thread) + .await + .map_err(|e| format!("ingest: {e}"))? + } + SourceKind::Document => { + let doc: DocumentInput = serde_json::from_value(payload) + .map_err(|e| format!("invalid document payload: {e}"))?; + do_ingest_document(config, &source_id, &owner, tags, doc) + .await + .map_err(|e| format!("ingest: {e}"))? + } + }; Ok(RpcOutcome::single_log( result, diff --git a/src/openhuman/memory/tree/schemas.rs b/src/openhuman/memory/tree/schemas.rs index 5ce660f45..6e5d59480 100644 --- a/src/openhuman/memory/tree/schemas.rs +++ b/src/openhuman/memory/tree/schemas.rs @@ -98,13 +98,19 @@ pub fn schemas(function: &str) -> ControllerSchema { FieldSchema { name: "chunks_written", ty: TypeSchema::U64, - comment: "Number of chunks persisted (including idempotent rewrites).", + comment: "Number of chunks persisted after admission.", + required: true, + }, + FieldSchema { + name: "chunks_dropped", + ty: TypeSchema::U64, + comment: "Number of chunks rejected by the admission gate.", required: true, }, FieldSchema { name: "chunk_ids", ty: TypeSchema::Array(Box::new(TypeSchema::String)), - comment: "IDs of all chunks produced.", + comment: "IDs of all chunks persisted after admission.", required: true, }, ], diff --git a/src/openhuman/memory/tree/score/extract/extractor.rs b/src/openhuman/memory/tree/score/extract/extractor.rs new file mode 100644 index 000000000..a915856fc --- /dev/null +++ b/src/openhuman/memory/tree/score/extract/extractor.rs @@ -0,0 +1,106 @@ +use async_trait::async_trait; + +use super::regex; +use super::types::ExtractedEntities; + +/// Interface for anything that can read a chunk's text and emit entities. +#[async_trait] +pub trait EntityExtractor: Send + Sync { + /// Human-readable name for logs and diagnostics. + fn name(&self) -> &'static str; + + /// Run extraction. Implementations should be idempotent per input. + async fn extract(&self, text: &str) -> anyhow::Result; +} + +/// Synchronous regex extractor adapted to the async [`EntityExtractor`] trait. +pub struct RegexEntityExtractor; + +#[async_trait] +impl EntityExtractor for RegexEntityExtractor { + fn name(&self) -> &'static str { + "regex" + } + + async fn extract(&self, text: &str) -> anyhow::Result { + Ok(regex::extract(text)) + } +} + +/// Runs a sequence of extractors and merges their results. +/// +/// An extractor returning an error is logged and skipped — one bad extractor +/// does not abort ingestion. +pub struct CompositeExtractor { + inner: Vec>, +} + +impl CompositeExtractor { + pub fn new(inner: Vec>) -> Self { + Self { inner } + } + + /// Convenience constructor: regex-only (the Phase 2 default). + pub fn regex_only() -> Self { + Self::new(vec![Box::new(RegexEntityExtractor)]) + } +} + +#[async_trait] +impl EntityExtractor for CompositeExtractor { + fn name(&self) -> &'static str { + "composite" + } + + async fn extract(&self, text: &str) -> anyhow::Result { + let mut out = ExtractedEntities::default(); + for ex in &self.inner { + match ex.extract(text).await { + Ok(batch) => out.merge(batch), + Err(e) => { + log::warn!( + "[memory_tree::extract] extractor `{}` failed: {e} — continuing", + ex.name() + ); + } + } + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::score::extract::EntityKind; + + #[tokio::test] + async fn regex_only_extractor_works() { + let c = CompositeExtractor::regex_only(); + let out = c.extract("hi @alice a@b.com #launch").await.unwrap(); + assert!(out.entities.iter().any(|e| e.kind == EntityKind::Handle)); + assert!(out.entities.iter().any(|e| e.kind == EntityKind::Email)); + assert!(out.entities.iter().any(|e| e.kind == EntityKind::Hashtag)); + } + + struct FailingExtractor; + #[async_trait] + impl EntityExtractor for FailingExtractor { + fn name(&self) -> &'static str { + "failing" + } + async fn extract(&self, _: &str) -> anyhow::Result { + Err(anyhow::anyhow!("boom")) + } + } + + #[tokio::test] + async fn composite_survives_one_failing_extractor() { + let c = CompositeExtractor::new(vec![ + Box::new(FailingExtractor), + Box::new(RegexEntityExtractor), + ]); + let out = c.extract("@alice").await.unwrap(); + assert!(out.entities.iter().any(|e| e.kind == EntityKind::Handle)); + } +} diff --git a/src/openhuman/memory/tree/score/extract/mod.rs b/src/openhuman/memory/tree/score/extract/mod.rs new file mode 100644 index 000000000..802d73a3c --- /dev/null +++ b/src/openhuman/memory/tree/score/extract/mod.rs @@ -0,0 +1,13 @@ +//! Entity extraction (Phase 2 / #708). +//! +//! Exposes [`EntityExtractor`] as a pluggable interface and a default +//! [`CompositeExtractor`] that runs a chain of extractors and merges their +//! output. Phase 2 ships with the mechanical regex extractor only; semantic +//! NER (GLiNER / LLM) plugs in later without changing any call sites. + +mod extractor; +pub mod regex; +pub mod types; + +pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor}; +pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; diff --git a/src/openhuman/memory/tree/score/extract/regex.rs b/src/openhuman/memory/tree/score/extract/regex.rs new file mode 100644 index 000000000..64308a49a --- /dev/null +++ b/src/openhuman/memory/tree/score/extract/regex.rs @@ -0,0 +1,209 @@ +//! Deterministic mechanical-entity extraction via regex. +//! +//! Catches the shapes regex handles cleanly and that are genuinely useful +//! as cross-platform identity anchors (email appearing in Slack + Gmail = +//! same person): +//! +//! - **Email** — RFC-ish pattern, boundary-guarded +//! - **URL** — `http(s)://…` up to whitespace or trailing punctuation +//! - **Handle** — `@alice`, `@alice.bsky.social`, or `alice#1234` +//! - **Hashtag** — `#launch-q2` +//! +//! Every match has `score = 1.0` (regex is deterministic). Spans are +//! char-offsets (not bytes) for UTF-8 safety. + +use once_cell::sync::Lazy; +use regex::Regex; + +use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; + +// ── Compiled regexes (once per process) ────────────────────────────────── + +static RE_EMAIL: Lazy = + Lazy::new(|| Regex::new(r"(?i)\b[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}\b").unwrap()); + +static RE_URL: Lazy = Lazy::new(|| { + // up-to trailing punctuation; avoids catastrophic backtracking + Regex::new(r"https?://[^\s<>\]\[()]+[^\s<>\]\[()\.\,;:\!\?]").unwrap() +}); + +static RE_HANDLE: Lazy = + Lazy::new(|| Regex::new(r"(?:^|[\s(])@([A-Za-z0-9_][A-Za-z0-9_.\-]{1,})").unwrap()); + +static RE_DISCRIM: Lazy = + Lazy::new(|| Regex::new(r"\b([A-Za-z0-9_.\-]{2,32})#\d{4}\b").unwrap()); + +static RE_HASHTAG: Lazy = + Lazy::new(|| Regex::new(r"(?:^|[\s(])#([A-Za-z][A-Za-z0-9_\-]{1,})").unwrap()); + +/// Extract all mechanical entities from `text`. +pub fn extract(text: &str) -> ExtractedEntities { + let mut entities: Vec = Vec::new(); + let mut topics: Vec = Vec::new(); + + for m in RE_EMAIL.find_iter(text) { + entities.push(to_entity(text, m.start(), m.end(), EntityKind::Email)); + } + for m in RE_URL.find_iter(text) { + entities.push(to_entity(text, m.start(), m.end(), EntityKind::Url)); + } + for cap in RE_HANDLE.captures_iter(text) { + if let Some(m) = cap.get(1) { + entities.push(to_entity(text, m.start(), m.end(), EntityKind::Handle)); + } + } + for cap in RE_DISCRIM.captures_iter(text) { + if let Some(m) = cap.get(0) { + entities.push(to_entity(text, m.start(), m.end(), EntityKind::Handle)); + } + } + for cap in RE_HASHTAG.captures_iter(text) { + if let Some(m) = cap.get(1) { + entities.push(to_entity(text, m.start(), m.end(), EntityKind::Hashtag)); + topics.push(ExtractedTopic { + label: text[m.start()..m.end()].to_lowercase(), + score: 1.0, + }); + } + } + + ExtractedEntities { entities, topics } +} + +fn to_entity(text: &str, start: usize, end: usize, kind: EntityKind) -> ExtractedEntity { + ExtractedEntity { + kind, + text: text[start..end].to_string(), + span_start: char_index(text, start), + span_end: char_index(text, end), + score: 1.0, + } +} + +fn char_index(s: &str, byte_idx: usize) -> u32 { + let byte_idx = byte_idx.min(s.len()); + s[..byte_idx].chars().count() as u32 +} + +#[cfg(test)] +mod tests { + use super::*; + + fn kinds(e: &ExtractedEntities) -> Vec { + let mut k: Vec<_> = e.entities.iter().map(|x| x.kind).collect(); + k.sort_by_key(|k| *k as u8); + k + } + + #[test] + fn email_basic() { + let o = extract("contact alice@example.com please"); + assert_eq!(o.entities.len(), 1); + assert_eq!(o.entities[0].kind, EntityKind::Email); + assert_eq!(o.entities[0].text, "alice@example.com"); + } + + #[test] + fn url_stops_at_trailing_punct() { + let o = extract("see https://example.com/x?y=1 now."); + let urls: Vec<_> = o + .entities + .iter() + .filter(|e| e.kind == EntityKind::Url) + .collect(); + assert_eq!(urls.len(), 1); + assert_eq!(urls[0].text, "https://example.com/x?y=1"); + } + + #[test] + fn handle_vs_email_boundary() { + let o = extract("@alice met alice@example.com and @bob"); + let handles: Vec<_> = o + .entities + .iter() + .filter(|e| e.kind == EntityKind::Handle) + .map(|e| e.text.as_str()) + .collect(); + let emails: Vec<_> = o + .entities + .iter() + .filter(|e| e.kind == EntityKind::Email) + .map(|e| e.text.as_str()) + .collect(); + assert_eq!(handles, vec!["alice", "bob"]); + assert_eq!(emails, vec!["alice@example.com"]); + } + + #[test] + fn discord_style_handle() { + let o = extract("ping alice#1234"); + let h: Vec<_> = o + .entities + .iter() + .filter(|e| e.kind == EntityKind::Handle) + .collect(); + assert_eq!(h.len(), 1); + assert_eq!(h[0].text, "alice#1234"); + } + + #[test] + fn hashtag_emits_topic() { + let o = extract("tracking #launch-q2 updates"); + assert_eq!( + o.entities + .iter() + .filter(|e| e.kind == EntityKind::Hashtag) + .count(), + 1 + ); + assert_eq!(o.topics.len(), 1); + assert_eq!(o.topics[0].label, "launch-q2"); + } + + #[test] + fn hashtag_requires_leading_letter() { + let o = extract("#123 no, #x1 yes"); + let tags: Vec<_> = o + .entities + .iter() + .filter(|e| e.kind == EntityKind::Hashtag) + .collect(); + assert_eq!(tags.len(), 1); + assert_eq!(tags[0].text, "x1"); + } + + #[test] + fn utf8_span_is_char_not_byte() { + let o = extract("中 a@b.com"); + let email = o + .entities + .iter() + .find(|e| e.kind == EntityKind::Email) + .unwrap(); + assert_eq!(email.span_start, 2); + } + + #[test] + fn all_mechanical_kinds_in_one_pass() { + let o = extract("email a@b.com, url https://x.com, @alice, #topic1"); + let k = kinds(&o); + assert!(k.contains(&EntityKind::Email)); + assert!(k.contains(&EntityKind::Url)); + assert!(k.contains(&EntityKind::Handle)); + assert!(k.contains(&EntityKind::Hashtag)); + } + + #[test] + fn scores_always_one() { + let o = extract("a@b.com #x @y https://q.com"); + for e in &o.entities { + assert!((e.score - 1.0).abs() < f32::EPSILON); + } + } + + #[test] + fn empty_input_no_matches() { + let o = extract("plain prose with no identifiers"); + assert!(o.entities.is_empty()); + } +} diff --git a/src/openhuman/memory/tree/score/extract/types.rs b/src/openhuman/memory/tree/score/extract/types.rs new file mode 100644 index 000000000..883aed584 --- /dev/null +++ b/src/openhuman/memory/tree/score/extract/types.rs @@ -0,0 +1,259 @@ +//! Types produced by entity extractors (Phase 2 / #708). +//! +//! The pipeline runs one or more [`super::EntityExtractor`] impls over each +//! admitted chunk and collects all their output into [`ExtractedEntities`]. + +use serde::{Deserialize, Serialize}; + +/// Classification of an extracted span. +/// +/// Split into two categories: +/// - **Mechanical** — regex finds these deterministically. Stable, high precision, +/// limited recall. These are "identifiers" (pointers), not "entities" +/// in the semantic sense. +/// - **Semantic** — model-based (future GLiNER / LLM). Named references to +/// real-world objects: Person, Organization, Location, Event, Product. +/// +/// Phase 2 ships with mechanical-only; semantic variants are populated in +/// Phase 3+ either at seal time by the summariser LLM or by a dedicated +/// per-chunk NER step if added later. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +#[non_exhaustive] +pub enum EntityKind { + // Mechanical + Email, + Url, + Handle, + Hashtag, + // Semantic (reserved — not emitted in Phase 2) + Person, + Organization, + Location, + Event, + Product, + Misc, +} + +impl EntityKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Email => "email", + Self::Url => "url", + Self::Handle => "handle", + Self::Hashtag => "hashtag", + Self::Person => "person", + Self::Organization => "organization", + Self::Location => "location", + Self::Event => "event", + Self::Product => "product", + Self::Misc => "misc", + } + } + + pub fn parse(s: &str) -> Result { + match s { + "email" => Ok(Self::Email), + "url" => Ok(Self::Url), + "handle" => Ok(Self::Handle), + "hashtag" => Ok(Self::Hashtag), + "person" => Ok(Self::Person), + "organization" => Ok(Self::Organization), + "location" => Ok(Self::Location), + "event" => Ok(Self::Event), + "product" => Ok(Self::Product), + "misc" => Ok(Self::Misc), + other => Err(format!("unknown entity kind: {other}")), + } + } + + /// Whether this kind comes from deterministic extraction. + pub fn is_mechanical(self) -> bool { + matches!(self, Self::Email | Self::Url | Self::Handle | Self::Hashtag) + } +} + +/// One extracted span from a chunk's content. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ExtractedEntity { + pub kind: EntityKind, + /// Surface form as it appears in the chunk. + pub text: String, + /// Character offsets `[start, end)` into the chunk text. + pub span_start: u32, + pub span_end: u32, + /// Extractor confidence `[0.0, 1.0]`. Regex = 1.0; model-based = output. + pub score: f32, +} + +/// Topic candidate (hashtag-style or summariser-labeled). +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ExtractedTopic { + /// Normalised topic text (lowercase, no leading `#`). + pub label: String, + pub score: f32, +} + +/// Aggregate output of one or more extractors on a single chunk. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ExtractedEntities { + pub entities: Vec, + pub topics: Vec, +} + +impl ExtractedEntities { + pub fn is_empty(&self) -> bool { + self.entities.is_empty() && self.topics.is_empty() + } + + /// Count of unique `(kind, text)` pairs, case-insensitive. Used as a scoring signal. + pub fn unique_entity_count(&self) -> usize { + use std::collections::BTreeSet; + self.entities + .iter() + .map(|e| (e.kind, e.text.to_lowercase())) + .collect::>() + .len() + } + + /// Merge another extractor's output into this one. + /// + /// Deduplicates by `(kind, normalised_text, span_start)` so the same + /// match from two extractors doesn't get double-counted. + pub fn merge(&mut self, other: ExtractedEntities) { + use std::collections::BTreeSet; + let mut seen: BTreeSet<(EntityKind, String, u32)> = self + .entities + .iter() + .map(|e| (e.kind, e.text.to_lowercase(), e.span_start)) + .collect(); + for e in other.entities { + let key = (e.kind, e.text.to_lowercase(), e.span_start); + if seen.insert(key) { + self.entities.push(e); + } + } + let mut topic_seen: BTreeSet = + self.topics.iter().map(|t| t.label.clone()).collect(); + for t in other.topics { + if topic_seen.insert(t.label.clone()) { + self.topics.push(t); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn entity_kind_round_trip() { + for k in [ + EntityKind::Email, + EntityKind::Url, + EntityKind::Handle, + EntityKind::Hashtag, + EntityKind::Person, + EntityKind::Organization, + EntityKind::Location, + EntityKind::Event, + EntityKind::Product, + EntityKind::Misc, + ] { + assert_eq!(EntityKind::parse(k.as_str()).unwrap(), k); + } + } + + #[test] + fn mechanical_classification() { + assert!(EntityKind::Email.is_mechanical()); + assert!(EntityKind::Url.is_mechanical()); + assert!(EntityKind::Handle.is_mechanical()); + assert!(EntityKind::Hashtag.is_mechanical()); + assert!(!EntityKind::Person.is_mechanical()); + } + + #[test] + fn unique_entity_count_dedups_case_insensitive() { + let e = ExtractedEntities { + entities: vec![ + ExtractedEntity { + kind: EntityKind::Person, + text: "Alice".into(), + span_start: 0, + span_end: 5, + score: 1.0, + }, + ExtractedEntity { + kind: EntityKind::Person, + text: "alice".into(), + span_start: 10, + span_end: 15, + score: 1.0, + }, + ], + topics: vec![], + }; + assert_eq!(e.unique_entity_count(), 1); + } + + #[test] + fn unique_entity_count_keeps_different_kinds_distinct() { + let e = ExtractedEntities { + entities: vec![ + ExtractedEntity { + kind: EntityKind::Handle, + text: "alice".into(), + span_start: 0, + span_end: 5, + score: 1.0, + }, + ExtractedEntity { + kind: EntityKind::Hashtag, + text: "alice".into(), + span_start: 10, + span_end: 15, + score: 1.0, + }, + ], + topics: vec![], + }; + assert_eq!(e.unique_entity_count(), 2); + } + + #[test] + fn merge_dedups_by_kind_text_span() { + let mut a = ExtractedEntities { + entities: vec![ExtractedEntity { + kind: EntityKind::Email, + text: "x@y.com".into(), + span_start: 0, + span_end: 7, + score: 1.0, + }], + topics: vec![], + }; + let b = ExtractedEntities { + entities: vec![ + ExtractedEntity { + kind: EntityKind::Email, + text: "x@y.com".into(), + span_start: 0, + span_end: 7, + score: 1.0, + }, // dup + ExtractedEntity { + kind: EntityKind::Email, + text: "x@y.com".into(), + span_start: 50, + span_end: 57, + score: 1.0, + }, // different span — keep + ], + topics: vec![], + }; + a.merge(b); + assert_eq!(a.entities.len(), 2); + } +} diff --git a/src/openhuman/memory/tree/score/mod.rs b/src/openhuman/memory/tree/score/mod.rs new file mode 100644 index 000000000..6ec12565c --- /dev/null +++ b/src/openhuman/memory/tree/score/mod.rs @@ -0,0 +1,305 @@ +//! Phase 2: scoring / admission / enrichment pipeline (#708). +//! +//! Wraps extraction, signal computation, admission gate, canonicalisation, +//! and persistence into one call per chunk. Phase 1 `_ingest_one_chunk` +//! passes each chunk through [`score_chunk`] after chunking and before +//! storing. + +pub mod extract; +pub mod resolver; +pub mod signals; +pub mod store; + +use std::sync::Arc; + +use anyhow::Result; +use chrono::Utc; +use futures_util::future::try_join_all; +use rusqlite::Transaction; +use serde::{Deserialize, Serialize}; + +use self::extract::{EntityExtractor, ExtractedEntities}; +use self::resolver::{canonicalise, CanonicalEntity}; +use self::signals::{ScoreSignals, SignalWeights}; +use crate::openhuman::memory::tree::types::{approx_token_count, Chunk, SourceKind}; + +/// Default drop threshold. Chunks with `total < DEFAULT_DROP_THRESHOLD` +/// are tombstoned and never reach the chunk store. +pub const DEFAULT_DROP_THRESHOLD: f32 = 0.3; + +/// Whole outcome of [`score_chunk`]. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ScoreResult { + pub chunk_id: String, + pub total: f32, + pub signals: ScoreSignals, + pub kept: bool, + pub drop_reason: Option, + pub extracted: ExtractedEntities, + pub canonical_entities: Vec, +} + +/// Configuration passed through the ingest pipeline for Phase 2 behaviour. +/// +/// Held as a struct (vs config struct fields) so callers can override per-run +/// without mutating global config — useful for tests and explicit threshold +/// tuning. +pub struct ScoringConfig { + pub extractor: Arc, + pub weights: SignalWeights, + pub drop_threshold: f32, +} + +impl ScoringConfig { + /// Phase 2 default: regex-only extractor, default weights, default threshold. + pub fn default_regex_only() -> Self { + Self { + extractor: Arc::new(extract::CompositeExtractor::regex_only()), + weights: SignalWeights::default(), + drop_threshold: DEFAULT_DROP_THRESHOLD, + } + } +} + +/// Compute the score for one chunk. +/// +/// Pure function — does not touch the store. Callers decide what to persist +/// based on [`ScoreResult::kept`]. +pub async fn score_chunk(chunk: &Chunk, cfg: &ScoringConfig) -> Result { + log::debug!( + "[memory_tree::score] score_chunk chunk_id={} tokens={}", + chunk.id, + chunk.token_count + ); + + let scoring_content = scoring_content_for_chunk(chunk); + let scoring_token_count = approx_token_count(&scoring_content); + + // 1. Extract entities (regex + any configured semantic extractors) + let extracted = cfg.extractor.extract(&scoring_content).await?; + + // 2. Compute signals + let signals = self::signals::compute( + &chunk.metadata, + &scoring_content, + scoring_token_count, + &extracted, + ); + + // 3. Weighted combine + let total = self::signals::combine(&signals, &cfg.weights); + + // 4. Admission gate. Source and interaction priors are deliberately + // non-zero, so guard against very short entity-free chatter being kept by + // metadata alone. + let tiny_entity_free = + scoring_token_count < self::signals::token_count::TOKEN_MIN && extracted.is_empty(); + let kept = !tiny_entity_free && total >= cfg.drop_threshold; + let drop_reason = if kept { + None + } else if tiny_entity_free { + Some(format!( + "token_count {} < minimum {} and no entities extracted", + scoring_token_count, + self::signals::token_count::TOKEN_MIN + )) + } else { + Some(format!( + "total {total:.3} < threshold {:.3}", + cfg.drop_threshold + )) + }; + + // 5. Canonicalise for indexing (only meaningful when kept — but we + // canonicalise unconditionally so the result is inspectable in tests) + let canonical_entities = canonicalise(&extracted); + + if !kept { + log::debug!( + "[memory_tree::score] drop chunk_id={} total={:.3} reason={:?}", + chunk.id, + total, + drop_reason + ); + } + + Ok(ScoreResult { + chunk_id: chunk.id.clone(), + total, + signals, + kept, + drop_reason, + extracted, + canonical_entities, + }) +} + +fn scoring_content_for_chunk(chunk: &Chunk) -> String { + if chunk.metadata.source_kind != SourceKind::Chat { + return chunk.content.clone(); + } + + chunk + .content + .lines() + .filter(|line| { + let trimmed = line.trim_start(); + !trimmed.starts_with("# Chat transcript") && !trimmed.starts_with("## ") + }) + .collect::>() + .join("\n") +} + +/// Score a batch of chunks. Errors from any single chunk fail the batch — +/// scoring is pure-ish (only the extractor may error) and a failure here is +/// a real bug, not a per-chunk issue to tolerate silently. +pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result> { + try_join_all(chunks.iter().map(|chunk| score_chunk(chunk, cfg))).await +} + +// ── Persistence helpers used by the ingest orchestrator ───────────────── + +/// Persist the score row + entity-index rows for one kept chunk. +/// +/// The caller is responsible for having already written the chunk itself +/// into `mem_tree_chunks` (so the FK-like relation is satisfied). Dropped +/// chunks still get a score row persisted for diagnostics — callers should +/// pass `None` for `tree_id` in that case, since the chunk won't appear in +/// a tree. +pub fn persist_score( + config: &crate::openhuman::config::Config, + result: &ScoreResult, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result<()> { + let row = score_row(result); + store::upsert_score(config, &row)?; + + if result.kept { + // Clear any stale entity-index rows for this chunk before re-indexing. + // INSERT OR REPLACE on (entity_id, node_id) never deletes rows whose + // entity_id is no longer present in the new extraction — so a re-score + // that drops an entity would otherwise leave a phantom index row. + store::clear_entity_index_for_node(config, &result.chunk_id)?; + if !result.canonical_entities.is_empty() { + store::index_entities( + config, + &result.canonical_entities, + &result.chunk_id, + "leaf", + timestamp_ms, + tree_id, + )?; + } + } + + Ok(()) +} + +pub(crate) fn persist_score_tx( + tx: &Transaction<'_>, + result: &ScoreResult, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result<()> { + let row = score_row(result); + store::upsert_score_tx(tx, &row)?; + + if result.kept { + // See persist_score for why we clear before re-indexing. + store::clear_entity_index_for_node_tx(tx, &result.chunk_id)?; + if !result.canonical_entities.is_empty() { + store::index_entities_tx( + tx, + &result.canonical_entities, + &result.chunk_id, + "leaf", + timestamp_ms, + tree_id, + )?; + } + } + + Ok(()) +} + +fn score_row(result: &ScoreResult) -> store::ScoreRow { + // Score rows keep wall-clock scoring time; the separate timestamp_ms + // argument used for entity indexes is the source/ingest ordering time. + store::ScoreRow { + chunk_id: result.chunk_id.clone(), + total: result.total, + signals: result.signals.clone(), + dropped: !result.kept, + reason: result.drop_reason.clone(), + computed_at_ms: Utc::now().timestamp_millis(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind}; + use chrono::Utc; + + fn test_chunk(content: &str) -> Chunk { + let meta = Metadata::point_in_time(SourceKind::Email, "t1", "alice", Utc::now()); + Chunk { + id: chunk_id(SourceKind::Email, "t1", 0), + content: content.to_string(), + token_count: crate::openhuman::memory::tree::types::approx_token_count(content), + metadata: meta, + seq_in_source: 0, + created_at: Utc::now(), + } + } + + #[tokio::test] + async fn substantive_chunk_is_kept() { + let c = test_chunk( + "We decided to ship Phoenix on Friday after reviewing \ + alice@example.com and the migration plan carefully. \ + @bob will coordinate and we discussed #launch-q2 details.", + ); + let cfg = ScoringConfig::default_regex_only(); + let r = score_chunk(&c, &cfg).await.unwrap(); + assert!(r.kept, "expected kept, got total={}", r.total); + assert!(r.drop_reason.is_none()); + assert!(!r.extracted.entities.is_empty()); + assert!(!r.canonical_entities.is_empty()); + } + + #[tokio::test] + async fn noise_chunk_is_dropped() { + // Very short — below TOKEN_MIN — and no entities. + let c = test_chunk("lol"); + let cfg = ScoringConfig::default_regex_only(); + let r = score_chunk(&c, &cfg).await.unwrap(); + assert!(!r.kept); + assert!(r.drop_reason.is_some()); + } + + #[tokio::test] + async fn threshold_override_respected() { + let c = test_chunk("just ok content, mid-signal"); + let mut cfg = ScoringConfig::default_regex_only(); + cfg.drop_threshold = 0.99; // unreasonably high + let r = score_chunk(&c, &cfg).await.unwrap(); + assert!(!r.kept); + } + + #[tokio::test] + async fn entities_are_canonicalised() { + let c = test_chunk("ping Alice@Example.com — she @alice replied to thread"); + let cfg = ScoringConfig::default_regex_only(); + let r = score_chunk(&c, &cfg).await.unwrap(); + // Email (lowercased) and handle canonical ids should both appear + let ids: Vec<_> = r + .canonical_entities + .iter() + .map(|e| e.canonical_id.as_str()) + .collect(); + assert!(ids.iter().any(|id| *id == "email:alice@example.com")); + assert!(ids.iter().any(|id| *id == "handle:alice")); + } +} diff --git a/src/openhuman/memory/tree/score/resolver.rs b/src/openhuman/memory/tree/score/resolver.rs new file mode 100644 index 000000000..c3b1d1fd2 --- /dev/null +++ b/src/openhuman/memory/tree/score/resolver.rs @@ -0,0 +1,136 @@ +//! Entity canonicalisation / cross-platform merge (Phase 2 / #708, V1). +//! +//! Exact-match only: normalises surface forms (lowercase emails, strip +//! leading `@` on handles) and assigns a canonical `entity_id` string. +//! +//! Fuzzy matching (alice-slack ≡ Alice-Discord by soft match) is deferred +//! until we have real entity-graph data — the current implementation +//! handles the mechanical cases cleanly without producing false merges. + +use serde::{Deserialize, Serialize}; + +use crate::openhuman::memory::tree::score::extract::{EntityKind, ExtractedEntities}; + +/// Canonicalised entity — same shape as [`ExtractedEntity`] plus a stable +/// `canonical_id` suitable for indexing. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CanonicalEntity { + pub canonical_id: String, + pub kind: EntityKind, + pub surface: String, + pub span_start: u32, + pub span_end: u32, + pub score: f32, +} + +/// Canonicalise a batch of extracted entities. +/// +/// Same surface form (after normalisation) → same `canonical_id` regardless +/// of how many times it appears in a chunk. Preserves source spans by +/// emitting one [`CanonicalEntity`] per occurrence. +pub fn canonicalise(extracted: &ExtractedEntities) -> Vec { + extracted + .entities + .iter() + .map(|e| CanonicalEntity { + canonical_id: canonical_id_for(e.kind, &e.text), + kind: e.kind, + surface: e.text.clone(), + span_start: e.span_start, + span_end: e.span_end, + score: e.score, + }) + .collect() +} + +/// Canonical id form per kind. Deterministic so the same surface always +/// maps to the same id. +/// +/// - Email: `email:lowercased` +/// - Handle: `handle:lowercased` with leading `@` stripped +/// - Hashtag: `hashtag:lowercased` with leading `#` stripped +/// - URL: `url:trimmed` with case preserved for path/query exact matching +/// - Semantic kinds: `kind:lowercased-surface` (V1; fuzzy merge deferred) +pub fn canonical_id_for(kind: EntityKind, surface: &str) -> String { + let trimmed = surface.trim(); + let clean = if kind == EntityKind::Url { + trimmed.to_string() + } else { + trimmed + .to_lowercase() + .trim_start_matches('@') + .trim_start_matches('#') + .to_string() + }; + format!("{}:{}", kind.as_str(), clean) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::score::extract::ExtractedEntity; + + fn entity(kind: EntityKind, text: &str) -> ExtractedEntity { + ExtractedEntity { + kind, + text: text.to_string(), + span_start: 0, + span_end: text.chars().count() as u32, + score: 1.0, + } + } + + #[test] + fn email_case_insensitive_canonicalises() { + let a = canonical_id_for(EntityKind::Email, "Alice@Example.com"); + let b = canonical_id_for(EntityKind::Email, "alice@example.com"); + assert_eq!(a, b); + assert_eq!(a, "email:alice@example.com"); + } + + #[test] + fn handle_strips_leading_at() { + let a = canonical_id_for(EntityKind::Handle, "@alice"); + let b = canonical_id_for(EntityKind::Handle, "alice"); + assert_eq!(a, b); + assert_eq!(a, "handle:alice"); + } + + #[test] + fn hashtag_strips_leading_hash() { + let a = canonical_id_for(EntityKind::Hashtag, "#launch"); + let b = canonical_id_for(EntityKind::Hashtag, "launch"); + assert_eq!(a, b); + } + + #[test] + fn url_preserves_case() { + let id = canonical_id_for(EntityKind::Url, " https://example.com/Path?Token=ABC "); + assert_eq!(id, "url:https://example.com/Path?Token=ABC"); + } + + #[test] + fn canonicalise_batch_preserves_spans() { + let ex = ExtractedEntities { + entities: vec![ + entity(EntityKind::Email, "Alice@Example.com"), + entity(EntityKind::Email, "alice@example.com"), + ], + topics: vec![], + }; + let out = canonicalise(&ex); + assert_eq!(out.len(), 2); + // Both map to the same canonical id (merge-equivalent) + assert_eq!(out[0].canonical_id, out[1].canonical_id); + // But surface forms remain distinct + assert_ne!(out[0].surface, out[1].surface); + } + + #[test] + fn different_kinds_produce_different_ids_for_same_text() { + assert_ne!( + canonical_id_for(EntityKind::Handle, "alice"), + canonical_id_for(EntityKind::Person, "alice") + ); + } +} diff --git a/src/openhuman/memory/tree/score/signals/interaction.rs b/src/openhuman/memory/tree/score/signals/interaction.rs new file mode 100644 index 000000000..1ba96b85e --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/interaction.rs @@ -0,0 +1,101 @@ +//! Interaction-weight signal — boosts chunks the user actively engaged with. +//! +//! Direct engagement is one of the strongest retention signals — "a message +//! you replied to" is almost always worth remembering, even if its content +//! looks noisy by other signals. +//! +//! Phase 2 infers engagement from a small set of reserved **tags**: +//! - `reply` — the user replied to this message/thread +//! - `sent` — the user authored this content +//! - `mention` — the user was @-mentioned +//! - `dm` — this arrived in a direct-message channel +//! +//! Ingest adapters can attach these tags during canonicalisation when the +//! upstream source supports the distinction. Absent tags → neutral score. + +use crate::openhuman::memory::tree::types::Metadata; + +pub const TAG_REPLY: &str = "reply"; +pub const TAG_SENT: &str = "sent"; +pub const TAG_MENTION: &str = "mention"; +pub const TAG_DM: &str = "dm"; + +/// Score in `[0.0, 1.0]` based on engagement tags present on the chunk. +/// +/// Multiple tags stack (capped at 1.0): +/// - `sent` → +0.6 (author) +/// - `reply` → +0.5 (active dialogue) +/// - `dm` → +0.3 (scoped audience) +/// - `mention` → +0.2 (addressed) +/// +/// Absent any of these → 0.5 (neutral — don't drop the chunk on this signal +/// alone since most content lacks explicit engagement tags). +pub fn score(meta: &Metadata) -> f32 { + let mut any_tag = false; + let mut total: f32 = 0.0; + for t in &meta.tags { + match t.as_str() { + TAG_SENT => { + total += 0.6; + any_tag = true; + } + TAG_REPLY => { + total += 0.5; + any_tag = true; + } + TAG_DM => { + total += 0.3; + any_tag = true; + } + TAG_MENTION => { + total += 0.2; + any_tag = true; + } + _ => {} + } + } + if !any_tag { + return 0.5; + } + total.clamp(0.0, 1.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::types::SourceKind; + use chrono::Utc; + + fn meta(tags: &[&str]) -> Metadata { + let mut m = Metadata::point_in_time(SourceKind::Chat, "x", "owner", Utc::now()); + m.tags = tags.iter().map(|s| s.to_string()).collect(); + m + } + + #[test] + fn no_tags_neutral() { + assert_eq!(score(&meta(&[])), 0.5); + assert_eq!(score(&meta(&["unrelated"])), 0.5); + } + + #[test] + fn sent_tag_high_score() { + assert!((score(&meta(&["sent"])) - 0.6).abs() < 1e-6); + } + + #[test] + fn stacking_capped_at_one() { + // sent (0.6) + reply (0.5) + mention (0.2) = 1.3 → clamp to 1.0 + assert!((score(&meta(&["sent", "reply", "mention"])) - 1.0).abs() < 1e-6); + } + + #[test] + fn reply_only() { + assert!((score(&meta(&["reply"])) - 0.5).abs() < 1e-6); + } + + #[test] + fn dm_plus_mention() { + assert!((score(&meta(&["dm", "mention"])) - 0.5).abs() < 1e-6); + } +} diff --git a/src/openhuman/memory/tree/score/signals/metadata_weight.rs b/src/openhuman/memory/tree/score/signals/metadata_weight.rs new file mode 100644 index 000000000..3b2a30265 --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/metadata_weight.rs @@ -0,0 +1,49 @@ +//! Metadata-weight signal — base weight from the source kind's grouping. +//! +//! The idea: a 1:1 email thread is inherently higher-signal than a broadcast +//! Slack channel, regardless of content. This signal captures the "shape" +//! of the interaction: how scoped is the audience? +//! +//! Phase 2 keeps this simple: one weight per `SourceKind`. Per-grouping +//! context (e.g., channel size, thread participant count) is a future +//! refinement when we actually have that metadata at ingest. + +use crate::openhuman::memory::tree::types::{Metadata, SourceKind}; + +/// Base weight for each source kind. +/// +/// Email threads are typically scoped (1:1 or small groups, directed). +/// Documents are single-author outputs — high intentionality per chunk. +/// Chats vary widely; base weight is lower because the channel could be +/// a 200-person broadcast or a tight DM — the interaction signal disambiguates. +pub fn score(meta: &Metadata) -> f32 { + match meta.source_kind { + SourceKind::Email => 0.8, + SourceKind::Document => 0.9, + SourceKind::Chat => 0.5, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + + fn meta(kind: SourceKind) -> Metadata { + Metadata::point_in_time(kind, "x", "owner", Utc::now()) + } + + #[test] + fn per_kind_weights() { + assert!(score(&meta(SourceKind::Document)) > score(&meta(SourceKind::Email))); + assert!(score(&meta(SourceKind::Email)) > score(&meta(SourceKind::Chat))); + } + + #[test] + fn bounded_zero_one() { + for k in [SourceKind::Chat, SourceKind::Email, SourceKind::Document] { + let s = score(&meta(k)); + assert!((0.0..=1.0).contains(&s)); + } + } +} diff --git a/src/openhuman/memory/tree/score/signals/mod.rs b/src/openhuman/memory/tree/score/signals/mod.rs new file mode 100644 index 000000000..321e13437 --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/mod.rs @@ -0,0 +1,20 @@ +//! Score signals + weighted combine (Phase 2 / #708). +//! +//! Each submodule computes one scoring signal in `[0.0, 1.0]`. [`combine`] +//! aggregates them into a total score using per-signal weights. The output +//! is still `[0.0, 1.0]` after normalisation by total weight. +//! +//! Storing per-signal values alongside the total (via [`ScoreSignals`]) is +//! what makes admission decisions debuggable — when a chunk is dropped, we +//! persist *which* signals fired at what values. + +pub mod interaction; +pub mod metadata_weight; +mod ops; +pub mod source_weight; +pub mod token_count; +mod types; +pub mod unique_words; + +pub use ops::{combine, compute, entity_density_score}; +pub use types::{ScoreSignals, SignalWeights}; diff --git a/src/openhuman/memory/tree/score/signals/ops.rs b/src/openhuman/memory/tree/score/signals/ops.rs new file mode 100644 index 000000000..900de969b --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/ops.rs @@ -0,0 +1,142 @@ +use super::{interaction, metadata_weight, source_weight, token_count, unique_words}; +use super::{ScoreSignals, SignalWeights}; +use crate::openhuman::memory::tree::score::extract::ExtractedEntities; +use crate::openhuman::memory::tree::types::Metadata; + +/// Compute all signals for a chunk. +pub fn compute( + meta: &Metadata, + content: &str, + token_count: u32, + ex: &ExtractedEntities, +) -> ScoreSignals { + ScoreSignals { + token_count: token_count::score(token_count), + unique_words: unique_words::score(content), + metadata_weight: metadata_weight::score(meta), + source_weight: source_weight::score(meta), + interaction: interaction::score(meta), + entity_density: entity_density_score(token_count, ex), + } +} + +/// Entity-density signal: entities per token, capped. +/// +/// More distinct entities per unit of content → more substantive. Calibrated +/// so ~1 entity per 100 tokens maxes out the signal. +pub fn entity_density_score(token_count: u32, ex: &ExtractedEntities) -> f32 { + let unique = ex.unique_entity_count() as f32; + if token_count == 0 { + return 0.0; + } + let per_token = unique / token_count as f32; + // cap at 0.01 entities/token = 1 entity per 100 tokens + (per_token / 0.01).min(1.0) +} + +/// Weighted sum of signals, normalised to `[0.0, 1.0]`. +pub fn combine(signals: &ScoreSignals, w: &SignalWeights) -> f32 { + let total_weight = w.token_count + + w.unique_words + + w.metadata_weight + + w.source_weight + + w.interaction + + w.entity_density; + if total_weight <= 0.0 { + return 0.0; + } + let weighted = signals.token_count * w.token_count + + signals.unique_words * w.unique_words + + signals.metadata_weight * w.metadata_weight + + signals.source_weight * w.source_weight + + signals.interaction * w.interaction + + signals.entity_density * w.entity_density; + (weighted / total_weight).clamp(0.0, 1.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::score::extract::{ + EntityKind, ExtractedEntities, ExtractedEntity, + }; + use crate::openhuman::memory::tree::types::SourceKind; + use chrono::Utc; + + fn meta(tags: &[&str], kind: SourceKind) -> Metadata { + let mut m = Metadata::point_in_time(kind, "x", "owner", Utc::now()); + m.tags = tags.iter().map(|s| s.to_string()).collect(); + m + } + + fn make_entities(n: usize) -> ExtractedEntities { + ExtractedEntities { + entities: (0..n) + .map(|i| ExtractedEntity { + kind: EntityKind::Email, + text: format!("user{i}@example.com"), + span_start: 0, + span_end: 10, + score: 1.0, + }) + .collect(), + topics: vec![], + } + } + + #[test] + fn combine_all_zeros_is_zero() { + let s = ScoreSignals::default(); + assert!(combine(&s, &SignalWeights::default()) < 0.01); + } + + #[test] + fn combine_all_ones_is_one() { + let s = ScoreSignals { + token_count: 1.0, + unique_words: 1.0, + metadata_weight: 1.0, + source_weight: 1.0, + interaction: 1.0, + entity_density: 1.0, + }; + assert!((combine(&s, &SignalWeights::default()) - 1.0).abs() < 1e-6); + } + + #[test] + fn weights_influence_total() { + let s = ScoreSignals { + token_count: 0.0, + unique_words: 0.0, + metadata_weight: 0.0, + source_weight: 0.0, + interaction: 1.0, + entity_density: 0.0, + }; + let total = combine(&s, &SignalWeights::default()); + assert!((total - (3.0 / 9.0)).abs() < 1e-6); + } + + #[test] + fn compute_wires_all_signals() { + let m = meta(&["reply"], SourceKind::Email); + let ex = make_entities(3); + let s = compute( + &m, + "Some substantive text about Phoenix launch planning.", + 12, + &ex, + ); + assert!(s.interaction > 0.0); + assert!(s.metadata_weight > 0.0); + assert!(s.source_weight > 0.0); + } + + #[test] + fn entity_density_scales() { + let ex = make_entities(1); + assert!((entity_density_score(100, &ex) - 1.0).abs() < 1e-6); + assert!((entity_density_score(1000, &ex) - 0.1).abs() < 1e-6); + assert_eq!(entity_density_score(0, &ex), 0.0); + } +} diff --git a/src/openhuman/memory/tree/score/signals/source_weight.rs b/src/openhuman/memory/tree/score/signals/source_weight.rs new file mode 100644 index 000000000..2339e0836 --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/source_weight.rs @@ -0,0 +1,110 @@ +//! Source-weight signal — per-provider base weight derived from the +//! `DataSource` when it can be inferred from a chunk's tags. +//! +//! Rationale from `Memory Architecture.md` (Step 2.3 "Source scoring"): +//! - High-intentionality messaging (direct DMs, personal emails) scores higher +//! - Broadcast/channel content scores lower +//! - Documents authored by the user score higher than shared-but-unmodified drops +//! +//! Phase 2 takes a conservative approach: per-[`DataSource`] base weight. +//! Finer distinction (DM vs channel on Slack specifically) requires richer +//! ingest-time metadata and is deferred. + +use crate::openhuman::memory::tree::types::{DataSource, Metadata, SourceKind}; + +const PROVIDER_PREFIX: &str = "provider:"; + +/// Best-effort map from `Metadata` to a [`DataSource`] — checks the `tags` +/// list for a stable `provider:` provider tag. If not present, +/// falls back to kind-based defaults. +/// +/// The ingestion pipeline can (and should) add a provider tag on the +/// canonicalised output so this signal fires deterministically. Until that's +/// wired everywhere, we fall back to the kind-level default. +pub fn infer_data_source(meta: &Metadata) -> Option { + for tag in &meta.tags { + let Some(provider) = tag.strip_prefix(PROVIDER_PREFIX) else { + continue; + }; + if let Ok(ds) = DataSource::parse(provider) { + return Some(ds); + } + } + None +} + +/// Score in `[0.0, 1.0]` for the chunk's originating provider. +pub fn score(meta: &Metadata) -> f32 { + if let Some(ds) = infer_data_source(meta) { + return weight_for(ds); + } + // Fallback: kind-level defaults consistent with per-provider averages. + match meta.source_kind { + SourceKind::Email => 0.75, + SourceKind::Document => 0.7, + SourceKind::Chat => 0.5, + } +} + +fn weight_for(ds: DataSource) -> f32 { + match ds { + // Personal email providers score high — typically small, directed audiences + DataSource::Gmail => 0.8, + DataSource::OtherEmail => 0.7, + // Chat providers differ: WhatsApp is typically DM-heavy, Discord + // can be broadcast-heavy, Telegram mixes both + DataSource::Whatsapp => 0.75, + DataSource::Telegram => 0.6, + DataSource::Discord => 0.5, + // Documents: Notion = structured, Drive = mixed, Meeting notes = high value + DataSource::Notion => 0.75, + DataSource::DriveDocs => 0.6, + DataSource::MeetingNotes => 0.85, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + + fn meta_with_tag(kind: SourceKind, tag: &str) -> Metadata { + let mut m = Metadata::point_in_time(kind, "x", "owner", Utc::now()); + m.tags.push(tag.to_string()); + m + } + + #[test] + fn data_source_inferred_from_tags() { + let m = meta_with_tag(SourceKind::Chat, "provider:whatsapp"); + assert_eq!(infer_data_source(&m), Some(DataSource::Whatsapp)); + } + + #[test] + fn plain_user_label_does_not_infer_provider() { + let m = meta_with_tag(SourceKind::Email, "notion"); + assert_eq!(infer_data_source(&m), None); + assert!((score(&m) - 0.75).abs() < 1e-6); + } + + #[test] + fn unknown_tag_falls_back_to_kind_default() { + let m = meta_with_tag(SourceKind::Email, "not-a-data-source"); + let s = score(&m); + assert!((s - 0.75).abs() < 1e-6); + } + + #[test] + fn provider_specific_weights_applied() { + let m = meta_with_tag(SourceKind::Document, "provider:meeting_notes"); + assert!((score(&m) - 0.85).abs() < 1e-6); + } + + #[test] + fn all_data_sources_bounded() { + for ds in DataSource::all() { + let w = weight_for(*ds); + assert!((0.0..=1.0).contains(&w)); + } + } +} diff --git a/src/openhuman/memory/tree/score/signals/token_count.rs b/src/openhuman/memory/tree/score/signals/token_count.rs new file mode 100644 index 000000000..a3b6e0cd7 --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/token_count.rs @@ -0,0 +1,79 @@ +//! Token-count signal — penalises very short or very long chunks. +//! +//! Rationale: "+1", "lol", "👍" are usually noise; multi-page walls of text +//! are often pasted logs or attachments that overwhelm summarisation. +//! The signal is strongest in a middle band that corresponds to substantive +//! prose/discussion. +//! +//! Output is a score in `[0.0, 1.0]` shaped as a plateau between +//! `TOKEN_MIN` and `TOKEN_MAX` with linear ramps on both sides. + +pub const TOKEN_MIN: u32 = 10; // below this → score 0 +pub const TOKEN_RAMP_LOW: u32 = 30; // 10..30 → linear 0→1 +pub const TOKEN_RAMP_HIGH: u32 = 3_000; // 3000..8000 → linear 1→0.5 +pub const TOKEN_MAX: u32 = 8_000; // above → score 0.5 (not zero — still has content) + +/// Score for a chunk's token count. See module docs for shape. +pub fn score(token_count: u32) -> f32 { + if token_count < TOKEN_MIN { + return 0.0; + } + if token_count <= TOKEN_RAMP_LOW { + // linear 0..1 over [MIN, RAMP_LOW] + let span = (TOKEN_RAMP_LOW - TOKEN_MIN) as f32; + return (token_count - TOKEN_MIN) as f32 / span; + } + if token_count <= TOKEN_RAMP_HIGH { + return 1.0; + } + if token_count <= TOKEN_MAX { + // linear 1.0..0.5 over [RAMP_HIGH, MAX] + let span = (TOKEN_MAX - TOKEN_RAMP_HIGH) as f32; + let t = (token_count - TOKEN_RAMP_HIGH) as f32 / span; + return 1.0 - 0.5 * t; + } + 0.5 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tiny_is_zero() { + assert_eq!(score(0), 0.0); + assert_eq!(score(5), 0.0); + assert_eq!(score(9), 0.0); + } + + #[test] + fn ramp_up_linear() { + // score(MIN) = 0, score(RAMP_LOW) = 1.0 + assert!((score(TOKEN_MIN) - 0.0).abs() < 1e-4); + assert!((score(TOKEN_RAMP_LOW) - 1.0).abs() < 1e-4); + // midpoint ~0.5 + let mid = TOKEN_MIN + (TOKEN_RAMP_LOW - TOKEN_MIN) / 2; + assert!((score(mid) - 0.5).abs() < 0.05); + } + + #[test] + fn plateau_is_one() { + assert_eq!(score(200), 1.0); + assert_eq!(score(1000), 1.0); + assert_eq!(score(TOKEN_RAMP_HIGH), 1.0); + } + + #[test] + fn ramp_down_to_half() { + assert!((score(TOKEN_MAX) - 0.5).abs() < 1e-4); + assert_eq!(score(TOKEN_MAX + 10_000), 0.5); + } + + #[test] + fn monotonic_in_bands() { + // Strictly increasing on the up-ramp + assert!(score(TOKEN_MIN + 1) < score(TOKEN_RAMP_LOW - 1)); + // Strictly decreasing on the down-ramp + assert!(score(TOKEN_RAMP_HIGH + 1) > score(TOKEN_MAX - 1)); + } +} diff --git a/src/openhuman/memory/tree/score/signals/types.rs b/src/openhuman/memory/tree/score/signals/types.rs new file mode 100644 index 000000000..7954fd510 --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/types.rs @@ -0,0 +1,37 @@ +use serde::{Deserialize, Serialize}; + +/// Per-signal score breakdown for one chunk. Persisted alongside the total +/// for diagnostics. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct ScoreSignals { + pub token_count: f32, + pub unique_words: f32, + pub metadata_weight: f32, + pub source_weight: f32, + pub interaction: f32, + pub entity_density: f32, +} + +/// Default weights applied to each signal in `combine`. +#[derive(Clone, Debug)] +pub struct SignalWeights { + pub token_count: f32, + pub unique_words: f32, + pub metadata_weight: f32, + pub source_weight: f32, + pub interaction: f32, + pub entity_density: f32, +} + +impl Default for SignalWeights { + fn default() -> Self { + Self { + token_count: 1.0, + unique_words: 1.0, + metadata_weight: 1.5, + source_weight: 1.5, + interaction: 3.0, // strongest signal — direct user engagement + entity_density: 1.0, + } + } +} diff --git a/src/openhuman/memory/tree/score/signals/unique_words.rs b/src/openhuman/memory/tree/score/signals/unique_words.rs new file mode 100644 index 000000000..846b7cb5e --- /dev/null +++ b/src/openhuman/memory/tree/score/signals/unique_words.rs @@ -0,0 +1,85 @@ +//! Unique-word-ratio signal — noise detector that fires on low-diversity text. +//! +//! Example: "yay yay yay yay lol lol lol" has high repetition = low diversity. +//! A substantive message has high type-token ratio (roughly, unique words / +//! total words). +//! +//! For very short messages the ratio is naturally ~1.0, so we require a +//! minimum total count before this signal contributes — otherwise "hi bob" +//! would score identically to a real message. + +pub const MIN_TOTAL_WORDS: usize = 5; + +/// Score in `[0.0, 1.0]` from the type-token ratio of `text`. +/// +/// - Too few total words → `0.5` (indeterminate — defer to other signals) +/// - Ratio < 0.3 (heavy repetition) → 0.0 +/// - Ratio >= 0.7 (substantive) → 1.0 +/// - Linear in between +pub fn score(text: &str) -> f32 { + let mut total: usize = 0; + let mut uniq: std::collections::BTreeSet = std::collections::BTreeSet::new(); + + for raw in text.split_whitespace() { + let w: String = raw + .trim_matches(|c: char| !c.is_alphanumeric()) + .to_lowercase(); + if w.is_empty() { + continue; + } + total += 1; + uniq.insert(w); + } + + if total < MIN_TOTAL_WORDS { + return 0.5; + } + + let ratio = uniq.len() as f32 / total as f32; + if ratio <= 0.3 { + 0.0 + } else if ratio >= 0.7 { + 1.0 + } else { + (ratio - 0.3) / 0.4 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn short_text_returns_neutral() { + assert_eq!(score(""), 0.5); + assert_eq!(score("hi bob"), 0.5); + } + + #[test] + fn high_repetition_scored_low() { + let noisy = "yay yay yay yay yay yay yay yay yay yay lol lol lol lol"; + assert!(score(noisy) < 0.2); + } + + #[test] + fn substantive_text_scored_high() { + let good = + "We decided to ship Phoenix on Friday after reviewing the migration plan carefully."; + assert!(score(good) >= 0.9); + } + + #[test] + fn medium_repetition_ramps() { + // ~50% unique ratio should score around 0.5 + let med = "alpha beta alpha beta gamma alpha delta beta gamma alpha"; + let s = score(med); + assert!(s > 0.2 && s < 0.8); + } + + #[test] + fn punctuation_stripped() { + let s1 = score("ship phoenix friday ship phoenix friday ship phoenix"); + let s2 = score("ship! phoenix, friday. ship! phoenix, friday. ship! phoenix."); + assert!((s1 - s2).abs() < 0.05); + } +} diff --git a/src/openhuman/memory/tree/score/store.rs b/src/openhuman/memory/tree/score/store.rs new file mode 100644 index 000000000..fb47306a2 --- /dev/null +++ b/src/openhuman/memory/tree/score/store.rs @@ -0,0 +1,473 @@ +//! Persistence for Phase 2 artefacts (#708): +//! +//! - `mem_tree_score` — per-chunk score rationale (which signals fired, why +//! dropped/kept) +//! - `mem_tree_entity_index` — inverted index `entity_id → node_id` so +//! retrieval can resolve entity-scoped queries in O(lookup) +//! +//! Schema is declared in `memory/tree/store.rs::SCHEMA`; this file only +//! owns the CRUD operations. + +use anyhow::Result; +use rusqlite::{params, Connection, OptionalExtension, Transaction}; +use serde::{Deserialize, Serialize}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::score::extract::EntityKind; +use crate::openhuman::memory::tree::score::resolver::CanonicalEntity; +use crate::openhuman::memory::tree::score::signals::ScoreSignals; +use crate::openhuman::memory::tree::store::with_connection; + +/// Serialized per-chunk score rationale. Mirrors the `mem_tree_score` row. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ScoreRow { + pub chunk_id: String, + pub total: f32, + pub signals: ScoreSignals, + pub dropped: bool, + pub reason: Option, + pub computed_at_ms: i64, +} + +/// Upsert one score rationale row, replacing any existing entry for `chunk_id`. +pub fn upsert_score(config: &Config, row: &ScoreRow) -> Result<()> { + with_connection(config, |conn| { + upsert_score_on_connection(conn, row)?; + Ok(()) + }) +} + +pub(crate) fn upsert_score_tx(tx: &Transaction<'_>, row: &ScoreRow) -> Result<()> { + tx.execute( + SCORE_UPSERT_SQL, + params![ + row.chunk_id, + row.total, + row.signals.token_count, + row.signals.unique_words, + row.signals.metadata_weight, + row.signals.source_weight, + row.signals.interaction, + row.signals.entity_density, + i32::from(row.dropped), + row.reason, + row.computed_at_ms, + ], + )?; + Ok(()) +} + +const SCORE_UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_score ( + chunk_id, total, + token_count_signal, unique_words_signal, + metadata_weight, source_weight, interaction_weight, entity_density, + dropped, reason, computed_at_ms + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)"; + +fn upsert_score_on_connection(conn: &Connection, row: &ScoreRow) -> Result<()> { + conn.execute( + SCORE_UPSERT_SQL, + params![ + row.chunk_id, + row.total, + row.signals.token_count, + row.signals.unique_words, + row.signals.metadata_weight, + row.signals.source_weight, + row.signals.interaction, + row.signals.entity_density, + i32::from(row.dropped), + row.reason, + row.computed_at_ms, + ], + )?; + Ok(()) +} + +/// Fetch one chunk's score rationale. +pub fn get_score(config: &Config, chunk_id: &str) -> Result> { + with_connection(config, |conn| { + conn.query_row( + "SELECT chunk_id, total, + token_count_signal, unique_words_signal, + metadata_weight, source_weight, interaction_weight, entity_density, + dropped, reason, computed_at_ms + FROM mem_tree_score WHERE chunk_id = ?1", + params![chunk_id], + |row| { + Ok(ScoreRow { + chunk_id: row.get(0)?, + total: row.get(1)?, + signals: ScoreSignals { + token_count: row.get(2)?, + unique_words: row.get(3)?, + metadata_weight: row.get(4)?, + source_weight: row.get(5)?, + interaction: row.get(6)?, + entity_density: row.get(7)?, + }, + dropped: row.get::<_, i32>(8)? != 0, + reason: row.get(9)?, + computed_at_ms: row.get(10)?, + }) + }, + ) + .optional() + .map_err(anyhow::Error::from) + }) +} + +/// Index one (entity, chunk) association. +/// +/// Idempotent on the composite primary key `(entity_id, node_id)` so +/// re-indexing the same association is a no-op update. +pub fn index_entity( + config: &Config, + entity: &CanonicalEntity, + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result<()> { + with_connection(config, |conn| { + conn.execute( + "INSERT OR REPLACE INTO mem_tree_entity_index ( + entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + entity.canonical_id, + node_id, + node_kind, + entity.kind.as_str(), + entity.surface, + entity.score, + timestamp_ms, + tree_id, + ], + )?; + Ok(()) + }) +} + +/// Batch index all entities extracted from a chunk. +pub fn index_entities( + config: &Config, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result { + if entities.is_empty() { + return Ok(0); + } + with_connection(config, |conn| { + let tx = conn.unchecked_transaction()?; + { + let mut stmt = tx.prepare( + "INSERT OR REPLACE INTO mem_tree_entity_index ( + entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + )?; + for e in entities { + stmt.execute(params![ + e.canonical_id, + node_id, + node_kind, + e.kind.as_str(), + e.surface, + e.score, + timestamp_ms, + tree_id, + ])?; + } + } + tx.commit()?; + Ok(entities.len()) + }) +} + +/// Remove all entity-index rows for a given node. Used before re-indexing +/// a re-scored chunk so entities dropped from the new extraction don't leak +/// through as stale `INSERT OR REPLACE` never deletes. +pub fn clear_entity_index_for_node(config: &Config, node_id: &str) -> Result { + with_connection(config, |conn| { + let n = conn.execute( + "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", + params![node_id], + )?; + Ok(n) + }) +} + +pub(crate) fn clear_entity_index_for_node_tx(tx: &Transaction<'_>, node_id: &str) -> Result { + let n = tx.execute( + "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", + params![node_id], + )?; + Ok(n) +} + +pub(crate) fn index_entities_tx( + tx: &Transaction<'_>, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result { + if entities.is_empty() { + return Ok(0); + } + let mut stmt = tx.prepare( + "INSERT OR REPLACE INTO mem_tree_entity_index ( + entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + )?; + for e in entities { + stmt.execute(params![ + e.canonical_id, + node_id, + node_kind, + e.kind.as_str(), + e.surface, + e.score, + timestamp_ms, + tree_id, + ])?; + } + Ok(entities.len()) +} + +/// Result row from [`lookup_entity`]. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct EntityHit { + pub entity_id: String, + pub node_id: String, + pub node_kind: String, + pub entity_kind: EntityKind, + pub surface: String, + pub score: f32, + pub timestamp_ms: i64, + pub tree_id: Option, +} + +/// Find all nodes indexed against `entity_id`, newest first. +pub fn lookup_entity( + config: &Config, + entity_id: &str, + limit: Option, +) -> Result> { + // Clamp to i64::MAX before casting so callers can't wrap a large usize + // into a negative LIMIT and bypass it. + let limit = limit.unwrap_or(100).min(i64::MAX as usize) as i64; + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id + FROM mem_tree_entity_index + WHERE entity_id = ?1 + ORDER BY timestamp_ms DESC + LIMIT ?2", + )?; + let rows = stmt + .query_map(params![entity_id, limit], |row| { + let kind_s: String = row.get(3)?; + let entity_kind = EntityKind::parse(&kind_s).map_err(|e| { + rusqlite::Error::FromSqlConversionFailure( + 3, + rusqlite::types::Type::Text, + e.into(), + ) + })?; + Ok(EntityHit { + entity_id: row.get(0)?, + node_id: row.get(1)?, + node_kind: row.get(2)?, + entity_kind, + surface: row.get(4)?, + score: row.get(5)?, + timestamp_ms: row.get(6)?, + tree_id: row.get(7)?, + }) + })? + .collect::>>()?; + Ok(rows) + }) +} + +/// Count rows in the entity index (for tests / diagnostics). +pub fn count_entity_index(config: &Config) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_entity_index", [], |r| { + r.get(0) + })?; + Ok(n.max(0) as u64) + }) +} + +/// Count score rows (for tests / diagnostics). +pub fn count_scores(config: &Config) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_score", [], |r| r.get(0))?; + Ok(n.max(0) as u64) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + 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(); + (tmp, cfg) + } + + fn sample_row(id: &str, dropped: bool) -> ScoreRow { + ScoreRow { + chunk_id: id.to_string(), + total: 0.7, + signals: ScoreSignals { + token_count: 1.0, + unique_words: 0.8, + metadata_weight: 0.9, + source_weight: 0.5, + interaction: 0.6, + entity_density: 0.3, + }, + dropped, + reason: if dropped { + Some("below threshold".into()) + } else { + None + }, + computed_at_ms: 1_700_000_000_000, + } + } + + fn sample_entity(id: &str) -> CanonicalEntity { + CanonicalEntity { + canonical_id: format!("email:{id}"), + kind: EntityKind::Email, + surface: format!("{id}@example.com"), + span_start: 0, + span_end: (id.len() + 12) as u32, + score: 1.0, + } + } + + #[test] + fn upsert_then_get_score() { + let (_tmp, cfg) = test_config(); + let row = sample_row("c1", false); + upsert_score(&cfg, &row).unwrap(); + let got = get_score(&cfg, "c1").unwrap().expect("row exists"); + assert_eq!(got.chunk_id, row.chunk_id); + assert!((got.total - row.total).abs() < 1e-6); + assert_eq!(got.dropped, row.dropped); + assert_eq!(got.reason, row.reason); + assert_eq!(got.computed_at_ms, row.computed_at_ms); + assert!((got.signals.token_count - row.signals.token_count).abs() < 1e-6); + } + + #[test] + fn upsert_score_idempotent() { + let (_tmp, cfg) = test_config(); + let r = sample_row("c1", false); + upsert_score(&cfg, &r).unwrap(); + upsert_score(&cfg, &r).unwrap(); + assert_eq!(count_scores(&cfg).unwrap(), 1); + } + + #[test] + fn dropped_flag_persists() { + let (_tmp, cfg) = test_config(); + let r = sample_row("c1", true); + upsert_score(&cfg, &r).unwrap(); + let got = get_score(&cfg, "c1").unwrap().unwrap(); + assert!(got.dropped); + assert_eq!(got.reason.as_deref(), Some("below threshold")); + } + + #[test] + fn get_missing_score_is_none() { + let (_tmp, cfg) = test_config(); + assert!(get_score(&cfg, "missing").unwrap().is_none()); + } + + #[test] + fn index_and_lookup_entity() { + let (_tmp, cfg) = test_config(); + let e = sample_entity("alice"); + index_entity(&cfg, &e, "chunk-1", "leaf", 1000, Some("source:chat")).unwrap(); + index_entity(&cfg, &e, "chunk-2", "leaf", 2000, Some("source:chat")).unwrap(); + + let hits = lookup_entity(&cfg, "email:alice", None).unwrap(); + assert_eq!(hits.len(), 2); + // newest first + assert_eq!(hits[0].node_id, "chunk-2"); + assert_eq!(hits[1].node_id, "chunk-1"); + } + + #[test] + fn index_batch() { + let (_tmp, cfg) = test_config(); + let entities = vec![sample_entity("a"), sample_entity("b"), sample_entity("c")]; + let n = index_entities(&cfg, &entities, "chunk-1", "leaf", 1000, None).unwrap(); + assert_eq!(n, 3); + assert_eq!(count_entity_index(&cfg).unwrap(), 3); + } + + #[test] + fn clear_entity_index_drops_stale_rows() { + let (_tmp, cfg) = test_config(); + let a = sample_entity("a"); + let b = sample_entity("b"); + index_entities(&cfg, &[a.clone(), b], "chunk-1", "leaf", 1000, None).unwrap(); + assert_eq!(count_entity_index(&cfg).unwrap(), 2); + + // Simulate a re-score that only keeps entity "a". + let cleared = clear_entity_index_for_node(&cfg, "chunk-1").unwrap(); + assert_eq!(cleared, 2); + index_entities(&cfg, &[a], "chunk-1", "leaf", 1000, None).unwrap(); + + let hits = lookup_entity(&cfg, "email:b", None).unwrap(); + assert!(hits.is_empty(), "stale entity should be removed"); + let hits = lookup_entity(&cfg, "email:a", None).unwrap(); + assert_eq!(hits.len(), 1); + } + + #[test] + fn index_idempotent_per_entity_node_pair() { + let (_tmp, cfg) = test_config(); + let e = sample_entity("alice"); + index_entity(&cfg, &e, "chunk-1", "leaf", 1000, None).unwrap(); + index_entity(&cfg, &e, "chunk-1", "leaf", 1000, None).unwrap(); + assert_eq!(count_entity_index(&cfg).unwrap(), 1); + } + + #[test] + fn lookup_limit_respected() { + let (_tmp, cfg) = test_config(); + let e = sample_entity("alice"); + for i in 0..5 { + index_entity( + &cfg, + &e, + &format!("chunk-{i}"), + "leaf", + 1000 + i as i64, + None, + ) + .unwrap(); + } + let hits = lookup_entity(&cfg, "email:alice", Some(2)).unwrap(); + assert_eq!(hits.len(), 2); + } +} diff --git a/src/openhuman/memory/tree/store.rs b/src/openhuman/memory/tree/store.rs index de19d6b36..bcb4a6ec1 100644 --- a/src/openhuman/memory/tree/store.rs +++ b/src/openhuman/memory/tree/store.rs @@ -9,7 +9,7 @@ use anyhow::{Context, Result}; use chrono::{DateTime, TimeZone, Utc}; -use rusqlite::{params, Connection, OptionalExtension}; +use rusqlite::{params, Connection, OptionalExtension, Transaction}; use std::time::Duration; use crate::openhuman::config::Config; @@ -48,6 +48,46 @@ CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_owner ON mem_tree_chunks(owner); CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_source_seq ON mem_tree_chunks(source_kind, source_id, seq_in_source); + +-- Phase 2 (#708): per-chunk score rationale for admission debugging. +CREATE TABLE IF NOT EXISTS mem_tree_score ( + chunk_id TEXT PRIMARY KEY, + total REAL NOT NULL, + token_count_signal REAL NOT NULL, + unique_words_signal REAL NOT NULL, + metadata_weight REAL NOT NULL, + source_weight REAL NOT NULL, + interaction_weight REAL NOT NULL, + entity_density REAL NOT NULL, + dropped INTEGER NOT NULL DEFAULT 0, + reason TEXT, + computed_at_ms INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_score_total + ON mem_tree_score(total); +CREATE INDEX IF NOT EXISTS idx_mem_tree_score_dropped + ON mem_tree_score(dropped); + +-- Phase 2 (#708): inverted index entity_id -> node_id for retrieval. +CREATE TABLE IF NOT EXISTS mem_tree_entity_index ( + entity_id TEXT NOT NULL, + node_id TEXT NOT NULL, + node_kind TEXT NOT NULL, + entity_kind TEXT NOT NULL, + surface TEXT NOT NULL, + score REAL NOT NULL, + timestamp_ms INTEGER NOT NULL, + tree_id TEXT, + PRIMARY KEY (entity_id, node_id) +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_entity + ON mem_tree_entity_index(entity_id); +CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_node + ON mem_tree_entity_index(node_id); +CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_timestamp + ON mem_tree_entity_index(timestamp_ms); "; /// Upsert a batch of chunks atomically. @@ -68,35 +108,85 @@ pub fn upsert_chunks(config: &Config, chunks: &[Chunk]) -> Result { let tx = conn.unchecked_transaction()?; { let mut stmt = tx.prepare( - "INSERT OR REPLACE INTO mem_tree_chunks ( + "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 - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)", + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) + 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", )?; - for chunk in chunks { - 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)?, - chunk.content, - chunk.token_count, - chunk.seq_in_source, - chunk.created_at.timestamp_millis(), - ])?; - } + upsert_chunks_with_statement(&mut stmt, chunks)?; } tx.commit()?; Ok(chunks.len()) }) } +/// Upsert chunks using an existing transaction, preserving previously stored embeddings. +pub(crate) fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result { + if chunks.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 + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) + 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", + )?; + upsert_chunks_with_statement(&mut stmt, chunks)?; + Ok(chunks.len()) +} + +fn upsert_chunks_with_statement( + stmt: &mut rusqlite::Statement<'_>, + chunks: &[Chunk], +) -> Result<()> { + for chunk in chunks { + 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)?, + chunk.content, + chunk.token_count, + chunk.seq_in_source, + chunk.created_at.timestamp_millis(), + ])?; + } + Ok(()) +} + /// Fetch one chunk by its id. pub fn get_chunk(config: &Config, id: &str) -> Result> { with_connection(config, |conn| { @@ -238,7 +328,14 @@ fn ms_to_utc(ms: i64) -> rusqlite::Result> { }) } -fn with_connection(config: &Config, f: impl FnOnce(&Connection) -> Result) -> Result { +/// Open the memory_tree SQLite DB and run a closure against it. +/// +/// Visible to sibling modules (e.g. `score::store`) so Phase 2 can reuse +/// the same connection setup / schema initialisation without duplication. +pub(crate) fn with_connection( + config: &Config, + f: impl FnOnce(&Connection) -> Result, +) -> Result { let dir = config.workspace_dir.join(DB_DIR); std::fs::create_dir_all(&dir) .with_context(|| format!("Failed to create memory_tree dir: {}", dir.display()))?; @@ -251,6 +348,8 @@ fn with_connection(config: &Config, f: impl FnOnce(&Connection) -> Result) .context("Failed to enable memory_tree WAL mode")?; conn.execute_batch(SCHEMA) .context("Failed to initialize memory_tree schema")?; + // Phase 2 migrations — additive, idempotent. + add_column_if_missing(&conn, "mem_tree_chunks", "embedding", "BLOB")?; f(&conn) } @@ -261,6 +360,69 @@ fn normalized_limit(requested: Option) -> i64 { i64::try_from(clamped).unwrap_or(MAX_LIST_LIMIT as i64) } +/// Idempotent `ALTER TABLE ADD COLUMN` — treats an existing column as success. +fn add_column_if_missing(conn: &Connection, table: &str, name: &str, sql_type: &str) -> Result<()> { + match conn.execute( + &format!("ALTER TABLE {table} ADD COLUMN {name} {sql_type}"), + [], + ) { + Ok(_) => { + log::debug!("[memory_tree::store] migration: added column {table}.{name} ({sql_type})"); + Ok(()) + } + Err(err) if err.to_string().contains("duplicate column name") => Ok(()), + Err(err) => Err(err).with_context(|| format!("Failed to add column {table}.{name}")), + } +} + +// ── Phase 2: embedding column accessors ───────────────────────────────── + +/// Store a chunk's embedding as a packed little-endian `f32` blob. +/// +/// Length is `embedding.len() * 4` bytes. The caller is responsible for +/// ensuring all embeddings in a given deployment share the same dimension. +pub fn set_chunk_embedding(config: &Config, chunk_id: &str, embedding: &[f32]) -> Result<()> { + let bytes: Vec = embedding.iter().flat_map(|f| f.to_le_bytes()).collect(); + with_connection(config, |conn| { + let changed = conn.execute( + "UPDATE mem_tree_chunks SET embedding = ?1 WHERE id = ?2", + rusqlite::params![bytes, chunk_id], + )?; + if changed == 0 { + log::warn!("[memory_tree::store] set_chunk_embedding: no row for chunk_id={chunk_id}"); + } + Ok(()) + }) +} + +/// Fetch a chunk's embedding, decoding the stored little-endian `f32` blob. +/// +/// Returns `Ok(None)` if the chunk doesn't exist or has no embedding stored. +pub fn get_chunk_embedding(config: &Config, chunk_id: &str) -> Result>> { + with_connection(config, |conn| { + let blob: Option>> = conn + .query_row( + "SELECT embedding FROM mem_tree_chunks WHERE id = ?1", + rusqlite::params![chunk_id], + |r| r.get::<_, Option>>(0), + ) + .optional()?; + match blob.flatten() { + None => Ok(None), + Some(bytes) => { + if !bytes.len().is_multiple_of(4) { + anyhow::bail!("embedding blob length {} not a multiple of 4", bytes.len()); + } + let floats: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + Ok(Some(floats)) + } + } + }) +} + #[cfg(test)] mod tests { use super::*; @@ -313,6 +475,24 @@ mod tests { assert_eq!(count_chunks(&cfg).unwrap(), 1); } + #[test] + fn reingest_preserves_existing_embedding() { + let (_tmp, cfg) = test_config(); + let mut c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + set_chunk_embedding(&cfg, &c.id, &[0.1, 0.2, 0.3]).unwrap(); + + c.content = "updated content".into(); + c.token_count = 99; + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + + let embedding = get_chunk_embedding(&cfg, &c.id).unwrap().unwrap(); + assert_eq!(embedding, vec![0.1, 0.2, 0.3]); + let got = get_chunk(&cfg, &c.id).unwrap().unwrap(); + assert_eq!(got.content, "updated content"); + assert_eq!(got.token_count, 99); + } + #[test] fn list_filters_by_source_kind() { let (_tmp, cfg) = test_config(); diff --git a/tests/json_rpc_e2e.rs b/tests/json_rpc_e2e.rs index b48830a06..7ce504c71 100644 --- a/tests/json_rpc_e2e.rs +++ b/tests/json_rpc_e2e.rs @@ -819,7 +819,7 @@ async fn json_rpc_memory_tree_end_to_end() { "payload": { "provider": "notion", "title": "Launch Plan", - "body": "Alpha\n\nBeta", + "body": "We decided to ship Phoenix on Friday after reviewing alice@example.com and the migration plan carefully. @bob will coordinate rollout, track #launch-q2 details, and update the Notion launch checklist with staging validation notes.", "modified_at": 1700000000000_i64, "source_ref": " notion://page/launch-plan " } @@ -833,6 +833,7 @@ async fn json_rpc_memory_tree_end_to_end() { Some(&json!("notion:launch-plan")) ); assert_eq!(ingest_result.get("chunks_written"), Some(&json!(1))); + assert_eq!(ingest_result.get("chunks_dropped"), Some(&json!(0))); let chunk_ids = ingest_result .get("chunk_ids") .and_then(Value::as_array)