mirror of
https://github.com/perstarkse/minne.git
synced 2026-06-25 11:26:17 +02:00
chore: centralize embedding errors, retrieval strategy, and test DB helpers.
Replace anyhow in embedding production code with EmbeddingError, move RetrievalStrategy into common config, and deduplicate Surreal test setup via common::test_utils.
This commit is contained in:
@@ -0,0 +1,96 @@
|
||||
//! Shared helpers for in-memory SurrealDB tests.
|
||||
#![cfg(any(test, feature = "test-utils"))]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::storage::{
|
||||
db::SurrealDbClient,
|
||||
indexes::{ensure_runtime, rebuild},
|
||||
types::{
|
||||
knowledge_entity_embedding::KnowledgeEntityEmbedding,
|
||||
system_settings::SystemSettings,
|
||||
text_chunk_embedding::TextChunkEmbedding,
|
||||
},
|
||||
};
|
||||
|
||||
const TEST_NAMESPACE: &str = "test_ns";
|
||||
|
||||
/// Starts an in-memory database, applies migrations, and returns a client.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if the database cannot be started or migrations fail.
|
||||
pub async fn setup_test_db() -> Result<SurrealDbClient> {
|
||||
let database = Uuid::new_v4().to_string();
|
||||
let db = SurrealDbClient::memory(TEST_NAMESPACE, &database)
|
||||
.await
|
||||
.context("start in-memory surrealdb")?;
|
||||
|
||||
db.apply_migrations()
|
||||
.await
|
||||
.context("apply migrations")?;
|
||||
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Updates singleton [`SystemSettings`] embedding dimensions for tests.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if settings cannot be loaded or updated.
|
||||
pub async fn configure_embedding_dimension(db: &SurrealDbClient, dimension: u32) -> Result<()> {
|
||||
let mut settings = SystemSettings::get_current(db).await?;
|
||||
settings.embedding_dimensions = dimension;
|
||||
SystemSettings::update(db, settings).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Starts a test database and sets the embedding dimension in system settings.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if setup or settings update fails.
|
||||
pub async fn setup_test_db_with_embedding_dimension(dimension: u32) -> Result<SurrealDbClient> {
|
||||
let db = setup_test_db().await?;
|
||||
configure_embedding_dimension(&db, dimension).await?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Prepares a database for text-chunk embedding tests at the given dimension.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if setup, settings update, or index redefinition fails.
|
||||
pub async fn prepare_text_chunk_test_db(dimension: u32) -> Result<SurrealDbClient> {
|
||||
let db = setup_test_db_with_embedding_dimension(dimension).await?;
|
||||
TextChunkEmbedding::redefine_hnsw_index(&db, dimension as usize)
|
||||
.await
|
||||
.with_context(|| format!("set text chunk index dimension to {dimension}"))?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Prepares a database for knowledge-entity embedding tests at the given dimension.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if setup, settings update, or index redefinition fails.
|
||||
pub async fn prepare_knowledge_entity_test_db(dimension: u32) -> Result<SurrealDbClient> {
|
||||
let db = setup_test_db_with_embedding_dimension(dimension).await?;
|
||||
KnowledgeEntityEmbedding::redefine_hnsw_index(&db, dimension as usize)
|
||||
.await
|
||||
.with_context(|| format!("set knowledge entity index dimension to {dimension}"))?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Starts a test database and ensures runtime FTS/HNSW indexes are ready.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if setup, index creation, or rebuild fails.
|
||||
pub async fn setup_test_db_with_runtime_indexes() -> Result<SurrealDbClient> {
|
||||
let db = setup_test_db().await?;
|
||||
ensure_runtime(&db, 1536).await?;
|
||||
rebuild(&db).await?;
|
||||
Ok(db)
|
||||
}
|
||||
Reference in New Issue
Block a user