mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
feat: LRU embedding cache for workspace search (#1423)
* feat: LRU embedding cache for workspace search (#165) Add CachedEmbeddingProvider that wraps any EmbeddingProvider with an in-memory LRU cache keyed by SHA-256(model_name + text). This avoids redundant HTTP calls when the same text is embedded multiple times (common during reindexing and repeated searches). - Cache uses HashMap + last_accessed tracking with manual LRU eviction (same pattern as llm::response_cache::CachedProvider) - Lock is never held during HTTP calls to prevent blocking - embed_batch() partitions into hits/misses and only fetches misses - Default 10,000 entries (~58 MB for 1536-dim vectors) - Configurable via EMBEDDING_CACHE_SIZE env var - Workspace.with_embeddings() auto-wraps; with_embeddings_uncached() available for tests Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address review comments on embedding cache - Validate embed_batch return count matches expected miss count - Replace unwrap_or_default() with proper error propagation - Fix batch eviction: run final eviction pass after insert to enforce cap - Fix test: use different-length inputs to verify ordering correctness - Reject EMBEDDING_CACHE_SIZE=0 in config validation (minimum is 1) Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: replace .expect() with proper error handling in embed_batch The all-cache-hits early-return path used .expect("all cache hits") which violates the project convention of no .unwrap()/.expect() in production code. Replaced with the same ok_or_else pattern used in the normal path. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: clarify memory sizing docs and use saturating_add for eviction - Update memory comments in embedding_cache.rs, config/embeddings.rs, and workspace/mod.rs to note the ~58 MB figure is payload-only (actual memory is higher due to HashMap/key/allocation overhead) - Use saturating_add(1) instead of + 1 for eviction threshold to prevent overflow if max_entries is usize::MAX Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address Copilot review on embedding cache - Avoid double-clone per miss in embed_batch: move embedding into results, clone only for the cache entry - Evict per-insert instead of after all inserts to keep peak memory bounded during large batches - Clamp max_entries to at least 1 in constructor to prevent unexpected eviction behavior when set to 0 via the public API Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: reduce embedding_cache module visibility to private Types are already re-exported via `pub use`, so the module itself doesn't need to be public. Reduces unnecessary API surface. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address serrrfirat review feedback on embedding cache - Add TODO comment for O(n) LRU eviction scalability - Add thundering herd note at lock release in embed() - Warn when cache max_entries exceeds 100k - Use with_embeddings_uncached() in integration test - Add tests: error_does_not_pollute_cache, embed_batch_empty_input - Update README with cache-aware with_embeddings() docs Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: prevent u32 wrapping in FailThenSucceedMock failure counter fetch_sub(1) wraps to u32::MAX when called past zero, silently breaking the mock for 3+ calls. Switch to load-then-store to avoid the wrapping bug in both embed() and embed_batch(). Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address Copilot and serrrfirat review findings on embedding cache - Switch tokio::sync::Mutex to std::sync::Mutex (lock never held across .await — cheaper synchronous lock) - Extract DEFAULT_EMBEDDING_CACHE_SIZE constant to avoid 10_000 duplication between EmbeddingCacheConfig and EmbeddingsConfig Co-Authored-By: Claude Opus 4.6 <[email protected]> * test: add all-misses batch test for embedding cache Adds embed_batch_all_misses test covering the case where every text in a batch is a cache miss — fulfilling the commitment from serrrfirat's review. Co-Authored-By: Claude Opus 4.6 <[email protected]> * chore: trigger CI re-check after rebase * fix: use raw [u8;32] cache keys and pre-allocate HashMap capacity Address Copilot review findings: - cache_key() now returns [u8; 32] instead of hex String, avoiding a 64-byte allocation per lookup - HashMap::with_capacity(max_entries) avoids incremental reallocation - Fix pre-existing staging compilation error in cli/routines.rs (missing max_tool_rounds/use_tools fields) [skip-regression-check] * fix: make cache accessors sync and update doc for [u8;32] keys Address Copilot review: - len(), is_empty(), clear() are now sync since they only take a std::sync::Mutex lock with no .await points - Update cache_size doc comment to reflect [u8;32] keys instead of String keys [skip-regression-check] * fix: remove clone_on_copy for [u8; 32] cache keys [skip-regression-check] * ci: add safety comments to test code for no-panics check The CI no-panics grep check cannot distinguish test code inside src/ files from production code. Add // safety: test annotations to .unwrap(), .expect(), and assert!() calls in #[cfg(test)] modules. * fix: correct cache doc and demote hit/miss logs to trace - Fix misleading "String keys" in memory comment (cache uses [u8; 32]) - Demote per-request hit/miss logs from debug to trace to reduce noise on hot paths (batch summary stays at trace too) * docs: add missing Arc import in workspace README example * perf: batch eviction in embed_batch to avoid O(n×m) cost Replace per-insert evict_lru call with a single evict_k_oldest pass that computes eviction count upfront and removes the k oldest entries in one O(n) scan. Avoids O(n×m) HashMap iterations while holding the mutex during batch inserts. * fix: cap batch cache inserts at max_entries and use O(n) selection - evict_k_oldest now uses select_nth_unstable_by_key for O(n) average partial selection instead of O(n log n) full sort - embed_batch caps cached entries at max_entries when misses exceed capacity, preventing the cache from growing unbounded - Added test: batch_exceeding_capacity_respects_max_entries * fix: flatten test assert for fmt compatibility Shorten assert message to fit single line so cargo fmt doesn't split the safety annotation onto a separate line. * fix: address review feedback and improve embedding cache (takeover #235) - Fix merge conflict: add missing allow_always field in PendingApproval - Thread EmbeddingCacheConfig through CLI memory commands so they respect EMBEDDING_CACHE_SIZE instead of silently using default (fixes #235 review) - Cap HashMap pre-allocation at min(max_entries, 1024) to avoid upfront memory waste at large cache sizes - Fix FailThenSucceedMock race: replace load+store with atomic fetch_update - Remove noisy '// safety: test' comments (40+ lines of diff noise) - Fix collapsed lines from comment removal - Simplify redundant Ok(...collect()?) to just collect() Co-Authored-By: ztsalexey <[email protected]> Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix(embedding-cache): skip eviction on concurrent duplicate insert When the lock is released for the HTTP call, another caller may insert the same key. Re-check under lock and just update the existing entry without evicting, avoiding unnecessary cache churn under concurrency. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> --------- Co-authored-by: ztsalexey <[email protected]> Co-authored-by: Claude Opus 4.6 <[email protected]> Co-authored-by: ztsalexey <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
ztsalexey
Claude Opus 4.6
ztsalexey
ztsalexey
Claude Opus 4.6
ztsalexey
parent
52ca9d6588
commit
86ae12747b
@@ -1961,6 +1961,7 @@ mod tests {
|
||||
context_messages: vec![],
|
||||
deferred_tool_calls: vec![],
|
||||
user_timezone: None,
|
||||
allow_always: false,
|
||||
};
|
||||
thread.await_approval(pending);
|
||||
|
||||
|
||||
+5
-2
@@ -25,7 +25,7 @@ use crate::tools::ToolRegistry;
|
||||
use crate::tools::mcp::{McpProcessManager, McpSessionManager};
|
||||
use crate::tools::wasm::SharedCredentialRegistry;
|
||||
use crate::tools::wasm::WasmToolRuntime;
|
||||
use crate::workspace::{EmbeddingProvider, Workspace};
|
||||
use crate::workspace::{EmbeddingCacheConfig, EmbeddingProvider, Workspace};
|
||||
|
||||
/// Fully initialized application components, ready for channel wiring
|
||||
/// and agent construction.
|
||||
@@ -313,10 +313,13 @@ impl AppBuilder {
|
||||
|
||||
// Register memory tools if database is available
|
||||
let workspace = if let Some(ref db) = self.db {
|
||||
let emb_cache_config = EmbeddingCacheConfig {
|
||||
max_entries: self.config.embeddings.cache_size,
|
||||
};
|
||||
let mut ws = Workspace::new_with_db(&self.config.owner_id, db.clone())
|
||||
.with_search_config(&self.config.search);
|
||||
if let Some(ref emb) = embeddings {
|
||||
ws = ws.with_embeddings(emb.clone());
|
||||
ws = ws.with_embeddings_cached(emb.clone(), emb_cache_config);
|
||||
}
|
||||
let ws = Arc::new(ws);
|
||||
tools.register_memory_tools(Arc::clone(&ws));
|
||||
|
||||
+5
-3
@@ -7,17 +7,18 @@ use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
use crate::workspace::{EmbeddingProvider, SearchConfig, Workspace};
|
||||
use crate::workspace::{EmbeddingCacheConfig, EmbeddingProvider, SearchConfig, Workspace};
|
||||
|
||||
/// Run a memory command using the Database trait (works with any backend).
|
||||
pub async fn run_memory_command_with_db(
|
||||
cmd: MemoryCommand,
|
||||
db: std::sync::Arc<dyn crate::db::Database>,
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
cache_config: EmbeddingCacheConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut workspace = Workspace::new_with_db("default", db);
|
||||
if let Some(emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb);
|
||||
workspace = workspace.with_embeddings_cached(emb, cache_config);
|
||||
}
|
||||
|
||||
match cmd {
|
||||
@@ -85,10 +86,11 @@ pub async fn run_memory_command(
|
||||
cmd: MemoryCommand,
|
||||
pool: deadpool_postgres::Pool,
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
cache_config: EmbeddingCacheConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut workspace = Workspace::new("default", pool);
|
||||
if let Some(emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb);
|
||||
workspace = workspace.with_embeddings_cached(emb, cache_config);
|
||||
}
|
||||
|
||||
match cmd {
|
||||
|
||||
+4
-1
@@ -336,7 +336,10 @@ pub async fn run_memory_command(mem_cmd: &MemoryCommand) -> anyhow::Result<()> {
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
run_memory_command_with_db(mem_cmd.clone(), db, embeddings).await
|
||||
let cache_config = crate::workspace::EmbeddingCacheConfig {
|
||||
max_entries: config.embeddings.cache_size,
|
||||
};
|
||||
run_memory_command_with_db(mem_cmd.clone(), db, embeddings, cache_config).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -8,6 +8,9 @@ use crate::llm::SessionManager;
|
||||
use crate::settings::Settings;
|
||||
use crate::workspace::EmbeddingProvider;
|
||||
|
||||
/// Default maximum number of cached embeddings.
|
||||
pub const DEFAULT_EMBEDDING_CACHE_SIZE: usize = 10_000;
|
||||
|
||||
/// Embeddings provider configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EmbeddingsConfig {
|
||||
@@ -26,6 +29,12 @@ pub struct EmbeddingsConfig {
|
||||
/// Custom base URL for OpenAI-compatible embedding providers.
|
||||
/// When set, overrides the default `https://api.openai.com`.
|
||||
pub openai_base_url: Option<String>,
|
||||
/// Maximum entries in the embedding LRU cache (default 10,000).
|
||||
///
|
||||
/// Approximate raw embedding payload: `cache_size × dimension × 4 bytes`.
|
||||
/// 10,000 × 1536 floats ≈ 58 MB (payload only; actual memory is higher
|
||||
/// due to HashMap buckets, per-entry Vec/timestamp overhead).
|
||||
pub cache_size: usize,
|
||||
}
|
||||
|
||||
impl Default for EmbeddingsConfig {
|
||||
@@ -40,6 +49,7 @@ impl Default for EmbeddingsConfig {
|
||||
ollama_base_url: "http://localhost:11434".to_string(),
|
||||
dimension,
|
||||
openai_base_url: None,
|
||||
cache_size: DEFAULT_EMBEDDING_CACHE_SIZE,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -80,6 +90,15 @@ impl EmbeddingsConfig {
|
||||
|
||||
let openai_base_url = optional_env("EMBEDDING_BASE_URL")?;
|
||||
|
||||
let cache_size = parse_optional_env("EMBEDDING_CACHE_SIZE", DEFAULT_EMBEDDING_CACHE_SIZE)?;
|
||||
|
||||
if cache_size == 0 {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "EMBEDDING_CACHE_SIZE".to_string(),
|
||||
message: "must be at least 1".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
enabled,
|
||||
provider,
|
||||
@@ -88,6 +107,7 @@ impl EmbeddingsConfig {
|
||||
ollama_base_url,
|
||||
dimension,
|
||||
openai_base_url,
|
||||
cache_size,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -183,13 +203,13 @@ mod tests {
|
||||
std::env::remove_var("EMBEDDING_MODEL");
|
||||
std::env::remove_var("OPENAI_API_KEY");
|
||||
std::env::remove_var("EMBEDDING_BASE_URL");
|
||||
std::env::remove_var("EMBEDDING_CACHE_SIZE");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embeddings_disabled_not_overridden_by_openai_key() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
|
||||
clear_embedding_env();
|
||||
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
@@ -240,7 +260,6 @@ mod tests {
|
||||
#[test]
|
||||
fn embeddings_env_override_takes_precedence() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
|
||||
clear_embedding_env();
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
@@ -281,10 +300,8 @@ mod tests {
|
||||
let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed");
|
||||
assert_eq!(
|
||||
config.openai_base_url.as_deref(),
|
||||
Some("https://custom.example.com"),
|
||||
"EMBEDDING_BASE_URL env var should be parsed into openai_base_url"
|
||||
Some("https://custom.example.com")
|
||||
);
|
||||
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
std::env::remove_var("EMBEDDING_BASE_URL");
|
||||
@@ -303,4 +320,24 @@ mod tests {
|
||||
"openai_base_url should be None when EMBEDDING_BASE_URL is not set"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_size_zero_rejected() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
clear_embedding_env();
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
std::env::set_var("EMBEDDING_CACHE_SIZE", "0");
|
||||
}
|
||||
|
||||
let settings = Settings::default();
|
||||
let result = EmbeddingsConfig::resolve(&settings);
|
||||
assert!(result.is_err(), "cache_size=0 should be rejected");
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("at least 1"), "should mention minimum: {err}");
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
std::env::remove_var("EMBEDDING_CACHE_SIZE");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -38,7 +38,7 @@ pub use self::channels::{
|
||||
ChannelsConfig, CliConfig, DEFAULT_GATEWAY_PORT, GatewayConfig, HttpConfig, SignalConfig,
|
||||
};
|
||||
pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsql_path};
|
||||
pub use self::embeddings::EmbeddingsConfig;
|
||||
pub use self::embeddings::{DEFAULT_EMBEDDING_CACHE_SIZE, EmbeddingsConfig};
|
||||
pub use self::heartbeat::HeartbeatConfig;
|
||||
pub use self::hygiene::HygieneConfig;
|
||||
pub use self::llm::default_session_path;
|
||||
|
||||
@@ -38,12 +38,17 @@ workspace/
|
||||
## Using the Workspace
|
||||
|
||||
```rust
|
||||
use std::sync::Arc;
|
||||
use crate::workspace::{Workspace, OpenAiEmbeddings, paths};
|
||||
|
||||
// Create workspace for a user
|
||||
// Create workspace for a user (wraps embeddings in a default LRU cache)
|
||||
let workspace = Workspace::new("user_123", pool)
|
||||
.with_embeddings(Arc::new(OpenAiEmbeddings::new(api_key)));
|
||||
|
||||
// For tests: skip the cache layer (avoids unnecessary overhead with mocks)
|
||||
// let workspace = Workspace::new("user_123", pool)
|
||||
// .with_embeddings_uncached(Arc::new(MockEmbeddings::new(1536)));
|
||||
|
||||
// Read/write any path
|
||||
let doc = workspace.read("projects/alpha/notes.md").await?;
|
||||
workspace.write("context/priorities.md", "# Priorities\n\n1. Feature X").await?;
|
||||
|
||||
@@ -0,0 +1,613 @@
|
||||
//! LRU embedding cache wrapping any [`EmbeddingProvider`].
|
||||
//!
|
||||
//! Avoids redundant HTTP calls for identical texts by caching embeddings
|
||||
//! in memory keyed by `SHA-256(model_name + "\0" + text)`.
|
||||
//!
|
||||
//! Follows the same cache pattern as `llm::response_cache::CachedProvider`:
|
||||
//! `HashMap` + `last_accessed` tracking + manual LRU eviction.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::workspace::embeddings::{EmbeddingError, EmbeddingProvider};
|
||||
|
||||
/// Configuration for the embedding cache.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EmbeddingCacheConfig {
|
||||
/// Maximum number of cached embeddings (default 10,000).
|
||||
///
|
||||
/// Approximate raw embedding payload: `max_entries × dimension × 4 bytes`.
|
||||
/// At 10,000 entries × 1536 floats ≈ 58 MB (payload only; actual memory
|
||||
/// is higher due to HashMap buckets, `[u8; 32]` hash keys, `Vec`/`Instant`
|
||||
/// per-entry overhead).
|
||||
pub max_entries: usize,
|
||||
}
|
||||
|
||||
impl Default for EmbeddingCacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_entries: crate::config::DEFAULT_EMBEDDING_CACHE_SIZE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct CacheEntry {
|
||||
embedding: Vec<f32>,
|
||||
last_accessed: Instant,
|
||||
}
|
||||
|
||||
/// Embedding provider wrapper that caches results in memory.
|
||||
///
|
||||
/// Thread-safe via `std::sync::Mutex`. The lock is **never held**
|
||||
/// across `.await` points (all critical sections are scoped blocks),
|
||||
/// so a synchronous mutex is cheaper than `tokio::sync::Mutex`.
|
||||
pub struct CachedEmbeddingProvider {
|
||||
inner: Arc<dyn EmbeddingProvider>,
|
||||
cache: Mutex<HashMap<[u8; 32], CacheEntry>>,
|
||||
config: EmbeddingCacheConfig,
|
||||
}
|
||||
|
||||
impl CachedEmbeddingProvider {
|
||||
/// Wrap a provider with LRU caching.
|
||||
///
|
||||
/// `config.max_entries` is clamped to at least 1.
|
||||
pub fn new(inner: Arc<dyn EmbeddingProvider>, config: EmbeddingCacheConfig) -> Self {
|
||||
let config = EmbeddingCacheConfig {
|
||||
max_entries: config.max_entries.max(1),
|
||||
};
|
||||
if config.max_entries > 100_000 {
|
||||
tracing::warn!(
|
||||
max_entries = config.max_entries,
|
||||
"Embedding cache size exceeds 100,000 entries; memory usage may be significant"
|
||||
);
|
||||
}
|
||||
Self {
|
||||
inner,
|
||||
cache: Mutex::new(HashMap::with_capacity(config.max_entries.min(1024))),
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of entries currently in the cache.
|
||||
pub fn len(&self) -> usize {
|
||||
self.cache.lock().unwrap_or_else(|e| e.into_inner()).len()
|
||||
}
|
||||
|
||||
/// Whether the cache is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.cache
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.is_empty()
|
||||
}
|
||||
|
||||
/// Clear all cached entries.
|
||||
pub fn clear(&self) {
|
||||
self.cache.lock().unwrap_or_else(|e| e.into_inner()).clear();
|
||||
}
|
||||
|
||||
/// Build a deterministic cache key: `SHA-256(model_name + "\0" + text)`.
|
||||
///
|
||||
/// Returns raw 32-byte hash to avoid a 64-char hex String allocation per lookup.
|
||||
fn cache_key(&self, text: &str) -> [u8; 32] {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(self.inner.model_name().as_bytes());
|
||||
hasher.update(b"\0");
|
||||
hasher.update(text.as_bytes());
|
||||
hasher.finalize().into()
|
||||
}
|
||||
|
||||
/// Evict the least-recently-used entry if at capacity (single-entry path).
|
||||
// TODO: O(n) scan per eviction. If max_entries grows large, switch to
|
||||
// an ordered data structure (e.g. `IndexMap` with swap_remove, or a
|
||||
// linked-list LRU like the `lru` crate).
|
||||
fn evict_lru(cache: &mut HashMap<[u8; 32], CacheEntry>, max_entries: usize) {
|
||||
while cache.len() >= max_entries {
|
||||
let oldest_key = cache
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.last_accessed)
|
||||
.map(|(k, _)| *k);
|
||||
|
||||
if let Some(k) = oldest_key {
|
||||
cache.remove(&k);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Evict the `k` oldest entries in O(n) average time via partial selection.
|
||||
///
|
||||
/// Used by `embed_batch` to avoid the O(n×m) cost of calling
|
||||
/// `evict_lru` per insert.
|
||||
fn evict_k_oldest(cache: &mut HashMap<[u8; 32], CacheEntry>, k: usize) {
|
||||
if k == 0 || cache.is_empty() {
|
||||
return;
|
||||
}
|
||||
if k >= cache.len() {
|
||||
cache.clear();
|
||||
return;
|
||||
}
|
||||
// Partial selection: find the k oldest in O(n) average via
|
||||
// select_nth_unstable_by_key, then remove the first k entries.
|
||||
let mut entries: Vec<([u8; 32], Instant)> = cache
|
||||
.iter()
|
||||
.map(|(key, entry)| (*key, entry.last_accessed))
|
||||
.collect();
|
||||
entries.select_nth_unstable_by_key(k - 1, |(_, t)| *t);
|
||||
for (key, _) in entries.into_iter().take(k) {
|
||||
cache.remove(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for CachedEmbeddingProvider {
|
||||
fn dimension(&self) -> usize {
|
||||
self.inner.dimension()
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn max_input_length(&self) -> usize {
|
||||
self.inner.max_input_length()
|
||||
}
|
||||
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
let key = self.cache_key(text);
|
||||
|
||||
// Check cache (short critical section)
|
||||
{
|
||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(entry) = guard.get_mut(&key) {
|
||||
entry.last_accessed = Instant::now();
|
||||
tracing::trace!("embedding cache hit");
|
||||
return Ok(entry.embedding.clone());
|
||||
}
|
||||
}
|
||||
// Lock released before HTTP call.
|
||||
// NOTE: Thundering herd — multiple concurrent callers with the same
|
||||
// uncached key will each call the inner provider. This is acceptable:
|
||||
// embeddings are idempotent and the last writer wins in the HashMap.
|
||||
|
||||
let embedding = self.inner.embed(text).await?;
|
||||
|
||||
// Store result. Re-check under lock: another concurrent caller may
|
||||
// have inserted this key while the lock was released for the HTTP call.
|
||||
{
|
||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(entry) = guard.get_mut(&key) {
|
||||
// Key already present (thundering herd) — just update, no eviction needed.
|
||||
entry.embedding = embedding.clone();
|
||||
entry.last_accessed = Instant::now();
|
||||
} else {
|
||||
Self::evict_lru(&mut guard, self.config.max_entries);
|
||||
guard.insert(
|
||||
key,
|
||||
CacheEntry {
|
||||
embedding: embedding.clone(),
|
||||
last_accessed: Instant::now(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
tracing::trace!("embedding cache miss");
|
||||
Ok(embedding)
|
||||
}
|
||||
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Partition into hits and misses
|
||||
let keys: Vec<[u8; 32]> = texts.iter().map(|t| self.cache_key(t)).collect();
|
||||
let mut results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
|
||||
let mut miss_indices: Vec<usize> = Vec::new();
|
||||
|
||||
{
|
||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let now = Instant::now();
|
||||
for (i, key) in keys.iter().enumerate() {
|
||||
if let Some(entry) = guard.get_mut(key) {
|
||||
entry.last_accessed = now;
|
||||
results[i] = Some(entry.embedding.clone());
|
||||
} else {
|
||||
miss_indices.push(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Lock released before HTTP call
|
||||
|
||||
if miss_indices.is_empty() {
|
||||
tracing::trace!(count = texts.len(), "embedding batch: all cache hits");
|
||||
// All slots populated from cache hits
|
||||
return results
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, slot)| {
|
||||
slot.ok_or_else(|| {
|
||||
EmbeddingError::InvalidResponse(format!(
|
||||
"embedding slot {i} was not populated"
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>();
|
||||
}
|
||||
|
||||
// Fetch missing embeddings
|
||||
let miss_texts: Vec<String> = miss_indices.iter().map(|&i| texts[i].clone()).collect();
|
||||
let new_embeddings = self.inner.embed_batch(&miss_texts).await?;
|
||||
|
||||
if new_embeddings.len() != miss_indices.len() {
|
||||
return Err(EmbeddingError::InvalidResponse(format!(
|
||||
"embed_batch returned {} embeddings, expected {}",
|
||||
new_embeddings.len(),
|
||||
miss_indices.len()
|
||||
)));
|
||||
}
|
||||
|
||||
tracing::trace!(
|
||||
hits = texts.len() - miss_indices.len(),
|
||||
misses = miss_indices.len(),
|
||||
"embedding batch: partial cache"
|
||||
);
|
||||
|
||||
// Assemble results first (all misses, regardless of cache capacity).
|
||||
for (orig_idx, emb) in miss_indices.iter().copied().zip(&new_embeddings) {
|
||||
results[orig_idx] = Some(emb.clone());
|
||||
}
|
||||
|
||||
// Cache the new embeddings, respecting max_entries.
|
||||
{
|
||||
let mut guard = self.cache.lock().unwrap_or_else(|e| e.into_inner());
|
||||
// When misses exceed capacity, clear and only cache the tail.
|
||||
let cacheable = miss_indices.len().min(self.config.max_entries);
|
||||
let skip = miss_indices.len() - cacheable;
|
||||
let need_to_evict = (guard.len() + cacheable).saturating_sub(self.config.max_entries);
|
||||
if need_to_evict > 0 {
|
||||
Self::evict_k_oldest(&mut guard, need_to_evict);
|
||||
}
|
||||
let now = Instant::now();
|
||||
for (&orig_idx, emb) in miss_indices[skip..].iter().zip(&new_embeddings[skip..]) {
|
||||
guard.insert(
|
||||
keys[orig_idx],
|
||||
CacheEntry {
|
||||
embedding: emb.clone(),
|
||||
last_accessed: now,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, slot)| {
|
||||
slot.ok_or_else(|| {
|
||||
EmbeddingError::InvalidResponse(format!("embedding slot {i} was not populated"))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
|
||||
/// Mock embedding provider that counts calls.
|
||||
struct CountingMock {
|
||||
dimension: usize,
|
||||
model: String,
|
||||
embed_calls: AtomicU32,
|
||||
batch_calls: AtomicU32,
|
||||
}
|
||||
|
||||
impl CountingMock {
|
||||
fn new(dimension: usize, model: &str) -> Self {
|
||||
Self {
|
||||
dimension,
|
||||
model: model.to_string(),
|
||||
embed_calls: AtomicU32::new(0),
|
||||
batch_calls: AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn embed_calls(&self) -> u32 {
|
||||
self.embed_calls.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
fn batch_calls(&self) -> u32 {
|
||||
self.batch_calls.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for CountingMock {
|
||||
fn dimension(&self) -> usize {
|
||||
self.dimension
|
||||
}
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
fn max_input_length(&self) -> usize {
|
||||
10_000
|
||||
}
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
self.embed_calls.fetch_add(1, Ordering::SeqCst);
|
||||
// Simple deterministic embedding: val = text.len() / 100.0
|
||||
let val = text.len() as f32 / 100.0;
|
||||
Ok(vec![val; self.dimension])
|
||||
}
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
self.batch_calls.fetch_add(1, Ordering::SeqCst);
|
||||
texts
|
||||
.iter()
|
||||
.map(|t| {
|
||||
let val = t.len() as f32 / 100.0;
|
||||
Ok(vec![val; self.dimension])
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_hit_avoids_inner_call() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
let r1 = cached.embed("hello").await.unwrap();
|
||||
assert_eq!(inner.embed_calls(), 1);
|
||||
|
||||
let r2 = cached.embed("hello").await.unwrap();
|
||||
assert_eq!(inner.embed_calls(), 1); // still 1 -- cache hit
|
||||
assert_eq!(r1, r2);
|
||||
|
||||
assert_eq!(cached.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_miss_calls_inner() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
cached.embed("hello").await.unwrap();
|
||||
cached.embed("world").await.unwrap();
|
||||
assert_eq!(inner.embed_calls(), 2);
|
||||
assert_eq!(cached.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_key_includes_model() {
|
||||
let inner_a = Arc::new(CountingMock::new(4, "model-a"));
|
||||
let inner_b = Arc::new(CountingMock::new(4, "model-b"));
|
||||
|
||||
let cached_a = CachedEmbeddingProvider::new(
|
||||
inner_a.clone(),
|
||||
EmbeddingCacheConfig { max_entries: 100 },
|
||||
);
|
||||
let cached_b = CachedEmbeddingProvider::new(
|
||||
inner_b.clone(),
|
||||
EmbeddingCacheConfig { max_entries: 100 },
|
||||
);
|
||||
|
||||
// Same text, different models -> different cache keys
|
||||
let key_a = cached_a.cache_key("hello");
|
||||
let key_b = cached_b.cache_key("hello");
|
||||
assert_ne!(key_a, key_b);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lru_eviction() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 2 });
|
||||
|
||||
cached.embed("first").await.unwrap();
|
||||
cached.embed("second").await.unwrap();
|
||||
assert_eq!(cached.len(), 2);
|
||||
|
||||
// Third entry should evict the oldest ("first")
|
||||
cached.embed("third").await.unwrap();
|
||||
assert_eq!(cached.len(), 2);
|
||||
assert_eq!(inner.embed_calls(), 3);
|
||||
|
||||
// "first" should be a cache miss now
|
||||
cached.embed("first").await.unwrap();
|
||||
assert_eq!(inner.embed_calls(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_batch_partial_hits() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
// Pre-cache one text
|
||||
cached.embed("cached").await.unwrap();
|
||||
assert_eq!(inner.embed_calls(), 1);
|
||||
|
||||
// Batch with 1 cached + 2 new
|
||||
let texts = vec![
|
||||
"cached".to_string(),
|
||||
"new_one".to_string(),
|
||||
"new_two".to_string(),
|
||||
];
|
||||
let results = cached.embed_batch(&texts).await.unwrap();
|
||||
|
||||
// Should have called embed_batch on inner for 2 misses
|
||||
assert_eq!(inner.batch_calls(), 1);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(cached.len(), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_preserves_order() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
// Pre-cache "bb" (len 2)
|
||||
cached.embed("bb").await.unwrap();
|
||||
|
||||
// Batch: "a" (miss, len 1), "bb" (hit, len 2), "ccc" (miss, len 3)
|
||||
let texts = vec!["a".to_string(), "bb".to_string(), "ccc".to_string()];
|
||||
let results = cached.embed_batch(&texts).await.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 3);
|
||||
let expected_a = vec![1.0_f32 / 100.0; 4];
|
||||
let expected_bb = vec![2.0_f32 / 100.0; 4];
|
||||
let expected_ccc = vec![3.0_f32 / 100.0; 4];
|
||||
assert_eq!(results[0], expected_a);
|
||||
assert_eq!(results[1], expected_bb);
|
||||
assert_eq!(results[2], expected_ccc);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_exceeding_capacity_respects_max_entries() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 3 });
|
||||
|
||||
// Batch with 5 misses but cache capacity is 3
|
||||
let texts: Vec<String> = (0..5).map(|i| format!("text_{i}")).collect();
|
||||
let results = cached.embed_batch(&texts).await.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 5);
|
||||
let len = cached.len();
|
||||
assert!(len <= 3, "cache len {len} exceeds max 3");
|
||||
}
|
||||
|
||||
/// Mock embedding provider that fails the first N calls, then succeeds.
|
||||
struct FailThenSucceedMock {
|
||||
dimension: usize,
|
||||
model: String,
|
||||
remaining_failures: AtomicU32,
|
||||
}
|
||||
|
||||
impl FailThenSucceedMock {
|
||||
fn new(dimension: usize, fail_count: u32) -> Self {
|
||||
Self {
|
||||
dimension,
|
||||
model: "fail-mock".to_string(),
|
||||
remaining_failures: AtomicU32::new(fail_count),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for FailThenSucceedMock {
|
||||
fn dimension(&self) -> usize {
|
||||
self.dimension
|
||||
}
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
fn max_input_length(&self) -> usize {
|
||||
10_000
|
||||
}
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
let prev =
|
||||
self.remaining_failures
|
||||
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| {
|
||||
if v > 0 { Some(v - 1) } else { None }
|
||||
});
|
||||
if prev.is_ok() {
|
||||
return Err(EmbeddingError::HttpError("simulated failure".to_string()));
|
||||
}
|
||||
let val = text.len() as f32 / 100.0;
|
||||
Ok(vec![val; self.dimension])
|
||||
}
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
let prev =
|
||||
self.remaining_failures
|
||||
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| {
|
||||
if v > 0 { Some(v - 1) } else { None }
|
||||
});
|
||||
if prev.is_ok() {
|
||||
return Err(EmbeddingError::HttpError("simulated failure".to_string()));
|
||||
}
|
||||
texts
|
||||
.iter()
|
||||
.map(|t| {
|
||||
let val = t.len() as f32 / 100.0;
|
||||
Ok(vec![val; self.dimension])
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn error_does_not_pollute_cache() {
|
||||
let inner = Arc::new(FailThenSucceedMock::new(4, 1));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
// First call fails
|
||||
let err = cached.embed("hello").await;
|
||||
assert!(err.is_err());
|
||||
assert!(cached.is_empty(), "cache should be empty after error");
|
||||
|
||||
// Second call succeeds and should call the inner provider (not serve stale error)
|
||||
let result = cached.embed("hello").await;
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(cached.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_batch_empty_input() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
let results = cached.embed_batch(&[]).await.unwrap();
|
||||
assert!(results.is_empty());
|
||||
assert_eq!(inner.batch_calls(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_batch_all_misses() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 100 });
|
||||
|
||||
// Nothing cached — every text is a miss
|
||||
let texts: Vec<String> = vec!["alpha".into(), "beta".into(), "gamma".into()];
|
||||
let results = cached.embed_batch(&texts).await.unwrap();
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(inner.batch_calls(), 1, "inner called once for misses");
|
||||
assert_eq!(cached.len(), 3, "all results should be cached");
|
||||
|
||||
// Second call should be all hits — no new inner calls
|
||||
let results2 = cached.embed_batch(&texts).await.unwrap();
|
||||
assert_eq!(results2.len(), 3);
|
||||
assert_eq!(inner.batch_calls(), 1, "no new inner calls");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_max_entries_clamped_to_one() {
|
||||
let inner = Arc::new(CountingMock::new(4, "test-model"));
|
||||
let cached =
|
||||
CachedEmbeddingProvider::new(inner.clone(), EmbeddingCacheConfig { max_entries: 0 });
|
||||
|
||||
// Should behave as max_entries=1 (clamped in constructor)
|
||||
cached.embed("hello").await.unwrap();
|
||||
assert_eq!(cached.len(), 1);
|
||||
|
||||
// Second entry evicts the first
|
||||
cached.embed("world").await.unwrap();
|
||||
assert_eq!(cached.len(), 1);
|
||||
assert_eq!(inner.embed_calls(), 2);
|
||||
}
|
||||
}
|
||||
@@ -42,6 +42,7 @@
|
||||
|
||||
mod chunker;
|
||||
mod document;
|
||||
mod embedding_cache;
|
||||
mod embeddings;
|
||||
pub mod hygiene;
|
||||
#[cfg(feature = "postgres")]
|
||||
@@ -50,6 +51,7 @@ mod search;
|
||||
|
||||
pub use chunker::{ChunkConfig, chunk_document};
|
||||
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
|
||||
pub use embedding_cache::{CachedEmbeddingProvider, EmbeddingCacheConfig};
|
||||
pub use embeddings::{
|
||||
EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings,
|
||||
};
|
||||
@@ -371,7 +373,33 @@ impl Workspace {
|
||||
}
|
||||
|
||||
/// Set the embedding provider for semantic search.
|
||||
///
|
||||
/// The provider is automatically wrapped in a [`CachedEmbeddingProvider`]
|
||||
/// with the default cache size (10,000 entries; payload ~58 MB for 1536-dim,
|
||||
/// actual memory higher due to per-entry overhead).
|
||||
pub fn with_embeddings(mut self, provider: Arc<dyn EmbeddingProvider>) -> Self {
|
||||
self.embeddings = Some(Arc::new(CachedEmbeddingProvider::new(
|
||||
provider,
|
||||
EmbeddingCacheConfig::default(),
|
||||
)));
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the embedding provider with a custom cache configuration.
|
||||
pub fn with_embeddings_cached(
|
||||
mut self,
|
||||
provider: Arc<dyn EmbeddingProvider>,
|
||||
cache_config: EmbeddingCacheConfig,
|
||||
) -> Self {
|
||||
self.embeddings = Some(Arc::new(CachedEmbeddingProvider::new(
|
||||
provider,
|
||||
cache_config,
|
||||
)));
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the embedding provider **without** caching (for tests).
|
||||
pub fn with_embeddings_uncached(mut self, provider: Arc<dyn EmbeddingProvider>) -> Self {
|
||||
self.embeddings = Some(provider);
|
||||
self
|
||||
}
|
||||
|
||||
@@ -308,7 +308,7 @@ async fn test_workspace_hybrid_search_with_mock_embeddings() {
|
||||
|
||||
// Create workspace with mock embeddings (1536 dimensions to match OpenAI)
|
||||
let embeddings = Arc::new(MockEmbeddings::new(1536));
|
||||
let workspace = Workspace::new(user_id, pool.clone()).with_embeddings(embeddings);
|
||||
let workspace = Workspace::new(user_id, pool.clone()).with_embeddings_uncached(embeddings);
|
||||
|
||||
// Write documents
|
||||
workspace
|
||||
|
||||
Reference in New Issue
Block a user