Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

### Changed

- Embedding calls go through one module, #1538 (part of #1543). Initialisation, the embedder lock and the retry now live in `embedding.rs` (`Uteke::embed`, `embed_once`, `embedder_max_seq_len`); the ten call sites (recall, explain, document search, document chunks, room recall, update, import, `embed_text`, writes, reembed) ask it, and the five `expect("embedder ensured above")` panics are gone. Behaviour change: query, document and room embedding now retry a transient failure like writes did (3 attempts, 200/400/800 ms) instead of failing at once; `reembed` takes the lock per row instead of for the whole loop and still makes one attempt per row, so a repair over a dead backend reports each failed row without waiting. A write that cannot generate a vector is still stored FTS5-only with a warning; an embedder that cannot start is still an error.
- Internal: tests can supply their own embedder through the public open path (#1542, part of #1543). `test_support` gains `open_with_embedder`, `open_with_embedder_at` and a `RecordingEmbedder` (records every text, can fail its first calls); the 33 core tests that were `#[ignore]`d because `Uteke::open` needs the ONNX runtime now run in CI with the deterministic keyword embedder (the document-chunk, wikilink-enrichment, room, dedup, fusion and verify/repair tests among them). No production code change.
- Internal: one `EngineSettings` and one apply step, #1515 part 3 (round 2 of the architecture deepening, #1523). `uteke_core::EngineSettings` holds `[tier]`, `[recall]`, `[dream]`, `[lifecycle]` and `[embed_fallback]`, and `FallbackSettings` is deserializable like the others; the CLI config and the server config embed it with `#[serde(flatten)]` instead of writing the five sections twice, and `Uteke::apply_settings` replaces the sequence of setters that the CLI and the server each repeated (the CLI no longer applies the fallback in a separate block; the server logs the dream and lifecycle settings on every start, not only when the section exists). MCP does not use it yet (next part). No behaviour change, checked by running the CLI and the server built from `develop` and from this branch on the same config with all five sections, a partial file and an invalid one.
- Internal: `[recall]` is read into one core type, #1515 part 2 (round 2 of the architecture deepening, #1523). `uteke_core::RecallSettings` owns the nine keys, their defaults and the `UTEKE_RECALL_*` / `UTEKE_GRAPH_*` overrides with their range checks (moved from the CLI, tested with an injected lookup); the CLI config uses it as its `[recall]` type and the server drops `RecallFileSection` and its `Option` chains. The server still does not read the `UTEKE_RECALL_*` variables: that is decided in the apply step of part 3. No behaviour change, checked by running the CLI and the server built from `develop` and from this branch with the same `[recall]` section and environment.
Expand Down
112 changes: 112 additions & 0 deletions crates/uteke-core/src/embedding.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
//! The one way the engine asks its embedder for a vector (#1538).
//!
//! Every embedding call goes through here, so lazy initialisation, the embedder
//! lock and the retry live in one place and no call site can reach the embedder
//! slot without having initialised it:
//!
//! - [`Uteke::embed`] initialises the embedder if needed and retries a transient
//! failure (3 attempts, 200/400/800 ms); the caller decides what a final
//! failure means (an error, or FTS5-only storage for a write).
//! - [`Uteke::embed_once`] is the same without the retry, for a bulk loop that
//! reports each failed row instead of waiting on a dead backend per row.
//!
//! The lock is held for one embed call and never across a loop.

use crate::{Error, Uteke};

impl Uteke {
/// Embed `text`, initialising the embedder on first use and retrying a
/// transient failure. A validation error is returned at once.
pub(crate) fn embed(&self, text: &str) -> Result<Vec<f32>, Error> {
self.ensure_embedder()?;
crate::operations::retry_embed(&self.embedder, text)
}

/// Embed `text` once: no retry, for loops that count failures per row.
pub(crate) fn embed_once(&self, text: &str) -> Result<Vec<f32>, Error> {
self.ensure_embedder()?;
self.with_embedder(|e| e.embed(text))
}

/// The embedder's maximum sequence length, initialising it if needed.
pub(crate) fn embedder_max_seq_len(&self) -> Result<usize, Error> {
self.ensure_embedder()?;
self.with_embedder(|e| Ok(e.max_seq_len()))
}

fn with_embedder<T>(
&self,
f: impl FnOnce(&dyn crate::embed::Embedder) -> Result<T, Error>,
) -> Result<T, Error> {
let guard = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock"))?;
let embedder = guard
.as_deref()
.ok_or_else(|| Error::embed_msg("No embedder configured"))?;
f(embedder)
}
}

#[cfg(test)]
mod tests {
use crate::test_support::{RecordingEmbedder, Seen, open_with_embedder};

fn store(failures: usize) -> (crate::Uteke, Seen) {
let (e, seen) = RecordingEmbedder::failing_first(failures);
(open_with_embedder(Box::new(e)), seen)
}

#[test]
fn a_query_retries_a_transient_failure_like_a_write() {
let (u, seen) = store(1);
u.recall("rust", 5, None, None, 0.0, None, None).unwrap();
assert_eq!(seen.lock().unwrap().len(), 2, "one failure, one retry");
}

#[test]
fn document_search_and_embed_text_retry_too() {
let (u, seen) = store(1);
u.doc_search_semantic("rust", 3).unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
let (u, seen) = store(1);
assert_eq!(u.embed_text("rust").unwrap().len(), 64);
assert_eq!(seen.lock().unwrap().len(), 2);
}

#[test]
fn a_persistent_failure_is_an_error_after_three_attempts() {
let (u, seen) = store(3);
assert!(u.embed_text("rust").is_err());
assert_eq!(seen.lock().unwrap().len(), 3);
}

#[test]
fn embed_once_does_not_retry() {
let (u, seen) = store(1);
assert!(u.embed_once("rust").is_err());
assert_eq!(seen.lock().unwrap().len(), 1);
}

#[test]
fn a_write_that_cannot_embed_is_fts5_only_with_a_warning_not_an_error() {
let (u, _) = store(3);
let out = u.remember_detailed("kept anyway", &[], None, None).unwrap();
assert!(!out.embedding_written);
assert!(out.warning.is_some());
}

#[test]
fn reembed_repairs_rows_with_one_attempt_each() {
let (u, seen) = store(0);
u.remember("one", &[], None, None).unwrap();
let before = seen.lock().unwrap().len();
u.graph_store()
.execute("UPDATE memories SET embedding = NULL", rusqlite::params![])
.unwrap();
let report = u.reembed_missing().unwrap();
assert_eq!(report.reembedded, 1);
assert_eq!(seen.lock().unwrap().len(), before + 1);
}
}
3 changes: 1 addition & 2 deletions crates/uteke-core/src/import_export.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,8 +124,7 @@ impl crate::Uteke {
let redactions = prepared.redactions;

// Re-embed the content with retry (consistent with remember path, #1005).
self.ensure_embedder()?;
let embedding = crate::operations::retry_embed(&self.embedder, &entry.content)?;
let embedding = self.embed(&entry.content)?;

// Dedup check: skip if cosine >= 0.95 to existing memory (#1005).
if let Some(_existing_id) = self.check_duplicate(&embedding, namespace)? {
Expand Down
49 changes: 8 additions & 41 deletions crates/uteke-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ pub mod consolidation_plan;
pub mod dream;
mod edges;
mod embed;
mod embedding;
#[cfg(test)]
mod engine_settings_tests;
mod error;
Expand Down Expand Up @@ -752,13 +753,7 @@ impl Uteke {
/// Used by the consolidation executor to embed consolidated records
/// before writing them so they stay vector-searchable.
pub fn embed_text(&self, text: &str) -> Result<Vec<f32>, Error> {
self.ensure_embedder()?;
self.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during embed_text"))?
.as_ref()
.expect("embedder ensured above")
.embed(text)
self.embed(text)
}

/// Add an already-persisted memory's embedding to the vector index.
Expand Down Expand Up @@ -1848,37 +1843,16 @@ impl Uteke {
content: &str,
tags: &[String],
) -> Result<Vec<(DocumentChunk, Vec<f32>)>, Error> {
self.ensure_embedder()?;
// The embedder mutex is shared with recall/remember (query embedding),
// so take it per call, never across the whole loop: a long document
// must not make every other embed wait for all of its chunks (#1323).
let max_chars = {
let guard = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during document chunking"))?;
guard
.as_ref()
.expect("embedder ensured above")
.max_seq_len()
.saturating_mul(4)
.max(1024)
};
// The embedder lock is taken per call, never across the whole loop: a
// long document must not make every other embed wait for all of its
// chunks (#1323).
let max_chars = self.embedder_max_seq_len()?.saturating_mul(4).max(1024);
let chunks = crate::chunker::chunk_markdown(content, max_chars);
chunks
.iter()
.enumerate()
.map(|(i, chunk)| {
let embedding = {
let guard = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during chunk embedding"))?;
guard
.as_ref()
.expect("embedder ensured above")
.embed(&chunk.content)?
};
let embedding = self.embed(&chunk.content)?;
Ok((
DocumentChunk {
id: uuid::Uuid::now_v7().to_string(),
Expand Down Expand Up @@ -2198,14 +2172,7 @@ impl Uteke {
) -> Result<Vec<crate::memory::documents::DocumentSearchResult>, Error> {
use crate::memory::documents::DocumentSearchResult;

self.ensure_embedder()?;
let query_embedding = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during doc search"))?
.as_ref()
.expect("embedder ensured above")
.embed(query)?;
let query_embedding = self.embed(query)?;

let index = self
.index
Expand Down
17 changes: 5 additions & 12 deletions crates/uteke-core/src/maintenance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -294,25 +294,18 @@ impl crate::Uteke {
"Re-embedding memories with missing vectors"
);

// Ensure embedder is initialized.
// An embedder that cannot start is an error, not one failed row each.
self.ensure_embedder()?;

let mut reembedded = 0usize;
let mut failed = 0usize;
let mut new_items: Vec<(String, Vec<f32>)> = Vec::new();

// Acquire embedder lock ONCE before the loop to avoid
// per-iteration mutex overhead (codecoradev/uteke#919 review).
let embedder_guard = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during reembed"))?;
let embedder = embedder_guard
.as_ref()
.ok_or_else(|| Error::embed("reembed", "embedder not initialized"))?;

// One attempt per row, and the embedder lock is taken per row: a repair
// over a dead backend reports each failed row instead of waiting on a
// retry for every one, and does not block recall for the whole loop.
for mem in &missing {
match embedder.embed(&mem.content) {
match self.embed_once(&mem.content) {
Ok(vec) => {
// Update memory in database.
let mut updated = (*mem).clone();
Expand Down
26 changes: 7 additions & 19 deletions crates/uteke-core/src/operations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,7 @@ pub(crate) fn retry_embed(
.lock()
.map_err(|_| Error::lock("embedder lock during remember"))?;
let Some(embedder) = lock.as_ref() else {
return Err(Error::embed_msg(
"No embedder configured; the memory was not embedded",
));
return Err(Error::embed_msg("No embedder configured"));
};
match embedder.embed(text) {
Ok(embedding) => return Ok(embedding),
Expand Down Expand Up @@ -331,10 +329,11 @@ impl crate::Uteke {
tracing::debug!("No embedding backend configured; storing memory FTS5-only (#1166)");
(Vec::new(), None)
} else {
// Failing to initialise the embedder is an error; failing to generate
// a vector, after the retry in `embed`, is FTS5-only with a warning
// (embedding failures silently drop vector entries, #621).
self.ensure_embedder()?;
// Retry embedding generation up to 3 times with exponential backoff.
// Embedding failures silently drop vector entries, causing desync (#621).
match self::retry_embed(&self.embedder, &embed_text) {
match self.embed(&embed_text) {
Ok(e) => (e, None),
Err(e) => {
tracing::warn!(
Expand Down Expand Up @@ -720,17 +719,7 @@ impl crate::Uteke {
) -> Result<Vec<SearchResult>, Error> {
// Embed query outside any lock — CPU-intensive (~50ms), no shared state needed.
// Only the embedder Mutex is held here, allowing concurrent index reads.
// Lazy-load embedder on first use.
self.ensure_embedder()?;
let query_embedding = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during recall"))?
.as_ref()
.ok_or_else(|| {
Error::embed_msg("No embedder configured; semantic recall is unavailable")
})?
.embed(query)?;
let query_embedding = self.embed(query)?;
// Embedder lock dropped here — other threads can embed or recall concurrently.

// Search usearch index with retry: if post-filtering removes too many
Expand Down Expand Up @@ -2051,8 +2040,7 @@ impl crate::Uteke {
} else {
c.to_string()
};
self.ensure_embedder()?;
let new_embedding = retry_embed(&self.embedder, &embed_text)?;
let new_embedding = self.embed(&embed_text)?;

// Update SQLite first, then vector index
let updated = self.store.update_fields(
Expand Down
10 changes: 1 addition & 9 deletions crates/uteke-core/src/recall_explain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,15 +113,7 @@ impl Uteke {
let query_embedding: Option<Vec<f32>> = if strategy == RecallStrategy::Fts5 {
None
} else {
self.ensure_embedder()?;
Some(
self.embedder
.lock()
.map_err(|_| Error::lock("embedder lock during recall_explained"))?
.as_ref()
.ok_or_else(|| Error::embed_msg("no embedder configured"))?
.embed(query)?,
)
Some(self.embed(query)?)
};

let boost_window = limit.saturating_mul(4).saturating_add(16);
Expand Down
9 changes: 1 addition & 8 deletions crates/uteke-core/src/room_ops.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,14 +188,7 @@ impl crate::Uteke {

if !missing_memories.is_empty() {
// Embed query and score each missing memory by cosine similarity.
self.ensure_embedder()?;
let query_emb = self
.embedder
.lock()
.map_err(|_| Error::lock("embedder in room fallback"))?
.as_ref()
.expect("embedder ensured above")
.embed(query)?;
let query_emb = self.embed(query)?;

for mem in missing_memories {
let score = crate::consolidate::cosine_similarity(&query_emb, &mem.embedding);
Expand Down
Loading