From 342f22317d6cd5aed5acc00c9270886dad16897b Mon Sep 17 00:00:00 2001 From: Austin Varnon Date: Mon, 6 Jul 2026 22:32:30 -0500 Subject: [PATCH 1/4] impolemented save batching --- CHANGELOG.md | 4 ++ benchmarks/longmemeval/src/ingest.rs | 27 +++++--- benchmarks/longmemeval/src/main.rs | 42 +++++++++---- examples/lexicon.rs | 6 +- src/builder.rs | 93 +++++++++++++++++++++------- src/engine.rs | 41 ++++++++++++ src/lib.rs | 47 ++++++++++++-- src/storage/turso_storage.rs | 32 ++++++++-- src/traits.rs | 19 ++++++ tests/lexicon_scoring.rs | 41 +++++++++++- 10 files changed, 293 insertions(+), 59 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 402f654..c3a357e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +### Breaking Changes + +- **Lexicon scoring is now opt-in.** `ContextForge::builder(...)` no longer auto-seeds `DefaultEnglishScorer`; by default the engine ranks on relevance (BM25, plus semantic when an embedding model is set) with no lexicon layer. Enable the English importance scorer explicitly with `.with_default_english_scorer()`, and note that `.with_persona_scorer(...)` no longer implies the English layer — call both to compose them. Rationale: the always-on, query-independent importance boost was shown to degrade factual retrieval (LongMemEval Recall@1 dropped 0.71 → 0.17 with the scorer forced on). Callers that want the old behavior — persona/importance use cases, e.g. a chat assistant — add `.with_default_english_scorer()`. + ### Features - **`OpenAiCompatDistiller::with_api_key`**: bearer token support for authenticated cloud endpoints. Pass any OpenAI-compatible API key (OpenAI, Groq, Together AI, GLM/ZhipuAI, etc.) and it is sent as `Authorization: Bearer ` on every request. The distiller now uses `rustls` and supports `https://` base URLs; the previous `http://`-only restriction is lifted. diff --git a/benchmarks/longmemeval/src/ingest.rs b/benchmarks/longmemeval/src/ingest.rs index 80401b0..4cc97b5 100644 --- a/benchmarks/longmemeval/src/ingest.rs +++ b/benchmarks/longmemeval/src/ingest.rs @@ -55,24 +55,31 @@ impl Ingest { instance: &Instance, scope: &str, ) -> anyhow::Result<()> { + // Collect all of the instance's turns and save them in one batch, so the + // tantivy index commits once per instance instead of once per turn. + let mut items: Vec<(String, String, SaveOptions)> = Vec::new(); for (session_id, turns) in instance.sessions() { for turn in turns { // Some LongMemEval turns have empty content; CF rejects empty - // entries, so skip them rather than fail the whole instance. + // entries, so skip them rather than fail the batch. if turn.content.trim().is_empty() { continue; } - let opts = SaveOptions { - session_id: Some(session_id.clone()), - scope: Some(scope.to_owned()), - ..Default::default() - }; - forge - .save(&turn.content, TURN_KIND, &opts) - .await - .map_err(|e| anyhow::anyhow!("save turn in session {session_id}: {e}"))?; + items.push(( + turn.content.clone(), + TURN_KIND.to_owned(), + SaveOptions { + session_id: Some(session_id.clone()), + scope: Some(scope.to_owned()), + ..Default::default() + }, + )); } } + forge + .save_batch(&items) + .await + .map_err(|e| anyhow::anyhow!("save batch for scope {scope}: {e}"))?; Ok(()) } diff --git a/benchmarks/longmemeval/src/main.rs b/benchmarks/longmemeval/src/main.rs index a7feb09..cf3e27c 100644 --- a/benchmarks/longmemeval/src/main.rs +++ b/benchmarks/longmemeval/src/main.rs @@ -8,10 +8,10 @@ //! Usage: //! cargo run -p longmemeval-bench --release -- [flags] //! Flags: -//! --pipeline bm25|lexicon|full (default bm25) -//! --ingest raw|distill (default raw) -//! --limit N (only the first N instances) -//! --embed-dir PATH (model cache dir for full; or CF_EMBED_MODEL_DIR) +//! --pipeline bm25|lexicon|semantic|full (default bm25) +//! --ingest raw|distill (default raw) +//! --limit N (only the first N instances) +//! --embed-dir PATH (model cache dir for semantic/full; or CF_EMBED_MODEL_DIR) mod dataset; mod ingest; @@ -37,9 +37,13 @@ const UNBOUNDED_BUDGET: usize = 100_000_000; enum Pipeline { /// `open`: BM25 + recency, no lexicon, no embedder. Bm25, - /// `builder().build()`: BM25 + `DefaultEnglishScorer` (builder always seeds it). + /// `builder().with_default_english_scorer().build()`: BM25 + lexicon. Lexicon, - /// `builder().with_embedding_model().build()`: BM25 + semantic + lexicon. + /// `builder().with_embedding_model().build()`: BM25 + semantic, no lexicon. + /// (Expressible now that the English scorer is opt-in.) + Semantic, + /// `builder().with_default_english_scorer().with_embedding_model().build()`: + /// BM25 + semantic + lexicon. Full, } @@ -48,6 +52,7 @@ impl Pipeline { match s { "bm25" => Ok(Pipeline::Bm25), "lexicon" => Ok(Pipeline::Lexicon), + "semantic" => Ok(Pipeline::Semantic), "full" => Ok(Pipeline::Full), other => Err(anyhow::anyhow!("unknown pipeline {other:?}")), } @@ -56,6 +61,7 @@ impl Pipeline { match self { Pipeline::Bm25 => "bm25", Pipeline::Lexicon => "lexicon", + Pipeline::Semantic => "semantic", Pipeline::Full => "full", } } @@ -122,15 +128,29 @@ async fn build_forge( let mut config = Config::default(); config.db_path = PathBuf::from(":memory:"); config.max_entries = 100_000_000; // never evict during a run + let require_dir = |embed_dir: &Option| { + embed_dir.clone().ok_or_else(|| { + anyhow::anyhow!("--embed-dir (or CF_EMBED_MODEL_DIR) required for this pipeline") + }) + }; let forge = match pipeline { Pipeline::Bm25 => ContextForge::open(config).await?, - Pipeline::Lexicon => ContextForge::builder(config).build().await?, + Pipeline::Lexicon => { + ContextForge::builder(config) + .with_default_english_scorer() + .build() + .await? + } + Pipeline::Semantic => { + ContextForge::builder(config) + .with_embedding_model(require_dir(embed_dir)?) + .build() + .await? + } Pipeline::Full => { - let dir = embed_dir.clone().ok_or_else(|| { - anyhow::anyhow!("--embed-dir (or CF_EMBED_MODEL_DIR) required for --pipeline full") - })?; ContextForge::builder(config) - .with_embedding_model(dir) + .with_default_english_scorer() + .with_embedding_model(require_dir(embed_dir)?) .build() .await? } diff --git a/examples/lexicon.rs b/examples/lexicon.rs index 51369ff..8443db8 100644 --- a/examples/lexicon.rs +++ b/examples/lexicon.rs @@ -57,14 +57,16 @@ patterns = [ // ── 3. Wire into ContextForge via the builder ───────────────────────────── // - // The builder always pre-seeds DefaultEnglishScorer. with_persona_scorer - // stacks the WH40k lexicon on top via CompositeLexiconScorer. + // Lexicon scoring is opt-in. `with_default_english_scorer` enables the + // built-in English importance markers; `with_persona_scorer` stacks the + // WH40k lexicon on top via CompositeLexiconScorer. let mut config = context_forge::Config::default(); // :memory: so the example leaves no files behind. config.db_path = std::path::PathBuf::from(":memory:"); let cf = ContextForge::builder(config) + .with_default_english_scorer() .with_persona_scorer(persona) .build() .await?; diff --git a/src/builder.rs b/src/builder.rs index 356490f..f55a633 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -13,11 +13,18 @@ use crate::ContextForge; /// Builder for [`ContextForge`]. /// -/// The builder always pre-seeds [`DefaultEnglishScorer`] so plain-English -/// importance signals ("confirmed", "we decided", "never mind", etc.) are -/// active without any additional configuration. A caller-provided persona -/// scorer is optional and stacks additively on top via -/// [`CompositeLexiconScorer`]. +/// **Lexicon scoring is opt-in.** By default the engine ranks on relevance +/// (BM25, plus semantic when an embedding model is set) with no lexicon layer. +/// Lexicon scoring applies a *query-independent* importance boost, which suits +/// persona/importance use cases but degrades pure relevance retrieval — so it +/// must be requested explicitly: +/// +/// - [`with_default_english_scorer`](Self::with_default_english_scorer) enables +/// the built-in [`DefaultEnglishScorer`] (plain-English commitment/ +/// confirmation/decision markers). +/// - [`with_persona_scorer`](Self::with_persona_scorer) adds a domain scorer. +/// +/// When both are set they compose additively via [`CompositeLexiconScorer`]. /// /// # Example /// @@ -30,19 +37,30 @@ use crate::ContextForge; /// let mut config = Config::default(); /// config.db_path = PathBuf::from("memory.db"); /// -/// // English scorer only (most common case): +/// // Relevance only, no lexicon (default): /// let cf = ContextForge::builder(config.clone()).build().await?; /// -/// // English + persona scorer loaded from a TOML string: +/// // Opt into the English importance scorer: +/// let cf = ContextForge::builder(config.clone()) +/// .with_default_english_scorer() +/// .build() +/// .await?; +/// +/// // English + a persona scorer loaded from a TOML string: /// let toml = "[affirmations]\npatterns = [\"for the emperor\"]"; /// let persona: ConfigLexiconScorer = toml.parse()?; -/// let cf = ContextForge::builder(config).with_persona_scorer(persona).build().await?; +/// let cf = ContextForge::builder(config) +/// .with_default_english_scorer() +/// .with_persona_scorer(persona) +/// .build() +/// .await?; /// /// Ok(()) /// } /// ``` pub struct ContextForgeBuilder { config: Config, + english_defaults: bool, persona_scorer: Option>, #[cfg(feature = "semantic")] embedding_cache_dir: Option, @@ -56,17 +74,36 @@ impl ContextForgeBuilder { pub fn new(config: Config) -> Self { Self { config, + english_defaults: false, persona_scorer: None, #[cfg(feature = "semantic")] embedding_cache_dir: None, } } - /// Stack a persona scorer on top of the always-on [`DefaultEnglishScorer`]. + /// Enable the built-in [`DefaultEnglishScorer`] (opt-in). /// - /// The persona scorer is typically a [`crate::ConfigLexiconScorer`] loaded - /// from a domain-specific TOML file. Its boosts are summed with the English - /// layer; the engine applies a `-1.0` floor after fusion. + /// Applies a query-independent importance boost to entries containing + /// plain-English commitment/confirmation/decision/correction markers + /// ("i'll fix it", "confirmed", "we decided", "never mind"). This is the + /// right signal for persona/importance use cases (surfacing what matters to + /// a user) but **hurts pure relevance retrieval** — a factual-QA benchmark + /// showed it lowering recall by burying evidence under important-*sounding* + /// distractors. Off by default for that reason; enable it when importance, + /// not just relevance, is what you want to rank on. + #[must_use] + pub fn with_default_english_scorer(mut self) -> Self { + self.english_defaults = true; + self + } + + /// Add a persona scorer (opt-in). + /// + /// Typically a [`crate::ConfigLexiconScorer`] loaded from a domain-specific + /// TOML file. If [`with_default_english_scorer`](Self::with_default_english_scorer) + /// was also called, the two compose additively via [`CompositeLexiconScorer`] + /// and the engine applies a `-1.0` floor after fusion; otherwise the persona + /// scorer is used alone. #[must_use] pub fn with_persona_scorer(mut self, scorer: impl LexiconScorer + 'static) -> Self { self.persona_scorer = Some(Arc::new(scorer)); @@ -87,11 +124,13 @@ impl ContextForgeBuilder { self } - /// Open the database and build a [`ContextForge`] with the configured scorer. + /// Open the database and build a [`ContextForge`] with the configured scorers. /// - /// The [`DefaultEnglishScorer`] is always included. If - /// [`Self::with_persona_scorer`] was called, both scorers are composed via - /// [`CompositeLexiconScorer`]. + /// Lexicon scoring is applied only if + /// [`with_default_english_scorer`](Self::with_default_english_scorer) and/or + /// [`with_persona_scorer`](Self::with_persona_scorer) were called; with both, + /// they compose via [`CompositeLexiconScorer`]. With neither, the engine ranks + /// on relevance (BM25 + semantic) only. /// /// # Errors /// @@ -103,15 +142,25 @@ impl ContextForgeBuilder { let (storage, searcher) = open_storage(Path::new(&db_path), max_entries).await?; - let english: Arc = Arc::new(DefaultEnglishScorer::default()); - let scorer: Arc = match self.persona_scorer { - Some(persona) => Arc::new(CompositeLexiconScorer::new(vec![english, persona])), - None => english, + // Compose only the opted-in scorers; default is none (relevance only). + let mut scorers: Vec> = Vec::new(); + if self.english_defaults { + scorers.push(Arc::new(DefaultEnglishScorer::default())); + } + if let Some(persona) = self.persona_scorer { + scorers.push(persona); + } + let scorer: Option> = match scorers.len() { + 0 => None, + 1 => scorers.pop(), + _ => Some(Arc::new(CompositeLexiconScorer::new(scorers))), }; #[cfg_attr(not(feature = "semantic"), allow(unused_mut))] - let mut engine = ContextEngine::new(Box::new(storage), Box::new(searcher), self.config) - .with_scorer(scorer); + let mut engine = ContextEngine::new(Box::new(storage), Box::new(searcher), self.config); + if let Some(scorer) = scorer { + engine = engine.with_scorer(scorer); + } #[cfg(feature = "semantic")] if let Some(cache_dir) = self.embedding_cache_dir { diff --git a/src/engine.rs b/src/engine.rs index 7615632..d625e95 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -317,6 +317,47 @@ impl ContextEngine { Ok(id) } + /// Save multiple snapshot entries with a **single** search-index commit for + /// the whole batch (see [`crate::traits::ContextStorage::save_batch`]), + /// instead of one commit per entry. Embeddings are still generated per entry + /// (non-fatal). Returns the generated ids in the same order as `items`. + /// + /// Each item is `(content, kind, options)`. Content must already be scrubbed + /// by the caller, as with [`Self::save_snapshot`]. + pub(crate) async fn save_snapshot_batch( + &self, + items: &[(String, String, SaveOptions)], + ) -> Result> { + if items.iter().any(|(content, _, _)| content.is_empty()) { + return Err(Error::InvalidEntry("content must not be empty".into())); + } + + let entries: Vec = items + .iter() + .map(|(content, kind, options)| ContextEntry { + id: Uuid::now_v7().to_string(), + content: content.clone(), + timestamp: current_timestamp(), + kind: kind.clone(), + scope: options.scope.clone(), + session_id: options.session_id.clone(), + token_count: Some(estimate_tokens(content)), + metadata: options.metadata.clone(), + }) + .collect(); + + self.storage.save_batch(&entries).await?; + tracing::trace!(count = %entries.len(), "entry batch saved"); + + // Embeddings are independent of the batched index commit (a separate + // turso UPDATE per entry); generate them after the batch is saved. + for entry in &entries { + self.embed_and_store(&entry.id, &entry.content).await; + } + + Ok(entries.into_iter().map(|e| e.id).collect()) + } + /// Generate and persist an embedding for an already-saved entry. /// /// Errors are logged and swallowed — the entry is always stored even when diff --git a/src/lib.rs b/src/lib.rs index 0348996..39df04d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -189,6 +189,32 @@ impl ContextForge { self.engine.save_snapshot(&scrubbed, kind, opts).await } + /// Save many entries in one batch, committing the search index **once** for + /// the whole batch instead of once per entry — much cheaper for bulk ingest. + /// + /// Each item is `(content, kind, options)`, handled exactly as + /// [`save`](Self::save) (content is scrubbed per item). Returns the generated + /// ids in the same order as `items`. + /// + /// # Errors + /// + /// Returns an error if any item's `content` is empty or if the underlying + /// storage write fails. On failure, entries written before the failure may + /// remain persisted. + pub async fn save_batch(&self, items: &[(String, String, SaveOptions)]) -> Result> { + let scrubbed: Vec<(String, String, SaveOptions)> = items + .iter() + .map(|(content, kind, opts)| { + ( + scrub_secrets(content, &self.scrub_config).into_owned(), + kind.clone(), + opts.clone(), + ) + }) + .collect(); + self.engine.save_snapshot_batch(&scrubbed).await + } + /// Distill `transcript` into a summary and durable facts, then save /// them as separate entries sharing `opts.scope` and /// `opts.session_id`. @@ -228,10 +254,15 @@ impl ContextForge { let memory = tokio::task::block_in_place(|| distiller.distill(&scrubbed_transcript))?; let memory = crate::distill::cap_distilled_memory(memory); - let mut ids = Vec::with_capacity(1 + memory.facts.len()); - - let summary_id = self.save(&memory.summary, kind::SUMMARY, opts).await?; - ids.push(summary_id); + // Build all entries up front, then persist them in one batch so the + // search index commits once (not once per summary + fact). Content is + // scrubbed here — the same scrubbing `save` applies — before storage. + let mut items: Vec<(String, String, SaveOptions)> = Vec::with_capacity(1 + memory.facts.len()); + items.push(( + scrub_secrets(&memory.summary, &self.scrub_config).into_owned(), + kind::SUMMARY.to_owned(), + opts.clone(), + )); for fact in &memory.facts { let fact_kind_str = match fact.kind { @@ -249,10 +280,14 @@ impl ContextForge { scope: opts.scope.clone(), metadata: Some(metadata), }; - let fact_id = self.save(&fact.text, kind::FACT, &fact_opts).await?; - ids.push(fact_id); + items.push(( + scrub_secrets(&fact.text, &self.scrub_config).into_owned(), + kind::FACT.to_owned(), + fact_opts, + )); } + let ids = self.engine.save_snapshot_batch(&items).await?; Ok(ids) } diff --git a/src/storage/turso_storage.rs b/src/storage/turso_storage.rs index bbb9d0c..cc1a1eb 100644 --- a/src/storage/turso_storage.rs +++ b/src/storage/turso_storage.rs @@ -176,9 +176,11 @@ fn get_int_opt(row: &turso::Row, idx: usize, field: &str) -> crate::Result crate::Result<()> { +impl TursoStorage { + /// Persist one entry to turso and stage it in the tantivy index **without** + /// committing the index. Callers commit afterward — [`save`](ContextStorage::save) + /// once per entry, [`save_batch`](ContextStorage::save_batch) once per batch. + async fn persist_one(&self, entry: &ContextEntry) -> crate::Result<()> { let conn = self.db.connect()?; conn.busy_timeout(Duration::from_secs(5))?; let max_entries = self.max_entries; @@ -259,16 +261,34 @@ impl ContextStorage for TursoStorage { return result; } - // Turso write committed — update tantivy. If this fails the entry is - // still persisted in turso; next startup rebuild will re-sync the index. + // Turso write committed — stage in tantivy (no commit; the caller + // commits). If this fails the entry is still persisted in turso; the + // next startup rebuild will re-sync the index. if let Some(id) = evicted_id { self.fts.remove(&id)?; } self.fts .add(&entry.id, &entry.content, entry.scope.as_deref())?; + + Ok(()) + } +} + +#[async_trait] +impl ContextStorage for TursoStorage { + async fn save(&self, entry: &ContextEntry) -> crate::Result<()> { + self.persist_one(entry).await?; self.fts.commit()?; + Ok(()) + } - result + async fn save_batch(&self, entries: &[ContextEntry]) -> crate::Result<()> { + for entry in entries { + self.persist_one(entry).await?; + } + // Single index commit for the whole batch — the point of this method. + self.fts.commit()?; + Ok(()) } async fn get_top_k(&self, k: usize) -> crate::Result> { diff --git a/src/traits.rs b/src/traits.rs index 68d73c8..2463320 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -26,6 +26,25 @@ pub trait ContextStorage: Send + Sync { /// Returns an error if the underlying storage write fails. async fn save(&self, entry: &ContextEntry) -> Result<()>; + /// Persist multiple entries, committing any derived search index **once** + /// for the whole batch. + /// + /// Semantically equivalent to calling [`save`](Self::save) for each entry, + /// but implementations with a separate index (e.g. a tantivy BM25 index) + /// should override this to avoid a per-entry index commit. The default + /// implementation simply loops [`save`](Self::save). + /// + /// # Errors + /// + /// Returns an error if any underlying storage write fails. On error, entries + /// already written before the failure may remain persisted. + async fn save_batch(&self, entries: &[ContextEntry]) -> Result<()> { + for entry in entries { + self.save(entry).await?; + } + Ok(()) + } + /// Return the top-k entries (most recent or highest priority). /// /// # Errors diff --git a/tests/lexicon_scoring.rs b/tests/lexicon_scoring.rs index 5df1094..6cea52a 100644 --- a/tests/lexicon_scoring.rs +++ b/tests/lexicon_scoring.rs @@ -20,8 +20,9 @@ fn in_memory_config() -> context_forge::Config { // ── builder wiring ──────────────────────────────────────────────────────────── #[tokio::test] -async fn builder_without_persona_scorer_still_applies_english_defaults() { +async fn opting_into_english_scorer_boosts_affirmation() { let cf = ContextForge::builder(in_memory_config()) + .with_default_english_scorer() .build() .await .unwrap(); @@ -40,11 +41,44 @@ async fn builder_without_persona_scorer_still_applies_english_defaults() { assert_eq!(hits.len(), 2); assert_eq!( hits[0].id, affirm_id, - "English affirmation entry should rank first even without a persona scorer" + "opted-in English affirmation entry should rank first" ); let _ = neutral_id; } +#[tokio::test] +async fn default_build_does_not_apply_english_defaults() { + // Regression guard for the opt-in change: a persona scorer WITHOUT + // `with_default_english_scorer` must not pull in English scoring. The persona + // boost (0.3) is deliberately smaller than the English "confirmed" boost + // (0.5) — if English were still auto-on, the English entry would win, so the + // persona entry winning proves English defaults are off. + let persona: ConfigLexiconScorer = "[terms]\n\"beacon\" = 0.3".parse().unwrap(); + let cf = ContextForge::builder(in_memory_config()) + .with_persona_scorer(persona) + .build() + .await + .unwrap(); + + let opts = SaveOptions::default(); + let english_id = cf + .save("confirmed, that is correct", kind::SNAPSHOT, &opts) + .await + .unwrap(); + let persona_id = cf + .save("light the beacon", kind::SNAPSHOT, &opts) + .await + .unwrap(); + + let hits = cf.query(MATCH_ALL_QUERY, None, 10_000).await.unwrap(); + assert_eq!(hits.len(), 2); + assert_eq!( + hits[0].id, persona_id, + "persona (0.3) must outrank the English marker (0.5-if-on), proving English defaults are off by default" + ); + let _ = english_id; +} + #[tokio::test] async fn builder_with_persona_scorer_stacks_on_english_baseline() { let persona_toml = r#" @@ -60,6 +94,7 @@ patterns = ["the emperor frowns upon this"] let persona: ConfigLexiconScorer = persona_toml.parse().unwrap(); let cf = ContextForge::builder(in_memory_config()) + .with_default_english_scorer() .with_persona_scorer(persona) .build() .await @@ -103,6 +138,7 @@ patterns = ["the emperor frowns upon this"] #[tokio::test] async fn english_scorer_boosts_commissive_over_neutral() { let cf = ContextForge::builder(in_memory_config()) + .with_default_english_scorer() .build() .await .unwrap(); @@ -130,6 +166,7 @@ async fn english_scorer_boosts_commissive_over_neutral() { async fn negation_window_suppresses_false_affirmation() { // "not confirmed" should not fire the "confirmed" affirmation. let cf = ContextForge::builder(in_memory_config()) + .with_default_english_scorer() .build() .await .unwrap(); From 6d211242d6367595d8238af99da4bf9eb9d24ddf Mon Sep 17 00:00:00 2001 From: Austin Varnon Date: Mon, 6 Jul 2026 23:14:27 -0500 Subject: [PATCH 2/4] =?UTF-8?q?-=20src/storage/turso=5Fstorage.rs=20?= =?UTF-8?q?=E2=80=94=20prepared-statement=20save=5Fbatch=20=20=20-=20src/s?= =?UTF-8?q?torage/mod.rs=20=E2=80=94=20two=20new=20tests:=20batched-evicti?= =?UTF-8?q?on=20correctness=20+=20=20=20=20=20batch=20persist/index=20=20?= =?UTF-8?q?=20-=20src/lib.rs=20+=20src/lexicon/mod.rs=20=E2=80=94=20fixed?= =?UTF-8?q?=20stale=20docs=20from=201a=20(the=20=20=20=20=20"builder=20alw?= =?UTF-8?q?ays=20pre-seeds=20DefaultEnglishScorer"=20lines,=20now=20false)?= =?UTF-8?q?;=20=20=20=20=20lib.rs=20also=20picked=20up=20a=20fmt=20fix=20o?= =?UTF-8?q?n=20a=20long=20line=20that=20would've=20failed=20=20=20=20=20CI?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/lexicon/mod.rs | 12 ++-- src/lib.rs | 14 +++-- src/storage/mod.rs | 43 +++++++++++++++ src/storage/turso_storage.rs | 103 ++++++++++++++++++++++++++++++++++- 4 files changed, 159 insertions(+), 13 deletions(-) diff --git a/src/lexicon/mod.rs b/src/lexicon/mod.rs index 474ad6e..d54e25f 100644 --- a/src/lexicon/mod.rs +++ b/src/lexicon/mod.rs @@ -12,11 +12,13 @@ //! //! ## Two-layer design //! -//! The builder always pre-seeds a [`DefaultEnglishScorer`] that recognizes -//! plain-English importance signals ("confirmed", "never mind", etc.). A -//! persona scorer ([`ConfigLexiconScorer`] loaded from a TOML file) is -//! optional and stacks on top via [`CompositeLexiconScorer`]. Both layers -//! are additive — the engine applies the `-1.0` floor clamp after fusion. +//! Lexicon scoring is opt-in on the builder. [`DefaultEnglishScorer`] recognizes +//! plain-English importance signals ("confirmed", "never mind", etc.) and is +//! enabled via `with_default_english_scorer`. A persona scorer +//! ([`ConfigLexiconScorer`] loaded from a TOML file) is added via +//! `with_persona_scorer`. When both are set they stack via +//! [`CompositeLexiconScorer`] — additive, with the engine applying the `-1.0` +//! floor clamp after fusion. pub use self::appender::{LexiconAppender, LexiconProposal}; pub use self::bootstrap::bootstrap_prompt; diff --git a/src/lib.rs b/src/lib.rs index 39df04d..5340409 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -140,11 +140,12 @@ impl ContextForge { /// Create a builder for `ContextForge`. /// - /// The builder always pre-seeds [`DefaultEnglishScorer`] so plain-English - /// importance signals are active without any configuration. Use - /// [`ContextForgeBuilder::with_persona_scorer`] to stack a domain-specific - /// scorer on top. See also [`Self::open`] for the lower-level path that - /// wires no scorer. + /// Lexicon scoring is **opt-in**: by default the engine ranks on relevance + /// (BM25, plus semantic when an embedding model is set), with no lexicon + /// layer. Enable it with + /// [`ContextForgeBuilder::with_default_english_scorer`] and/or + /// [`ContextForgeBuilder::with_persona_scorer`]. [`Self::open`] is the + /// lower-level path that likewise wires no scorer. #[must_use] pub fn builder(config: Config) -> ContextForgeBuilder { ContextForgeBuilder::new(config) @@ -257,7 +258,8 @@ impl ContextForge { // Build all entries up front, then persist them in one batch so the // search index commits once (not once per summary + fact). Content is // scrubbed here — the same scrubbing `save` applies — before storage. - let mut items: Vec<(String, String, SaveOptions)> = Vec::with_capacity(1 + memory.facts.len()); + let mut items: Vec<(String, String, SaveOptions)> = + Vec::with_capacity(1 + memory.facts.len()); items.push(( scrub_secrets(&memory.summary, &self.scrub_config).into_owned(), kind::SUMMARY.to_owned(), diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 159b84d..dd84c8b 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -214,6 +214,49 @@ mod tests { assert!(ids.contains(&"e3")); } + #[tokio::test] + async fn test_save_batch_evicts_oldest_over_capacity() { + // A single batch that overflows capacity must evict the oldest rows once, + // keeping the newest `max_entries` — the batched-eviction path. + let (storage, _) = open_storage(Path::new(":memory:"), 3).await.unwrap(); + let batch = vec![ + make_entry("e1", "oldest", 100, kind::MANUAL), + make_entry("e2", "second", 200, kind::MANUAL), + make_entry("e3", "third", 300, kind::MANUAL), + make_entry("e4", "fourth", 400, kind::MANUAL), + make_entry("e5", "newest", 500, kind::MANUAL), + ]; + storage.save_batch(&batch).await.unwrap(); + + assert_eq!(storage.count().await.unwrap(), 3); + let all = storage.get_all().await.unwrap(); + let ids: Vec<&str> = all.iter().map(|e| e.id.as_str()).collect(); + assert!(!ids.contains(&"e1"), "oldest should be evicted"); + assert!(!ids.contains(&"e2"), "second-oldest should be evicted"); + assert!(ids.contains(&"e3")); + assert!(ids.contains(&"e4")); + assert!(ids.contains(&"e5")); + } + + #[tokio::test] + async fn test_save_batch_persists_and_indexes() { + // save_batch must persist every entry and commit the FTS index so the + // batch is immediately searchable — equivalent to a sequence of save(). + let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap(); + let batch = vec![ + make_entry("b1", "rust programming language", 100, kind::MANUAL), + make_entry("b2", "python scripting tools", 200, kind::MANUAL), + ]; + storage.save_batch(&batch).await.unwrap(); + + assert_eq!(storage.count().await.unwrap(), 2); + let hits = searcher.search("rust", None, 10).await.unwrap(); + assert!( + hits.iter().any(|e| e.entry.id == "b1"), + "batch-saved entry should be FTS-searchable" + ); + } + #[tokio::test] async fn test_fts_search() { let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap(); diff --git a/src/storage/turso_storage.rs b/src/storage/turso_storage.rs index cc1a1eb..e0b79c4 100644 --- a/src/storage/turso_storage.rs +++ b/src/storage/turso_storage.rs @@ -283,10 +283,109 @@ impl ContextStorage for TursoStorage { } async fn save_batch(&self, entries: &[ContextEntry]) -> crate::Result<()> { + if entries.is_empty() { + return Ok(()); + } + + let conn = self.db.connect()?; + conn.busy_timeout(Duration::from_secs(5))?; + let max_entries = self.max_entries; + let mut evicted_ids: Vec = Vec::new(); + + let result = async { + conn.execute("BEGIN IMMEDIATE", ()).await?; + + // Prepare the INSERT once and reuse it for every row: all rows are the + // same statement, only the bound values differ, so this avoids the + // per-row SQL re-parse that `conn.execute(sql, ...)` incurs. INSERT OR + // REPLACE is idempotent, so no per-row existence check is needed. + let mut insert = conn + .prepare_cached( + "INSERT OR REPLACE INTO entries \ + (id, content, timestamp, kind, scope, session_id, token_count, metadata) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + ) + .await?; + + for entry in entries { + let metadata_json = entry + .metadata + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|e| { + crate::Error::InvalidEntry(format!("metadata is not valid JSON: {e}")) + })?; + insert + .execute(( + entry.id.clone(), + entry.content.clone(), + entry.timestamp, + entry.kind.clone(), + entry.scope.clone(), + entry.session_id.clone(), + entry + .token_count + .map(|n| i64::try_from(n).unwrap_or(i64::MAX)), + metadata_json, + )) + .await?; + } + + // Eviction is handled once for the whole batch: if the inserts pushed + // the store over capacity, drop the oldest overflow rows. The batch's + // own entries carry current timestamps, so they are never the ones + // evicted here. + let mut count_rows = conn.query("SELECT COUNT(*) FROM entries", ()).await?; + let count: i64 = match count_rows.next().await? { + Some(row) => match row.get_value(0)? { + turso::Value::Integer(n) => n, + _ => 0, + }, + None => 0, + }; + drop(count_rows); + + let cap = i64::try_from(max_entries).unwrap_or(i64::MAX); + let overflow = count - cap; + if overflow > 0 { + let mut id_rows = conn + .query( + "SELECT id FROM entries ORDER BY timestamp ASC LIMIT ?1", + (overflow,), + ) + .await?; + while let Some(row) = id_rows.next().await? { + if let turso::Value::Text(id) = row.get_value(0)? { + evicted_ids.push(id); + } + } + drop(id_rows); + for id in &evicted_ids { + conn.execute("DELETE FROM entries WHERE id = ?1", (id.clone(),)) + .await?; + } + } + + conn.execute("COMMIT", ()).await?; + Ok(()) + } + .await; + + if result.is_err() { + let _ = conn.execute("ROLLBACK", ()).await; + return result; + } + + // Turso batch committed — sync tantivy: remove evicted docs, stage all + // new entries, and commit the index once for the whole batch. + for id in &evicted_ids { + self.fts.remove(id)?; + } for entry in entries { - self.persist_one(entry).await?; + self.fts + .add(&entry.id, &entry.content, entry.scope.as_deref())?; } - // Single index commit for the whole batch — the point of this method. self.fts.commit()?; Ok(()) } From 478cfad17e2f3988dfdcd376f6ec84611046a823 Mon Sep 17 00:00:00 2001 From: Austin Varnon Date: Mon, 6 Jul 2026 23:31:02 -0500 Subject: [PATCH 3/4] bench to test if queries was bottleneck or not --- benchmarks/longmemeval/src/main.rs | 51 ++++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 7 deletions(-) diff --git a/benchmarks/longmemeval/src/main.rs b/benchmarks/longmemeval/src/main.rs index cf3e27c..04eff78 100644 --- a/benchmarks/longmemeval/src/main.rs +++ b/benchmarks/longmemeval/src/main.rs @@ -19,6 +19,7 @@ mod metrics; use std::collections::{BTreeMap, HashSet}; use std::path::PathBuf; +use std::time::{Duration, Instant}; use context_forge::{Config, ContextForge}; @@ -208,10 +209,16 @@ async fn main() -> anyhow::Result<()> { let mut overall = Agg::default(); let mut by_type: BTreeMap = BTreeMap::new(); let mut errors = 0usize; + let mut total = Timing::default(); for (i, inst) in instances.iter().enumerate() { match run_instance(&args, inst).await { - Ok(result) => accumulate(&mut overall, &mut by_type, inst, result), + Ok((result, timing)) => { + total.build += timing.build; + total.ingest += timing.ingest; + total.query += timing.query; + accumulate(&mut overall, &mut by_type, inst, result); + } Err(e) => { errors += 1; eprintln!(" [!] instance {i} ({}) failed: {e}", inst.question_id); @@ -223,6 +230,12 @@ async fn main() -> anyhow::Result<()> { } print_report(&overall, &by_type, errors); + println!( + "\nphase totals: build={:.1}s ingest={:.1}s query={:.1}s", + total.build.as_secs_f64(), + total.ingest.as_secs_f64(), + total.query.as_secs_f64(), + ); Ok(()) } @@ -233,11 +246,27 @@ struct InstanceResult { budget_sessions: BTreeMap>, } -async fn run_instance(args: &Args, inst: &Instance) -> anyhow::Result { - let forge = build_forge(args.pipeline, &args.embed_dir).await?; +/// Wall-clock split of one instance across the three phases, so we can see +/// where the run's time actually goes (setup vs. ingest vs. query). +#[derive(Default, Clone, Copy)] +struct Timing { + build: Duration, + ingest: Duration, + query: Duration, +} + +async fn run_instance(args: &Args, inst: &Instance) -> anyhow::Result<(InstanceResult, Timing)> { let scope = inst.question_id.as_str(); + + let t = Instant::now(); + let forge = build_forge(args.pipeline, &args.embed_dir).await?; + let build = t.elapsed(); + + let t = Instant::now(); args.ingest.run(&forge, inst, scope).await?; + let ingest = t.elapsed(); + let t = Instant::now(); // Unbounded query → ranked session order for Recall@k / NDCG@k. let ranked = forge .query(&inst.question, Some(scope), UNBOUNDED_BUDGET) @@ -251,11 +280,19 @@ async fn run_instance(args: &Args, inst: &Instance) -> anyhow::Result = entries.into_iter().filter_map(|e| e.session_id).collect(); budget_sessions.insert(budget, sessions); } + let query = t.elapsed(); - Ok(InstanceResult { - ranked_sessions, - budget_sessions, - }) + Ok(( + InstanceResult { + ranked_sessions, + budget_sessions, + }, + Timing { + build, + ingest, + query, + }, + )) } fn accumulate( From f977cdf323544def2f49cd4574e13200fa9849d9 Mon Sep 17 00:00:00 2001 From: Austin Varnon Date: Tue, 7 Jul 2026 00:18:48 -0500 Subject: [PATCH 4/4] 1c pluggable embedder implemented --- .gitignore | 4 ++ CHANGELOG.md | 2 + benchmarks/longmemeval/src/main.rs | 45 ++++++++++++---- src/builder.rs | 45 ++++++++++++++-- src/engine.rs | 13 ++--- src/lib.rs | 3 ++ tests/embedder_injection.rs | 86 ++++++++++++++++++++++++++++++ 7 files changed, 178 insertions(+), 20 deletions(-) create mode 100644 tests/embedder_injection.rs diff --git a/.gitignore b/.gitignore index bc282c0..30a55d1 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,7 @@ spikes/ # Benchmark datasets (large, fetched manually — see benchmarks/*/README.md) benchmarks/*/data/ + +# Embedding model cache (downloaded ONNX weights + tokenizer, ~90 MB) +/models/ +.fastembed_cache/ diff --git a/CHANGELOG.md b/CHANGELOG.md index c3a357e..8c50fa8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ All notable changes to this project will be documented in this file. ### Features +- **Pluggable embedder — `ContextForgeBuilder::with_embedder`.** Inject any `Embedder` implementation as an `Arc`, so one loaded model can be shared across many `ContextForge` instances (load once, clone the `Arc`) or an alternate/remote backend can be supplied. `with_embedding_model(dir)` remains as a convenience that builds a `FasEmbedder`. `Embedder` (and `FasEmbedder`, under the `semantic` feature) are now re-exported at the crate root. `ContextEngine::with_embedder` is now `pub` and takes `Arc`. +- **Batched writes — `ContextForge::save_batch`.** Persist many entries with a single search-index commit and a single prepared-statement transaction, instead of per-entry commits and per-entry SQL re-parses. `distill_and_save` now uses this path internally. Measured ~45% faster bulk ingest on the LongMemEval benchmark; no change to retrieval results. - **`OpenAiCompatDistiller::with_api_key`**: bearer token support for authenticated cloud endpoints. Pass any OpenAI-compatible API key (OpenAI, Groq, Together AI, GLM/ZhipuAI, etc.) and it is sent as `Authorization: Bearer ` on every request. The distiller now uses `rustls` and supports `https://` base URLs; the previous `http://`-only restriction is lifted. ## [0.8.1] - 2026-07-05 diff --git a/benchmarks/longmemeval/src/main.rs b/benchmarks/longmemeval/src/main.rs index 04eff78..e7216f3 100644 --- a/benchmarks/longmemeval/src/main.rs +++ b/benchmarks/longmemeval/src/main.rs @@ -19,9 +19,10 @@ mod metrics; use std::collections::{BTreeMap, HashSet}; use std::path::PathBuf; +use std::sync::Arc; use std::time::{Duration, Instant}; -use context_forge::{Config, ContextForge}; +use context_forge::{Config, ContextForge, Embedder, FasEmbedder}; use dataset::Instance; use ingest::Ingest; @@ -121,18 +122,22 @@ fn next(it: &mut impl Iterator, flag: &str) -> anyhow::Result, + shared_embedder: &Option>, ) -> anyhow::Result { // Config is #[non_exhaustive] — build from Default, then set fields. let mut config = Config::default(); config.db_path = PathBuf::from(":memory:"); config.max_entries = 100_000_000; // never evict during a run - let require_dir = |embed_dir: &Option| { - embed_dir.clone().ok_or_else(|| { - anyhow::anyhow!("--embed-dir (or CF_EMBED_MODEL_DIR) required for this pipeline") - }) + let require_embedder = || { + shared_embedder + .clone() + .ok_or_else(|| anyhow::anyhow!("this pipeline requires an embedder (--embed-dir)")) }; let forge = match pipeline { Pipeline::Bm25 => ContextForge::open(config).await?, @@ -144,14 +149,14 @@ async fn build_forge( } Pipeline::Semantic => { ContextForge::builder(config) - .with_embedding_model(require_dir(embed_dir)?) + .with_embedder(require_embedder()?) .build() .await? } Pipeline::Full => { ContextForge::builder(config) .with_default_english_scorer() - .with_embedding_model(require_dir(embed_dir)?) + .with_embedder(require_embedder()?) .build() .await? } @@ -206,13 +211,27 @@ async fn main() -> anyhow::Result<()> { args.ingest.label(), ); + // Load the embedding model ONCE for the whole run (not per instance) and + // inject it into every ContextForge via with_embedder. Requires 1c's public + // Embedder injection; before it, each instance reloaded the model. + let shared_embedder: Option> = match args.pipeline { + Pipeline::Semantic | Pipeline::Full => { + let dir = args.embed_dir.clone().ok_or_else(|| { + anyhow::anyhow!("--embed-dir (or CF_EMBED_MODEL_DIR) required for this pipeline") + })?; + eprintln!(" loading embedding model once from {}...", dir.display()); + Some(Arc::new(FasEmbedder::new(&dir)?)) + } + Pipeline::Bm25 | Pipeline::Lexicon => None, + }; + let mut overall = Agg::default(); let mut by_type: BTreeMap = BTreeMap::new(); let mut errors = 0usize; let mut total = Timing::default(); for (i, inst) in instances.iter().enumerate() { - match run_instance(&args, inst).await { + match run_instance(&args, inst, &shared_embedder).await { Ok((result, timing)) => { total.build += timing.build; total.ingest += timing.ingest; @@ -255,11 +274,15 @@ struct Timing { query: Duration, } -async fn run_instance(args: &Args, inst: &Instance) -> anyhow::Result<(InstanceResult, Timing)> { +async fn run_instance( + args: &Args, + inst: &Instance, + shared_embedder: &Option>, +) -> anyhow::Result<(InstanceResult, Timing)> { let scope = inst.question_id.as_str(); let t = Instant::now(); - let forge = build_forge(args.pipeline, &args.embed_dir).await?; + let forge = build_forge(args.pipeline, shared_embedder).await?; let build = t.elapsed(); let t = Instant::now(); diff --git a/src/builder.rs b/src/builder.rs index f55a633..0f8c793 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -7,6 +7,8 @@ use crate::config::Config; use crate::engine::ContextEngine; use crate::lexicon::{CompositeLexiconScorer, DefaultEnglishScorer, LexiconScorer}; use crate::scrub::ScrubConfig; +#[cfg(feature = "semantic")] +use crate::semantic::{Embedder, FasEmbedder}; use crate::storage::open_storage; use crate::traits::Result; use crate::ContextForge; @@ -64,6 +66,8 @@ pub struct ContextForgeBuilder { persona_scorer: Option>, #[cfg(feature = "semantic")] embedding_cache_dir: Option, + #[cfg(feature = "semantic")] + embedder: Option>, } impl ContextForgeBuilder { @@ -78,6 +82,8 @@ impl ContextForgeBuilder { persona_scorer: None, #[cfg(feature = "semantic")] embedding_cache_dir: None, + #[cfg(feature = "semantic")] + embedder: None, } } @@ -116,7 +122,15 @@ impl ContextForgeBuilder { /// (~22 MB). The model is downloaded automatically on first use; subsequent /// starts load from the local cache. /// + /// This is a convenience wrapper over + /// [`with_embedder`](Self::with_embedder) that constructs a [`FasEmbedder`] + /// at [`build`](Self::build) time. To reuse one already-loaded model across + /// many `ContextForge` instances, or to plug in a different backend, use + /// [`with_embedder`](Self::with_embedder) directly. + /// /// Requires the `semantic` Cargo feature. + /// + /// [`FasEmbedder`]: crate::semantic::FasEmbedder #[cfg(feature = "semantic")] #[must_use] pub fn with_embedding_model(mut self, cache_dir: impl AsRef) -> Self { @@ -124,6 +138,21 @@ impl ContextForgeBuilder { self } + /// Inject an [`Embedder`](crate::semantic::Embedder) for semantic search. + /// + /// Accepts any implementation as an `Arc`, so a single loaded + /// model can be shared across many builds (load once, clone the `Arc`), or a + /// custom/remote backend can be supplied. Takes precedence over + /// [`with_embedding_model`](Self::with_embedding_model) if both are set. + /// + /// Requires the `semantic` Cargo feature. + #[cfg(feature = "semantic")] + #[must_use] + pub fn with_embedder(mut self, embedder: Arc) -> Self { + self.embedder = Some(embedder); + self + } + /// Open the database and build a [`ContextForge`] with the configured scorers. /// /// Lexicon scoring is applied only if @@ -162,10 +191,20 @@ impl ContextForgeBuilder { engine = engine.with_scorer(scorer); } + // A directly-injected embedder wins; otherwise construct a FasEmbedder + // from the configured cache dir if one was set. #[cfg(feature = "semantic")] - if let Some(cache_dir) = self.embedding_cache_dir { - let embedder = crate::semantic::FasEmbedder::new(&cache_dir)?; - engine = engine.with_embedder(Arc::new(embedder)); + { + let embedder: Option> = match self.embedder { + Some(e) => Some(e), + None => match self.embedding_cache_dir { + Some(cache_dir) => Some(Arc::new(FasEmbedder::new(&cache_dir)?)), + None => None, + }, + }; + if let Some(embedder) = embedder { + engine = engine.with_embedder(embedder); + } } Ok(ContextForge::from_parts(engine, scrub_config)) diff --git a/src/engine.rs b/src/engine.rs index d625e95..e89ea11 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,10 +8,9 @@ use crate::config::{Config, DEFAULT_RECENCY_HALF_LIFE_SECS}; use crate::entry::ContextEntry; use crate::error::Error; use crate::lexicon::LexiconScorer; -use crate::traits::{ContextStorage, Result, Searcher}; -// Bring the Embedder trait into scope for method dispatch on Arc. #[cfg(feature = "semantic")] -use crate::semantic::Embedder as _; +use crate::semantic::Embedder; +use crate::traits::{ContextStorage, Result, Searcher}; /// Default candidate limit when fetching search results for assembly. const DEFAULT_SEARCH_LIMIT: usize = 50; @@ -59,7 +58,7 @@ pub struct ContextEngine { config: Config, scorer: Option>, #[cfg(feature = "semantic")] - embedder: Option>, + embedder: Option>, } impl ContextEngine { @@ -105,10 +104,12 @@ impl ContextEngine { /// /// When set, [`Self::save_snapshot`] generates and stores an embedding for /// each new entry, and [`Self::assemble`] blends BM25 and semantic - /// candidates via Reciprocal Rank Fusion (RRF, k=60). + /// candidates via Reciprocal Rank Fusion (RRF, k=60). Accepts any + /// [`Embedder`](crate::semantic::Embedder) implementation, so callers can + /// inject a shared instance (load the model once) or an alternate backend. #[cfg(feature = "semantic")] #[must_use] - pub(crate) fn with_embedder(mut self, embedder: Arc) -> Self { + pub fn with_embedder(mut self, embedder: Arc) -> Self { self.embedder = Some(embedder); self } diff --git a/src/lib.rs b/src/lib.rs index 5340409..463b9e3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -113,6 +113,9 @@ pub use lexicon::{ LexiconAppender, LexiconConfig, LexiconPatterns, LexiconProposal, LexiconScorer, }; pub use scrub::{scrub_secrets, ScrubConfig}; +pub use semantic::Embedder; +#[cfg(feature = "semantic")] +pub use semantic::FasEmbedder; pub use session::{group_entries_by_session, SessionGroup}; pub use storage::{open_storage, TursoSearcher, TursoStorage}; pub use traits::{ContextStorage, Result, Searcher}; diff --git a/tests/embedder_injection.rs b/tests/embedder_injection.rs new file mode 100644 index 0000000..dded215 --- /dev/null +++ b/tests/embedder_injection.rs @@ -0,0 +1,86 @@ +//! Integration test for the public `Embedder` injection path (Phase 1c). +//! +//! Verifies that a caller-defined [`Embedder`] can be injected via +//! `ContextForgeBuilder::with_embedder` — the API that lets a single loaded +//! model be shared across many builds, or an alternate backend be plugged in. +//! Semantic-gated: the injection API and semantic search require the feature. +#![cfg(feature = "semantic")] + +use std::path::PathBuf; +use std::sync::Arc; + +use context_forge::{kind, Config, ContextForge, Embedder, SaveOptions}; + +/// A trivial, deterministic embedder — proves an arbitrary `Embedder` impl (not +/// just the built-in `FasEmbedder`) is accepted and driven end to end. +struct ConstantEmbedder { + dims: usize, +} + +impl Embedder for ConstantEmbedder { + fn embed(&self, text: &str) -> context_forge::Result> { + // Deterministic vector from the text so distinct inputs differ slightly. + let seed = (text.len() % 13) as f32 / 13.0; + Ok(vec![seed; self.dims]) + } +} + +fn in_memory_config() -> Config { + let mut cfg = Config::default(); + cfg.db_path = PathBuf::from(":memory:"); + cfg +} + +#[tokio::test] +async fn builder_accepts_injected_custom_embedder() { + // all-MiniLM-L6-v2 is 384-dim; the schema's F32_BLOB is sized for that. + let embedder: Arc = Arc::new(ConstantEmbedder { dims: 384 }); + + let cf = ContextForge::builder(in_memory_config()) + .with_embedder(embedder) + .build() + .await + .expect("build with injected embedder"); + + // Saving drives the injected embedder (embed_and_store) without error. + let opts = SaveOptions::default(); + cf.save("the deploy used a canary rollout", kind::SNAPSHOT, &opts) + .await + .expect("save with injected embedder"); + cf.save("lunch options near the office", kind::SNAPSHOT, &opts) + .await + .expect("save second entry"); + + // Query runs BM25 + semantic (via the injected embedder) and returns without + // error — proving the injected embedder is wired through the whole path. + let hits = cf + .query("deploy rollout", None, 4096) + .await + .expect("query with injected embedder"); + assert!( + hits.iter().any(|e| e.content.contains("canary")), + "the relevant entry should be retrieved" + ); +} + +#[tokio::test] +async fn injected_embedder_shared_across_builds() { + // The core 1c use case: one Arc shared across multiple ContextForge builds + // (load once, clone the Arc) rather than reconstructing per instance. + let shared: Arc = Arc::new(ConstantEmbedder { dims: 384 }); + + for _ in 0..3 { + let cf = ContextForge::builder(in_memory_config()) + .with_embedder(Arc::clone(&shared)) + .build() + .await + .expect("build sharing one embedder"); + cf.save( + "shared embedder entry", + kind::SNAPSHOT, + &SaveOptions::default(), + ) + .await + .expect("save"); + } +}