From 3b7c4cdc387a0f5bda7c8b623e4eea4c33f5da2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EB=AF=BC=EC=9E=AC?= Date: Mon, 18 May 2026 10:09:37 +0900 Subject: [PATCH] feat(memory): add per-model embedding tables (#2010) Co-authored-by: honor2030 <19909783+honor2030@users.noreply.github.com> --- src/openhuman/memory/store/unified/events.rs | 85 +++++++++++++- .../memory/store/unified/events_tests.rs | 49 ++++++++ .../memory/store/unified/segments.rs | 85 ++++++++++++++ .../memory/store/unified/segments_tests.rs | 38 +++++++ src/openhuman/memory/tree/store.rs | 106 ++++++++++++++++++ src/openhuman/memory/tree/store_tests.rs | 33 ++++++ .../memory/tree/tree_source/store.rs | 89 +++++++++++++++ .../memory/tree/tree_source/store_tests.rs | 44 ++++++++ 8 files changed, 528 insertions(+), 1 deletion(-) diff --git a/src/openhuman/memory/store/unified/events.rs b/src/openhuman/memory/store/unified/events.rs index 46c294c68..6bf993e45 100644 --- a/src/openhuman/memory/store/unified/events.rs +++ b/src/openhuman/memory/store/unified/events.rs @@ -7,7 +7,7 @@ //! - Tier B (local LLM): runs on segment close if local AI is enabled. use parking_lot::Mutex; -use rusqlite::{params, Connection}; +use rusqlite::{params, Connection, OptionalExtension}; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -62,6 +62,21 @@ CREATE TRIGGER IF NOT EXISTS event_au AFTER UPDATE ON event_log BEGIN INSERT INTO event_fts(rowid, content, subject, event_type) VALUES (new.rowid, new.content, new.subject, new.event_type); END; + +-- Per-(event, embedding model) vectors (#1574). The legacy event_log.embedding +-- column stays available during the dual-write migration; this table records +-- vector-space provenance for safe provider/model switches. +CREATE TABLE IF NOT EXISTS event_embeddings ( + event_id TEXT NOT NULL REFERENCES event_log(event_id) ON DELETE CASCADE, + model_signature TEXT NOT NULL, + vector BLOB NOT NULL, + dim INTEGER NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (event_id, model_signature) +); + +CREATE INDEX IF NOT EXISTS idx_event_embeddings_model + ON event_embeddings(model_signature); "#; /// Event types extracted from conversations. @@ -153,6 +168,54 @@ pub fn event_insert(conn: &Arc>, event: &EventRecord) -> anyho Ok(()) } +/// Store an event embedding for a specific provider/model/dimension signature. +/// +/// This writes only the per-model table introduced for #1574. The legacy +/// `event_log.embedding` column remains available for dual-read fallback. +pub fn event_embedding_upsert( + conn: &Arc>, + event_id: &str, + model_signature: &str, + embedding: &[f32], + created_at: f64, +) -> anyhow::Result<()> { + let bytes = vec_to_bytes(embedding); + let dim = i64::try_from(embedding.len())?; + let conn = conn.lock(); + conn.execute( + "INSERT INTO event_embeddings (event_id, model_signature, vector, dim, created_at) + VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(event_id, model_signature) DO UPDATE SET + vector = excluded.vector, + dim = excluded.dim, + created_at = excluded.created_at", + params![event_id, model_signature, bytes, dim, created_at], + )?; + Ok(()) +} + +/// Fetch an event embedding for exactly one provider/model/dimension signature. +pub fn event_embedding_get( + conn: &Arc>, + event_id: &str, + model_signature: &str, +) -> anyhow::Result>> { + let conn = conn.lock(); + let row: Option<(Vec, i64)> = conn + .query_row( + "SELECT vector, dim + FROM event_embeddings + WHERE event_id = ?1 AND model_signature = ?2", + params![event_id, model_signature], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .optional()?; + match row { + None => Ok(None), + Some((bytes, dim)) => decode_embedding_row(&bytes, dim), + } +} + /// Search events via FTS5, scoped to a namespace. pub fn event_search_fts( conn: &Arc>, @@ -398,6 +461,26 @@ fn bytes_to_vec(bytes: &[u8]) -> Vec { .collect() } +fn decode_embedding_row(bytes: &[u8], dim: i64) -> anyhow::Result>> { + if dim < 0 { + anyhow::bail!("event embedding has negative dimension {dim}"); + } + if !bytes.len().is_multiple_of(4) { + anyhow::bail!( + "event embedding blob length {} not a multiple of 4", + bytes.len() + ); + } + let vector = bytes_to_vec(bytes); + if vector.len() != dim as usize { + anyhow::bail!( + "event embedding dimension mismatch: dim column says {dim}, blob contains {} floats", + vector.len() + ); + } + Ok(Some(vector)) +} + #[cfg(test)] #[path = "events_tests.rs"] mod tests; diff --git a/src/openhuman/memory/store/unified/events_tests.rs b/src/openhuman/memory/store/unified/events_tests.rs index ce700b3b0..72156871d 100644 --- a/src/openhuman/memory/store/unified/events_tests.rs +++ b/src/openhuman/memory/store/unified/events_tests.rs @@ -210,3 +210,52 @@ fn event_fts_matches_subject_field() { assert_eq!(by_subject.len(), 1, "FTS should match on subject field"); assert_eq!(by_subject[0].event_id, "evt-subj"); } + +#[test] +fn event_embeddings_are_scoped_by_model_signature() { + let conn = setup_db(); + let event = EventRecord { + event_id: "evt-embed".into(), + segment_id: "seg-1".into(), + session_id: "s1".into(), + namespace: "global".into(), + event_type: EventType::Fact, + content: "The user prefers Korean summaries".into(), + subject: Some("language preference".into()), + timestamp_ref: None, + confidence: 0.9, + embedding: None, + source_turn_ids: None, + created_at: 1000.0, + }; + event_insert(&conn, &event).unwrap(); + + event_embedding_upsert( + &conn, + "evt-embed", + "openai/text-embedding-3-small@1536", + &[0.1, 0.2], + 1001.0, + ) + .unwrap(); + event_embedding_upsert( + &conn, + "evt-embed", + "local/bge-small@384", + &[0.3, 0.4, 0.5], + 1002.0, + ) + .unwrap(); + + assert_eq!( + event_embedding_get(&conn, "evt-embed", "openai/text-embedding-3-small@1536").unwrap(), + Some(vec![0.1, 0.2]) + ); + assert_eq!( + event_embedding_get(&conn, "evt-embed", "local/bge-small@384").unwrap(), + Some(vec![0.3, 0.4, 0.5]) + ); + assert!(event_embedding_get(&conn, "evt-embed", "missing/model@1") + .unwrap() + .is_none()); +} diff --git a/src/openhuman/memory/store/unified/segments.rs b/src/openhuman/memory/store/unified/segments.rs index 4ff9fc332..46e63ce10 100644 --- a/src/openhuman/memory/store/unified/segments.rs +++ b/src/openhuman/memory/store/unified/segments.rs @@ -37,6 +37,22 @@ CREATE INDEX IF NOT EXISTS idx_segments_namespace CREATE INDEX IF NOT EXISTS idx_segments_status ON conversation_segments(status, session_id); + +-- Per-model segment embeddings for #1574. The legacy +-- `conversation_segments.embedding` column stays in place during staged +-- migration; this table lets provider/model switches become query-time +-- filters instead of destructive rewrites. +CREATE TABLE IF NOT EXISTS segment_embeddings ( + segment_id TEXT NOT NULL REFERENCES conversation_segments(segment_id) ON DELETE CASCADE, + model_signature TEXT NOT NULL, + vector BLOB NOT NULL, + dim INTEGER NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (segment_id, model_signature) +); + +CREATE INDEX IF NOT EXISTS idx_segment_embeddings_model + ON segment_embeddings(model_signature); "#; /// Segment status lifecycle: open → closed → summarised. @@ -258,6 +274,55 @@ pub fn segment_set_embedding( Ok(()) } +/// Store a segment embedding for a specific provider/model/dimension signature. +/// +/// This writes only the per-model table introduced for #1574. The legacy +/// `conversation_segments.embedding` column remains available for dual-read +/// fallback while query paths migrate. +pub fn segment_embedding_upsert( + conn: &Arc>, + segment_id: &str, + model_signature: &str, + embedding: &[f32], + created_at: f64, +) -> anyhow::Result<()> { + let bytes = vec_to_bytes(embedding); + let dim = i64::try_from(embedding.len())?; + let conn = conn.lock(); + conn.execute( + "INSERT INTO segment_embeddings (segment_id, model_signature, vector, dim, created_at) + VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(segment_id, model_signature) DO UPDATE SET + vector = excluded.vector, + dim = excluded.dim, + created_at = excluded.created_at", + params![segment_id, model_signature, bytes, dim, created_at], + )?; + Ok(()) +} + +/// Fetch a segment embedding for exactly one provider/model/dimension signature. +pub fn segment_embedding_get( + conn: &Arc>, + segment_id: &str, + model_signature: &str, +) -> anyhow::Result>> { + let conn = conn.lock(); + let row: Option<(Vec, i64)> = conn + .query_row( + "SELECT vector, dim + FROM segment_embeddings + WHERE segment_id = ?1 AND model_signature = ?2", + params![segment_id, model_signature], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .optional()?; + match row { + None => Ok(None), + Some((bytes, dim)) => decode_embedding_row(&bytes, dim), + } +} + /// Store topic keywords for the segment. pub fn segment_set_keywords( conn: &Arc>, @@ -502,6 +567,26 @@ fn bytes_to_vec(bytes: &[u8]) -> Vec { .collect() } +fn decode_embedding_row(bytes: &[u8], dim: i64) -> anyhow::Result>> { + if dim < 0 { + anyhow::bail!("segment embedding has negative dimension {dim}"); + } + if !bytes.len().is_multiple_of(4) { + anyhow::bail!( + "segment embedding blob length {} not a multiple of 4", + bytes.len() + ); + } + let floats = bytes_to_vec(bytes); + if floats.len() != dim as usize { + anyhow::bail!( + "segment embedding dimension mismatch: dim column says {dim}, blob contains {} floats", + floats.len() + ); + } + Ok(Some(floats)) +} + #[cfg(test)] #[path = "segments_tests.rs"] mod tests; diff --git a/src/openhuman/memory/store/unified/segments_tests.rs b/src/openhuman/memory/store/unified/segments_tests.rs index abd356811..468439fc4 100644 --- a/src/openhuman/memory/store/unified/segments_tests.rs +++ b/src/openhuman/memory/store/unified/segments_tests.rs @@ -21,6 +21,44 @@ fn create_and_get_segment() { assert_eq!(seg.status, SegmentStatus::Open); } +#[test] +fn segment_embeddings_are_scoped_by_model_signature() { + let conn = setup_db(); + segment_create(&conn, "seg-embed", "s1", "global", 1, 1000.0, 1000.0).unwrap(); + + segment_embedding_upsert( + &conn, + "seg-embed", + "openai/text-embedding-3-small@1536", + &[0.1, 0.2], + 1001.0, + ) + .unwrap(); + segment_embedding_upsert( + &conn, + "seg-embed", + "local/bge-small@384", + &[0.3, 0.4, 0.5], + 1002.0, + ) + .unwrap(); + + assert_eq!( + segment_embedding_get(&conn, "seg-embed", "openai/text-embedding-3-small@1536").unwrap(), + Some(vec![0.1, 0.2]) + ); + assert_eq!( + segment_embedding_get(&conn, "seg-embed", "local/bge-small@384").unwrap(), + Some(vec![0.3, 0.4, 0.5]) + ); + assert!(segment_embedding_get(&conn, "seg-embed", "missing/model@1") + .unwrap() + .is_none()); + + let legacy_segment = segment_get(&conn, "seg-embed").unwrap().unwrap(); + assert!(legacy_segment.embedding.is_none()); +} + #[test] fn append_and_close_segment() { let conn = setup_db(); diff --git a/src/openhuman/memory/tree/store.rs b/src/openhuman/memory/tree/store.rs index a1dced01d..f0c8a5067 100644 --- a/src/openhuman/memory/tree/store.rs +++ b/src/openhuman/memory/tree/store.rs @@ -71,6 +71,21 @@ CREATE INDEX IF NOT EXISTS idx_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); +-- Per-(chunk, embedding model) vectors (#1574). The legacy +-- mem_tree_chunks.embedding column remains in place during the dual-write +-- migration; this table lets multiple vector spaces coexist safely. +CREATE TABLE IF NOT EXISTS mem_tree_chunk_embeddings ( + chunk_id TEXT NOT NULL REFERENCES mem_tree_chunks(id) ON DELETE CASCADE, + model_signature TEXT NOT NULL, + vector BLOB NOT NULL, + dim INTEGER NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (chunk_id, model_signature) +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_chunk_embeddings_model + ON mem_tree_chunk_embeddings(model_signature); + -- Phase 2 (#708): per-chunk score rationale for admission debugging. CREATE TABLE IF NOT EXISTS mem_tree_score ( chunk_id TEXT PRIMARY KEY, @@ -163,6 +178,21 @@ CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_sealed_at CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_deleted ON mem_tree_summaries(deleted); +-- Per-(summary, embedding model) vectors (#1574). Kept separate from the +-- legacy mem_tree_summaries.embedding column so provider/model switches can +-- be query-time filters instead of destructive rewrites. +CREATE TABLE IF NOT EXISTS mem_tree_summary_embeddings ( + summary_id TEXT NOT NULL REFERENCES mem_tree_summaries(id) ON DELETE CASCADE, + model_signature TEXT NOT NULL, + vector BLOB NOT NULL, + dim INTEGER NOT NULL, + created_at REAL NOT NULL, + PRIMARY KEY (summary_id, model_signature) +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_summary_embeddings_model + ON mem_tree_summary_embeddings(model_signature); + -- `mem_tree_buffers` holds the unsealed frontier per (tree, level). One row -- per active level per tree; deleted when the buffer seals (clears) in the -- same transaction as the new summary node row. @@ -933,6 +963,58 @@ pub fn set_chunk_embedding(config: &Config, chunk_id: &str, embedding: &[f32]) - }) } +/// Store a chunk embedding for a specific provider/model/dimension signature. +/// +/// This is the Stage-1 per-model table write path for #1574. The legacy +/// `mem_tree_chunks.embedding` column is intentionally left untouched by this +/// helper so callers can dual-write while query paths migrate. +pub fn set_chunk_embedding_for_signature( + config: &Config, + chunk_id: &str, + model_signature: &str, + embedding: &[f32], +) -> Result<()> { + let bytes = embedding_to_blob(embedding); + let dim = i64::try_from(embedding.len()).context("embedding dimension does not fit i64")?; + let created_at = Utc::now().timestamp_millis() as f64 / 1000.0; + with_connection(config, |conn| { + conn.execute( + "INSERT INTO mem_tree_chunk_embeddings + (chunk_id, model_signature, vector, dim, created_at) + VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(chunk_id, model_signature) DO UPDATE SET + vector = excluded.vector, + dim = excluded.dim, + created_at = excluded.created_at", + rusqlite::params![chunk_id, model_signature, bytes, dim, created_at], + )?; + Ok(()) + }) +} + +/// Fetch a chunk embedding for exactly one provider/model/dimension signature. +pub fn get_chunk_embedding_for_signature( + config: &Config, + chunk_id: &str, + model_signature: &str, +) -> Result>> { + with_connection(config, |conn| { + let row: Option<(Vec, i64)> = conn + .query_row( + "SELECT vector, dim + FROM mem_tree_chunk_embeddings + WHERE chunk_id = ?1 AND model_signature = ?2", + rusqlite::params![chunk_id, model_signature], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .optional()?; + match row { + None => Ok(None), + Some((bytes, dim)) => embedding_from_blob(&bytes, dim, "chunk embedding"), + } + }) +} + /// 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. @@ -961,6 +1043,30 @@ pub fn get_chunk_embedding(config: &Config, chunk_id: &str) -> Result Vec { + embedding.iter().flat_map(|f| f.to_le_bytes()).collect() +} + +fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result>> { + if dim < 0 { + anyhow::bail!("{label} has negative dimension {dim}"); + } + if !bytes.len().is_multiple_of(4) { + anyhow::bail!("{label} blob length {} not a multiple of 4", bytes.len()); + } + let floats: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + if floats.len() != dim as usize { + anyhow::bail!( + "{label} dimension mismatch: dim column says {dim}, blob contains {} floats", + floats.len() + ); + } + Ok(Some(floats)) +} + #[cfg(test)] #[path = "store_tests.rs"] mod tests; diff --git a/src/openhuman/memory/tree/store_tests.rs b/src/openhuman/memory/tree/store_tests.rs index c4fbe5051..cb70b7dda 100644 --- a/src/openhuman/memory/tree/store_tests.rs +++ b/src/openhuman/memory/tree/store_tests.rs @@ -71,6 +71,39 @@ fn reingest_preserves_existing_embedding() { assert_eq!(got.token_count, 99); } +#[test] +fn chunk_embeddings_are_scoped_by_model_signature() { + let (_tmp, cfg) = test_config(); + let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + + set_chunk_embedding_for_signature( + &cfg, + &c.id, + "openai/text-embedding-3-small@1536", + &[0.1, 0.2], + ) + .unwrap(); + set_chunk_embedding_for_signature(&cfg, &c.id, "local/bge-small@384", &[0.3, 0.4, 0.5]) + .unwrap(); + + assert_eq!( + get_chunk_embedding_for_signature(&cfg, &c.id, "openai/text-embedding-3-small@1536") + .unwrap(), + Some(vec![0.1, 0.2]) + ); + assert_eq!( + get_chunk_embedding_for_signature(&cfg, &c.id, "local/bge-small@384").unwrap(), + Some(vec![0.3, 0.4, 0.5]) + ); + assert!( + get_chunk_embedding_for_signature(&cfg, &c.id, "missing/model@1") + .unwrap() + .is_none() + ); + assert!(get_chunk_embedding(&cfg, &c.id).unwrap().is_none()); +} + #[test] fn list_filters_by_source_kind() { let (_tmp, cfg) = test_config(); diff --git a/src/openhuman/memory/tree/tree_source/store.rs b/src/openhuman/memory/tree/tree_source/store.rs index f31b43ee2..2ecd8721f 100644 --- a/src/openhuman/memory/tree/tree_source/store.rs +++ b/src/openhuman/memory/tree/tree_source/store.rs @@ -290,6 +290,68 @@ pub fn get_summary_embedding(config: &Config, summary_id: &str) -> Result