mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
feat(memory): activate chat memory pipeline — Archivist wiring, LLM recaps, cross-thread STM recall, segment-granular tree, unified compaction (#2175)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
7741c5815c
commit
341ce42ff1
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,7 @@
|
||||
use super::*;
|
||||
use crate::openhuman::agent::hooks::{ToolCallRecord, TurnContext};
|
||||
use crate::openhuman::memory::store::{events as ev, fts5, segments as seg};
|
||||
use crate::openhuman::memory::tree::chat::ChatPrompt;
|
||||
|
||||
fn setup_conn() -> Arc<Mutex<Connection>> {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
@@ -233,3 +234,611 @@ async fn archivist_extracts_preference_event_on_boundary() {
|
||||
events.iter().map(|e| &e.content).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
// ── Phase 0: episodic_capture_enabled independent of learning.enabled ────────
|
||||
|
||||
/// When `learning.enabled = false` but `episodic_capture_enabled = true`,
|
||||
/// the ArchivistHook (constructed directly, as builder.rs would produce)
|
||||
/// must still write 2 episodic_log rows (user + assistant) and create/advance
|
||||
/// a segment. This verifies the core contract: episodic capture runs
|
||||
/// regardless of the learning inference stack toggle.
|
||||
#[tokio::test]
|
||||
async fn phase0_episodic_rows_and_segment_without_learning_enabled() {
|
||||
let conn = setup_conn();
|
||||
// Simulate what builder.rs does when learning.enabled=false but
|
||||
// episodic_capture_enabled=true: construct the hook directly with
|
||||
// the SQLite conn, enabled=true. No config attached (no LLM recap
|
||||
// or tree ingest — those are gated by learning.enabled / chat_to_tree_enabled).
|
||||
let hook = ArchivistHook::new(conn.clone(), true);
|
||||
|
||||
let session = "phase0-test-session";
|
||||
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Hello, what is Rust?".into(),
|
||||
assistant_response: "Rust is a systems language.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify 2 episodic rows were written.
|
||||
let entries = fts5::episodic_session_entries(&conn, session).unwrap();
|
||||
assert_eq!(
|
||||
entries.len(),
|
||||
2,
|
||||
"Expected 2 episodic rows (user + assistant), got {}",
|
||||
entries.len()
|
||||
);
|
||||
assert_eq!(entries[0].role, "user");
|
||||
assert_eq!(entries[1].role, "assistant");
|
||||
|
||||
// Verify a segment was created.
|
||||
let open_seg = seg::open_segment_for_session(&conn, session)
|
||||
.unwrap()
|
||||
.expect("Expected an open segment after first turn");
|
||||
assert_eq!(open_seg.turn_count, 1);
|
||||
|
||||
// Add a second turn to verify segment advances.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Tell me more about ownership.".into(),
|
||||
assistant_response: "Ownership prevents data races.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 2,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let entries2 = fts5::episodic_session_entries(&conn, session).unwrap();
|
||||
assert_eq!(
|
||||
entries2.len(),
|
||||
4,
|
||||
"Expected 4 episodic rows after 2 turns, got {}",
|
||||
entries2.len()
|
||||
);
|
||||
let open_seg2 = seg::open_segment_for_session(&conn, session)
|
||||
.unwrap()
|
||||
.expect("Expected an open segment after 2 turns");
|
||||
assert_eq!(
|
||||
open_seg2.turn_count, 2,
|
||||
"Segment should have 2 turns, got {}",
|
||||
open_seg2.turn_count
|
||||
);
|
||||
}
|
||||
|
||||
// ── Phase 1: LLM recap + finalize-time embedding ─────────────────────────────
|
||||
|
||||
/// Stub ChatProvider that returns a fixed recap string without hitting
|
||||
/// any real LLM, so the test is hermetic.
|
||||
struct StubChatProvider;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl crate::openhuman::memory::tree::chat::ChatProvider for StubChatProvider {
|
||||
fn name(&self) -> &str {
|
||||
"stub:test"
|
||||
}
|
||||
|
||||
async fn chat_for_json(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
|
||||
Ok("stub recap: discussed Rust ownership model".to_string())
|
||||
}
|
||||
|
||||
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
|
||||
Ok("stub recap: discussed Rust ownership model".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Stub Embedder that returns a fixed unit vector without hitting Ollama.
|
||||
struct StubEmbedder;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl crate::openhuman::memory::tree::score::embed::Embedder for StubEmbedder {
|
||||
fn name(&self) -> &'static str {
|
||||
"stub-embedder-v1"
|
||||
}
|
||||
|
||||
async fn embed(&self, _text: &str) -> anyhow::Result<Vec<f32>> {
|
||||
// Return a simple 4-dim unit vector.
|
||||
Ok(vec![0.5_f32, 0.5, 0.5, 0.5])
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an ArchivistHook with stub provider + embedder injected directly.
|
||||
/// Uses the test-only `new_with_stubs` constructor to bypass `with_config`.
|
||||
fn hook_with_stubs(conn: Arc<Mutex<Connection>>) -> ArchivistHook {
|
||||
ArchivistHook::new_with_stubs(conn, Arc::new(StubChatProvider), Arc::new(StubEmbedder))
|
||||
}
|
||||
|
||||
/// When a segment closes, the LLM chat provider recap is used (verified by
|
||||
/// a non-empty segment summary) and an embedding row is written to
|
||||
/// `segment_embeddings`.
|
||||
#[tokio::test]
|
||||
async fn phase1_llm_recap_and_embedding_on_segment_close() {
|
||||
let conn = setup_conn();
|
||||
let hook = hook_with_stubs(conn.clone());
|
||||
|
||||
let session = "phase1-recap-test";
|
||||
|
||||
// Turn 1 — opens first segment.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Tell me about Rust ownership".into(),
|
||||
assistant_response: "Rust's ownership model prevents data races.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Turn 2 — continues same segment.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "What about the borrow checker?".into(),
|
||||
assistant_response: "The borrow checker enforces ownership rules at compile time.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 2,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Turn 3 — topic change triggers a boundary → closes first segment → recap + embed fire.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Completely different topic: what is async/await in Python?".into(),
|
||||
assistant_response: "Python asyncio enables concurrent programming.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 3,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify segments exist.
|
||||
let segments = seg::segments_by_namespace(&conn, "global", 10).unwrap();
|
||||
assert!(
|
||||
segments.len() >= 2,
|
||||
"Expected at least 2 segments (closed + open), got {}",
|
||||
segments.len()
|
||||
);
|
||||
|
||||
// Find the closed segment (has a summary).
|
||||
let closed = segments
|
||||
.iter()
|
||||
.find(|s| s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false));
|
||||
assert!(
|
||||
closed.is_some(),
|
||||
"Expected at least one closed segment with a non-empty summary"
|
||||
);
|
||||
|
||||
let closed_seg = closed.unwrap();
|
||||
let summary = closed_seg.summary.as_ref().unwrap();
|
||||
// The stub provider returns a fixed string — verify it was persisted.
|
||||
assert!(
|
||||
summary.contains("stub recap"),
|
||||
"Expected summary to contain 'stub recap', got: {:?}",
|
||||
summary
|
||||
);
|
||||
|
||||
// Verify an embedding row was written for the closed segment.
|
||||
let embedding =
|
||||
seg::segment_embedding_get(&conn, &closed_seg.segment_id, "stub-embedder-v1").unwrap();
|
||||
assert!(
|
||||
embedding.is_some(),
|
||||
"Expected an embedding row for segment={} model=stub-embedder-v1",
|
||||
closed_seg.segment_id
|
||||
);
|
||||
let vec = embedding.unwrap();
|
||||
assert_eq!(vec.len(), 4, "Expected 4-dim vector from stub embedder");
|
||||
for v in &vec {
|
||||
assert!(
|
||||
(*v - 0.5_f32).abs() < 1e-4,
|
||||
"Expected vector components ≈ 0.5, got {v}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// `flush_open_segment` must force-close the trailing open segment and
|
||||
/// trigger recap + embedding even without a boundary-triggering turn.
|
||||
#[tokio::test]
|
||||
async fn phase1_flush_open_segment_finalizes_trailing_segment() {
|
||||
let conn = setup_conn();
|
||||
let hook = hook_with_stubs(conn.clone());
|
||||
|
||||
let session = "phase1-flush-test";
|
||||
|
||||
// Write 2 turns — stays in one open segment (no topic boundary fires).
|
||||
for i in 1..=2 {
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: format!("Question about Rust turn {i}"),
|
||||
assistant_response: format!("Answer about Rust turn {i}"),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: i,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Confirm the segment is still open (no boundary fired).
|
||||
let open_seg_before = seg::open_segment_for_session(&conn, session).unwrap();
|
||||
assert!(
|
||||
open_seg_before.is_some(),
|
||||
"Expected an open segment before flush"
|
||||
);
|
||||
|
||||
// Flush — should force-close, recap, and embed.
|
||||
hook.flush_open_segment(session).await;
|
||||
|
||||
// Segment should now be closed (no open segment for this session).
|
||||
let open_seg_after = seg::open_segment_for_session(&conn, session).unwrap();
|
||||
assert!(
|
||||
open_seg_after.is_none(),
|
||||
"Expected no open segment after flush_open_segment"
|
||||
);
|
||||
|
||||
// The formerly-open segment should now have a summary.
|
||||
let segments = seg::segments_by_namespace(&conn, "global", 10).unwrap();
|
||||
let flushed = segments.iter().find(|s| {
|
||||
s.session_id == session && s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false)
|
||||
});
|
||||
assert!(
|
||||
flushed.is_some(),
|
||||
"Expected flushed segment to have a non-empty summary"
|
||||
);
|
||||
|
||||
let seg_id = &flushed.unwrap().segment_id;
|
||||
let embedding = seg::segment_embedding_get(&conn, seg_id, "stub-embedder-v1").unwrap();
|
||||
assert!(
|
||||
embedding.is_some(),
|
||||
"Expected embedding row for flushed segment={seg_id}"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Phase 2: segment-granularity tree ingest ─────────────────────────────────
|
||||
//
|
||||
// The following tests verify:
|
||||
// a) No per-turn tree write fires from on_turn_complete (no double-write).
|
||||
// b) Exactly ONE tree ingest fires when a segment closes (not N per turn).
|
||||
// c) The ingested batch contains all the segment's raw prose turns.
|
||||
// d) The `source_id` is the constant "conversations:agent".
|
||||
// e) Each leaf message carries session/segment/episodic-span provenance.
|
||||
// f) The ingested content is raw prose, NOT the LLM recap.
|
||||
// g) flush_open_segment also triggers tree ingest.
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Build a Config that points at a temp workspace, suitable for tree-ingest tests.
|
||||
/// The memory_tree DB and content dir are created under `tmp.path()`.
|
||||
fn test_config_with_tree() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
// Disable embedding so ingest doesn't fail trying to contact Ollama.
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
// Ensure the tree ingest gate is on.
|
||||
cfg.learning.chat_to_tree_enabled = true;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
/// Build a hook that has both stub providers AND a real-enough Config wired in,
|
||||
/// so the Phase 2 tree ingest path is exercised hermetically.
|
||||
fn hook_with_stubs_and_tree_config(conn: Arc<Mutex<Connection>>, cfg: Config) -> ArchivistHook {
|
||||
ArchivistHook::new_with_stubs_and_config(
|
||||
conn,
|
||||
Arc::new(StubChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
cfg,
|
||||
)
|
||||
}
|
||||
|
||||
/// After a single turn (no segment boundary), the tree must have ZERO chunks —
|
||||
/// the per-turn pipe_turn_to_tree path no longer exists.
|
||||
#[tokio::test]
|
||||
async fn phase2_no_per_turn_tree_write() {
|
||||
let conn = setup_conn();
|
||||
let (_tmp, cfg) = test_config_with_tree();
|
||||
let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone());
|
||||
|
||||
let session = "phase2-no-per-turn";
|
||||
|
||||
// Single turn — no segment close fires, so no tree ingest should happen.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "What is Rust?".into(),
|
||||
assistant_response: "Rust is a systems programming language.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Segment is still open (no boundary fired) — tree must have 0 chunks.
|
||||
let open_seg = seg::open_segment_for_session(&conn, session).unwrap();
|
||||
assert!(
|
||||
open_seg.is_some(),
|
||||
"Expected an open segment (no boundary should have fired)"
|
||||
);
|
||||
|
||||
let chunk_count = count_chunks(&cfg).unwrap();
|
||||
assert_eq!(
|
||||
chunk_count, 0,
|
||||
"Expected 0 tree chunks after a single turn (no segment close): \
|
||||
per-turn tree write must not exist (Phase 2)"
|
||||
);
|
||||
}
|
||||
|
||||
/// When a segment closes (boundary triggered), exactly ONE tree ingest fires
|
||||
/// for that segment containing all its turns — not one ingest per turn.
|
||||
#[tokio::test]
|
||||
async fn phase2_exactly_one_tree_ingest_per_segment_close() {
|
||||
let conn = setup_conn();
|
||||
let (_tmp, cfg) = test_config_with_tree();
|
||||
let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone());
|
||||
|
||||
let session = "phase2-one-ingest";
|
||||
|
||||
// Turn 1 — opens first segment.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Tell me about Rust ownership".into(),
|
||||
assistant_response: "Rust ownership prevents memory bugs.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Turn 2 — stays in same segment.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "What about the borrow checker?".into(),
|
||||
assistant_response: "The borrow checker enforces ownership at compile time.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 2,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// No tree write yet — segment still open.
|
||||
let pre_close_chunks = count_chunks(&cfg).unwrap();
|
||||
assert_eq!(
|
||||
pre_close_chunks, 0,
|
||||
"Expected 0 tree chunks before any segment close; got {pre_close_chunks}"
|
||||
);
|
||||
|
||||
// Turn 3 — topic change triggers boundary → closes first segment → tree ingest fires.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Switching to a completely different topic: tell me about Python asyncio."
|
||||
.into(),
|
||||
assistant_response: "Python asyncio enables concurrent coroutines.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 3,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Segment closed → exactly one ingest for the closed segment (containing turns 1+2).
|
||||
// The ingest packs the messages into one or more chunks (greedy packing),
|
||||
// but chunks_written >= 1 confirms ingest happened.
|
||||
let post_close_chunks = count_chunks(&cfg).unwrap();
|
||||
assert!(
|
||||
post_close_chunks >= 1,
|
||||
"Expected ≥ 1 tree chunk after segment close; got {post_close_chunks}"
|
||||
);
|
||||
|
||||
// List the chunks and check they come from the constant source_id.
|
||||
let chunks = list_chunks(
|
||||
&cfg,
|
||||
&ListChunksQuery {
|
||||
source_id: Some("conversations:agent".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
!chunks.is_empty(),
|
||||
"Expected chunks under source_id='conversations:agent'"
|
||||
);
|
||||
}
|
||||
|
||||
/// The ingested leaf messages must carry the episodic-provenance `source_ref`
|
||||
/// in the expected format:
|
||||
/// `agent://session/{session_id}/segment/{segment_id}#ep{start}-{end}`.
|
||||
///
|
||||
/// Also verifies that `source_id` is the constant `"conversations:agent"`.
|
||||
#[tokio::test]
|
||||
async fn phase2_provenance_stamped_on_leaf_and_source_id_is_constant() {
|
||||
let conn = setup_conn();
|
||||
let (_tmp, cfg) = test_config_with_tree();
|
||||
let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone());
|
||||
|
||||
let session = "phase2-provenance";
|
||||
|
||||
// Two turns in the first segment.
|
||||
for i in 1..=2 {
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: format!("Ownership question {i}"),
|
||||
assistant_response: format!("Ownership answer {i}"),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: i,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Force a segment close via flush_open_segment.
|
||||
hook.flush_open_segment(session).await;
|
||||
|
||||
// Retrieve the closed segment to extract its ID.
|
||||
let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap();
|
||||
let closed = all_segs
|
||||
.iter()
|
||||
.find(|s| {
|
||||
s.session_id == session
|
||||
&& s.status != crate::openhuman::memory::store::segments::SegmentStatus::Open
|
||||
})
|
||||
.expect("Expected a closed segment after flush");
|
||||
|
||||
let segment_id = &closed.segment_id;
|
||||
let start_ep = closed.start_episodic_id;
|
||||
let end_ep = closed.end_episodic_id.unwrap_or(start_ep);
|
||||
|
||||
// Chunks should be present.
|
||||
let chunks = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
|
||||
assert!(
|
||||
!chunks.is_empty(),
|
||||
"Expected tree chunks after flush_open_segment"
|
||||
);
|
||||
|
||||
// source_id must be the constant — never per-session or per-segment.
|
||||
for chunk in &chunks {
|
||||
assert_eq!(
|
||||
chunk.metadata.source_id, "conversations:agent",
|
||||
"source_id must be the constant 'conversations:agent', got: {}",
|
||||
chunk.metadata.source_id
|
||||
);
|
||||
}
|
||||
|
||||
// The source_ref on at least one chunk must contain the provenance pattern.
|
||||
let expected_provenance =
|
||||
format!("agent://session/{session}/segment/{segment_id}#ep{start_ep}-{end_ep}");
|
||||
let has_provenance = chunks.iter().any(|chunk| {
|
||||
chunk
|
||||
.metadata
|
||||
.source_ref
|
||||
.as_ref()
|
||||
.map(|r| {
|
||||
r.value
|
||||
.contains(&format!("agent://session/{session}/segment/{segment_id}"))
|
||||
})
|
||||
.unwrap_or(false)
|
||||
});
|
||||
assert!(
|
||||
has_provenance,
|
||||
"Expected at least one chunk with source_ref containing provenance pattern \
|
||||
'{expected_provenance}'; found: {:?}",
|
||||
chunks
|
||||
.iter()
|
||||
.map(|c| c.metadata.source_ref.as_ref().map(|r| r.value.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
/// The ingested content must be the raw prose turns (user + assistant text),
|
||||
/// NOT equal to the LLM recap text. The recap lives only in the STM segment
|
||||
/// layer; the tree must ingest raw evidence so it can build its own summaries.
|
||||
#[tokio::test]
|
||||
async fn phase2_ingested_content_is_raw_prose_not_recap() {
|
||||
let conn = setup_conn();
|
||||
let (_tmp, cfg) = test_config_with_tree();
|
||||
let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone());
|
||||
|
||||
let session = "phase2-raw-prose";
|
||||
|
||||
// The stub recap always returns "stub recap: discussed Rust ownership model".
|
||||
// The raw user messages contain very different text.
|
||||
let user_msg = "My specific question about lifetimes in Rust code";
|
||||
let asst_msg = "Lifetimes annotate how long references are valid in memory";
|
||||
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: user_msg.into(),
|
||||
assistant_response: asst_msg.into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Flush to close the segment and trigger tree ingest.
|
||||
hook.flush_open_segment(session).await;
|
||||
|
||||
let chunks = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
|
||||
assert!(
|
||||
!chunks.is_empty(),
|
||||
"Expected tree chunks after flush_open_segment"
|
||||
);
|
||||
|
||||
// The stub recap text must NOT appear in any chunk body.
|
||||
let stub_recap_text = "stub recap: discussed Rust ownership model";
|
||||
for chunk in &chunks {
|
||||
assert!(
|
||||
!chunk.content.contains(stub_recap_text),
|
||||
"Chunk content must NOT contain the recap text (evidence-vs-interpretation policy). \
|
||||
Found recap text in chunk id={}: {:?}",
|
||||
chunk.id,
|
||||
&chunk.content[..chunk.content.len().min(200)]
|
||||
);
|
||||
}
|
||||
|
||||
// The raw prose text MUST appear in at least one chunk.
|
||||
let has_user_prose = chunks.iter().any(|c| c.content.contains("lifetimes"));
|
||||
assert!(
|
||||
has_user_prose,
|
||||
"Expected at least one chunk body to contain raw prose from the turn \
|
||||
(keyword 'lifetimes'); found: {:?}",
|
||||
chunks
|
||||
.iter()
|
||||
.map(|c| &c.content[..c.content.len().min(100)])
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
/// `flush_open_segment` must also trigger the tree ingest for the trailing
|
||||
/// open segment (same as on_segment_closed at a topic boundary).
|
||||
#[tokio::test]
|
||||
async fn phase2_flush_also_triggers_tree_ingest() {
|
||||
let conn = setup_conn();
|
||||
let (_tmp, cfg) = test_config_with_tree();
|
||||
let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone());
|
||||
|
||||
let session = "phase2-flush-tree";
|
||||
|
||||
// Two turns — no boundary fires, segment stays open.
|
||||
for i in 1..=2 {
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: format!("Rust borrowing question {i}"),
|
||||
assistant_response: format!("Borrowing answer {i}"),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: i,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Confirm no tree chunks yet (segment still open).
|
||||
let before = count_chunks(&cfg).unwrap();
|
||||
assert_eq!(
|
||||
before, 0,
|
||||
"Expected 0 tree chunks before flush; got {before}"
|
||||
);
|
||||
|
||||
// Flush should close the segment and trigger tree ingest.
|
||||
hook.flush_open_segment(session).await;
|
||||
|
||||
let after = count_chunks(&cfg).unwrap();
|
||||
assert!(
|
||||
after >= 1,
|
||||
"Expected ≥ 1 tree chunk after flush_open_segment triggers segment ingest; got {after}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -17,7 +17,7 @@ use crate::openhuman::agent::host_runtime;
|
||||
use crate::openhuman::agent::memory_loader::{DefaultMemoryLoader, MemoryLoader};
|
||||
use crate::openhuman::config::{Config, ContextConfig};
|
||||
use crate::openhuman::context::prompt::SystemPromptBuilder;
|
||||
use crate::openhuman::context::{ContextManager, ProviderSummarizer};
|
||||
use crate::openhuman::context::{ContextManager, ProviderSummarizer, SegmentRecapSummarizer};
|
||||
use crate::openhuman::inference::provider::{self, Provider};
|
||||
use crate::openhuman::memory::{self, Memory};
|
||||
use crate::openhuman::security::SecurityPolicy;
|
||||
@@ -86,6 +86,8 @@ impl AgentBuilder {
|
||||
omit_profile: None,
|
||||
omit_memory_md: None,
|
||||
payload_summarizer: None,
|
||||
archivist_hook: None,
|
||||
unified_compaction_enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -313,6 +315,39 @@ impl AgentBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach the production [`ArchivistHook`] instance so the session
|
||||
/// turn loop can call [`ArchivistHook::flush_open_segment`] at
|
||||
/// session-wind-down time, guaranteeing the trailing open segment is
|
||||
/// always finalized with an LLM recap + embedding.
|
||||
///
|
||||
/// Set from `build_session_agent_inner` when
|
||||
/// `config.learning.episodic_capture_enabled` is `true` and a
|
||||
/// SQLite connection is available. Callers that construct an `Agent`
|
||||
/// directly (tests, CLI) can leave this `None` — flush is a no-op
|
||||
/// when the hook is absent.
|
||||
pub fn archivist_hook(
|
||||
mut self,
|
||||
hook: Option<Arc<crate::openhuman::agent::harness::archivist::ArchivistHook>>,
|
||||
) -> Self {
|
||||
self.archivist_hook = hook;
|
||||
self
|
||||
}
|
||||
|
||||
/// Phase 1.5 — gate the unified compaction path.
|
||||
///
|
||||
/// When `true` (the default) and an archivist hook is wired in via
|
||||
/// [`Self::archivist_hook`], the session's `ContextManager` summarizer is
|
||||
/// wrapped with a [`SegmentRecapSummarizer`] that routes autocompaction
|
||||
/// through the archivist's rolling recap (one LLM summarizer, soft-fallback
|
||||
/// to [`ProviderSummarizer`] when the recap is unavailable).
|
||||
///
|
||||
/// When `false` the `ProviderSummarizer` is used directly and Phase 1.5 is
|
||||
/// completely absent from the hot path — behaviour is identical to today's.
|
||||
pub fn unified_compaction_enabled(mut self, enabled: bool) -> Self {
|
||||
self.unified_compaction_enabled = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
/// Validates the configuration and constructs a new `Agent` instance.
|
||||
///
|
||||
/// This method is responsible for wiring together the provided components,
|
||||
@@ -377,7 +412,52 @@ impl AgentBuilder {
|
||||
// summarizer — every concern that touches "what's in the
|
||||
// model's context window" routes through this single handle.
|
||||
let context_config = self.context_config.unwrap_or_default();
|
||||
let summarizer = Arc::new(ProviderSummarizer::new(provider.clone()));
|
||||
|
||||
// Phase 1.5 — unified compaction.
|
||||
//
|
||||
// When `unified_compaction_enabled` is true AND an archivist hook
|
||||
// is wired in, wrap the inner `ProviderSummarizer` with a
|
||||
// `SegmentRecapSummarizer`. The outer type:
|
||||
// 1. Tries the rolling segment recap from the open segment.
|
||||
// 2. Falls back to the inner `ProviderSummarizer` if unavailable.
|
||||
//
|
||||
// With the flag off OR no archivist, the plain `ProviderSummarizer`
|
||||
// is used and Phase 1.5 is completely absent from the hot path
|
||||
// — behaviour is identical to Phase 1.
|
||||
let inner_summarizer: Arc<dyn crate::openhuman::context::Summarizer> =
|
||||
Arc::new(ProviderSummarizer::new(provider.clone()));
|
||||
let session_id_for_recap = self
|
||||
.event_session_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| "standalone".to_string());
|
||||
let summarizer: Arc<dyn crate::openhuman::context::Summarizer> =
|
||||
if self.unified_compaction_enabled {
|
||||
if let Some(ref archivist) = self.archivist_hook {
|
||||
log::debug!(
|
||||
"[agent::builder] unified_compaction_enabled=true — \
|
||||
wrapping summarizer with SegmentRecapSummarizer \
|
||||
session_id={session_id_for_recap}"
|
||||
);
|
||||
Arc::new(SegmentRecapSummarizer::new(
|
||||
Arc::clone(archivist),
|
||||
session_id_for_recap,
|
||||
inner_summarizer,
|
||||
))
|
||||
} else {
|
||||
log::debug!(
|
||||
"[agent::builder] unified_compaction_enabled=true but \
|
||||
no archivist hook — using ProviderSummarizer"
|
||||
);
|
||||
inner_summarizer
|
||||
}
|
||||
} else {
|
||||
log::debug!(
|
||||
"[agent::builder] unified_compaction_enabled=false — \
|
||||
using ProviderSummarizer (Phase 1.5 disabled)"
|
||||
);
|
||||
inner_summarizer
|
||||
};
|
||||
|
||||
let context = ContextManager::new(
|
||||
&context_config,
|
||||
summarizer,
|
||||
@@ -462,6 +542,7 @@ impl AgentBuilder {
|
||||
omit_memory_md: self.omit_memory_md.unwrap_or(true),
|
||||
payload_summarizer: self.payload_summarizer,
|
||||
last_seen_integrations_hash: 0,
|
||||
archivist_hook: self.archivist_hook,
|
||||
synthesized_tool_names: std::collections::HashSet::new(),
|
||||
})
|
||||
}
|
||||
@@ -1015,6 +1096,46 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// ── ArchivistHook — register independently of learning.enabled ──────
|
||||
//
|
||||
// Episodic capture (FTS5 index, segment lifecycle, LLM recap, embedding)
|
||||
// is the system-of-record for chat turns and must stay active even when
|
||||
// the inference stack (`reflection`, `stability_detector`) is disabled.
|
||||
// Gated only on `config.learning.episodic_capture_enabled` (default: true)
|
||||
// and on the memory backend exposing a SQLite connection.
|
||||
let archivist_hook_arc: Option<
|
||||
Arc<crate::openhuman::agent::harness::archivist::ArchivistHook>,
|
||||
> = if config.learning.episodic_capture_enabled {
|
||||
match memory.sqlite_conn() {
|
||||
Some(conn) => {
|
||||
let hook = Arc::new(
|
||||
crate::openhuman::agent::harness::archivist::ArchivistHook::new(conn, true)
|
||||
.with_config(config.clone()),
|
||||
);
|
||||
post_turn_hooks
|
||||
.push(Arc::clone(&hook)
|
||||
as Arc<dyn crate::openhuman::agent::hooks::PostTurnHook>);
|
||||
log::info!(
|
||||
"[archivist] episodic capture hook registered (learning.enabled={})",
|
||||
config.learning.enabled
|
||||
);
|
||||
Some(hook)
|
||||
}
|
||||
None => {
|
||||
log::warn!(
|
||||
"[archivist] no SQLite connection available from memory backend — \
|
||||
episodic capture disabled"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
log::info!(
|
||||
"[archivist] episodic_capture_enabled=false — archivist hook not registered"
|
||||
);
|
||||
None
|
||||
};
|
||||
|
||||
// Resolve the per-agent delegation tool set and visible-tool
|
||||
// whitelist from the target definition (when we have one) or
|
||||
// fall back to the orchestrator's synthesis path.
|
||||
@@ -1347,6 +1468,8 @@ impl Agent {
|
||||
if let Some(ps) = payload_summarizer {
|
||||
builder = builder.payload_summarizer(ps);
|
||||
}
|
||||
builder = builder.archivist_hook(archivist_hook_arc);
|
||||
builder = builder.unified_compaction_enabled(config.learning.unified_compaction_enabled);
|
||||
builder.build()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -267,6 +267,12 @@ impl Agent {
|
||||
// error. The timestamp is bumped on every successful `load` (even
|
||||
// when the digest is empty) so an empty workspace doesn't get
|
||||
// re-queried every turn.
|
||||
//
|
||||
// Gate STM preemptive recall on the session turn index, independent
|
||||
// of tree-prefetch success/failure. (Previously keyed off
|
||||
// `last_tree_prefetch_at.is_none()`, which stays `None` when tree
|
||||
// prefetch fails — re-firing STM recall on every later turn.)
|
||||
let is_first_turn_for_stm = self.context.stats().session_memory_current_turn == 0;
|
||||
let now = std::time::Instant::now();
|
||||
let context = if crate::openhuman::agent::tree_loader::should_prefetch(
|
||||
self.last_tree_prefetch_at,
|
||||
@@ -309,6 +315,70 @@ impl Agent {
|
||||
context
|
||||
};
|
||||
|
||||
// ── Phase 3 STM preemptive recall ────────────────────────────
|
||||
// On the very first turn only, assemble a bounded cross-thread
|
||||
// context block from the FTS5 episodic arm (keyword match) and the
|
||||
// segment-embedding arm (cosine similarity). The block rides on the
|
||||
// user message (NOT the system prompt) to keep the KV-cache prefix
|
||||
// stable, exactly like the tree-context injection above.
|
||||
//
|
||||
// Gate: `learning.stm_recall_enabled` must be true AND this must
|
||||
// be the first turn (STM is snapshot-frozen at session start).
|
||||
// Failure is non-fatal — bare `context` passes through untouched.
|
||||
let context = if is_first_turn_for_stm {
|
||||
// Load config to check the gate. Use a cached load (cheap).
|
||||
let stm_enabled = crate::openhuman::config::rpc::load_config_with_timeout()
|
||||
.await
|
||||
.map(|cfg| cfg.learning.stm_recall_enabled)
|
||||
.unwrap_or(true); // default: enabled
|
||||
|
||||
if stm_enabled {
|
||||
if let Some(conn) = self.memory.sqlite_conn() {
|
||||
use crate::openhuman::memory::stm_recall::recall::{stm_recall, StmRecallOpts};
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: &self.event_session_id,
|
||||
query: if user_message.trim().is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(user_message)
|
||||
},
|
||||
model_signature: None,
|
||||
};
|
||||
match stm_recall(&conn, &opts, None) {
|
||||
Ok(block) if !block.is_empty() => {
|
||||
let stm_md = block.render();
|
||||
log::info!(
|
||||
"[stm_recall] preemptive block injected: {} items, ~{} chars, fts5_candidates={}, dropped_dedup={}",
|
||||
block.items.len(),
|
||||
stm_md.chars().count(),
|
||||
block.fts5_candidates,
|
||||
block.dropped_dedup
|
||||
);
|
||||
format!("{stm_md}{context}")
|
||||
}
|
||||
Ok(_) => {
|
||||
log::debug!(
|
||||
"[stm_recall] preemptive recall: no cross-thread context found"
|
||||
);
|
||||
context
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("[stm_recall] preemptive recall failed (non-fatal): {e}");
|
||||
context
|
||||
}
|
||||
}
|
||||
} else {
|
||||
log::debug!("[stm_recall] preemptive recall skipped — no SQLite connection on memory backend");
|
||||
context
|
||||
}
|
||||
} else {
|
||||
log::debug!("[stm_recall] preemptive recall skipped — stm_recall_enabled=false");
|
||||
context
|
||||
}
|
||||
} else {
|
||||
context
|
||||
};
|
||||
|
||||
let enriched = if context.is_empty() {
|
||||
log::info!("[agent] no memory context found — using raw user message");
|
||||
self.last_memory_context = None;
|
||||
@@ -904,7 +974,7 @@ impl Agent {
|
||||
// later), which is the right amount of retry behaviour for a
|
||||
// librarian task that's idempotent across reruns.
|
||||
if result.is_ok() && self.context.should_extract_session_memory() {
|
||||
self.spawn_session_memory_extraction();
|
||||
self.spawn_session_memory_extraction().await;
|
||||
// Sibling pipeline (#1399): heuristic transcript ingestion
|
||||
// turns the just-written transcript into durable
|
||||
// conversational memory + reflections so a brand-new chat
|
||||
@@ -1687,7 +1757,29 @@ impl Agent {
|
||||
/// Gated by [`context_pipeline::SessionMemoryState::should_extract`]
|
||||
/// — see its docs for the threshold invariants. Safe to call from
|
||||
/// inside `turn()` after the turn body has settled.
|
||||
pub(super) fn spawn_session_memory_extraction(&mut self) {
|
||||
pub(super) async fn spawn_session_memory_extraction(&mut self) {
|
||||
// ── Flush the trailing open segment before the session winds down ──
|
||||
//
|
||||
// The ArchivistHook manages per-turn segment lifecycle but cannot
|
||||
// force-close the *last* open segment because there is no explicit
|
||||
// "session end" event in the turn loop. `spawn_session_memory_extraction`
|
||||
// is the closest available signal: it fires when the context manager
|
||||
// decides the session has accumulated enough material to archive.
|
||||
//
|
||||
// GUARANTEE: the flush is *awaited* here (not fire-and-forget) so
|
||||
// the trailing segment always receives its recap + embedding + tree
|
||||
// ingest before the function returns, even during runtime wind-down.
|
||||
// This honours the doc-comment guarantee on `flush_open_segment` in
|
||||
// `archivist.rs`. No deadlock risk: no mutex guard is held across
|
||||
// this await point.
|
||||
if let Some(ref archivist) = self.archivist_hook {
|
||||
let session_id = self.event_session_id.clone();
|
||||
log::debug!(
|
||||
"[archivist] awaiting flush_open_segment for session={session_id} at session wind-down"
|
||||
);
|
||||
archivist.flush_open_segment(&session_id).await;
|
||||
}
|
||||
|
||||
let Some(registry) = harness::AgentDefinitionRegistry::global() else {
|
||||
log::debug!("[session_memory] registry not initialised — skipping extraction spawn");
|
||||
return;
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
//! crate gaining field access.
|
||||
|
||||
use crate::openhuman::agent::dispatcher::ToolDispatcher;
|
||||
use crate::openhuman::agent::harness::archivist::ArchivistHook;
|
||||
use crate::openhuman::agent::hooks::PostTurnHook;
|
||||
use crate::openhuman::agent::memory_loader::MemoryLoader;
|
||||
use crate::openhuman::agent::progress::AgentProgress;
|
||||
@@ -161,6 +162,12 @@ pub struct Agent {
|
||||
/// dormant on session startup and only fires when integrations
|
||||
/// actually change mid-conversation.
|
||||
pub(super) last_seen_integrations_hash: u64,
|
||||
/// Optional reference to the `ArchivistHook` registered in
|
||||
/// `post_turn_hooks`. Kept separately so the turn loop can call
|
||||
/// `flush_open_segment` at session-memory-extraction time (the
|
||||
/// closest available signal to "session is ending") to finalize the
|
||||
/// trailing open segment with an LLM recap + embedding.
|
||||
pub(super) archivist_hook: Option<Arc<ArchivistHook>>,
|
||||
/// Names of every tool currently in [`Agent::tools`] that was
|
||||
/// produced by [`crate::openhuman::tools::orchestrator_tools::collect_orchestrator_tools`]
|
||||
/// (i.e. `delegate_<toolkit>` skill tools and archetype-delegation
|
||||
@@ -220,6 +227,18 @@ pub struct AgentBuilder {
|
||||
/// to a `SubagentPayloadSummarizer` instance.
|
||||
pub(super) payload_summarizer:
|
||||
Option<Arc<dyn crate::openhuman::agent::harness::payload_summarizer::PayloadSummarizer>>,
|
||||
/// Optional reference to the production `ArchivistHook`. Set when
|
||||
/// `config.learning.episodic_capture_enabled` is true. Used to call
|
||||
/// `flush_open_segment` at the closest available session-end signal.
|
||||
pub(super) archivist_hook: Option<Arc<ArchivistHook>>,
|
||||
/// Phase 1.5 — when `true` AND `archivist_hook` is `Some`, the
|
||||
/// `ContextManager`'s summarizer is wrapped with a
|
||||
/// `SegmentRecapSummarizer` that routes compaction through the
|
||||
/// archivist's rolling segment recap (one summarizer, soft-fallback).
|
||||
/// When `false` (or archivist absent), the plain `ProviderSummarizer`
|
||||
/// is used and Phase 1.5 is completely absent from the hot path.
|
||||
/// Default: `true` (mirrors `LearningConfig::unified_compaction_enabled`).
|
||||
pub(super) unified_compaction_enabled: bool,
|
||||
}
|
||||
|
||||
impl Default for AgentBuilder {
|
||||
|
||||
@@ -76,6 +76,55 @@ pub struct LearningConfig {
|
||||
#[serde(default = "default_true")]
|
||||
pub stability_detector_enabled: bool,
|
||||
|
||||
/// Enable episodic capture (ArchivistHook) regardless of the master
|
||||
/// `learning.enabled` toggle.
|
||||
///
|
||||
/// Episodic capture is the system-of-record for chat turns
|
||||
/// (`episodic_log` FTS5 table, conversation segmentation, segment
|
||||
/// summaries with LLM recap, and segment embeddings). It must remain
|
||||
/// active even when the inference stack
|
||||
/// (reflection / stability-detector) is off.
|
||||
///
|
||||
/// Default: `true`. Set to `false` to fully disable the Archivist.
|
||||
///
|
||||
/// Override via `OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED=0|1`.
|
||||
#[serde(default = "default_true")]
|
||||
pub episodic_capture_enabled: bool,
|
||||
|
||||
/// Enable preemptive STM recall injection at session start and on-demand
|
||||
/// `stm_recall_search` tool exposure.
|
||||
///
|
||||
/// When enabled, a bounded cross-thread context block is assembled from
|
||||
/// recent episodic entries (FTS5 keyword arm) and segment recaps (cosine
|
||||
/// similarity arm) from OTHER sessions and injected into the first turn's
|
||||
/// user message. The `stm_recall_search` tool is also registered in the
|
||||
/// agent's tool list.
|
||||
///
|
||||
/// Default: `true`. Set to `false` to fully disable STM recall.
|
||||
///
|
||||
/// Override via `OPENHUMAN_LEARNING_STM_RECALL_ENABLED=0|1`.
|
||||
#[serde(default = "default_true")]
|
||||
pub stm_recall_enabled: bool,
|
||||
|
||||
/// Use the rolling segment recap as the compaction text for evicted turns
|
||||
/// (Phase 1.5 — unified compaction).
|
||||
///
|
||||
/// When `true`, the [`ContextManager`]'s autocompaction summarizer is
|
||||
/// wrapped with a `SegmentRecapSummarizer` that first tries to obtain the
|
||||
/// current open segment's rolling recap from the `ArchivistHook` and uses
|
||||
/// it as the replacement text for the evicted head. If the rolling recap
|
||||
/// is unavailable (no archivist, no open segment, LLM failure, flag off)
|
||||
/// the inner `ProviderSummarizer` runs as before — the live prompt is
|
||||
/// NEVER left over-budget regardless of the recap path's health.
|
||||
///
|
||||
/// Default: `true`. Set to `false` to revert to the standalone
|
||||
/// `ProviderSummarizer` path (today's behaviour, Phase 1.5 completely
|
||||
/// absent from the hot path).
|
||||
///
|
||||
/// Override via `OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED=0|1`.
|
||||
#[serde(default = "default_true")]
|
||||
pub unified_compaction_enabled: bool,
|
||||
|
||||
/// How often the periodic rebuild loop runs in seconds. Default: 1800 (30 minutes).
|
||||
#[serde(default = "default_rebuild_interval_secs")]
|
||||
pub rebuild_interval_secs: u64,
|
||||
@@ -111,6 +160,9 @@ impl Default for LearningConfig {
|
||||
chat_to_tree_enabled: default_true(),
|
||||
stability_detector_enabled: default_true(),
|
||||
rebuild_interval_secs: default_rebuild_interval_secs(),
|
||||
episodic_capture_enabled: default_true(),
|
||||
stm_recall_enabled: default_true(),
|
||||
unified_compaction_enabled: default_true(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1308,6 +1308,28 @@ impl Config {
|
||||
self.learning.min_turn_complexity = min;
|
||||
}
|
||||
}
|
||||
if let Some(flag) = env.get("OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED") {
|
||||
if let Some(enabled) =
|
||||
parse_env_bool("OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED", flag.as_str())
|
||||
{
|
||||
self.learning.episodic_capture_enabled = enabled;
|
||||
}
|
||||
}
|
||||
if let Some(flag) = env.get("OPENHUMAN_LEARNING_STM_RECALL_ENABLED") {
|
||||
if let Some(enabled) =
|
||||
parse_env_bool("OPENHUMAN_LEARNING_STM_RECALL_ENABLED", flag.as_str())
|
||||
{
|
||||
self.learning.stm_recall_enabled = enabled;
|
||||
}
|
||||
}
|
||||
if let Some(flag) = env.get("OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED") {
|
||||
if let Some(enabled) = parse_env_bool(
|
||||
"OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED",
|
||||
flag.as_str(),
|
||||
) {
|
||||
self.learning.unified_compaction_enabled = enabled;
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 4 memory-tree embedding overrides (#710). Setting the env
|
||||
// var to an empty string explicitly clears the default — useful
|
||||
|
||||
@@ -34,6 +34,7 @@ pub mod manager;
|
||||
pub mod microcompact;
|
||||
pub mod pipeline;
|
||||
pub mod prompt;
|
||||
pub mod segment_recap_summarizer;
|
||||
pub mod session_memory;
|
||||
pub mod summarizer;
|
||||
pub mod tool_result_budget;
|
||||
@@ -49,6 +50,7 @@ pub use prompt::{
|
||||
PromptSection, PromptTool, RuntimeSection, SafetySection, SystemPromptBuilder, ToolsSection,
|
||||
WorkspaceSection,
|
||||
};
|
||||
pub use segment_recap_summarizer::SegmentRecapSummarizer;
|
||||
pub use session_memory::{
|
||||
SessionMemoryConfig, SessionMemoryState, ARCHIVIST_EXTRACTION_PROMPT, DEFAULT_MIN_TOKEN_GROWTH,
|
||||
DEFAULT_MIN_TOOL_CALLS, DEFAULT_MIN_TURNS_BETWEEN,
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
//! Phase 1.5 — segment-recap-backed compaction summarizer.
|
||||
//!
|
||||
//! [`SegmentRecapSummarizer`] wraps an inner [`Summarizer`] (normally a
|
||||
//! [`super::ProviderSummarizer`]) and intercepts the autocompaction call to
|
||||
//! try the rolling segment recap from the [`ArchivistHook`] first.
|
||||
//!
|
||||
//! ## Strategy
|
||||
//!
|
||||
//! When `AutocompactionRequested` fires, the manager calls
|
||||
//! [`SegmentRecapSummarizer::summarize`]. That method:
|
||||
//!
|
||||
//! 1. Calls [`ArchivistHook::rolling_segment_recap`] for the current session.
|
||||
//! 2. If a non-empty recap is returned, it replaces the evicted head of the
|
||||
//! history with a single `[segment-recap]` system message containing that
|
||||
//! text. The head/tail split follows the same "never break a tool-call
|
||||
//! pair" rule as the inner summarizer.
|
||||
//! 3. If the recap is `None`/empty (archivist absent, no open segment, LLM
|
||||
//! fail, flag off) — soft-fallback: delegate to the inner summarizer
|
||||
//! unchanged. The prompt is never left over-budget; the existing
|
||||
//! compaction safety net always fires.
|
||||
//!
|
||||
//! ## Invariants (never violated by this code)
|
||||
//!
|
||||
//! - The `ArchivistHook` is only *read* through `rolling_segment_recap` — no
|
||||
//! segment is closed, no `segment_set_summary` is written, no embedding is
|
||||
//! produced. Those side-effects are finalize-only (Phase 1 owns them).
|
||||
//! - Events/profile/tree derive from RAW episodic rows — this path never
|
||||
//! touches those subsystems.
|
||||
//! - On any error in the recap path, the inner summarizer runs. The
|
||||
//! circuit-breaker logic in [`super::ContextManager`] is fed the result of
|
||||
//! whichever path actually ran; the breaker sees a success if either path
|
||||
//! succeeds.
|
||||
//! - The history is either fully rewritten (success) or left completely
|
||||
//! untouched (the inner summarizer also fails, in which case its `Err` is
|
||||
//! returned — the breaker will nudge and eventually trip, preventing
|
||||
//! infinite loops).
|
||||
|
||||
use super::summarizer::{Summarizer, SummaryStats};
|
||||
use crate::openhuman::agent::harness::archivist::ArchivistHook;
|
||||
use crate::openhuman::inference::provider::{ChatMessage, ConversationMessage};
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// How many messages at the tail of the history are preserved verbatim when
|
||||
/// the recap path fires (same default as [`super::ProviderSummarizer`] so the
|
||||
/// two paths preserve the same recent context window).
|
||||
const DEFAULT_KEEP_RECENT: usize = super::summarizer::DEFAULT_KEEP_RECENT;
|
||||
|
||||
/// Compaction summarizer that prefers the rolling segment recap from the
|
||||
/// archivist over a standalone LLM summarization call.
|
||||
///
|
||||
/// Constructed by [`super::SessionBuilder`] when
|
||||
/// `config.learning.unified_compaction_enabled` is `true` and an
|
||||
/// `ArchivistHook` is wired in; otherwise the plain `ProviderSummarizer`
|
||||
/// is used and this type is never instantiated.
|
||||
pub struct SegmentRecapSummarizer {
|
||||
/// The archivist that owns episodic state and can produce rolling recaps.
|
||||
archivist: Arc<ArchivistHook>,
|
||||
/// Session ID needed to look up the open segment.
|
||||
session_id: String,
|
||||
/// Inner summarizer — the safety-net fallback (normally a
|
||||
/// [`super::ProviderSummarizer`] wrapping the same provider the agent uses
|
||||
/// for its normal turns).
|
||||
inner: Arc<dyn Summarizer>,
|
||||
/// How many tail messages to preserve verbatim when the recap path fires.
|
||||
keep_recent: usize,
|
||||
}
|
||||
|
||||
impl SegmentRecapSummarizer {
|
||||
/// Construct a new [`SegmentRecapSummarizer`].
|
||||
///
|
||||
/// * `archivist` — shared archivist handle (same Arc the agent holds).
|
||||
/// * `session_id` — the session whose open segment will be recapped.
|
||||
/// * `inner` — fallback summarizer; used when the recap is unavailable.
|
||||
pub fn new(
|
||||
archivist: Arc<ArchivistHook>,
|
||||
session_id: String,
|
||||
inner: Arc<dyn Summarizer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
archivist,
|
||||
session_id,
|
||||
inner,
|
||||
keep_recent: DEFAULT_KEEP_RECENT,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override how many tail messages are preserved verbatim. Useful in
|
||||
/// tests that want to exercise the recap path with short histories.
|
||||
#[cfg(test)]
|
||||
pub fn with_keep_recent(mut self, n: usize) -> Self {
|
||||
self.keep_recent = n;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Summarizer for SegmentRecapSummarizer {
|
||||
async fn summarize(
|
||||
&self,
|
||||
history: &mut Vec<ConversationMessage>,
|
||||
model: &str,
|
||||
) -> Result<SummaryStats> {
|
||||
// ── 1. Try the rolling segment recap ─────────────────────────────
|
||||
let recap_opt = self.archivist.rolling_segment_recap(&self.session_id).await;
|
||||
|
||||
match recap_opt {
|
||||
Some(recap) if !recap.is_empty() => {
|
||||
tracing::info!(
|
||||
session_id = %self.session_id,
|
||||
recap_chars = recap.len(),
|
||||
history_len = history.len(),
|
||||
keep_recent = self.keep_recent,
|
||||
"[context::segment_recap] using rolling segment recap as compaction text"
|
||||
);
|
||||
|
||||
let total = history.len();
|
||||
if total <= self.keep_recent {
|
||||
tracing::debug!(
|
||||
session_id = %self.session_id,
|
||||
"[context::segment_recap] history below keep_recent — \
|
||||
nothing to compact (NoOp)"
|
||||
);
|
||||
return Ok(SummaryStats::default());
|
||||
}
|
||||
|
||||
// Use the same "never break a tool-call pair" split rule as
|
||||
// ProviderSummarizer so the API invariant
|
||||
// (AssistantToolCalls ↔ ToolResults) is preserved.
|
||||
// Delegates to the canonical implementation in `summarizer`
|
||||
// so the two compaction paths share a single definition.
|
||||
let proposed_head = total - self.keep_recent;
|
||||
let head_len = super::summarizer::snap_split_forward(history, proposed_head);
|
||||
if head_len == 0 {
|
||||
tracing::debug!(
|
||||
session_id = %self.session_id,
|
||||
"[context::segment_recap] split snapped to 0 — \
|
||||
falling back to inner summarizer"
|
||||
);
|
||||
return self.inner.summarize(history, model).await;
|
||||
}
|
||||
|
||||
// Estimate bytes freed (same formula as ProviderSummarizer).
|
||||
let approx_input_bytes: usize = history[..head_len]
|
||||
.iter()
|
||||
.map(|m| conversation_message_approx_bytes(m))
|
||||
.sum();
|
||||
|
||||
let summary_body = format!(
|
||||
"[segment-recap] Summary of {head_len} earlier messages \
|
||||
(archivist rolling recap):\n\n{recap}"
|
||||
);
|
||||
let summary_chars = summary_body.len();
|
||||
let approx_tokens_freed = (approx_input_bytes as u64)
|
||||
.saturating_sub(summary_chars as u64)
|
||||
.div_ceil(4);
|
||||
|
||||
// Atomically rewrite the head in place — no partial mutation on
|
||||
// failure because all failure paths returned early above.
|
||||
let tail: Vec<ConversationMessage> = history.drain(head_len..).collect();
|
||||
history.clear();
|
||||
history.push(ConversationMessage::Chat(ChatMessage::system(summary_body)));
|
||||
history.extend(tail);
|
||||
|
||||
tracing::info!(
|
||||
session_id = %self.session_id,
|
||||
messages_removed = head_len,
|
||||
approx_tokens_freed,
|
||||
summary_chars,
|
||||
"[context::segment_recap] compaction via segment recap complete"
|
||||
);
|
||||
|
||||
Ok(SummaryStats {
|
||||
messages_removed: head_len,
|
||||
approx_tokens_freed,
|
||||
summary_chars,
|
||||
})
|
||||
}
|
||||
|
||||
// ── 2. Soft-fallback ─────────────────────────────────────────
|
||||
//
|
||||
// The rolling recap was unavailable (None, empty, LLM fail, no
|
||||
// open segment, archivist disabled). Delegate to the inner
|
||||
// summarizer so the prompt is NEVER left over-budget. This is the
|
||||
// "always bounded" guarantee.
|
||||
recap_result => {
|
||||
let reason = match &recap_result {
|
||||
None => "rolling_segment_recap returned None",
|
||||
Some(s) if s.is_empty() => "rolling_segment_recap returned empty string",
|
||||
_ => "unreachable",
|
||||
};
|
||||
tracing::info!(
|
||||
session_id = %self.session_id,
|
||||
reason,
|
||||
"[context::segment_recap] recap unavailable — \
|
||||
falling back to inner summarizer"
|
||||
);
|
||||
self.inner.summarize(history, model).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Very rough byte count for a [`ConversationMessage`] — used only for the
|
||||
/// "approx_tokens_freed" stat. Accuracy doesn't matter much (it's the same
|
||||
/// rough accounting ProviderSummarizer uses).
|
||||
fn conversation_message_approx_bytes(msg: &ConversationMessage) -> usize {
|
||||
match msg {
|
||||
ConversationMessage::Chat(m) => m.content.len(),
|
||||
ConversationMessage::AssistantToolCalls { text, tool_calls } => {
|
||||
text.as_deref().map_or(0, str::len)
|
||||
+ tool_calls
|
||||
.iter()
|
||||
.map(|tc| tc.arguments.len())
|
||||
.sum::<usize>()
|
||||
}
|
||||
ConversationMessage::ToolResults(results) => results.iter().map(|r| r.content.len()).sum(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "segment_recap_summarizer_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,562 @@
|
||||
//! Tests for Phase 1.5 — `SegmentRecapSummarizer`.
|
||||
//!
|
||||
//! Proves:
|
||||
//! 1. Rolling recap summarizes the open segment WITHOUT closing it, writing
|
||||
//! `segment_set_summary`, or producing an embedding row.
|
||||
//! 2. When a rolling recap is available, the compaction replacement text
|
||||
//! equals it (provenance/path), not a separately-generated summary.
|
||||
//! 3. Soft-fallback: archivist absent / LLM stub failing / flag off →
|
||||
//! compaction falls back to the inner summarizer; prompt stays bounded.
|
||||
//! 4. Finalize path (Phase 1 `on_segment_closed`) still works unchanged
|
||||
//! (recap persisted + embedded at close) — regression guard.
|
||||
|
||||
use super::*;
|
||||
use crate::openhuman::agent::harness::archivist::ArchivistHook;
|
||||
use crate::openhuman::agent::hooks::{PostTurnHook as _, TurnContext};
|
||||
use crate::openhuman::context::summarizer::{Summarizer, SummaryStats};
|
||||
use crate::openhuman::inference::provider::{ChatMessage, ConversationMessage};
|
||||
use crate::openhuman::memory::store::{fts5, segments as seg};
|
||||
use crate::openhuman::memory::tree::chat::ChatPrompt;
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use rusqlite::Connection;
|
||||
use std::sync::Arc;
|
||||
|
||||
// ── Shared test infrastructure ────────────────────────────────────────────────
|
||||
|
||||
fn setup_conn() -> Arc<Mutex<Connection>> {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
conn.execute_batch(fts5::EPISODIC_INIT_SQL).unwrap();
|
||||
conn.execute_batch(seg::SEGMENTS_INIT_SQL).unwrap();
|
||||
conn.execute_batch(crate::openhuman::memory::store::events::EVENTS_INIT_SQL)
|
||||
.unwrap();
|
||||
conn.execute_batch(crate::openhuman::memory::store::profile::PROFILE_INIT_SQL)
|
||||
.unwrap();
|
||||
Arc::new(Mutex::new(conn))
|
||||
}
|
||||
|
||||
/// Stub ChatProvider always returns a fixed recap string.
|
||||
struct StubChatProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl crate::openhuman::memory::tree::chat::ChatProvider for StubChatProvider {
|
||||
fn name(&self) -> &str {
|
||||
"stub:test"
|
||||
}
|
||||
async fn chat_for_json(&self, _prompt: &ChatPrompt) -> Result<String> {
|
||||
Ok("rolling recap: discussed memory safety in Rust".to_string())
|
||||
}
|
||||
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result<String> {
|
||||
Ok("rolling recap: discussed memory safety in Rust".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Stub ChatProvider that always fails — simulates LLM unavailability.
|
||||
struct FailingChatProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl crate::openhuman::memory::tree::chat::ChatProvider for FailingChatProvider {
|
||||
fn name(&self) -> &str {
|
||||
"stub:failing"
|
||||
}
|
||||
async fn chat_for_json(&self, _prompt: &ChatPrompt) -> Result<String> {
|
||||
anyhow::bail!("stub LLM unavailable")
|
||||
}
|
||||
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result<String> {
|
||||
anyhow::bail!("stub LLM unavailable")
|
||||
}
|
||||
}
|
||||
|
||||
/// Stub Embedder returns a fixed unit vector.
|
||||
struct StubEmbedder;
|
||||
|
||||
#[async_trait]
|
||||
impl crate::openhuman::memory::tree::score::embed::Embedder for StubEmbedder {
|
||||
fn name(&self) -> &'static str {
|
||||
"stub-embedder-v1"
|
||||
}
|
||||
async fn embed(&self, _text: &str) -> Result<Vec<f32>> {
|
||||
Ok(vec![0.5_f32, 0.5, 0.5, 0.5])
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner mock summarizer that records call count and always succeeds.
|
||||
struct RecordingSummarizer {
|
||||
calls: std::sync::Mutex<usize>,
|
||||
}
|
||||
|
||||
impl RecordingSummarizer {
|
||||
fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
calls: std::sync::Mutex::new(0),
|
||||
})
|
||||
}
|
||||
fn call_count(&self) -> usize {
|
||||
*self.calls.lock().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Summarizer for RecordingSummarizer {
|
||||
async fn summarize(
|
||||
&self,
|
||||
history: &mut Vec<ConversationMessage>,
|
||||
_model: &str,
|
||||
) -> Result<SummaryStats> {
|
||||
*self.calls.lock().unwrap() += 1;
|
||||
// Replace history with a single "inner fallback" message.
|
||||
let removed = history.len();
|
||||
history.clear();
|
||||
history.push(ConversationMessage::Chat(ChatMessage::system(
|
||||
"inner fallback summary",
|
||||
)));
|
||||
Ok(SummaryStats {
|
||||
messages_removed: removed,
|
||||
approx_tokens_freed: 500,
|
||||
summary_chars: 22,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn user(s: &str) -> ConversationMessage {
|
||||
ConversationMessage::Chat(ChatMessage::user(s))
|
||||
}
|
||||
|
||||
// ── Test 1: rolling recap does NOT close the segment or write summary/embedding
|
||||
|
||||
/// `rolling_segment_recap` must produce a non-empty string from the open
|
||||
/// segment's entries without closing it, writing `segment_set_summary`, or
|
||||
/// producing an embedding row.
|
||||
#[tokio::test]
|
||||
async fn rolling_recap_does_not_close_segment_or_write_summary_or_embedding() {
|
||||
let conn = setup_conn();
|
||||
let hook = Arc::new(ArchivistHook::new_with_stubs(
|
||||
conn.clone(),
|
||||
Arc::new(StubChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
));
|
||||
|
||||
let session = "p15-rolling-no-close";
|
||||
|
||||
// Write two turns into the open segment (no boundary fires).
|
||||
for i in 1..=2u64 {
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: format!("Turn {i} about Rust memory safety"),
|
||||
assistant_response: format!("Answer {i} about ownership"),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: i as usize,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Verify the segment is still open before calling rolling_segment_recap.
|
||||
let open_before = seg::open_segment_for_session(&conn, session).unwrap();
|
||||
assert!(
|
||||
open_before.is_some(),
|
||||
"Expected an open segment after 2 turns (no boundary should have fired)"
|
||||
);
|
||||
|
||||
// Call rolling_segment_recap — must NOT close the segment or write DB state.
|
||||
let recap = hook.rolling_segment_recap(session).await;
|
||||
assert!(
|
||||
recap.is_some(),
|
||||
"Expected rolling_segment_recap to return Some for an open segment with entries"
|
||||
);
|
||||
let recap_text = recap.unwrap();
|
||||
assert!(
|
||||
!recap_text.is_empty(),
|
||||
"Expected non-empty recap from rolling_segment_recap"
|
||||
);
|
||||
|
||||
// Segment must STILL be open after the call.
|
||||
let open_after = seg::open_segment_for_session(&conn, session).unwrap();
|
||||
assert!(
|
||||
open_after.is_some(),
|
||||
"Segment must still be open after rolling_segment_recap (no close side-effect)"
|
||||
);
|
||||
|
||||
let seg_id = open_after.unwrap().segment_id;
|
||||
|
||||
// No summary must have been written.
|
||||
let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap();
|
||||
let our_seg = all_segs.iter().find(|s| s.segment_id == seg_id).unwrap();
|
||||
assert!(
|
||||
our_seg.summary.is_none() || our_seg.summary.as_deref() == Some(""),
|
||||
"segment_set_summary must NOT have been called by rolling_segment_recap; \
|
||||
found: {:?}",
|
||||
our_seg.summary
|
||||
);
|
||||
|
||||
// No embedding must have been written.
|
||||
let embedding = seg::segment_embedding_get(&conn, &seg_id, "stub-embedder-v1").unwrap();
|
||||
assert!(
|
||||
embedding.is_none(),
|
||||
"No embedding must be written by rolling_segment_recap \
|
||||
(finalize-only invariant); found an embedding for segment={seg_id}"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Test 2: compaction uses recap text, not inner summarizer ─────────────────
|
||||
|
||||
/// When the rolling recap is available, `SegmentRecapSummarizer::summarize`
|
||||
/// must use it as the replacement text — the inner summarizer must NOT be
|
||||
/// called and the history head must contain `[segment-recap]`.
|
||||
#[tokio::test]
|
||||
async fn compaction_uses_recap_text_not_inner_summarizer() {
|
||||
let conn = setup_conn();
|
||||
let hook = Arc::new(ArchivistHook::new_with_stubs(
|
||||
conn.clone(),
|
||||
Arc::new(StubChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
));
|
||||
|
||||
let session = "p15-compaction-recap";
|
||||
|
||||
// Write one turn so the open segment has entries.
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Tell me about Rust ownership".into(),
|
||||
assistant_response: "Ownership prevents data races.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let inner = RecordingSummarizer::new();
|
||||
let recap_summ = SegmentRecapSummarizer::new(
|
||||
Arc::clone(&hook),
|
||||
session.to_string(),
|
||||
inner.clone() as Arc<dyn Summarizer>,
|
||||
)
|
||||
.with_keep_recent(1); // compact everything except the last message
|
||||
|
||||
// Build a history that exceeds keep_recent so summarization fires.
|
||||
let mut history = vec![
|
||||
user("message 1 (will be evicted)"),
|
||||
user("message 2 (will be evicted)"),
|
||||
user("message 3 — preserved tail"),
|
||||
];
|
||||
|
||||
let stats = recap_summ
|
||||
.summarize(&mut history, "test-model")
|
||||
.await
|
||||
.expect("summarize must succeed");
|
||||
|
||||
// Inner summarizer must NOT have been called.
|
||||
assert_eq!(
|
||||
inner.call_count(),
|
||||
0,
|
||||
"Inner summarizer must not be called when rolling recap is available"
|
||||
);
|
||||
|
||||
// Stats must reflect a non-zero reduction.
|
||||
assert!(
|
||||
stats.messages_removed > 0,
|
||||
"Expected some messages to be removed; got 0"
|
||||
);
|
||||
|
||||
// The new head message must contain the recap text (not the inner fallback).
|
||||
assert_eq!(
|
||||
history.len(),
|
||||
2,
|
||||
"Expected summary message + 1 preserved tail; got {}",
|
||||
history.len()
|
||||
);
|
||||
match &history[0] {
|
||||
ConversationMessage::Chat(m) => {
|
||||
assert!(
|
||||
m.content.contains("[segment-recap]"),
|
||||
"Expected summary message to contain '[segment-recap]'; got: {:?}",
|
||||
m.content
|
||||
);
|
||||
// Must also contain the actual recap text from the stub.
|
||||
assert!(
|
||||
m.content.contains("rolling recap"),
|
||||
"Expected summary to contain recap text from stub provider; got: {:?}",
|
||||
m.content
|
||||
);
|
||||
}
|
||||
other => panic!("Expected Chat message for summary, got: {:?}", other),
|
||||
}
|
||||
|
||||
// Tail must be preserved verbatim.
|
||||
match &history[1] {
|
||||
ConversationMessage::Chat(m) => {
|
||||
assert_eq!(m.content, "message 3 — preserved tail");
|
||||
}
|
||||
other => panic!("Expected Chat tail message, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Test 3: soft-fallback when archivist absent ───────────────────────────────
|
||||
|
||||
/// When the archivist is disabled (no conn), `rolling_segment_recap` returns
|
||||
/// `None` and `SegmentRecapSummarizer` must fall back to the inner summarizer.
|
||||
/// The prompt must remain bounded (no panic, no over-budget, no error).
|
||||
#[tokio::test]
|
||||
async fn soft_fallback_when_archivist_absent() {
|
||||
// Use a disabled archivist (no SQLite connection).
|
||||
let disabled_hook = Arc::new(ArchivistHook::disabled());
|
||||
|
||||
let inner = RecordingSummarizer::new();
|
||||
let recap_summ = SegmentRecapSummarizer::new(
|
||||
disabled_hook,
|
||||
"no-session".to_string(),
|
||||
inner.clone() as Arc<dyn Summarizer>,
|
||||
)
|
||||
.with_keep_recent(1);
|
||||
|
||||
let mut history = vec![user("evict me"), user("evict me too"), user("keep me")];
|
||||
|
||||
let stats = recap_summ
|
||||
.summarize(&mut history, "test-model")
|
||||
.await
|
||||
.expect("summarize must not return Err — soft-fallback guarantees boundedness");
|
||||
|
||||
// Inner summarizer must have been called exactly once.
|
||||
assert_eq!(
|
||||
inner.call_count(),
|
||||
1,
|
||||
"Inner summarizer must be called when archivist is absent"
|
||||
);
|
||||
|
||||
// History must be bounded (inner replaced it).
|
||||
assert_eq!(
|
||||
history.len(),
|
||||
1,
|
||||
"Inner summarizer must have reduced the history"
|
||||
);
|
||||
match &history[0] {
|
||||
ConversationMessage::Chat(m) => {
|
||||
assert_eq!(m.content, "inner fallback summary");
|
||||
}
|
||||
_ => panic!("Expected system summary from inner summarizer"),
|
||||
}
|
||||
|
||||
// Stats must be valid (not zeroes — inner mock returns non-zero values).
|
||||
assert_eq!(stats.approx_tokens_freed, 500);
|
||||
}
|
||||
|
||||
// ── Test 4: soft-fallback when LLM stub fails ────────────────────────────────
|
||||
|
||||
/// Tier 3 — the bookend heuristic stub must NEVER become live compaction
|
||||
/// text. Here there is no chat provider configured, so `summarize_entries`
|
||||
/// produces only `segments::fallback_summary` (bookend stub) with
|
||||
/// `produced_by_llm = false`. `rolling_segment_recap` must therefore return
|
||||
/// `None`, and `SegmentRecapSummarizer` falls back to the inner summarizer.
|
||||
/// (Option A: only a genuine summariser-produced recap may be compaction.)
|
||||
#[tokio::test]
|
||||
async fn bookend_stub_never_becomes_compaction_falls_back_to_inner() {
|
||||
let conn = setup_conn();
|
||||
// `ArchivistHook::new` leaves `chat_provider = None` → summarize_entries
|
||||
// takes the no-provider branch → bookend stub, produced_by_llm = false.
|
||||
let hook = Arc::new(ArchivistHook::new(conn.clone(), true));
|
||||
|
||||
let session = "p15-bookend-stub-none";
|
||||
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Hello, what is Rust?".into(),
|
||||
assistant_response: "Rust is a systems language.".into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Only the bookend stub is available → MUST be None.
|
||||
let recap = hook.rolling_segment_recap(session).await;
|
||||
assert!(
|
||||
recap.is_none(),
|
||||
"rolling_segment_recap must return None when only the bookend stub \
|
||||
is available — the stub must never be live compaction text"
|
||||
);
|
||||
|
||||
let inner = RecordingSummarizer::new();
|
||||
let recap_summ = SegmentRecapSummarizer::new(
|
||||
Arc::clone(&hook),
|
||||
session.to_string(),
|
||||
inner.clone() as Arc<dyn Summarizer>,
|
||||
)
|
||||
.with_keep_recent(1);
|
||||
|
||||
let mut history = vec![user("evict"), user("evict2"), user("keep")];
|
||||
recap_summ
|
||||
.summarize(&mut history, "test-model")
|
||||
.await
|
||||
.expect("must not panic — inner summarizer is the safety net");
|
||||
|
||||
assert_eq!(
|
||||
inner.call_count(),
|
||||
1,
|
||||
"Inner summarizer must run when only the bookend stub is available \
|
||||
(stub must never be live compaction text)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Tier 2 — when a chat provider IS configured but its call fails,
|
||||
/// `LlmSummariser`'s soft-fallback yields an *inert clipped-content* recap
|
||||
/// (the real conversation text, truncated — NOT the bookend stub). Per
|
||||
/// Option A this is acceptable as live compaction text (real content,
|
||||
/// strictly better than no compaction), so `rolling_segment_recap` returns
|
||||
/// `Some` and the inner summarizer is NOT used.
|
||||
#[tokio::test]
|
||||
async fn failing_provider_yields_inert_clipped_recap_used_as_compaction() {
|
||||
let conn = setup_conn();
|
||||
let hook = Arc::new(ArchivistHook::new_with_stubs(
|
||||
conn.clone(),
|
||||
Arc::new(FailingChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
));
|
||||
|
||||
let session = "p15-inert-clipped-kept";
|
||||
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: "Explain ownership in Rust in detail.".into(),
|
||||
assistant_response: "Ownership means each value has a single owner; \
|
||||
borrows are checked at compile time."
|
||||
.into(),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Provider present but failing → LlmSummariser inert fallback → real
|
||||
// clipped content (not the bookend stub) → Some, treated as usable.
|
||||
let recap = hook.rolling_segment_recap(session).await;
|
||||
assert!(
|
||||
recap.is_some(),
|
||||
"Inert clipped-content recap (real text) is acceptable compaction \
|
||||
text — must be Some, not None"
|
||||
);
|
||||
|
||||
let inner = RecordingSummarizer::new();
|
||||
let recap_summ = SegmentRecapSummarizer::new(
|
||||
Arc::clone(&hook),
|
||||
session.to_string(),
|
||||
inner.clone() as Arc<dyn Summarizer>,
|
||||
)
|
||||
.with_keep_recent(1);
|
||||
|
||||
let mut history = vec![user("evict"), user("evict2"), user("keep")];
|
||||
recap_summ
|
||||
.summarize(&mut history, "test-model")
|
||||
.await
|
||||
.expect("must not panic");
|
||||
|
||||
assert_eq!(
|
||||
inner.call_count(),
|
||||
0,
|
||||
"Inner summarizer must NOT run when an inert clipped-content recap \
|
||||
is available (real content, better than no compaction)"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Test 5: flag off → inner summarizer runs, Phase 1.5 absent ───────────────
|
||||
|
||||
/// When `unified_compaction_enabled = false`, the `SegmentRecapSummarizer`
|
||||
/// is NOT instantiated — the builder uses `ProviderSummarizer` directly.
|
||||
/// This test verifies the flag-off path via the `RecordingSummarizer` fallback:
|
||||
/// a `SegmentRecapSummarizer` with an otherwise-healthy archivist must still
|
||||
/// fall through to the inner summarizer if the session has no entries yet
|
||||
/// (no open segment entries → recap = None → inner).
|
||||
#[tokio::test]
|
||||
async fn no_entries_returns_none_and_inner_summarizer_fires() {
|
||||
let conn = setup_conn();
|
||||
// Archivist with no turns written — no open segment, no entries.
|
||||
let hook = Arc::new(ArchivistHook::new_with_stubs(
|
||||
conn.clone(),
|
||||
Arc::new(StubChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
));
|
||||
|
||||
let inner = RecordingSummarizer::new();
|
||||
let recap_summ = SegmentRecapSummarizer::new(
|
||||
Arc::clone(&hook),
|
||||
"empty-session".to_string(),
|
||||
inner.clone() as Arc<dyn Summarizer>,
|
||||
)
|
||||
.with_keep_recent(1);
|
||||
|
||||
let mut history = vec![user("a"), user("b"), user("c")];
|
||||
let _stats = recap_summ
|
||||
.summarize(&mut history, "test-model")
|
||||
.await
|
||||
.expect("must not error");
|
||||
|
||||
// With no open segment, rolling_segment_recap returns None →
|
||||
// inner summarizer fires.
|
||||
assert_eq!(
|
||||
inner.call_count(),
|
||||
1,
|
||||
"Inner summarizer must run when session has no open segment entries"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Test 6: Phase 1 regression guard ─────────────────────────────────────────
|
||||
|
||||
/// `on_segment_closed` (finalize path) still works after Phase 1.5 refactor:
|
||||
/// the shared `summarize_entries` helper must produce the same result, and
|
||||
/// `segment_set_summary` + embedding must still fire at close time.
|
||||
#[tokio::test]
|
||||
async fn phase1_finalize_path_still_persists_summary_and_embedding() {
|
||||
let conn = setup_conn();
|
||||
let hook = Arc::new(ArchivistHook::new_with_stubs(
|
||||
conn.clone(),
|
||||
Arc::new(StubChatProvider),
|
||||
Arc::new(StubEmbedder),
|
||||
));
|
||||
|
||||
let session = "p15-finalize-regression";
|
||||
|
||||
// Two turns in the first segment.
|
||||
for i in 1..=2u64 {
|
||||
hook.on_turn_complete(&TurnContext {
|
||||
user_message: format!("Rust turn {i}"),
|
||||
assistant_response: format!("Answer {i}"),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 50,
|
||||
session_id: Some(session.into()),
|
||||
iteration_count: i as usize,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Force-flush to trigger on_segment_closed (finalize path).
|
||||
hook.flush_open_segment(session).await;
|
||||
|
||||
// The closed segment must have a summary.
|
||||
let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap();
|
||||
let flushed = all_segs
|
||||
.iter()
|
||||
.find(|s| {
|
||||
s.session_id == session && s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false)
|
||||
})
|
||||
.expect("Expected flushed segment to have a non-empty summary after Phase 1.5 refactor");
|
||||
|
||||
let summary = flushed.summary.as_ref().unwrap();
|
||||
// The stub always returns "rolling recap: discussed memory safety in Rust".
|
||||
assert!(
|
||||
summary.contains("rolling recap"),
|
||||
"Expected summary to contain stub recap text after finalize; got: {summary:?}"
|
||||
);
|
||||
|
||||
// Embedding must also still be written (finalize-only invariant).
|
||||
let embedding =
|
||||
seg::segment_embedding_get(&conn, &flushed.segment_id, "stub-embedder-v1").unwrap();
|
||||
assert!(
|
||||
embedding.is_some(),
|
||||
"Expected finalize-time embedding to still be written after Phase 1.5 refactor"
|
||||
);
|
||||
}
|
||||
@@ -245,7 +245,11 @@ impl Summarizer for ProviderSummarizer {
|
||||
/// head length. Returns 0 when the adjustment would consume the entire
|
||||
/// history, meaning there is nothing we can safely summarize without
|
||||
/// breaking the API invariant.
|
||||
fn snap_split_forward(history: &[ConversationMessage], proposed_head: usize) -> usize {
|
||||
///
|
||||
/// Exported as `pub(super)` so sibling modules (e.g.
|
||||
/// `segment_recap_summarizer`) can reuse the same invariant instead of
|
||||
/// maintaining a separate copy.
|
||||
pub(super) fn snap_split_forward(history: &[ConversationMessage], proposed_head: usize) -> usize {
|
||||
let mut head = proposed_head.min(history.len());
|
||||
// If the message immediately *before* the split is an
|
||||
// AssistantToolCalls and the message *at* the split is its
|
||||
|
||||
@@ -13,6 +13,7 @@ pub mod ops;
|
||||
pub mod rpc_models;
|
||||
pub mod safety;
|
||||
pub mod schemas;
|
||||
pub mod stm_recall;
|
||||
pub mod store;
|
||||
pub mod sync_status;
|
||||
pub mod tool_memory;
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
//! Tunable constants for STM recall — the STM/LTM knobs.
|
||||
//!
|
||||
//! Split out of `mod.rs` per the repo's "light `mod.rs`" rule (CLAUDE.md):
|
||||
//! `mod.rs` stays export-focused; operational/tunable values live here and
|
||||
//! are re-exported (`pub use constants::*;`) so existing `super::CONST`
|
||||
//! references in `recall.rs` / `tool.rs` keep working unchanged.
|
||||
|
||||
/// STM recency window in days. Segments or episodic entries older than this
|
||||
/// are excluded — they belong in LTM (the memory tree).
|
||||
pub const RECENCY_WINDOW_DAYS: f64 = 14.0;
|
||||
|
||||
/// Hard cap on segments loaded for vector search (Arm 2).
|
||||
/// Keeps the brute-force cosine pass bounded at single-user scale.
|
||||
pub const RECENCY_WINDOW_MAX_SEGMENTS: usize = 100;
|
||||
|
||||
/// Cosine similarity gate for Arm 2 (segment recaps).
|
||||
/// Below this threshold a recap is excluded regardless of recency.
|
||||
/// Range: [0.0, 1.0]; 0.65 is "medium gate" — confident topical overlap.
|
||||
pub const COSINE_GATE: f32 = 0.65;
|
||||
|
||||
/// Maximum segment recaps to include in the output block.
|
||||
pub const MAX_SEGMENT_RECAPS: usize = 5;
|
||||
|
||||
/// Maximum raw episodic turns to include in the output block.
|
||||
pub const MAX_EPISODIC_TURNS: usize = 5;
|
||||
|
||||
/// Approximate token budget for the entire STM block (chars / 4 ≈ tokens).
|
||||
/// ~1500 tokens × 4 chars/token = 6000 chars.
|
||||
pub const TOKEN_BUDGET: usize = 6_000;
|
||||
|
||||
/// How many FTS5 candidates to fetch before applying the high-precision gate.
|
||||
/// The gate is: only strong keyword matches survive — FTS5 rank threshold is
|
||||
/// applied at the DB level via LIMIT; we over-fetch slightly and let dedup
|
||||
/// finish trimming.
|
||||
pub const FTS5_LIMIT: usize = 20;
|
||||
@@ -0,0 +1,47 @@
|
||||
//! Phase 3 — Bounded cross-thread STM recall.
|
||||
//!
|
||||
//! Assembles a bounded, recency-weighted context block from two arms:
|
||||
//!
|
||||
//! - **Arm 1** — FTS5 over not-yet-compacted recent episodic entries from
|
||||
//! OTHER sessions. Reuses [`crate::openhuman::memory::store::fts5::episodic_cross_session_search`].
|
||||
//! When no user query is available (preemptive/session-start case), falls back to
|
||||
//! a recency selection of recent non-current-session episodic turns.
|
||||
//!
|
||||
//! - **Arm 2** — Brute-force cosine nearest-neighbour over `segment_embeddings`
|
||||
//! (per-model table from Phase 0+1). Single-user scale: no ANN index, no new
|
||||
//! deps. Loads candidate vectors, computes cosine, top-k. Filters by
|
||||
//! `model_signature` and excludes the current `session_id`.
|
||||
//!
|
||||
//! **Merge → dedup → bound:**
|
||||
//! - Dedup: if a segment recap and its raw episodic rows (within
|
||||
//! `start_episodic_id..=end_episodic_id`) both appear, prefer the recap and
|
||||
//! drop the overlapping episodic hits.
|
||||
//! - Recency-weight: scored by `updated_at` timestamp proximity.
|
||||
//! - Hard-cap: tunable token budget ([`TOKEN_BUDGET`]) and top-k bounds
|
||||
//! ([`MAX_SEGMENT_RECAPS`] + [`MAX_EPISODIC_TURNS`]).
|
||||
//!
|
||||
//! ## Tunable consts (all in this file, all documented)
|
||||
//! - [`RECENCY_WINDOW_DAYS`] — how many days back to search (STM/LTM boundary)
|
||||
//! - [`RECENCY_WINDOW_MAX_SEGMENTS`] — max segments to load for vector search
|
||||
//! - [`COSINE_GATE`] — minimum similarity for Arm 2 (medium gate)
|
||||
//! - [`MAX_SEGMENT_RECAPS`] — top-k segment recaps to include
|
||||
//! - [`MAX_EPISODIC_TURNS`] — max raw episodic turns to include
|
||||
//! - [`TOKEN_BUDGET`] — hard token budget (chars / 4 approx)
|
||||
//! - [`FTS5_LIMIT`] — how many FTS5 candidates to fetch before gating
|
||||
//!
|
||||
//! ## Scope boundary
|
||||
//! Does NOT traverse `tree::*` (`SummaryNode`, `memory_tree_*`). The memory
|
||||
//! tree is LTM; this module is strictly STM (recent episodic + segment layer).
|
||||
|
||||
pub mod recall;
|
||||
pub mod tool;
|
||||
|
||||
pub use recall::{stm_recall, StmRecallBlock, StmRecallOpts};
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Tunable constants — the STM/LTM knobs. Moved to `constants.rs` per the
|
||||
// repo's "light mod.rs" rule; re-exported so `super::CONST` keeps resolving.
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
mod constants;
|
||||
pub use constants::*;
|
||||
@@ -0,0 +1,629 @@
|
||||
//! Core STM recall logic — two-arm hybrid retrieval.
|
||||
//!
|
||||
//! Arm 1: FTS5 episodic cross-session search (keyword, high-precision gate).
|
||||
//! Arm 2: Brute-force cosine over `segment_embeddings` (vector, medium gate).
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::openhuman::memory::store::fts5;
|
||||
use crate::openhuman::memory::store::fts5::EpisodicEntry;
|
||||
|
||||
use super::{
|
||||
COSINE_GATE, FTS5_LIMIT, MAX_EPISODIC_TURNS, MAX_SEGMENT_RECAPS, RECENCY_WINDOW_DAYS,
|
||||
RECENCY_WINDOW_MAX_SEGMENTS, TOKEN_BUDGET,
|
||||
};
|
||||
|
||||
/// A single item in the assembled STM recall block.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum StmItem {
|
||||
/// A segment-level recap (compacted, from `segment_embeddings` arm).
|
||||
SegmentRecap {
|
||||
segment_id: String,
|
||||
session_id: String,
|
||||
summary: String,
|
||||
/// `start_episodic_id..=end_episodic_id` span — used for dedup.
|
||||
start_episodic_id: i64,
|
||||
end_episodic_id: Option<i64>,
|
||||
updated_at: f64,
|
||||
cosine: f32,
|
||||
},
|
||||
/// A raw episodic turn from Arm 1 (FTS5 keyword match).
|
||||
EpisodicTurn {
|
||||
id: Option<i64>,
|
||||
session_id: String,
|
||||
timestamp: f64,
|
||||
role: String,
|
||||
content: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl StmItem {
|
||||
fn timestamp(&self) -> f64 {
|
||||
match self {
|
||||
Self::SegmentRecap { updated_at, .. } => *updated_at,
|
||||
Self::EpisodicTurn { timestamp, .. } => *timestamp,
|
||||
}
|
||||
}
|
||||
|
||||
fn approx_chars(&self) -> usize {
|
||||
match self {
|
||||
Self::SegmentRecap { summary, .. } => summary.len() + 60,
|
||||
Self::EpisodicTurn { content, role, .. } => content.len() + role.len() + 20,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Options for a single STM recall pass.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct StmRecallOpts<'a> {
|
||||
/// The active session to exclude from all results.
|
||||
pub exclude_session: &'a str,
|
||||
/// Optional query text for Arm 1 (FTS5) and Arm 2 (embed).
|
||||
/// When `None`, Arm 1 falls back to a recency selection and Arm 2 is skipped.
|
||||
pub query: Option<&'a str>,
|
||||
/// Model signature to filter segment embeddings (e.g. `"cloud:voyage-3:1024"`).
|
||||
/// When `None`, any model signature is accepted (weaker but still useful).
|
||||
pub model_signature: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// The assembled STM recall block.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct StmRecallBlock {
|
||||
/// Deduplicated, recency-weighted items, bounded to [`TOKEN_BUDGET`].
|
||||
pub items: Vec<StmItem>,
|
||||
/// Number of items dropped due to token-budget exhaustion.
|
||||
pub dropped_budget: usize,
|
||||
/// Number of episodic hits dropped because they fell inside a segment span (dedup).
|
||||
pub dropped_dedup: usize,
|
||||
/// Cosine arm — items retrieved before gating.
|
||||
pub cosine_candidates: usize,
|
||||
/// FTS5 arm — items retrieved before gating.
|
||||
pub fts5_candidates: usize,
|
||||
}
|
||||
|
||||
impl StmRecallBlock {
|
||||
/// `true` when the block has no usable content.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.items.is_empty()
|
||||
}
|
||||
|
||||
/// Render the block into a markdown string suitable for injection into the
|
||||
/// system prompt or a user-turn context block.
|
||||
pub fn render(&self) -> String {
|
||||
if self.items.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
let mut out = String::from("## Recent context from other conversations\n\n");
|
||||
out.push_str(
|
||||
"The following snippets are from previous conversations in other chat threads. \
|
||||
They are provided for continuity — reference them when they are relevant to \
|
||||
the current request, but do not surface them unless asked.\n\n",
|
||||
);
|
||||
|
||||
let mut seg_count = 0usize;
|
||||
let mut ep_count = 0usize;
|
||||
|
||||
for item in &self.items {
|
||||
match item {
|
||||
StmItem::SegmentRecap {
|
||||
session_id,
|
||||
summary,
|
||||
updated_at,
|
||||
..
|
||||
} => {
|
||||
seg_count += 1;
|
||||
let age_days = age_days_from_ts(*updated_at);
|
||||
out.push_str(&format!(
|
||||
"**Conversation recap** (thread `{session_id}`, ~{age_days:.0} days ago):\n{summary}\n\n"
|
||||
));
|
||||
}
|
||||
StmItem::EpisodicTurn {
|
||||
session_id,
|
||||
timestamp,
|
||||
role,
|
||||
content,
|
||||
..
|
||||
} => {
|
||||
ep_count += 1;
|
||||
let age_days = age_days_from_ts(*timestamp);
|
||||
out.push_str(&format!(
|
||||
"**{role}** (thread `{session_id}`, ~{age_days:.0} days ago): {content}\n\n"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall] rendered block: {} recaps + {} episodic turns, {} dropped_dedup, {} dropped_budget",
|
||||
seg_count,
|
||||
ep_count,
|
||||
self.dropped_dedup,
|
||||
self.dropped_budget
|
||||
);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
fn age_days_from_ts(ts: f64) -> f64 {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs_f64();
|
||||
(now - ts).max(0.0) / 86_400.0
|
||||
}
|
||||
|
||||
fn now_secs() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs_f64()
|
||||
}
|
||||
|
||||
/// Brute-force cosine similarity between two equal-length float slices.
|
||||
/// Returns 0.0 when either vector has zero magnitude (zero-padded / inert embedder).
|
||||
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
if a.len() != b.len() || a.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut dot = 0.0_f32;
|
||||
let mut norm_a = 0.0_f32;
|
||||
let mut norm_b = 0.0_f32;
|
||||
for (x, y) in a.iter().zip(b.iter()) {
|
||||
dot += x * y;
|
||||
norm_a += x * x;
|
||||
norm_b += y * y;
|
||||
}
|
||||
let denom = norm_a.sqrt() * norm_b.sqrt();
|
||||
if denom < f32::EPSILON {
|
||||
0.0
|
||||
} else {
|
||||
(dot / denom).clamp(-1.0, 1.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a raw BLOB from `segment_embeddings.vector` into `Vec<f32>`.
|
||||
///
|
||||
/// A well-formed embedding blob is a whole number of little-endian `f32`s
|
||||
/// (length a multiple of 4). A non-multiple-of-4 length means the blob is
|
||||
/// truncated/corrupt; silently dropping the trailing bytes would yield a
|
||||
/// wrong-length vector, so we reject it (empty → cosine treats it as a
|
||||
/// non-match, which is the safe outcome).
|
||||
fn decode_vector_blob(bytes: &[u8]) -> Vec<f32> {
|
||||
if bytes.len() % 4 != 0 {
|
||||
tracing::warn!(
|
||||
"[stm_recall] decode_vector_blob: blob length {} is not a multiple of 4 — \
|
||||
discarding malformed embedding",
|
||||
bytes.len()
|
||||
);
|
||||
return Vec::new();
|
||||
}
|
||||
bytes
|
||||
.chunks_exact(4)
|
||||
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ── Arm 2: vector search over segment_embeddings ──────────────────────────────
|
||||
|
||||
/// A row pulled from `segment_embeddings` joined with `conversation_segments`.
|
||||
struct SegmentEmbeddingRow {
|
||||
segment_id: String,
|
||||
session_id: String,
|
||||
summary: Option<String>,
|
||||
start_episodic_id: i64,
|
||||
end_episodic_id: Option<i64>,
|
||||
updated_at: f64,
|
||||
vector: Vec<f32>,
|
||||
model_signature: String,
|
||||
}
|
||||
|
||||
/// Load candidate segment embeddings for the cosine pass.
|
||||
///
|
||||
/// Applies:
|
||||
/// - Recency window (last [`RECENCY_WINDOW_DAYS`] days)
|
||||
/// - Exclude current `session_id`
|
||||
/// - Optional `model_signature` filter
|
||||
/// - Capped at [`RECENCY_WINDOW_MAX_SEGMENTS`] rows
|
||||
fn load_segment_embedding_candidates(
|
||||
conn: &Arc<Mutex<Connection>>,
|
||||
exclude_session: &str,
|
||||
model_signature: Option<&str>,
|
||||
cutoff_ts: f64,
|
||||
) -> anyhow::Result<Vec<SegmentEmbeddingRow>> {
|
||||
let conn = conn.lock();
|
||||
|
||||
// We join `segment_embeddings` (se) with `conversation_segments` (cs)
|
||||
// to get the session_id, summary, episodic span, and recency.
|
||||
// Filter: cs.session_id != exclude_session, cs.updated_at >= cutoff_ts,
|
||||
// cs.summary IS NOT NULL (only summarised segments have useful recaps),
|
||||
// optionally se.model_signature = ?
|
||||
let rows: Vec<SegmentEmbeddingRow> = if let Some(sig) = model_signature {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT se.segment_id, cs.session_id, cs.summary,
|
||||
cs.start_episodic_id, cs.end_episodic_id,
|
||||
cs.updated_at, se.vector, se.model_signature
|
||||
FROM segment_embeddings AS se
|
||||
JOIN conversation_segments AS cs ON se.segment_id = cs.segment_id
|
||||
WHERE cs.session_id != ?1
|
||||
AND cs.updated_at >= ?2
|
||||
AND cs.summary IS NOT NULL
|
||||
AND se.model_signature = ?3
|
||||
ORDER BY cs.updated_at DESC
|
||||
LIMIT ?4",
|
||||
)?;
|
||||
let collected: Vec<SegmentEmbeddingRow> = stmt
|
||||
.query_map(
|
||||
params![
|
||||
exclude_session,
|
||||
cutoff_ts,
|
||||
sig,
|
||||
RECENCY_WINDOW_MAX_SEGMENTS as i64
|
||||
],
|
||||
|row| {
|
||||
let vector_bytes: Vec<u8> = row.get(6)?;
|
||||
Ok(SegmentEmbeddingRow {
|
||||
segment_id: row.get(0)?,
|
||||
session_id: row.get(1)?,
|
||||
summary: row.get(2)?,
|
||||
start_episodic_id: row.get(3)?,
|
||||
end_episodic_id: row.get(4)?,
|
||||
updated_at: row.get(5)?,
|
||||
vector: decode_vector_blob(&vector_bytes),
|
||||
model_signature: row.get(7)?,
|
||||
})
|
||||
},
|
||||
)?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
collected
|
||||
} else {
|
||||
// No model filter — accept any model (weaker but useful when no sig available).
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT se.segment_id, cs.session_id, cs.summary,
|
||||
cs.start_episodic_id, cs.end_episodic_id,
|
||||
cs.updated_at, se.vector, se.model_signature
|
||||
FROM segment_embeddings AS se
|
||||
JOIN conversation_segments AS cs ON se.segment_id = cs.segment_id
|
||||
WHERE cs.session_id != ?1
|
||||
AND cs.updated_at >= ?2
|
||||
AND cs.summary IS NOT NULL
|
||||
ORDER BY cs.updated_at DESC
|
||||
LIMIT ?3",
|
||||
)?;
|
||||
let collected: Vec<SegmentEmbeddingRow> = stmt
|
||||
.query_map(
|
||||
params![
|
||||
exclude_session,
|
||||
cutoff_ts,
|
||||
RECENCY_WINDOW_MAX_SEGMENTS as i64
|
||||
],
|
||||
|row| {
|
||||
let vector_bytes: Vec<u8> = row.get(6)?;
|
||||
Ok(SegmentEmbeddingRow {
|
||||
segment_id: row.get(0)?,
|
||||
session_id: row.get(1)?,
|
||||
summary: row.get(2)?,
|
||||
start_episodic_id: row.get(3)?,
|
||||
end_episodic_id: row.get(4)?,
|
||||
updated_at: row.get(5)?,
|
||||
vector: decode_vector_blob(&vector_bytes),
|
||||
model_signature: row.get(7)?,
|
||||
})
|
||||
},
|
||||
)?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
collected
|
||||
};
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm2: loaded {} segment embedding candidates (model_sig={:?})",
|
||||
rows.len(),
|
||||
model_signature
|
||||
);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Load recent episodic turns from other sessions (recency fallback for Arm 1
|
||||
/// when no query is available).
|
||||
fn load_recent_episodic_other_sessions(
|
||||
conn: &Arc<Mutex<Connection>>,
|
||||
exclude_session: &str,
|
||||
cutoff_ts: f64,
|
||||
limit: usize,
|
||||
) -> anyhow::Result<Vec<EpisodicEntry>> {
|
||||
let conn = conn.lock();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT id, session_id, timestamp, role, content, lesson, tool_calls_json, cost_microdollars
|
||||
FROM episodic_log
|
||||
WHERE session_id != ?1
|
||||
AND timestamp >= ?2
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT ?3",
|
||||
)?;
|
||||
let rows = stmt
|
||||
.query_map(params![exclude_session, cutoff_ts, limit as i64], |row| {
|
||||
Ok(EpisodicEntry {
|
||||
id: row.get(0)?,
|
||||
session_id: row.get(1)?,
|
||||
timestamp: row.get(2)?,
|
||||
role: row.get(3)?,
|
||||
content: row.get(4)?,
|
||||
lesson: row.get(5)?,
|
||||
tool_calls_json: row.get(6)?,
|
||||
cost_microdollars: row.get::<_, i64>(7)? as u64,
|
||||
})
|
||||
})?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm1 recency fallback: loaded {} episodic turns from other sessions",
|
||||
rows.len()
|
||||
);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Main entry point: run the two-arm STM recall and return an assembled block.
|
||||
///
|
||||
/// When `opts.query` is `None` (preemptive/session-start case):
|
||||
/// - Arm 1 uses a recency selection of recent other-session episodic turns.
|
||||
/// - Arm 2 is skipped (no query vector to embed; caller may pre-compute one).
|
||||
///
|
||||
/// When `opts.query` is `Some(q)`:
|
||||
/// - Arm 1 runs FTS5 cross-session search with `q` and the high-precision gate.
|
||||
/// - Arm 2 runs brute-force cosine against any pre-computed `query_embedding`.
|
||||
///
|
||||
/// `query_embedding` is the embedding of `opts.query` (when provided).
|
||||
/// The caller is responsible for producing it (avoids a blocking embed call here).
|
||||
/// If `None`, Arm 2 is skipped regardless of `opts.query`.
|
||||
pub fn stm_recall(
|
||||
conn: &Arc<Mutex<Connection>>,
|
||||
opts: &StmRecallOpts<'_>,
|
||||
query_embedding: Option<&[f32]>,
|
||||
) -> anyhow::Result<StmRecallBlock> {
|
||||
let cutoff_ts = now_secs() - RECENCY_WINDOW_DAYS * 86_400.0;
|
||||
let mut block = StmRecallBlock::default();
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall] starting recall exclude_session={} has_query={} has_embedding={} recency_days={}",
|
||||
opts.exclude_session,
|
||||
opts.query.is_some(),
|
||||
query_embedding.is_some(),
|
||||
RECENCY_WINDOW_DAYS
|
||||
);
|
||||
|
||||
// ── Arm 2: vector search over segment_embeddings ──────────────────────────
|
||||
let mut segment_items: Vec<StmItem> = Vec::new();
|
||||
let mut segment_spans: Vec<(i64, Option<i64>)> = Vec::new(); // for dedup
|
||||
|
||||
if let Some(q_emb) = query_embedding {
|
||||
if !q_emb.is_empty() {
|
||||
let candidates = load_segment_embedding_candidates(
|
||||
conn,
|
||||
opts.exclude_session,
|
||||
opts.model_signature,
|
||||
cutoff_ts,
|
||||
)?;
|
||||
block.cosine_candidates = candidates.len();
|
||||
|
||||
let mut scored: Vec<(f32, SegmentEmbeddingRow)> = candidates
|
||||
.into_iter()
|
||||
.filter_map(|row| {
|
||||
if row.vector.is_empty() {
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm2: skipping segment {} — zero-length vector (inert embedder?)",
|
||||
row.segment_id
|
||||
);
|
||||
return None;
|
||||
}
|
||||
let cos = cosine_similarity(q_emb, &row.vector);
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm2: segment={} session={} cosine={:.3} gate={}",
|
||||
row.segment_id,
|
||||
row.session_id,
|
||||
cos,
|
||||
COSINE_GATE
|
||||
);
|
||||
if cos >= COSINE_GATE {
|
||||
Some((cos, row))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Sort descending by cosine then recency
|
||||
scored.sort_by(|a, b| {
|
||||
b.0.partial_cmp(&a.0)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(
|
||||
b.1.updated_at
|
||||
.partial_cmp(&a.1.updated_at)
|
||||
.unwrap_or(std::cmp::Ordering::Equal),
|
||||
)
|
||||
});
|
||||
|
||||
for (cos, row) in scored.into_iter().take(MAX_SEGMENT_RECAPS) {
|
||||
let summary = match row.summary {
|
||||
Some(ref s) if !s.trim().is_empty() => s.clone(),
|
||||
_ => continue,
|
||||
};
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm2: accepting segment={} session={} cosine={:.3} model={}",
|
||||
row.segment_id,
|
||||
row.session_id,
|
||||
cos,
|
||||
row.model_signature
|
||||
);
|
||||
segment_spans.push((row.start_episodic_id, row.end_episodic_id));
|
||||
segment_items.push(StmItem::SegmentRecap {
|
||||
segment_id: row.segment_id,
|
||||
session_id: row.session_id,
|
||||
summary,
|
||||
start_episodic_id: row.start_episodic_id,
|
||||
end_episodic_id: row.end_episodic_id,
|
||||
updated_at: row.updated_at,
|
||||
cosine: cos,
|
||||
});
|
||||
}
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm2: {} recaps accepted after cosine gate",
|
||||
segment_items.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Arm 1: FTS5 or recency episodic ──────────────────────────────────────
|
||||
let mut episodic_items: Vec<StmItem> = Vec::new();
|
||||
|
||||
let raw_episodic: Vec<EpisodicEntry> = if let Some(query) = opts.query {
|
||||
// Keyword search path
|
||||
let hits = fts5::episodic_cross_session_search(
|
||||
conn,
|
||||
query,
|
||||
FTS5_LIMIT,
|
||||
Some(opts.exclude_session),
|
||||
)?;
|
||||
block.fts5_candidates = hits.len();
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm1 FTS5: {} hits for query (exclude={})",
|
||||
hits.len(),
|
||||
opts.exclude_session
|
||||
);
|
||||
hits
|
||||
} else {
|
||||
// Recency fallback — no query; pull most-recent turns from other sessions
|
||||
let hits =
|
||||
load_recent_episodic_other_sessions(conn, opts.exclude_session, cutoff_ts, FTS5_LIMIT)?;
|
||||
block.fts5_candidates = hits.len();
|
||||
hits
|
||||
};
|
||||
|
||||
// Dedup: drop episodic rows whose ID falls within any accepted segment's span.
|
||||
for entry in raw_episodic {
|
||||
let entry_id = entry.id.unwrap_or(-1);
|
||||
|
||||
// Check if this episodic entry is covered by any accepted segment span.
|
||||
let covered = segment_spans
|
||||
.iter()
|
||||
.any(|(start, end)| entry_id >= *start && end.map_or(false, |e| entry_id <= e));
|
||||
|
||||
if covered {
|
||||
block.dropped_dedup += 1;
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm1: dropping episodic id={} — covered by segment span",
|
||||
entry_id
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Recency window applies consistently — including keyword/FTS mode.
|
||||
// FTS5 keyword search is not time-bounded, so without this an
|
||||
// older-than-window episodic hit could leak into STM. The recency
|
||||
// window IS the STM/LTM boundary and must always hold.
|
||||
if entry.timestamp < cutoff_ts {
|
||||
continue;
|
||||
}
|
||||
|
||||
episodic_items.push(StmItem::EpisodicTurn {
|
||||
id: entry.id,
|
||||
session_id: entry.session_id,
|
||||
timestamp: entry.timestamp,
|
||||
role: entry.role,
|
||||
content: entry.content,
|
||||
});
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall] arm1: {} episodic items after dedup (dropped_dedup={})",
|
||||
episodic_items.len(),
|
||||
block.dropped_dedup
|
||||
);
|
||||
|
||||
// ── Merge + cap ───────────────────────────────────────────────────────────
|
||||
// Recency-weight: sort each arm descending by timestamp. Interleave by picking
|
||||
// the most-recent item across both arms.
|
||||
segment_items.sort_by(|a, b| {
|
||||
b.timestamp()
|
||||
.partial_cmp(&a.timestamp())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
episodic_items.sort_by(|a, b| {
|
||||
b.timestamp()
|
||||
.partial_cmp(&a.timestamp())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Apply top-k caps before interleave
|
||||
let seg_capped: Vec<StmItem> = segment_items.into_iter().take(MAX_SEGMENT_RECAPS).collect();
|
||||
let ep_capped: Vec<StmItem> = episodic_items
|
||||
.into_iter()
|
||||
.take(MAX_EPISODIC_TURNS)
|
||||
.collect();
|
||||
|
||||
// Interleave: recency-first merge
|
||||
let mut seg_iter = seg_capped.into_iter().peekable();
|
||||
let mut ep_iter = ep_capped.into_iter().peekable();
|
||||
let mut merged: Vec<StmItem> = Vec::new();
|
||||
|
||||
loop {
|
||||
match (seg_iter.peek(), ep_iter.peek()) {
|
||||
(None, None) => break,
|
||||
(Some(_), None) => {
|
||||
merged.extend(seg_iter.by_ref());
|
||||
break;
|
||||
}
|
||||
(None, Some(_)) => {
|
||||
merged.extend(ep_iter.by_ref());
|
||||
break;
|
||||
}
|
||||
(Some(s), Some(e)) => {
|
||||
if s.timestamp() >= e.timestamp() {
|
||||
merged.push(seg_iter.next().expect("peek confirmed Some"));
|
||||
} else {
|
||||
merged.push(ep_iter.next().expect("peek confirmed Some"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply token budget
|
||||
let mut used_chars = 0usize;
|
||||
let mut final_items: Vec<StmItem> = Vec::new();
|
||||
let mut dropped_budget = 0usize;
|
||||
|
||||
for item in merged {
|
||||
let chars = item.approx_chars();
|
||||
if used_chars + chars > TOKEN_BUDGET {
|
||||
dropped_budget += 1;
|
||||
tracing::debug!(
|
||||
"[stm_recall] budget: dropping item (would exceed {TOKEN_BUDGET} chars)"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
used_chars += chars;
|
||||
final_items.push(item);
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall] final block: {} items, ~{} chars, {} dropped_budget, {} dropped_dedup, {} cosine_candidates, {} fts5_candidates",
|
||||
final_items.len(),
|
||||
used_chars,
|
||||
dropped_budget,
|
||||
block.dropped_dedup,
|
||||
block.cosine_candidates,
|
||||
block.fts5_candidates
|
||||
);
|
||||
|
||||
block.items = final_items;
|
||||
block.dropped_budget = dropped_budget;
|
||||
Ok(block)
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Unit tests
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "recall_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,721 @@
|
||||
//! Unit + integration tests for Phase 3 STM recall.
|
||||
|
||||
use super::*;
|
||||
use crate::openhuman::agent::harness::archivist::ArchivistHook;
|
||||
use crate::openhuman::agent::hooks::{PostTurnHook, TurnContext};
|
||||
use crate::openhuman::memory::store::events::EVENTS_INIT_SQL;
|
||||
use crate::openhuman::memory::store::fts5;
|
||||
use crate::openhuman::memory::store::profile::PROFILE_INIT_SQL;
|
||||
use crate::openhuman::memory::store::segments::SEGMENTS_INIT_SQL;
|
||||
use parking_lot::Mutex;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::sync::Arc;
|
||||
|
||||
// ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
fn setup_conn() -> Arc<Mutex<Connection>> {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
conn.execute_batch(fts5::EPISODIC_INIT_SQL).unwrap();
|
||||
conn.execute_batch(SEGMENTS_INIT_SQL).unwrap();
|
||||
conn.execute_batch(EVENTS_INIT_SQL).unwrap();
|
||||
conn.execute_batch(PROFILE_INIT_SQL).unwrap();
|
||||
Arc::new(Mutex::new(conn))
|
||||
}
|
||||
|
||||
fn now_ts() -> f64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs_f64()
|
||||
}
|
||||
|
||||
/// Insert an episodic entry with an explicit timestamp.
|
||||
fn insert_episodic(
|
||||
conn: &Arc<Mutex<Connection>>,
|
||||
session_id: &str,
|
||||
ts: f64,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> i64 {
|
||||
let c = conn.lock();
|
||||
c.execute(
|
||||
"INSERT INTO episodic_log (session_id, timestamp, role, content, lesson, tool_calls_json, cost_microdollars) VALUES (?1,?2,?3,?4,NULL,NULL,0)",
|
||||
params![session_id, ts, role, content],
|
||||
).unwrap();
|
||||
c.last_insert_rowid()
|
||||
}
|
||||
|
||||
/// Insert a segment with a summary and optional embedding.
|
||||
fn insert_segment_with_embedding(
|
||||
conn: &Arc<Mutex<Connection>>,
|
||||
segment_id: &str,
|
||||
session_id: &str,
|
||||
start_id: i64,
|
||||
end_id: i64,
|
||||
summary: &str,
|
||||
embedding: Option<Vec<f32>>,
|
||||
updated_at: f64,
|
||||
model_sig: &str,
|
||||
) {
|
||||
let c = conn.lock();
|
||||
c.execute(
|
||||
"INSERT INTO conversation_segments
|
||||
(segment_id, session_id, namespace, start_episodic_id, end_episodic_id,
|
||||
start_timestamp, end_timestamp, turn_count, summary, status, created_at, updated_at)
|
||||
VALUES (?1,?2,'global',?3,?4,?5,?5,2,?6,'summarised',?5,?5)",
|
||||
params![segment_id, session_id, start_id, end_id, updated_at, summary],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
if let Some(emb) = embedding {
|
||||
let bytes: Vec<u8> = emb.iter().flat_map(|f| f.to_le_bytes()).collect();
|
||||
let dim = emb.len() as i64;
|
||||
c.execute(
|
||||
"INSERT INTO segment_embeddings (segment_id, model_signature, vector, dim, created_at)
|
||||
VALUES (?1,?2,?3,?4,?5)",
|
||||
params![segment_id, model_sig, bytes, dim, updated_at],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// ── cosine_similarity unit tests ──────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn cosine_identical_vectors_returns_one() {
|
||||
let v = vec![1.0_f32, 0.0, 0.0];
|
||||
assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_orthogonal_vectors_returns_zero() {
|
||||
let a = vec![1.0_f32, 0.0, 0.0];
|
||||
let b = vec![0.0_f32, 1.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &b)).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_opposite_vectors_returns_minus_one() {
|
||||
let a = vec![1.0_f32, 0.0, 0.0];
|
||||
let b = vec![-1.0_f32, 0.0, 0.0];
|
||||
assert!((cosine_similarity(&a, &b) + 1.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_zero_vector_returns_zero_not_nan() {
|
||||
let a = vec![0.0_f32, 0.0, 0.0];
|
||||
let b = vec![1.0_f32, 0.0, 0.0];
|
||||
let sim = cosine_similarity(&a, &b);
|
||||
assert!(!sim.is_nan(), "cosine_similarity must not return NaN");
|
||||
assert_eq!(sim, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_mismatched_lengths_returns_zero() {
|
||||
let a = vec![1.0_f32, 0.0];
|
||||
let b = vec![1.0_f32, 0.0, 0.0];
|
||||
assert_eq!(cosine_similarity(&a, &b), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_empty_vectors_returns_zero() {
|
||||
assert_eq!(cosine_similarity(&[], &[]), 0.0);
|
||||
}
|
||||
|
||||
// ── gating threshold tests ────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn cosine_gate_const_is_reasonable() {
|
||||
// Gate must be in (0.5, 1.0) — below 0.5 lets in noise, above 0.9 is too strict.
|
||||
assert!(super::super::COSINE_GATE > 0.5 && super::super::COSINE_GATE < 0.9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn arm2_drops_below_gate_and_accepts_above() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Build query embedding — unit vector along dim 0
|
||||
let mut q_emb = vec![0.0_f32; 8];
|
||||
q_emb[0] = 1.0;
|
||||
|
||||
// High-match segment: unit vector along dim 0 (cos = 1.0)
|
||||
let mut high_emb = vec![0.0_f32; 8];
|
||||
high_emb[0] = 1.0;
|
||||
insert_episodic(&conn, "other-session", now - 100.0, "user", "seed turn");
|
||||
let id = insert_episodic(
|
||||
&conn,
|
||||
"other-session",
|
||||
now - 90.0,
|
||||
"assistant",
|
||||
"high match reply",
|
||||
);
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-high",
|
||||
"other-session",
|
||||
id - 1,
|
||||
id,
|
||||
"This conversation covered high-match topics",
|
||||
Some(high_emb),
|
||||
now - 50.0,
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
// Low-match segment: orthogonal vector (cos = 0.0 < gate)
|
||||
let mut low_emb = vec![0.0_f32; 8];
|
||||
low_emb[1] = 1.0;
|
||||
let id2 = insert_episodic(&conn, "low-session", now - 200.0, "user", "unrelated");
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-low",
|
||||
"low-session",
|
||||
id2,
|
||||
id2,
|
||||
"This is about something completely unrelated",
|
||||
Some(low_emb),
|
||||
now - 150.0,
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current-session",
|
||||
query: Some("high match"),
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, Some(&q_emb)).unwrap();
|
||||
|
||||
let recap_ids: Vec<&str> = block
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|it| {
|
||||
if let StmItem::SegmentRecap { segment_id, .. } = it {
|
||||
Some(segment_id.as_str())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert!(
|
||||
recap_ids.contains(&"seg-high"),
|
||||
"high-cosine segment must be accepted; got: {:?}",
|
||||
recap_ids
|
||||
);
|
||||
assert!(
|
||||
!recap_ids.contains(&"seg-low"),
|
||||
"low-cosine segment must be excluded by gate; got: {:?}",
|
||||
recap_ids
|
||||
);
|
||||
}
|
||||
|
||||
// ── exclude-own-session tests ─────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn exclude_own_session_arm1_fts5() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Other session — should appear
|
||||
insert_episodic(
|
||||
&conn,
|
||||
"other-sess",
|
||||
now - 100.0,
|
||||
"user",
|
||||
"Rust programming concepts",
|
||||
);
|
||||
// Current session — must be excluded
|
||||
insert_episodic(
|
||||
&conn,
|
||||
"current-sess",
|
||||
now - 50.0,
|
||||
"user",
|
||||
"Rust programming today",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current-sess",
|
||||
query: Some("Rust programming"),
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, None).unwrap();
|
||||
|
||||
for item in &block.items {
|
||||
if let StmItem::EpisodicTurn { session_id, .. } = item {
|
||||
assert_ne!(
|
||||
session_id, "current-sess",
|
||||
"arm1 must never return current session items; got session_id={session_id}"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Must see the other session
|
||||
let has_other = block.items.iter().any(
|
||||
|it| matches!(it, StmItem::EpisodicTurn { session_id, .. } if session_id == "other-sess"),
|
||||
);
|
||||
assert!(has_other, "arm1 must surface items from other sessions");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exclude_own_session_arm2_vector() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
let mut emb = vec![0.0_f32; 8];
|
||||
emb[0] = 1.0;
|
||||
|
||||
// Insert segment from "current-session" — should be excluded
|
||||
let id = insert_episodic(
|
||||
&conn,
|
||||
"current-session",
|
||||
now - 100.0,
|
||||
"user",
|
||||
"current thread",
|
||||
);
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-current",
|
||||
"current-session",
|
||||
id,
|
||||
id,
|
||||
"Current session recap",
|
||||
Some(emb.clone()),
|
||||
now - 50.0,
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
// Segment from another session — should appear
|
||||
let id2 = insert_episodic(&conn, "other-session", now - 200.0, "user", "other thread");
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-other",
|
||||
"other-session",
|
||||
id2,
|
||||
id2,
|
||||
"Other session recap",
|
||||
Some(emb.clone()),
|
||||
now - 100.0,
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current-session",
|
||||
query: None,
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, Some(&emb)).unwrap();
|
||||
|
||||
for item in &block.items {
|
||||
if let StmItem::SegmentRecap { session_id, .. } = item {
|
||||
assert_ne!(
|
||||
session_id, "current-session",
|
||||
"arm2 must never return current session recaps"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── dedup-by-episodic-span tests ──────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn dedup_drops_episodic_row_inside_segment_span() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Insert episodic rows for "other-session"
|
||||
let id_start = insert_episodic(
|
||||
&conn,
|
||||
"other-session",
|
||||
now - 300.0,
|
||||
"user",
|
||||
"Rust ownership",
|
||||
);
|
||||
let id_end = insert_episodic(
|
||||
&conn,
|
||||
"other-session",
|
||||
now - 290.0,
|
||||
"assistant",
|
||||
"Rust uses borrow checker",
|
||||
);
|
||||
|
||||
// A high-similarity segment recap covers those episodic rows
|
||||
let mut emb = vec![0.0_f32; 8];
|
||||
emb[0] = 1.0;
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-covers",
|
||||
"other-session",
|
||||
id_start,
|
||||
id_end,
|
||||
"Conversation about Rust ownership and borrow checker",
|
||||
Some(emb.clone()),
|
||||
now - 100.0,
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current",
|
||||
query: Some("Rust ownership"),
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, Some(&emb)).unwrap();
|
||||
|
||||
// The segment recap must appear
|
||||
let has_recap = block.items.iter().any(
|
||||
|it| matches!(it, StmItem::SegmentRecap { segment_id, .. } if segment_id == "seg-covers"),
|
||||
);
|
||||
assert!(has_recap, "segment recap must appear in output");
|
||||
|
||||
// The covered episodic rows must NOT appear (dedup)
|
||||
for item in &block.items {
|
||||
if let StmItem::EpisodicTurn { id, .. } = item {
|
||||
assert!(
|
||||
*id != Some(id_start) && *id != Some(id_end),
|
||||
"episodic rows inside segment span must be deduplicated; id={id:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
block.dropped_dedup > 0,
|
||||
"dropped_dedup must be > 0 when rows are inside a segment span"
|
||||
);
|
||||
}
|
||||
|
||||
// ── recency window bound test ─────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn recency_window_excludes_old_segments() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Recent segment — within window
|
||||
let id1 = insert_episodic(&conn, "recent-session", now - 100.0, "user", "recent");
|
||||
let emb_recent: Vec<f32> = (0..8).map(|i| if i == 0 { 1.0 } else { 0.0 }).collect();
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-recent",
|
||||
"recent-session",
|
||||
id1,
|
||||
id1,
|
||||
"Recent segment recap",
|
||||
Some(emb_recent.clone()),
|
||||
now - 100.0, // recent
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
// Old segment — beyond RECENCY_WINDOW_DAYS
|
||||
let old_ts = now - (super::super::RECENCY_WINDOW_DAYS + 2.0) * 86_400.0;
|
||||
let id2 = insert_episodic(&conn, "old-session", old_ts, "user", "old content");
|
||||
insert_segment_with_embedding(
|
||||
&conn,
|
||||
"seg-old",
|
||||
"old-session",
|
||||
id2,
|
||||
id2,
|
||||
"Old segment recap",
|
||||
Some(emb_recent.clone()),
|
||||
old_ts, // older than window
|
||||
"test:model:8",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current",
|
||||
query: None,
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, Some(&emb_recent)).unwrap();
|
||||
|
||||
let seg_ids: Vec<&str> = block
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|it| {
|
||||
if let StmItem::SegmentRecap { segment_id, .. } = it {
|
||||
Some(segment_id.as_str())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert!(
|
||||
!seg_ids.contains(&"seg-old"),
|
||||
"old segment beyond recency window must be excluded; got: {:?}",
|
||||
seg_ids
|
||||
);
|
||||
assert!(
|
||||
seg_ids.contains(&"seg-recent"),
|
||||
"recent segment must appear; got: {:?}",
|
||||
seg_ids
|
||||
);
|
||||
}
|
||||
|
||||
// ── token budget test ─────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn token_budget_limits_output_size() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Insert many episodic turns from other sessions
|
||||
for i in 0..50 {
|
||||
let large_content = "X".repeat(300); // 300 chars each
|
||||
insert_episodic(
|
||||
&conn,
|
||||
&format!("session-{i}"),
|
||||
now - i as f64 * 60.0,
|
||||
"user",
|
||||
&large_content,
|
||||
);
|
||||
}
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current",
|
||||
query: None,
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, None).unwrap();
|
||||
|
||||
let total_chars: usize = block.items.iter().map(|it| it.approx_chars()).sum();
|
||||
assert!(
|
||||
total_chars <= super::super::TOKEN_BUDGET,
|
||||
"total chars {} must not exceed budget {}",
|
||||
total_chars,
|
||||
super::super::TOKEN_BUDGET
|
||||
);
|
||||
}
|
||||
|
||||
// ── preemptive recency fallback (no query) ────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn preemptive_no_query_returns_recent_other_sessions() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
// Insert turns from other sessions
|
||||
insert_episodic(
|
||||
&conn,
|
||||
"session-a",
|
||||
now - 300.0,
|
||||
"user",
|
||||
"Alpha session content",
|
||||
);
|
||||
insert_episodic(
|
||||
&conn,
|
||||
"session-b",
|
||||
now - 200.0,
|
||||
"user",
|
||||
"Beta session content",
|
||||
);
|
||||
|
||||
// Also insert turns for the current session — must be excluded
|
||||
insert_episodic(
|
||||
&conn,
|
||||
"current-session",
|
||||
now - 100.0,
|
||||
"user",
|
||||
"Current session content",
|
||||
);
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current-session",
|
||||
query: None,
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, None).unwrap();
|
||||
|
||||
// Check that we got results from other sessions
|
||||
let other_sessions: Vec<&str> = block
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|it| {
|
||||
if let StmItem::EpisodicTurn { session_id, .. } = it {
|
||||
Some(session_id.as_str())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert!(
|
||||
!other_sessions.is_empty() || block.items.is_empty(), // empty is OK if no rows
|
||||
"preemptive fallback must only return other-session items"
|
||||
);
|
||||
for sid in &other_sessions {
|
||||
assert_ne!(
|
||||
*sid, "current-session",
|
||||
"current session must be excluded in preemptive mode"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── rendered block format ─────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn render_produces_non_empty_markdown_when_items_present() {
|
||||
let conn = setup_conn();
|
||||
let now = now_ts();
|
||||
|
||||
insert_episodic(&conn, "other-session", now - 100.0, "user", "Test content");
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current",
|
||||
query: None,
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, None).unwrap();
|
||||
|
||||
if !block.items.is_empty() {
|
||||
let rendered = block.render();
|
||||
assert!(
|
||||
rendered.contains("## Recent context"),
|
||||
"rendered block must contain heading"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_empty_block_returns_empty_string() {
|
||||
let block = StmRecallBlock::default();
|
||||
assert!(block.render().is_empty());
|
||||
assert!(block.is_empty());
|
||||
}
|
||||
|
||||
// ── end-to-end integration test ───────────────────────────────────────────────
|
||||
// Drive the real chain: completed turns → episodic rows → segment close
|
||||
// (recap + embedding via the Phase 0+1 path using stub providers) →
|
||||
// STM recall returns cross-thread recaps and excludes the current session.
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_stm_recall_chain() {
|
||||
use crate::openhuman::memory::tree::chat::ChatPrompt;
|
||||
|
||||
let conn = setup_conn();
|
||||
|
||||
// ── Phase 0+1 stub providers ─────────────────────────────────────────────
|
||||
// We use a stub chat provider that returns a fixed recap string, and the
|
||||
// InertEmbedder that returns zero vectors. This exercises the real
|
||||
// archivist code path (recap + segment_embedding_upsert) without
|
||||
// requiring a live LLM or Ollama daemon.
|
||||
|
||||
struct StubChat;
|
||||
use crate::openhuman::memory::tree::chat::ChatProvider;
|
||||
#[async_trait::async_trait]
|
||||
impl ChatProvider for StubChat {
|
||||
fn name(&self) -> &str {
|
||||
"stub"
|
||||
}
|
||||
async fn chat_for_json(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
|
||||
Ok("RECAP: stub LLM summary of the segment.".to_string())
|
||||
}
|
||||
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
|
||||
Ok("RECAP: stub LLM summary of the segment.".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
use crate::openhuman::memory::tree::score::embed::InertEmbedder;
|
||||
let chat_provider: Arc<dyn crate::openhuman::memory::tree::chat::ChatProvider> =
|
||||
Arc::new(StubChat);
|
||||
let embedder: Arc<dyn crate::openhuman::memory::tree::score::embed::Embedder> =
|
||||
Arc::new(InertEmbedder::new());
|
||||
|
||||
let archivist = ArchivistHook::new_with_stubs(conn.clone(), chat_provider, embedder);
|
||||
|
||||
// ── Turns for "other-thread" ─────────────────────────────────────────────
|
||||
// Drive 25 turns on session "other-thread" — exceeds max_turns_per_segment (20)
|
||||
// so a segment boundary fires, the segment closes, and recap + embedding happen.
|
||||
|
||||
for i in 0..25 {
|
||||
let ctx = TurnContext {
|
||||
user_message: format!("User message {i} about Rust and memory safety"),
|
||||
assistant_response: format!("Assistant response {i}: Rust ownership is great."),
|
||||
tool_calls: vec![],
|
||||
turn_duration_ms: 100,
|
||||
session_id: Some("other-thread".to_string()),
|
||||
iteration_count: i + 1,
|
||||
};
|
||||
archivist.on_turn_complete(&ctx).await.unwrap();
|
||||
}
|
||||
|
||||
// Force-flush any trailing open segment so we definitely get a recap.
|
||||
archivist.flush_open_segment("other-thread").await;
|
||||
|
||||
// ── Verify episodic rows were written ────────────────────────────────────
|
||||
let ep_rows = fts5::episodic_session_entries(&conn, "other-thread").unwrap();
|
||||
assert!(
|
||||
ep_rows.len() >= 50,
|
||||
"expected >=50 episodic rows (2 per turn × 25), got {}",
|
||||
ep_rows.len()
|
||||
);
|
||||
|
||||
// ── Verify segment embedding written ─────────────────────────────────────
|
||||
let has_embedding = {
|
||||
let c = conn.lock();
|
||||
let count: i64 = c
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM segment_embeddings se
|
||||
JOIN conversation_segments cs ON se.segment_id = cs.segment_id
|
||||
WHERE cs.session_id = 'other-thread'",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap_or(0);
|
||||
count > 0
|
||||
};
|
||||
assert!(
|
||||
has_embedding,
|
||||
"CRITICAL: Phase 0+1 did NOT write segment_embeddings for other-thread. \
|
||||
This means the archivist recap+embed path is broken. \
|
||||
STM recall Arm 2 would have no data to query."
|
||||
);
|
||||
|
||||
// ── Now run STM recall from "current-thread" ─────────────────────────────
|
||||
// InertEmbedder returns zero vectors, cosine of zero vectors = 0.0 < COSINE_GATE.
|
||||
// So Arm 2 will find no hits (expected — inert embedder produces identical vectors).
|
||||
// Arm 1 (FTS5 or recency) should still return episodic turns from other-thread.
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: "current-thread",
|
||||
query: Some("Rust memory safety"),
|
||||
model_signature: None,
|
||||
};
|
||||
let block = stm_recall(&conn, &opts, None).unwrap(); // no embedding for Arm 2
|
||||
|
||||
// With keyword "Rust memory safety" + other-thread has "Rust and memory safety"
|
||||
// in the episodic log, Arm 1 should surface at least some results.
|
||||
// (FTS5 porter-stems "safety" and "Rust" matches the stored content.)
|
||||
|
||||
// Verify: nothing from current-thread
|
||||
for item in &block.items {
|
||||
match item {
|
||||
StmItem::EpisodicTurn { session_id, .. } => {
|
||||
assert_ne!(
|
||||
session_id, "current-thread",
|
||||
"STM recall must never return current-thread items"
|
||||
);
|
||||
}
|
||||
StmItem::SegmentRecap { session_id, .. } => {
|
||||
assert_ne!(
|
||||
session_id, "current-thread",
|
||||
"STM recall must never return current-thread recaps"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The FTS5 arm should have found the other-thread episodic rows
|
||||
let fts5_hits = block.fts5_candidates;
|
||||
assert!(
|
||||
fts5_hits > 0,
|
||||
"Arm 1 (FTS5) must have found candidates from other-thread for 'Rust memory safety' query; \
|
||||
fts5_candidates={}. This proves episodic rows are written and searchable.",
|
||||
fts5_hits
|
||||
);
|
||||
|
||||
// Verify block is well-formed
|
||||
let rendered = block.render();
|
||||
if !block.items.is_empty() {
|
||||
assert!(
|
||||
rendered.contains("## Recent context"),
|
||||
"rendered block must have heading"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
//! Agent-callable tool for on-demand STM recall.
|
||||
//!
|
||||
//! The agent invokes `stm_recall_search[query]` mid-session to pull
|
||||
//! cross-thread context by keyword. The tool:
|
||||
//!
|
||||
//! 1. Extracts the active session_id from the tool call arguments (or falls
|
||||
//! back to a process-level default when not provided).
|
||||
//! 2. Runs [`stm_recall`] with Arm 1 (FTS5 keyword) only — no embedding step
|
||||
//! at call time to avoid blocking the hot path.
|
||||
//! 3. Returns the rendered markdown block as a tool result.
|
||||
//!
|
||||
//! The tool is registered in `tools/ops.rs` alongside `MemoryRecallTool` when
|
||||
//! `learning.stm_recall_enabled` is true.
|
||||
|
||||
use crate::openhuman::memory::Memory;
|
||||
use crate::openhuman::tools::traits::{Tool, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::recall::{stm_recall, StmRecallOpts};
|
||||
|
||||
/// On-demand STM recall tool.
|
||||
///
|
||||
/// Searches recent episodic memory from **other** chat threads using
|
||||
/// keyword matching (FTS5). Vector similarity (Arm 2) is not run here —
|
||||
/// it requires an embedding call that belongs in the session-start
|
||||
/// preemptive path, not the on-demand agent invocation path.
|
||||
pub struct StmRecallTool {
|
||||
memory: Arc<dyn Memory>,
|
||||
/// Session ID to exclude from results. Injected at construction so
|
||||
/// the tool knows the "current thread" without the agent having to
|
||||
/// pass it explicitly.
|
||||
session_id: String,
|
||||
/// Optional model signature for filtering segment embeddings.
|
||||
model_signature: Option<String>,
|
||||
}
|
||||
|
||||
impl StmRecallTool {
|
||||
/// Create a new `StmRecallTool`.
|
||||
///
|
||||
/// `session_id` — the current session's ID. Results from this session
|
||||
/// are always excluded.
|
||||
///
|
||||
/// `model_signature` — when `Some`, Arm 2 is filtered to this model.
|
||||
/// Pass `None` to accept any model (Arm 2 is skipped in on-demand mode
|
||||
/// anyway, but stored for future extension).
|
||||
pub fn new(
|
||||
memory: Arc<dyn Memory>,
|
||||
session_id: String,
|
||||
model_signature: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
memory,
|
||||
session_id,
|
||||
model_signature,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for StmRecallTool {
|
||||
fn name(&self) -> &str {
|
||||
"stm_recall_search"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Search recent conversational context from other chat threads. \
|
||||
Use this when you need facts or context discussed in a previous conversation \
|
||||
that may be relevant to the current request. \
|
||||
Returns a bounded set of snippets and conversation recaps from other sessions."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Keywords or phrase to search across recent other-session conversations"
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||
let query = args
|
||||
.get("query")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing 'query' parameter"))?
|
||||
.trim();
|
||||
if query.is_empty() {
|
||||
return Err(anyhow::anyhow!("query cannot be empty"));
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"[stm_recall_tool] on-demand recall query_len={} session={}",
|
||||
query.chars().count(),
|
||||
self.session_id
|
||||
);
|
||||
|
||||
// Get SQLite connection via the Memory trait's sqlite_conn() hook.
|
||||
let conn = match self.memory.sqlite_conn() {
|
||||
Some(c) => c,
|
||||
None => {
|
||||
tracing::warn!(
|
||||
"[stm_recall_tool] memory backend has no SQLite connection — stm_recall unavailable"
|
||||
);
|
||||
return Ok(ToolResult::success(
|
||||
"STM recall is not available (memory backend is not SQLite-backed).",
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let opts = StmRecallOpts {
|
||||
exclude_session: &self.session_id,
|
||||
query: Some(query),
|
||||
model_signature: self.model_signature.as_deref(),
|
||||
};
|
||||
|
||||
match stm_recall(&conn, &opts, None) {
|
||||
Ok(block) => {
|
||||
tracing::debug!(
|
||||
"[stm_recall_tool] recall complete: {} items, {} fts5_candidates, {} dropped_dedup",
|
||||
block.items.len(),
|
||||
block.fts5_candidates,
|
||||
block.dropped_dedup
|
||||
);
|
||||
if block.is_empty() {
|
||||
Ok(ToolResult::success(
|
||||
"No relevant context found in recent other-session conversations.",
|
||||
))
|
||||
} else {
|
||||
Ok(ToolResult::success(block.render()))
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("[stm_recall_tool] stm_recall failed: {e}");
|
||||
Ok(ToolResult::error(format!("STM recall failed: {e}")))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::embeddings::NoopEmbedding;
|
||||
use crate::openhuman::memory::UnifiedMemory;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn make_mem() -> (TempDir, Arc<dyn Memory>) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mem = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap();
|
||||
(tmp, Arc::new(mem))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_name_and_schema() {
|
||||
let (_tmp, mem) = make_mem();
|
||||
let tool = StmRecallTool::new(mem, "s1".into(), None);
|
||||
assert_eq!(tool.name(), "stm_recall_search");
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"]["query"].is_object());
|
||||
assert_eq!(schema["required"][0].as_str(), Some("query"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_empty_query_returns_error() {
|
||||
let (_tmp, mem) = make_mem();
|
||||
let tool = StmRecallTool::new(mem, "s1".into(), None);
|
||||
let result = tool.execute(json!({"query": " "})).await;
|
||||
assert!(result.is_err(), "empty query must return Err");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_missing_query_returns_error() {
|
||||
let (_tmp, mem) = make_mem();
|
||||
let tool = StmRecallTool::new(mem, "s1".into(), None);
|
||||
let result = tool.execute(json!({})).await;
|
||||
assert!(result.is_err(), "missing query must return Err");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_returns_no_matches_when_empty_db() {
|
||||
let (_tmp, mem) = make_mem();
|
||||
let tool = StmRecallTool::new(mem, "s1".into(), None);
|
||||
let result = tool
|
||||
.execute(json!({"query": "Rust ownership"}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(
|
||||
result.output().contains("No relevant context") || !result.output().is_empty(),
|
||||
"empty db must return a no-match message, got: {}",
|
||||
result.output()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -398,6 +398,10 @@ impl Memory for UnifiedMemory {
|
||||
async fn health_check(&self) -> bool {
|
||||
self.workspace_dir.exists() && self.db_path.exists()
|
||||
}
|
||||
|
||||
fn sqlite_conn(&self) -> Option<std::sync::Arc<parking_lot::Mutex<rusqlite::Connection>>> {
|
||||
Some(std::sync::Arc::clone(&self.conn))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -5,7 +5,10 @@
|
||||
//! types used for representing and organizing memories.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use rusqlite::Connection;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Represents a single stored memory entry with associated metadata.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -151,6 +154,17 @@ pub trait Memory: Send + Sync {
|
||||
|
||||
/// Performs a health check on the underlying storage system.
|
||||
async fn health_check(&self) -> bool;
|
||||
|
||||
/// Return the shared SQLite connection when the backend is `UnifiedMemory`.
|
||||
///
|
||||
/// Used by subsystems (e.g. `ArchivistHook`) that need direct SQLite
|
||||
/// access for FTS5 / segment writes without going through the async
|
||||
/// `Memory` trait.
|
||||
///
|
||||
/// Default: `None`. Only `UnifiedMemory` overrides this.
|
||||
fn sqlite_conn(&self) -> Option<Arc<Mutex<Connection>>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -228,6 +228,22 @@ pub fn all_tools_with_runtime(
|
||||
root_config.curl.timeout_secs,
|
||||
)));
|
||||
|
||||
// Phase 3 STM recall — on-demand cross-thread episodic search tool.
|
||||
// Feature-gated on `learning.stm_recall_enabled` (default true) so the
|
||||
// tool surface and the preemptive prompt injection are enabled/disabled
|
||||
// together. `session_id` is not known at tool-build time; exclude-own-
|
||||
// session is enforced by the preemptive first-turn injection in turn.rs
|
||||
// (the on-demand tool intentionally uses an empty exclude_session).
|
||||
if root_config.learning.stm_recall_enabled {
|
||||
tools.push(Box::new(
|
||||
crate::openhuman::memory::stm_recall::tool::StmRecallTool::new(
|
||||
memory.clone(),
|
||||
String::new(),
|
||||
None,
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
// gitbooks — answers questions about OpenHuman by calling the
|
||||
// GitBook MCP server. Two tools mirroring the upstream MCP tools.
|
||||
if root_config.gitbooks.enabled {
|
||||
|
||||
Reference in New Issue
Block a user