fix: load embedding dimensions once per persist and trim vector search select.

This commit is contained in:
Per Stark
2026-06-12 13:54:51 +02:00
parent 28e8ede478
commit 4e8a58fff1
9 changed files with 90 additions and 52 deletions
+11 -14
View File
@@ -250,13 +250,10 @@ impl KnowledgeEntity {
pub async fn store_with_embedding(
entity: KnowledgeEntity,
embedding: Vec<f32>,
embedding_dimensions: usize,
db: &SurrealDbClient,
) -> Result<(), AppError> {
let settings = SystemSettings::get_current(db).await?;
KnowledgeEntityEmbedding::validate_dimension(
&embedding,
settings.embedding_dimensions as usize,
)?;
KnowledgeEntityEmbedding::validate_dimension(&embedding, embedding_dimensions)?;
let entity_id = entity.id.clone();
let emb = KnowledgeEntityEmbedding::new(
@@ -722,13 +719,13 @@ mod tests {
);
let emb = vec![0.1, 0.2, 0.3, 0.4, 0.5];
KnowledgeEntity::store_with_embedding(entity1.clone(), emb.clone(), &db)
KnowledgeEntity::store_with_embedding(entity1.clone(), emb.clone(), 5, &db)
.await
.with_context(|| "Failed to store entity 1".to_string())?;
KnowledgeEntity::store_with_embedding(entity2.clone(), emb.clone(), &db)
KnowledgeEntity::store_with_embedding(entity2.clone(), emb.clone(), 5, &db)
.await
.with_context(|| "Failed to store entity 2".to_string())?;
KnowledgeEntity::store_with_embedding(different_entity.clone(), emb.clone(), &db)
KnowledgeEntity::store_with_embedding(different_entity.clone(), emb.clone(), 5, &db)
.await
.with_context(|| "Failed to store different entity".to_string())?;
@@ -819,10 +816,10 @@ mod tests {
user_id,
);
KnowledgeEntity::store_with_embedding(entity1, vec![0.1, 0.2, 0.3], &db)
KnowledgeEntity::store_with_embedding(entity1, vec![0.1, 0.2, 0.3], 3, &db)
.await
.expect("store entity1");
KnowledgeEntity::store_with_embedding(entity2, vec![0.3, 0.2, 0.1], &db)
KnowledgeEntity::store_with_embedding(entity2, vec![0.3, 0.2, 0.1], 3, &db)
.await
.expect("store entity2");
@@ -895,7 +892,7 @@ mod tests {
user_id.clone(),
);
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.1, 0.2, 0.3], &db)
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store entity with embedding".to_string())?;
@@ -970,10 +967,10 @@ mod tests {
user_id.clone(),
);
KnowledgeEntity::store_with_embedding(e1.clone(), vec![1.0, 0.0, 0.0], &db)
KnowledgeEntity::store_with_embedding(e1.clone(), vec![1.0, 0.0, 0.0], 3, &db)
.await
.with_context(|| "store e1".to_string())?;
KnowledgeEntity::store_with_embedding(e2.clone(), vec![0.0, 1.0, 0.0], &db)
KnowledgeEntity::store_with_embedding(e2.clone(), vec![0.0, 1.0, 0.0], 3, &db)
.await
.with_context(|| "store e2".to_string())?;
@@ -1062,7 +1059,7 @@ mod tests {
user_id.clone(),
);
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.1, 0.2, 0.3], &db)
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store entity with embedding".to_string())?;
@@ -205,7 +205,7 @@ mod tests {
let embedding_vec = vec![0.11_f32, 0.22, 0.33];
let entity = build_knowledge_entity_with_id(entity_key, source_id, user_id);
KnowledgeEntity::store_with_embedding(entity.clone(), embedding_vec.clone(), &db)
KnowledgeEntity::store_with_embedding(entity.clone(), embedding_vec.clone(), 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
@@ -234,7 +234,7 @@ mod tests {
let entity = build_knowledge_entity_with_id(entity_key, source_id, user_id);
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.5_f32, 0.6, 0.7], &db)
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.5_f32, 0.6, 0.7], 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
@@ -266,7 +266,7 @@ mod tests {
let entity = build_knowledge_entity_with_id("entity-store", source_id, user_id);
KnowledgeEntity::store_with_embedding(entity.clone(), embedding.clone(), &db)
KnowledgeEntity::store_with_embedding(entity.clone(), embedding.clone(), 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
@@ -295,7 +295,7 @@ mod tests {
let db = prepare_knowledge_entity_test_db(3).await?;
let entity = build_knowledge_entity_with_id("entity-dim", "source-dim", "user-dim");
let result = KnowledgeEntity::store_with_embedding(entity, vec![0.1, 0.2], &db).await;
let result = KnowledgeEntity::store_with_embedding(entity, vec![0.1, 0.2], 3, &db).await;
assert!(matches!(result, Err(AppError::Validation(_))));
@@ -313,13 +313,13 @@ mod tests {
let entity2 = build_knowledge_entity_with_id("entity-s2", source_id, user_id);
let entity_other = build_knowledge_entity_with_id("entity-other", other_source, user_id);
KnowledgeEntity::store_with_embedding(entity1.clone(), vec![1.0_f32, 1.1, 1.2], &db)
KnowledgeEntity::store_with_embedding(entity1.clone(), vec![1.0_f32, 1.1, 1.2], 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
KnowledgeEntity::store_with_embedding(entity2.clone(), vec![2.0_f32, 2.1, 2.2], &db)
KnowledgeEntity::store_with_embedding(entity2.clone(), vec![2.0_f32, 2.1, 2.2], 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
KnowledgeEntity::store_with_embedding(entity_other.clone(), vec![3.0_f32, 3.1, 3.2], &db)
KnowledgeEntity::store_with_embedding(entity_other.clone(), vec![3.0_f32, 3.1, 3.2], 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
@@ -403,7 +403,7 @@ mod tests {
let source_id = "source-fetch";
let entity = build_knowledge_entity_with_id(entity_key, source_id, user_id);
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.7_f32, 0.8, 0.9], &db)
KnowledgeEntity::store_with_embedding(entity.clone(), vec![0.7_f32, 0.8, 0.9], 3, &db)
.await
.with_context(|| "Failed to store entity with embedding".to_string())?;
@@ -441,7 +441,7 @@ mod tests {
let source_id = "source-upsert";
let entity = build_knowledge_entity_with_id("entity-upsert", source_id, user_id);
KnowledgeEntity::store_with_embedding(entity.clone(), vec![1.0_f32, 0.0, 0.0], &db)
KnowledgeEntity::store_with_embedding(entity.clone(), vec![1.0_f32, 0.0, 0.0], 3, &db)
.await
.with_context(|| "initial store".to_string())?;
+1 -1
View File
@@ -691,7 +691,7 @@ mod tests {
"user1".into(),
);
TextChunk::store_with_embedding(initial_chunk.clone(), vec![0.1; 1536], &db)
TextChunk::store_with_embedding(initial_chunk.clone(), vec![0.1; 1536], 1536, &db)
.await
.with_context(|| "Failed to store initial chunk with embedding".to_string())?;
+16 -17
View File
@@ -3,7 +3,6 @@ use std::collections::HashMap;
use std::fmt::Write;
use crate::storage::indexes::hnsw_index_overwrite_sql;
use crate::storage::types::system_settings::SystemSettings;
use crate::storage::types::text_chunk_embedding::TextChunkEmbedding;
use crate::utils::embedding::RE_EMBED_BATCH_SIZE;
use crate::{error::AppError, storage::db::SurrealDbClient, stored_object};
@@ -67,10 +66,10 @@ impl TextChunk {
pub async fn store_with_embedding(
chunk: TextChunk,
embedding: Vec<f32>,
embedding_dimensions: usize,
db: &SurrealDbClient,
) -> Result<(), AppError> {
let settings = SystemSettings::get_current(db).await?;
TextChunkEmbedding::validate_dimension(&embedding, settings.embedding_dimensions as usize)?;
TextChunkEmbedding::validate_dimension(&embedding, embedding_dimensions)?;
let chunk_id = chunk.id.clone();
let emb = TextChunkEmbedding::new(
@@ -104,7 +103,7 @@ impl TextChunk {
Ok(())
}
/// Vector search over text chunks using the embedding table, fetching full chunk rows and embeddings.
/// Vector search over text chunks using the embedding table, fetching full chunk rows and scores.
pub async fn vector_search(
take: usize,
query_embedding: &[f32],
@@ -122,7 +121,6 @@ impl TextChunk {
r#"
SELECT
chunk_id,
embedding,
vector::similarity::cosine(embedding, $embedding) AS score
FROM {emb_table}
WHERE user_id = $user_id
@@ -466,15 +464,16 @@ mod tests {
user_id.clone(),
);
TextChunk::store_with_embedding(chunk1.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], &db)
TextChunk::store_with_embedding(chunk1.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], 5, &db)
.await
.with_context(|| "store chunk1".to_string())?;
TextChunk::store_with_embedding(chunk2.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], &db)
TextChunk::store_with_embedding(chunk2.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], 5, &db)
.await
.with_context(|| "store chunk2".to_string())?;
TextChunk::store_with_embedding(
different_chunk.clone(),
vec![0.1, 0.2, 0.3, 0.4, 0.5],
5,
&db,
)
.await
@@ -536,7 +535,7 @@ mod tests {
"user123".to_string(),
);
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], &db)
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], 5, &db)
.await
.with_context(|| "store chunk".to_string())?;
@@ -584,10 +583,10 @@ mod tests {
"user123".to_string(),
);
TextChunk::store_with_embedding(chunk1.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], &db)
TextChunk::store_with_embedding(chunk1.clone(), vec![0.1, 0.2, 0.3, 0.4, 0.5], 5, &db)
.await
.expect("store chunk1");
TextChunk::store_with_embedding(chunk2.clone(), vec![0.5, 0.4, 0.3, 0.2, 0.1], &db)
TextChunk::store_with_embedding(chunk2.clone(), vec![0.5, 0.4, 0.3, 0.2, 0.1], 5, &db)
.await
.expect("store chunk2");
@@ -632,7 +631,7 @@ mod tests {
.await
.with_context(|| "redefine index".to_string())?;
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], &db)
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store with embedding".to_string())?;
@@ -683,7 +682,7 @@ mod tests {
"runtime_user".to_string(),
);
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], &db)
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store with embedding".to_string())?;
@@ -755,7 +754,7 @@ mod tests {
user_id.clone(),
);
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], &db)
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store".to_string())?;
@@ -792,10 +791,10 @@ mod tests {
let chunk1 = TextChunk::new("s1".to_string(), "chunk one".to_string(), user_id.clone());
let chunk2 = TextChunk::new("s2".to_string(), "chunk two".to_string(), user_id.clone());
TextChunk::store_with_embedding(chunk1.clone(), vec![1.0, 0.0, 0.0], &db)
TextChunk::store_with_embedding(chunk1.clone(), vec![1.0, 0.0, 0.0], 3, &db)
.await
.with_context(|| "store chunk1".to_string())?;
TextChunk::store_with_embedding(chunk2.clone(), vec![0.0, 1.0, 0.0], &db)
TextChunk::store_with_embedding(chunk2.clone(), vec![0.0, 1.0, 0.0], 3, &db)
.await
.with_context(|| "store chunk2".to_string())?;
@@ -948,7 +947,7 @@ mod tests {
let chunk = TextChunk::new("src".to_string(), "body".to_string(), "user".to_string());
let err = TextChunk::store_with_embedding(chunk, vec![0.1, 0.2], &db)
let err = TextChunk::store_with_embedding(chunk, vec![0.1, 0.2], 3, &db)
.await
.expect_err("expected dimension validation failure");
assert!(matches!(err, AppError::Validation(_)));
@@ -978,7 +977,7 @@ mod tests {
user_id.clone(),
);
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], &db)
TextChunk::store_with_embedding(chunk.clone(), vec![0.1, 0.2, 0.3], 3, &db)
.await
.with_context(|| "store chunk with embedding".to_string())?;