diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c67d811d0..f0f056bf2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -63,5 +63,14 @@ jobs: - uses: Swatinem/rust-cache@v2 + # The memory tests share one embedding model cached at this path; + # without the cache every run downloads it from Hugging Face, which + # fails intermittently and turns unrelated PRs red. + - name: Cache embedding model + uses: actions/cache@v4 + with: + path: /tmp/spacebot-test-embedding-cache + key: fastembed-onnx-v1 + - name: cargo test (lib) run: cargo test --lib diff --git a/src/memory/embedding.rs b/src/memory/embedding.rs index 82b5588b7..ee6269be1 100644 --- a/src/memory/embedding.rs +++ b/src/memory/embedding.rs @@ -65,3 +65,18 @@ impl EmbeddingModel { pub async fn embed_text(model: &Arc, text: &str) -> Result> { model.embed_one(text).await } + +/// Shared embedding model for tests. Downloads once into a stable cache +/// directory and is reused by every test in the binary, so CI can cache the +/// model files across runs instead of fetching them from Hugging Face on +/// each test that needs embeddings. +#[cfg(test)] +pub(crate) fn shared_test_model() -> std::sync::Arc { + use std::sync::{Arc, OnceLock}; + static MODEL: OnceLock> = OnceLock::new(); + Arc::clone(MODEL.get_or_init(|| { + let cache_dir = std::env::temp_dir().join("spacebot-test-embedding-cache"); + std::fs::create_dir_all(&cache_dir).expect("failed to create embedding cache dir"); + Arc::new(EmbeddingModel::new(&cache_dir).expect("failed to initialize embedding model")) + })) +} diff --git a/src/memory/maintenance.rs b/src/memory/maintenance.rs index 67b935497..c5e4a58b5 100644 --- a/src/memory/maintenance.rs +++ b/src/memory/maintenance.rs @@ -488,20 +488,12 @@ pub struct MaintenanceReport { mod tests { use super::*; use crate::memory::{Association, RelationType}; - use std::sync::{Arc, OnceLock}; + use std::sync::Arc; use tempfile::tempdir; use tokio::time::Duration; fn shared_embedding_model() -> Arc { - static MODEL: OnceLock> = OnceLock::new(); - Arc::clone(MODEL.get_or_init(|| { - let cache_dir = std::env::temp_dir().join("spacebot-test-embedding-cache"); - std::fs::create_dir_all(&cache_dir).expect("failed to create embedding cache dir"); - Arc::new( - crate::memory::EmbeddingModel::new(&cache_dir) - .expect("failed to initialize embedding model"), - ) - })) + crate::memory::embedding::shared_test_model() } async fn create_memory_with_embedding( diff --git a/src/memory/search.rs b/src/memory/search.rs index 296686c81..8c4419498 100644 --- a/src/memory/search.rs +++ b/src/memory/search.rs @@ -580,7 +580,7 @@ mod tests { .await .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); - let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); + let embedding_model = crate::memory::embedding::shared_test_model(); let search = MemorySearch::new(store, embedding_table, embedding_model); let config = SearchConfig { @@ -608,7 +608,7 @@ mod tests { .await .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); - let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); + let embedding_model = crate::memory::embedding::shared_test_model(); let search = MemorySearch::new(store, embedding_table, embedding_model); let config = SearchConfig { @@ -633,7 +633,7 @@ mod tests { .await .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); - let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); + let embedding_model = crate::memory::embedding::shared_test_model(); let search = MemorySearch::new(store, embedding_table, embedding_model); let config = SearchConfig { @@ -658,7 +658,7 @@ mod tests { .await .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); - let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); + let embedding_model = crate::memory::embedding::shared_test_model(); let search = MemorySearch::new(store, embedding_table, embedding_model); let config = SearchConfig {