mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
* 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 <workspace>/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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>
This commit is contained in:
sanil-23
co-authored by
Claude Opus 4.7
Steven Enamakel
Claude Opus 4.7
Steven Enamakel
parent
0ce1253fe1
commit
5fff5568d3
Generated
+1
-1
@@ -4420,7 +4420,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
|
||||
|
||||
[[package]]
|
||||
name = "openhuman"
|
||||
version = "0.52.26"
|
||||
version = "0.52.27"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"anyhow",
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
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<IngestResult> {
|
||||
// 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<Chunk> = Vec::new();
|
||||
let mut all_results: Vec<(ScoreResult, i64)> = Vec::new();
|
||||
for (chunk, result) in chunks.iter().zip(scores.into_iter()) {
|
||||
let ts_ms = chunk.metadata.timestamp.timestamp_millis();
|
||||
if result.kept {
|
||||
kept_chunks.push(chunk.clone());
|
||||
}
|
||||
all_results.push((result, ts_ms));
|
||||
}
|
||||
|
||||
let 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<usize> {
|
||||
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"
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<IngestResult> {
|
||||
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,
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
],
|
||||
|
||||
@@ -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<ExtractedEntities>;
|
||||
}
|
||||
|
||||
/// 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<ExtractedEntities> {
|
||||
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<Box<dyn EntityExtractor>>,
|
||||
}
|
||||
|
||||
impl CompositeExtractor {
|
||||
pub fn new(inner: Vec<Box<dyn EntityExtractor>>) -> 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<ExtractedEntities> {
|
||||
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<ExtractedEntities> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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<Regex> =
|
||||
Lazy::new(|| Regex::new(r"(?i)\b[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}\b").unwrap());
|
||||
|
||||
static RE_URL: Lazy<Regex> = Lazy::new(|| {
|
||||
// up-to trailing punctuation; avoids catastrophic backtracking
|
||||
Regex::new(r"https?://[^\s<>\]\[()]+[^\s<>\]\[()\.\,;:\!\?]").unwrap()
|
||||
});
|
||||
|
||||
static RE_HANDLE: Lazy<Regex> =
|
||||
Lazy::new(|| Regex::new(r"(?:^|[\s(])@([A-Za-z0-9_][A-Za-z0-9_.\-]{1,})").unwrap());
|
||||
|
||||
static RE_DISCRIM: Lazy<Regex> =
|
||||
Lazy::new(|| Regex::new(r"\b([A-Za-z0-9_.\-]{2,32})#\d{4}\b").unwrap());
|
||||
|
||||
static RE_HASHTAG: Lazy<Regex> =
|
||||
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<ExtractedEntity> = Vec::new();
|
||||
let mut topics: Vec<ExtractedTopic> = 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<EntityKind> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -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<Self, String> {
|
||||
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<ExtractedEntity>,
|
||||
pub topics: Vec<ExtractedTopic>,
|
||||
}
|
||||
|
||||
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::<BTreeSet<_>>()
|
||||
.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<String> =
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
pub extracted: ExtractedEntities,
|
||||
pub canonical_entities: Vec<CanonicalEntity>,
|
||||
}
|
||||
|
||||
/// 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<dyn EntityExtractor>,
|
||||
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<ScoreResult> {
|
||||
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::<Vec<_>>()
|
||||
.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<Vec<ScoreResult>> {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
@@ -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<CanonicalEntity> {
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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:<snake_case>` 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<DataSource> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<String> = 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);
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
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<Option<ScoreRow>> {
|
||||
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<usize> {
|
||||
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<usize> {
|
||||
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<usize> {
|
||||
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<usize> {
|
||||
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<String>,
|
||||
}
|
||||
|
||||
/// Find all nodes indexed against `entity_id`, newest first.
|
||||
pub fn lookup_entity(
|
||||
config: &Config,
|
||||
entity_id: &str,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<EntityHit>> {
|
||||
// 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::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(rows)
|
||||
})
|
||||
}
|
||||
|
||||
/// Count rows in the entity index (for tests / diagnostics).
|
||||
pub fn count_entity_index(config: &Config) -> Result<u64> {
|
||||
with_connection(config, |conn| {
|
||||
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<u64> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<usize> {
|
||||
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<usize> {
|
||||
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<Option<Chunk>> {
|
||||
with_connection(config, |conn| {
|
||||
@@ -238,7 +328,14 @@ fn ms_to_utc(ms: i64) -> rusqlite::Result<DateTime<Utc>> {
|
||||
})
|
||||
}
|
||||
|
||||
fn with_connection<T>(config: &Config, f: impl FnOnce(&Connection) -> Result<T>) -> Result<T> {
|
||||
/// 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<T>(
|
||||
config: &Config,
|
||||
f: impl FnOnce(&Connection) -> Result<T>,
|
||||
) -> Result<T> {
|
||||
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<T>(config: &Config, f: impl FnOnce(&Connection) -> Result<T>)
|
||||
.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<usize>) -> 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<u8> = 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<Option<Vec<f32>>> {
|
||||
with_connection(config, |conn| {
|
||||
let blob: Option<Option<Vec<u8>>> = conn
|
||||
.query_row(
|
||||
"SELECT embedding FROM mem_tree_chunks WHERE id = ?1",
|
||||
rusqlite::params![chunk_id],
|
||||
|r| r.get::<_, Option<Vec<u8>>>(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<f32> = 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();
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user