mirror of
https://github.com/perstarkse/minne.git
synced 2026-04-25 02:08:30 +02:00
tidying stuff up, dto for search
This commit is contained in:
2
Cargo.lock
generated
2
Cargo.lock
generated
@@ -2194,7 +2194,7 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "eval"
|
||||
name = "evaluations"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
|
||||
10
Cargo.toml
10
Cargo.toml
@@ -7,7 +7,7 @@ members = [
|
||||
"ingestion-pipeline",
|
||||
"retrieval-pipeline",
|
||||
"json-stream-parser",
|
||||
"eval"
|
||||
"evaluations"
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
@@ -80,7 +80,7 @@ implicit_clone = "warn"
|
||||
redundant_clone = "warn"
|
||||
|
||||
# Security-focused lints
|
||||
integer_arithmetic = "warn"
|
||||
arithmetic_side_effects = "warn"
|
||||
indexing_slicing = "warn"
|
||||
unwrap_used = "warn"
|
||||
expect_used = "warn"
|
||||
@@ -90,7 +90,7 @@ todo = "warn"
|
||||
|
||||
# Async/Network lints
|
||||
async_yields_async = "warn"
|
||||
await_holding_invalid_state = "warn"
|
||||
await_holding_invalid_type = "warn"
|
||||
rc_buffer = "warn"
|
||||
|
||||
# Maintainability-focused lints
|
||||
@@ -109,6 +109,8 @@ wildcard_dependencies = "warn"
|
||||
missing_docs_in_private_items = "warn"
|
||||
|
||||
# Allow noisy lints that don't add value for this project
|
||||
manual_must_use = "allow"
|
||||
needless_raw_string_hashes = "allow"
|
||||
multiple_bound_locations = "allow"
|
||||
cargo_common_metadata = "allow"
|
||||
multiple-crate-versions = "allow"
|
||||
module_name_repetition = "allow"
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
-- Add embedding_backend field to system_settings for visibility of active backend
|
||||
|
||||
DEFINE FIELD IF NOT EXISTS embedding_backend ON system_settings TYPE option<string>;
|
||||
|
||||
-- Set default to 'openai' for existing installs to preserve backward compatibility
|
||||
UPDATE system_settings:current SET
|
||||
embedding_backend = 'openai'
|
||||
WHERE embedding_backend == NONE;
|
||||
@@ -428,6 +428,103 @@ impl KnowledgeEntity {
|
||||
info!("Re-embedding process for knowledge entities completed successfully.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Re-creates embeddings for all knowledge entities using an `EmbeddingProvider`.
|
||||
///
|
||||
/// This variant uses the application's configured embedding provider (FastEmbed, OpenAI, etc.)
|
||||
/// instead of directly calling OpenAI. Used during startup when embedding configuration changes.
|
||||
pub async fn update_all_embeddings_with_provider(
|
||||
db: &SurrealDbClient,
|
||||
provider: &crate::utils::embedding::EmbeddingProvider,
|
||||
) -> Result<(), AppError> {
|
||||
let new_dimensions = provider.dimension();
|
||||
info!(
|
||||
dimensions = new_dimensions,
|
||||
backend = provider.backend_label(),
|
||||
"Starting re-embedding process for all knowledge entities"
|
||||
);
|
||||
|
||||
// Fetch all entities first
|
||||
let all_entities: Vec<KnowledgeEntity> = db.select(Self::table_name()).await?;
|
||||
let total_entities = all_entities.len();
|
||||
if total_entities == 0 {
|
||||
info!("No knowledge entities to update. Just updating the index.");
|
||||
KnowledgeEntityEmbedding::redefine_hnsw_index(db, new_dimensions).await?;
|
||||
return Ok(());
|
||||
}
|
||||
info!(entities = total_entities, "Found entities to process");
|
||||
|
||||
// Generate all new embeddings in memory
|
||||
let mut new_embeddings: HashMap<String, (Vec<f32>, String)> = HashMap::new();
|
||||
info!("Generating new embeddings for all entities...");
|
||||
|
||||
for (i, entity) in all_entities.iter().enumerate() {
|
||||
if i > 0 && i % 100 == 0 {
|
||||
info!(progress = i, total = total_entities, "Re-embedding progress");
|
||||
}
|
||||
|
||||
let embedding_input = format!(
|
||||
"name: {}, description: {}, type: {:?}",
|
||||
entity.name, entity.description, entity.entity_type
|
||||
);
|
||||
|
||||
let embedding = provider
|
||||
.embed(&embedding_input)
|
||||
.await
|
||||
.map_err(|e| AppError::InternalError(format!("Embedding failed: {e}")))?;
|
||||
|
||||
// Safety check: ensure the generated embedding has the correct dimension.
|
||||
if embedding.len() != new_dimensions {
|
||||
let err_msg = format!(
|
||||
"CRITICAL: Generated embedding for entity {} has incorrect dimension ({}). Expected {}. Aborting.",
|
||||
entity.id, embedding.len(), new_dimensions
|
||||
);
|
||||
error!("{}", err_msg);
|
||||
return Err(AppError::InternalError(err_msg));
|
||||
}
|
||||
new_embeddings.insert(entity.id.clone(), (embedding, entity.user_id.clone()));
|
||||
}
|
||||
info!("Successfully generated all new embeddings.");
|
||||
|
||||
// Perform DB updates in a single transaction
|
||||
info!("Applying embedding updates in a transaction...");
|
||||
let mut transaction_query = String::from("BEGIN TRANSACTION;");
|
||||
|
||||
for (id, (embedding, user_id)) in new_embeddings {
|
||||
let embedding_str = format!(
|
||||
"[{}]",
|
||||
embedding
|
||||
.iter()
|
||||
.map(|f| f.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
transaction_query.push_str(&format!(
|
||||
"UPSERT type::thing('knowledge_entity_embedding', '{id}') SET \
|
||||
entity_id = type::thing('knowledge_entity', '{id}'), \
|
||||
embedding = {embedding}, \
|
||||
user_id = '{user_id}', \
|
||||
created_at = IF created_at != NONE THEN created_at ELSE time::now() END, \
|
||||
updated_at = time::now();",
|
||||
id = id,
|
||||
embedding = embedding_str,
|
||||
user_id = user_id
|
||||
));
|
||||
}
|
||||
|
||||
transaction_query.push_str(&format!(
|
||||
"DEFINE INDEX OVERWRITE idx_embedding_knowledge_entity_embedding ON TABLE knowledge_entity_embedding FIELDS embedding HNSW DIMENSION {};",
|
||||
new_dimensions
|
||||
));
|
||||
|
||||
transaction_query.push_str("COMMIT TRANSACTION;");
|
||||
|
||||
// Execute the entire atomic operation
|
||||
db.query(transaction_query).await?;
|
||||
|
||||
info!("Re-embedding process for knowledge entities completed successfully.");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -13,6 +13,9 @@ pub struct SystemSettings {
|
||||
pub processing_model: String,
|
||||
pub embedding_model: String,
|
||||
pub embedding_dimensions: u32,
|
||||
/// Active embedding backend ("openai", "fastembed", "hashed"). Read-only, synced from config.
|
||||
#[serde(default)]
|
||||
pub embedding_backend: Option<String>,
|
||||
pub query_system_prompt: String,
|
||||
pub ingestion_system_prompt: String,
|
||||
pub image_processing_model: String,
|
||||
@@ -49,6 +52,57 @@ impl SystemSettings {
|
||||
"Something went wrong updating the settings".into(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Syncs SystemSettings with the active embedding provider's properties.
|
||||
/// Updates embedding_backend, embedding_model, and embedding_dimensions if they differ.
|
||||
/// Returns true if any settings were changed.
|
||||
pub async fn sync_from_embedding_provider(
|
||||
db: &SurrealDbClient,
|
||||
provider: &crate::utils::embedding::EmbeddingProvider,
|
||||
) -> Result<(Self, bool), AppError> {
|
||||
let mut settings = Self::get_current(db).await?;
|
||||
let mut needs_update = false;
|
||||
|
||||
let backend_label = provider.backend_label().to_string();
|
||||
let provider_dimensions = provider.dimension() as u32;
|
||||
let provider_model = provider.model_code();
|
||||
|
||||
// Sync backend label
|
||||
if settings.embedding_backend.as_deref() != Some(&backend_label) {
|
||||
settings.embedding_backend = Some(backend_label);
|
||||
needs_update = true;
|
||||
}
|
||||
|
||||
// Sync dimensions
|
||||
if settings.embedding_dimensions != provider_dimensions {
|
||||
tracing::info!(
|
||||
old_dimensions = settings.embedding_dimensions,
|
||||
new_dimensions = provider_dimensions,
|
||||
"Embedding dimensions changed, updating SystemSettings"
|
||||
);
|
||||
settings.embedding_dimensions = provider_dimensions;
|
||||
needs_update = true;
|
||||
}
|
||||
|
||||
// Sync model if provider has one
|
||||
if let Some(model) = provider_model {
|
||||
if settings.embedding_model != model {
|
||||
tracing::info!(
|
||||
old_model = %settings.embedding_model,
|
||||
new_model = %model,
|
||||
"Embedding model changed, updating SystemSettings"
|
||||
);
|
||||
settings.embedding_model = model;
|
||||
needs_update = true;
|
||||
}
|
||||
}
|
||||
|
||||
if needs_update {
|
||||
settings = Self::update(db, settings).await?;
|
||||
}
|
||||
|
||||
Ok((settings, needs_update))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -323,6 +323,106 @@ impl TextChunk {
|
||||
info!("Re-embedding process for text chunks completed successfully.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Re-creates embeddings for all text chunks using an `EmbeddingProvider`.
|
||||
///
|
||||
/// This variant uses the application's configured embedding provider (FastEmbed, OpenAI, etc.)
|
||||
/// instead of directly calling OpenAI. Used during startup when embedding configuration changes.
|
||||
pub async fn update_all_embeddings_with_provider(
|
||||
db: &SurrealDbClient,
|
||||
provider: &crate::utils::embedding::EmbeddingProvider,
|
||||
) -> Result<(), AppError> {
|
||||
let new_dimensions = provider.dimension();
|
||||
info!(
|
||||
dimensions = new_dimensions,
|
||||
backend = provider.backend_label(),
|
||||
"Starting re-embedding process for all text chunks"
|
||||
);
|
||||
|
||||
// Fetch all chunks first
|
||||
let all_chunks: Vec<TextChunk> = db.select(Self::table_name()).await?;
|
||||
let total_chunks = all_chunks.len();
|
||||
if total_chunks == 0 {
|
||||
info!("No text chunks to update. Just updating the index.");
|
||||
TextChunkEmbedding::redefine_hnsw_index(db, new_dimensions).await?;
|
||||
return Ok(());
|
||||
}
|
||||
info!(chunks = total_chunks, "Found chunks to process");
|
||||
|
||||
// Generate all new embeddings in memory
|
||||
let mut new_embeddings: HashMap<String, (Vec<f32>, String, String)> = HashMap::new();
|
||||
info!("Generating new embeddings for all chunks...");
|
||||
|
||||
for (i, chunk) in all_chunks.iter().enumerate() {
|
||||
if i > 0 && i % 100 == 0 {
|
||||
info!(progress = i, total = total_chunks, "Re-embedding progress");
|
||||
}
|
||||
|
||||
let embedding = provider
|
||||
.embed(&chunk.chunk)
|
||||
.await
|
||||
.map_err(|e| AppError::InternalError(format!("Embedding failed: {e}")))?;
|
||||
|
||||
// Safety check: ensure the generated embedding has the correct dimension.
|
||||
if embedding.len() != new_dimensions {
|
||||
let err_msg = format!(
|
||||
"CRITICAL: Generated embedding for chunk {} has incorrect dimension ({}). Expected {}. Aborting.",
|
||||
chunk.id, embedding.len(), new_dimensions
|
||||
);
|
||||
error!("{}", err_msg);
|
||||
return Err(AppError::InternalError(err_msg));
|
||||
}
|
||||
new_embeddings.insert(
|
||||
chunk.id.clone(),
|
||||
(embedding, chunk.user_id.clone(), chunk.source_id.clone()),
|
||||
);
|
||||
}
|
||||
info!("Successfully generated all new embeddings.");
|
||||
|
||||
// Perform DB updates in a single transaction against the embedding table
|
||||
info!("Applying embedding updates in a transaction...");
|
||||
let mut transaction_query = String::from("BEGIN TRANSACTION;");
|
||||
|
||||
for (id, (embedding, user_id, source_id)) in new_embeddings {
|
||||
let embedding_str = format!(
|
||||
"[{}]",
|
||||
embedding
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
write!(
|
||||
&mut transaction_query,
|
||||
"UPSERT type::thing('text_chunk_embedding', '{id}') SET \
|
||||
chunk_id = type::thing('text_chunk', '{id}'), \
|
||||
source_id = '{source_id}', \
|
||||
embedding = {embedding}, \
|
||||
user_id = '{user_id}', \
|
||||
created_at = IF created_at != NONE THEN created_at ELSE time::now() END, \
|
||||
updated_at = time::now();",
|
||||
id = id,
|
||||
embedding = embedding_str,
|
||||
user_id = user_id,
|
||||
source_id = source_id
|
||||
)
|
||||
.map_err(|e| AppError::InternalError(e.to_string()))?;
|
||||
}
|
||||
|
||||
write!(
|
||||
&mut transaction_query,
|
||||
"DEFINE INDEX OVERWRITE idx_embedding_text_chunk_embedding ON TABLE text_chunk_embedding FIELDS embedding HNSW DIMENSION {};",
|
||||
new_dimensions
|
||||
)
|
||||
.map_err(|e| AppError::InternalError(e.to_string()))?;
|
||||
|
||||
transaction_query.push_str("COMMIT TRANSACTION;");
|
||||
|
||||
db.query(transaction_query).await?;
|
||||
|
||||
info!("Re-embedding process for text chunks completed successfully.");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -2,6 +2,19 @@ use config::{Config, ConfigError, Environment, File};
|
||||
use serde::Deserialize;
|
||||
use std::env;
|
||||
|
||||
/// Selects the embedding backend for vector generation.
|
||||
#[derive(Clone, Deserialize, Debug, Default, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum EmbeddingBackend {
|
||||
/// Use OpenAI-compatible API for embeddings.
|
||||
OpenAI,
|
||||
/// Use FastEmbed local embeddings (default).
|
||||
#[default]
|
||||
FastEmbed,
|
||||
/// Use deterministic hashed embeddings (for testing).
|
||||
Hashed,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Debug, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum StorageKind {
|
||||
@@ -60,6 +73,8 @@ pub struct AppConfig {
|
||||
pub fastembed_max_length: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub retrieval_strategy: Option<String>,
|
||||
#[serde(default)]
|
||||
pub embedding_backend: EmbeddingBackend,
|
||||
}
|
||||
|
||||
/// Default data directory for persisted assets.
|
||||
@@ -127,6 +142,7 @@ impl Default for AppConfig {
|
||||
fastembed_show_download_progress: None,
|
||||
fastembed_max_length: None,
|
||||
retrieval_strategy: None,
|
||||
embedding_backend: EmbeddingBackend::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -235,6 +235,34 @@ impl EmbeddingProvider {
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Creates an embedding provider based on application configuration.
|
||||
///
|
||||
/// Dispatches to the appropriate constructor based on `config.embedding_backend`:
|
||||
/// - `OpenAI`: Requires a valid OpenAI client
|
||||
/// - `FastEmbed`: Uses local embedding model
|
||||
/// - `Hashed`: Uses deterministic hashed embeddings (for testing)
|
||||
pub async fn from_config(
|
||||
config: &crate::utils::config::AppConfig,
|
||||
openai_client: Option<Arc<Client<async_openai::config::OpenAIConfig>>>,
|
||||
) -> Result<Self> {
|
||||
use crate::utils::config::EmbeddingBackend;
|
||||
|
||||
match config.embedding_backend {
|
||||
EmbeddingBackend::OpenAI => {
|
||||
let client = openai_client.ok_or_else(|| {
|
||||
anyhow!("OpenAI embedding backend requires an OpenAI client")
|
||||
})?;
|
||||
// Use defaults that match SystemSettings initial values
|
||||
Self::new_openai(client, "text-embedding-3-small".to_string(), 1536)
|
||||
}
|
||||
EmbeddingBackend::FastEmbed => {
|
||||
// Use nomic-embed-text-v1.5 as the default FastEmbed model
|
||||
Self::new_fastembed(Some("nomic-ai/nomic-embed-text-v1.5".to_string())).await
|
||||
}
|
||||
EmbeddingBackend::Hashed => Self::new_hashed(384),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions for hashed embeddings
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use common::storage::{db::SurrealDbClient, store::StorageManager};
|
||||
use common::utils::embedding::EmbeddingProvider;
|
||||
use common::utils::template_engine::{ProvidesTemplateEngine, TemplateEngine};
|
||||
use common::{create_template_engine, storage::db::ProvidesDb, utils::config::AppConfig};
|
||||
use retrieval_pipeline::{reranking::RerankerPool, RetrievalStrategy};
|
||||
@@ -16,6 +17,7 @@ pub struct HtmlState {
|
||||
pub config: AppConfig,
|
||||
pub storage: StorageManager,
|
||||
pub reranker_pool: Option<Arc<RerankerPool>>,
|
||||
pub embedding_provider: Arc<EmbeddingProvider>,
|
||||
}
|
||||
|
||||
impl HtmlState {
|
||||
@@ -26,6 +28,7 @@ impl HtmlState {
|
||||
storage: StorageManager,
|
||||
config: AppConfig,
|
||||
reranker_pool: Option<Arc<RerankerPool>>,
|
||||
embedding_provider: Arc<EmbeddingProvider>,
|
||||
) -> Result<Self, Box<dyn std::error::Error>> {
|
||||
let template_engine = create_template_engine!("templates");
|
||||
debug!("Template engine created for html_router.");
|
||||
@@ -38,6 +41,7 @@ impl HtmlState {
|
||||
config,
|
||||
storage,
|
||||
reranker_pool,
|
||||
embedding_provider,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -151,21 +151,46 @@ pub async fn update_model_settings(
|
||||
|
||||
let current_settings = SystemSettings::get_current(&state.db).await?;
|
||||
|
||||
// Determine if re-embedding is required
|
||||
let reembedding_needed = input
|
||||
.embedding_dimensions
|
||||
.is_some_and(|new_dims| new_dims != current_settings.embedding_dimensions);
|
||||
// Check if using FastEmbed - if so, embedding model/dimensions cannot be changed via UI
|
||||
let uses_local_embeddings = current_settings
|
||||
.embedding_backend
|
||||
.as_deref()
|
||||
.is_some_and(|b| b == "fastembed" || b == "hashed");
|
||||
|
||||
// For local embeddings, ignore any embedding model/dimension changes from the form
|
||||
let (final_embedding_model, final_embedding_dimensions, reembedding_needed) =
|
||||
if uses_local_embeddings {
|
||||
// Keep current values - they're controlled by config, not the admin UI
|
||||
info!(
|
||||
backend = ?current_settings.embedding_backend,
|
||||
"Embedding model/dimensions controlled by config, ignoring form input"
|
||||
);
|
||||
(
|
||||
current_settings.embedding_model.clone(),
|
||||
current_settings.embedding_dimensions,
|
||||
false,
|
||||
)
|
||||
} else {
|
||||
// OpenAI backend - allow changes from form
|
||||
let reembedding_needed = input
|
||||
.embedding_dimensions
|
||||
.is_some_and(|new_dims| new_dims != current_settings.embedding_dimensions);
|
||||
(
|
||||
input.embedding_model,
|
||||
input
|
||||
.embedding_dimensions
|
||||
.unwrap_or(current_settings.embedding_dimensions),
|
||||
reembedding_needed,
|
||||
)
|
||||
};
|
||||
|
||||
let new_settings = SystemSettings {
|
||||
query_model: input.query_model,
|
||||
processing_model: input.processing_model,
|
||||
image_processing_model: input.image_processing_model,
|
||||
voice_processing_model: input.voice_processing_model,
|
||||
embedding_model: input.embedding_model,
|
||||
// Use new dimensions if provided, otherwise retain the current ones.
|
||||
embedding_dimensions: input
|
||||
.embedding_dimensions
|
||||
.unwrap_or(current_settings.embedding_dimensions),
|
||||
embedding_model: final_embedding_model,
|
||||
embedding_dimensions: final_embedding_dimensions,
|
||||
..current_settings.clone()
|
||||
};
|
||||
|
||||
|
||||
@@ -132,6 +132,7 @@ pub async fn get_response_stream(
|
||||
let retrieval_result = match retrieval_pipeline::retrieve_entities(
|
||||
&state.db,
|
||||
&state.openai_client,
|
||||
Some(&*state.embedding_provider),
|
||||
&user_message.content,
|
||||
&user.id,
|
||||
config,
|
||||
|
||||
@@ -288,6 +288,7 @@ pub async fn suggest_knowledge_relationships(
|
||||
retrieval_pipeline::retrieve_entities(
|
||||
&state.db,
|
||||
&state.openai_client,
|
||||
Some(&*state.embedding_provider),
|
||||
&query,
|
||||
&user.id,
|
||||
config,
|
||||
|
||||
@@ -1,13 +1,21 @@
|
||||
use std::{fmt, str::FromStr, time::Duration};
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
fmt, str::FromStr,
|
||||
};
|
||||
|
||||
use axum::{
|
||||
extract::{Query, State},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use common::storage::types::{conversation::Conversation, user::User};
|
||||
use common::storage::types::{
|
||||
conversation::Conversation,
|
||||
text_content::{deserialize_flexible_id, TextContent},
|
||||
user::User,
|
||||
StoredObject,
|
||||
};
|
||||
use retrieval_pipeline::{RetrievalConfig, SearchResult, SearchTarget, StrategyOutput};
|
||||
use serde::{de, Deserialize, Deserializer, Serialize};
|
||||
use tokio::time::error::Elapsed;
|
||||
use surrealdb::RecordId;
|
||||
|
||||
use crate::{
|
||||
html_state::HtmlState,
|
||||
@@ -31,6 +39,113 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn source_id_suffix(source_id: &str) -> String {
|
||||
let start = source_id.len().saturating_sub(8);
|
||||
source_id[start..].to_string()
|
||||
}
|
||||
|
||||
fn truncate_label(value: &str, max_chars: usize) -> String {
|
||||
let mut end = None;
|
||||
let mut count = 0;
|
||||
for (idx, _) in value.char_indices() {
|
||||
if count == max_chars {
|
||||
end = Some(idx);
|
||||
break;
|
||||
}
|
||||
count += 1;
|
||||
}
|
||||
|
||||
match end {
|
||||
Some(idx) => format!("{}...", &value[..idx]),
|
||||
None => value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn first_non_empty_line(text: &str, max_chars: usize) -> Option<String> {
|
||||
for line in text.lines() {
|
||||
let trimmed = line.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return Some(truncate_label(trimmed, max_chars));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct UrlInfoLabel {
|
||||
#[serde(default)]
|
||||
title: String,
|
||||
#[serde(default)]
|
||||
url: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FileInfoLabel {
|
||||
#[serde(default)]
|
||||
file_name: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SourceLabelRow {
|
||||
#[serde(deserialize_with = "deserialize_flexible_id")]
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
url_info: Option<UrlInfoLabel>,
|
||||
#[serde(default)]
|
||||
file_info: Option<FileInfoLabel>,
|
||||
#[serde(default)]
|
||||
context: Option<String>,
|
||||
#[serde(default)]
|
||||
category: String,
|
||||
#[serde(default)]
|
||||
text: String,
|
||||
}
|
||||
|
||||
fn build_source_label(row: &SourceLabelRow) -> String {
|
||||
const MAX_LABEL_CHARS: usize = 80;
|
||||
|
||||
if let Some(url_info) = row.url_info.as_ref() {
|
||||
let title = url_info.title.trim();
|
||||
if !title.is_empty() {
|
||||
return title.to_string();
|
||||
}
|
||||
|
||||
let url = url_info.url.trim();
|
||||
if !url.is_empty() {
|
||||
return url.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(file_info) = row.file_info.as_ref() {
|
||||
let name = file_info.file_name.trim();
|
||||
if !name.is_empty() {
|
||||
return name.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(context) = row.context.as_ref() {
|
||||
let trimmed = context.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return truncate_label(trimmed, MAX_LABEL_CHARS);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(text_label) = first_non_empty_line(&row.text, MAX_LABEL_CHARS) {
|
||||
return text_label;
|
||||
}
|
||||
|
||||
let category = row.category.trim();
|
||||
if !category.is_empty() {
|
||||
return truncate_label(category, MAX_LABEL_CHARS);
|
||||
}
|
||||
|
||||
format!("Text snippet: {}", source_id_suffix(&row.id))
|
||||
}
|
||||
|
||||
fn fallback_source_label(source_id: &str) -> String {
|
||||
format!("Text snippet: {}", source_id_suffix(source_id))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SearchParams {
|
||||
#[serde(default, deserialize_with = "empty_string_as_none")]
|
||||
@@ -42,6 +157,7 @@ pub struct SearchParams {
|
||||
struct TextChunkForTemplate {
|
||||
id: String,
|
||||
source_id: String,
|
||||
source_label: String,
|
||||
chunk: String,
|
||||
score: f32,
|
||||
}
|
||||
@@ -54,6 +170,7 @@ struct KnowledgeEntityForTemplate {
|
||||
description: String,
|
||||
entity_type: String,
|
||||
source_id: String,
|
||||
source_label: String,
|
||||
score: f32,
|
||||
}
|
||||
|
||||
@@ -89,14 +206,21 @@ pub async fn search_result_handler(
|
||||
} else {
|
||||
// Use retrieval pipeline Search strategy
|
||||
let config = RetrievalConfig::for_search(SearchTarget::Both);
|
||||
|
||||
// Checkout a reranker lease if pool is available
|
||||
let reranker_lease = match &state.reranker_pool {
|
||||
Some(pool) => Some(pool.checkout().await),
|
||||
None => None,
|
||||
};
|
||||
|
||||
let result = retrieval_pipeline::pipeline::run_pipeline(
|
||||
&state.db,
|
||||
&state.openai_client,
|
||||
None, // No embedding provider in HtmlState
|
||||
Some(&state.embedding_provider),
|
||||
trimmed_query,
|
||||
&user.id,
|
||||
config,
|
||||
None, // No reranker for now
|
||||
reranker_lease,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -105,17 +229,74 @@ pub async fn search_result_handler(
|
||||
_ => SearchResult::new(vec![], vec![]),
|
||||
};
|
||||
|
||||
let mut source_ids = HashSet::new();
|
||||
for chunk_result in &search_result.chunks {
|
||||
source_ids.insert(chunk_result.chunk.source_id.clone());
|
||||
}
|
||||
for entity_result in &search_result.entities {
|
||||
source_ids.insert(entity_result.entity.source_id.clone());
|
||||
}
|
||||
|
||||
let source_label_map = if source_ids.is_empty() {
|
||||
HashMap::new()
|
||||
} else {
|
||||
let record_ids: Vec<RecordId> = source_ids
|
||||
.iter()
|
||||
.filter_map(|id| {
|
||||
if id.contains(':') {
|
||||
RecordId::from_str(id).ok()
|
||||
} else {
|
||||
Some(RecordId::from_table_key(TextContent::table_name(), id))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let mut response = state
|
||||
.db
|
||||
.client
|
||||
.query(
|
||||
"SELECT id, url_info, file_info, context, category, text FROM type::table($table_name) WHERE user_id = $user_id AND id INSIDE $record_ids",
|
||||
)
|
||||
.bind(("table_name", TextContent::table_name()))
|
||||
.bind(("user_id", user.id.clone()))
|
||||
.bind(("record_ids", record_ids))
|
||||
.await?;
|
||||
let contents: Vec<SourceLabelRow> = response.take(0)?;
|
||||
|
||||
tracing::debug!(
|
||||
source_id_count = source_ids.len(),
|
||||
label_row_count = contents.len(),
|
||||
"Resolved search source labels"
|
||||
);
|
||||
|
||||
let mut labels = HashMap::new();
|
||||
for content in contents {
|
||||
let label = build_source_label(&content);
|
||||
labels.insert(content.id.clone(), label.clone());
|
||||
labels.insert(
|
||||
format!("{}:{}", TextContent::table_name(), content.id),
|
||||
label,
|
||||
);
|
||||
}
|
||||
|
||||
labels
|
||||
};
|
||||
|
||||
let mut combined_results: Vec<SearchResultForTemplate> =
|
||||
Vec::with_capacity(search_result.chunks.len() + search_result.entities.len());
|
||||
|
||||
// Add chunk results
|
||||
for chunk_result in search_result.chunks {
|
||||
let source_label = source_label_map
|
||||
.get(&chunk_result.chunk.source_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| fallback_source_label(&chunk_result.chunk.source_id));
|
||||
combined_results.push(SearchResultForTemplate {
|
||||
result_type: "text_chunk".to_string(),
|
||||
score: chunk_result.score,
|
||||
text_chunk: Some(TextChunkForTemplate {
|
||||
id: chunk_result.chunk.id,
|
||||
source_id: chunk_result.chunk.source_id,
|
||||
source_label,
|
||||
chunk: chunk_result.chunk.chunk,
|
||||
score: chunk_result.score,
|
||||
}),
|
||||
@@ -125,6 +306,10 @@ pub async fn search_result_handler(
|
||||
|
||||
// Add entity results
|
||||
for entity_result in search_result.entities {
|
||||
let source_label = source_label_map
|
||||
.get(&entity_result.entity.source_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| fallback_source_label(&entity_result.entity.source_id));
|
||||
combined_results.push(SearchResultForTemplate {
|
||||
result_type: "knowledge_entity".to_string(),
|
||||
score: entity_result.score,
|
||||
@@ -135,6 +320,7 @@ pub async fn search_result_handler(
|
||||
description: entity_result.entity.description,
|
||||
entity_type: format!("{:?}", entity_result.entity.entity_type),
|
||||
source_id: entity_result.entity.source_id,
|
||||
source_label,
|
||||
score: entity_result.score,
|
||||
}),
|
||||
});
|
||||
|
||||
@@ -36,9 +36,12 @@
|
||||
<div id="system_prompt_section" class="nb-panel p-4">
|
||||
<div class="text-sm font-semibold mb-3">System Prompts</div>
|
||||
<div class="flex gap-2 flex-col sm:flex-row">
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-query-prompt" hx-target="#modal" hx-swap="innerHTML">Edit Query Prompt</button>
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-ingestion-prompt" hx-target="#modal" hx-swap="innerHTML">Edit Ingestion Prompt</button>
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-image-prompt" hx-target="#modal" hx-swap="innerHTML">Edit Image Prompt</button>
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-query-prompt" hx-target="#modal"
|
||||
hx-swap="innerHTML">Edit Query Prompt</button>
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-ingestion-prompt" hx-target="#modal"
|
||||
hx-swap="innerHTML">Edit Ingestion Prompt</button>
|
||||
<button type="button" class="nb-btn btn-sm" hx-get="/edit-image-prompt" hx-target="#modal"
|
||||
hx-swap="innerHTML">Edit Image Prompt</button>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -52,7 +55,8 @@
|
||||
<div class="text-sm opacity-80 mb-1">Query Model</div>
|
||||
<select name="query_model" class="nb-select w-full">
|
||||
{% for model in available_models.data %}
|
||||
<option value="{{model.id}}" {% if settings.query_model==model.id %} selected {% endif %}>{{model.id}}</option>
|
||||
<option value="{{model.id}}" {% if settings.query_model==model.id %} selected {% endif %}>{{model.id}}
|
||||
</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.query_model}}</span></p>
|
||||
@@ -63,7 +67,8 @@
|
||||
<div class="text-sm opacity-80 mb-1">Processing Model</div>
|
||||
<select name="processing_model" class="nb-select w-full">
|
||||
{% for model in available_models.data %}
|
||||
<option value="{{model.id}}" {% if settings.processing_model==model.id %} selected {% endif %}>{{model.id}}</option>
|
||||
<option value="{{model.id}}" {% if settings.processing_model==model.id %} selected {% endif %}>
|
||||
{{model.id}}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.processing_model}}</span></p>
|
||||
@@ -74,10 +79,12 @@
|
||||
<div class="text-sm opacity-80 mb-1">Image Processing Model</div>
|
||||
<select name="image_processing_model" class="nb-select w-full">
|
||||
{% for model in available_models.data %}
|
||||
<option value="{{model.id}}" {% if settings.image_processing_model==model.id %} selected {% endif %}>{{model.id}}</option>
|
||||
<option value="{{model.id}}" {% if settings.image_processing_model==model.id %} selected {% endif %}>
|
||||
{{model.id}}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.image_processing_model}}</span></p>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span
|
||||
class="font-mono">{{settings.image_processing_model}}</span></p>
|
||||
</div>
|
||||
|
||||
<!-- Voice Processing Model -->
|
||||
@@ -85,39 +92,66 @@
|
||||
<div class="text-sm opacity-80 mb-1">Voice Processing Model</div>
|
||||
<select name="voice_processing_model" class="nb-select w-full">
|
||||
{% for model in available_models.data %}
|
||||
<option value="{{model.id}}" {% if settings.voice_processing_model==model.id %} selected {% endif %}>{{model.id}}</option>
|
||||
<option value="{{model.id}}" {% if settings.voice_processing_model==model.id %} selected {% endif %}>
|
||||
{{model.id}}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.voice_processing_model}}</span></p>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span
|
||||
class="font-mono">{{settings.voice_processing_model}}</span></p>
|
||||
</div>
|
||||
|
||||
<!-- Embedding Model -->
|
||||
<div>
|
||||
<div class="text-sm opacity-80 mb-1">Embedding Model</div>
|
||||
{% if settings.embedding_backend == "fastembed" or settings.embedding_backend == "hashed" %}
|
||||
<input type="text" name="embedding_model" class="nb-input w-full opacity-60 cursor-not-allowed"
|
||||
value="{{settings.embedding_model}}" disabled />
|
||||
<p class="text-xs opacity-70 mt-1">Model: <span class="font-mono">{{settings.embedding_model}}
|
||||
({{settings.embedding_dimensions}} dims)</span></p>
|
||||
<p class="text-xs text-info mt-1">ℹ️ Embedding model is controlled by config when using <span
|
||||
class="font-mono">{{settings.embedding_backend}}</span> backend.</p>
|
||||
{% else %}
|
||||
<select name="embedding_model" class="nb-select w-full">
|
||||
{% for model in available_models.data %}
|
||||
<option value="{{model.id}}" {% if settings.embedding_model==model.id %} selected {% endif %}>{{model.id}}</option>
|
||||
<option value="{{model.id}}" {% if settings.embedding_model==model.id %} selected {% endif %}>{{model.id}}
|
||||
</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.embedding_model}} ({{settings.embedding_dimensions}} dims)</span></p>
|
||||
<p class="text-xs opacity-70 mt-1">Current: <span class="font-mono">{{settings.embedding_model}}
|
||||
({{settings.embedding_dimensions}} dims)</span></p>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<!-- Embedding Dimensions -->
|
||||
<div>
|
||||
<div class="text-sm opacity-80 mb-1" for="embedding_dimensions">Embedding Dimensions</div>
|
||||
<input type="number" id="embedding_dimensions" name="embedding_dimensions" class="nb-input w-full" value="{{ settings.embedding_dimensions }}" required />
|
||||
{% if settings.embedding_backend == "fastembed" or settings.embedding_backend == "hashed" %}
|
||||
<input type="number" id="embedding_dimensions" name="embedding_dimensions"
|
||||
class="nb-input w-full opacity-60 cursor-not-allowed" value="{{ settings.embedding_dimensions }}"
|
||||
disabled />
|
||||
<p class="text-xs text-info mt-1">ℹ️ Dimensions are fixed for <span
|
||||
class="font-mono">{{settings.embedding_backend}}</span> backend. Set <span
|
||||
class="font-mono">EMBEDDING_BACKEND=openai</span> to use OpenAI embeddings.</p>
|
||||
{% else %}
|
||||
<input type="number" id="embedding_dimensions" name="embedding_dimensions" class="nb-input w-full"
|
||||
value="{{ settings.embedding_dimensions }}" required />
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<!-- Alert -->
|
||||
{% if settings.embedding_backend != "fastembed" and settings.embedding_backend != "hashed" %}
|
||||
<div id="embedding-change-alert" class="nb-panel p-3 bg-warning/20 hidden">
|
||||
<div class="text-sm"><strong>Warning:</strong> Changing dimensions will require re-creating all embeddings. Look up your model's required dimensions or use a model that allows specifying them.</div>
|
||||
<div class="text-sm"><strong>Warning:</strong> Changing dimensions will require re-creating all embeddings.
|
||||
Look up your model's required dimensions or use a model that allows specifying them.</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="flex justify-end">
|
||||
<button type="submit" class="nb-btn nb-cta btn-sm">Save Model Settings</button>
|
||||
</div>
|
||||
</form>
|
||||
|
||||
{% if settings.embedding_backend != "fastembed" and settings.embedding_backend != "hashed" %}
|
||||
<script>
|
||||
// Rebind after HTMX swaps
|
||||
(() => {
|
||||
@@ -135,6 +169,7 @@
|
||||
}
|
||||
})();
|
||||
</script>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
</div>
|
||||
|
||||
@@ -143,7 +178,8 @@
|
||||
<label class="flex items-center gap-3">
|
||||
{% block registration_status_input %}
|
||||
<form hx-patch="/toggle-registrations" hx-swap="outerHTML" hx-trigger="change">
|
||||
<input name="registration_open" type="checkbox" class="nb-checkbox" {% if settings.registrations_enabled %}checked{% endif %} />
|
||||
<input name="registration_open" type="checkbox" class="nb-checkbox" {% if settings.registrations_enabled
|
||||
%}checked{% endif %} />
|
||||
</form>
|
||||
{% endblock %}
|
||||
<span class="text-sm">Enable Registrations</span>
|
||||
@@ -153,4 +189,4 @@
|
||||
</section>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
{% endblock %}
|
||||
@@ -22,7 +22,7 @@
|
||||
</div>
|
||||
|
||||
<div class="flex-1 min-w-0">
|
||||
<h3 class="text-lg font-extrabold mb-1 leading-snug">
|
||||
<h3 class="text-lg font-extrabold mb-2 leading-snug">
|
||||
<a hx-get="/content/{{ tc.id }}/read" hx-target="#modal" hx-swap="innerHTML" class="nb-link">
|
||||
{% set title_text = tc.highlighted_url_title
|
||||
| default(tc.url_info.title if tc.url_info else none, true)
|
||||
@@ -72,6 +72,44 @@
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
{% elif result.result_type == "text_chunk" %}
|
||||
{% set chunk = result.text_chunk %}
|
||||
<div
|
||||
class="w-10 h-10 flex-shrink-0 self-start mt-1 grid place-items-center border-2 border-neutral bg-base-100 shadow-[4px_4px_0_0_#000]">
|
||||
<div class="tooltip tooltip-right" data-tip="Text Chunk">
|
||||
{% include "icons/bars_icon.html" %}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="flex-1 min-w-0">
|
||||
<h3 class="text-lg font-extrabold mb-1 leading-snug">
|
||||
<a hx-get="/content/{{ chunk.source_id }}/read" hx-target="#modal" hx-swap="innerHTML" class="nb-link">
|
||||
Chunk from {{ chunk.source_label }}
|
||||
</a>
|
||||
</h3>
|
||||
|
||||
<div class="markdown-content prose-tufte-compact text-base-content/80 mb-4 overflow-hidden line-clamp-6"
|
||||
data-content="{{ chunk.chunk | escape }}">
|
||||
{{ chunk.chunk | escape }}
|
||||
</div>
|
||||
|
||||
<div class="text-xs flex flex-wrap gap-x-4 gap-y-2 items-center">
|
||||
<a hx-get="/content/{{ chunk.source_id }}/read" hx-target="#modal" hx-swap="innerHTML"
|
||||
class="nb-link uppercase tracking-wide">
|
||||
View full document
|
||||
</a>
|
||||
|
||||
<span class="inline-flex items-center min-w-0">
|
||||
<span class="uppercase tracking-wide opacity-60 mr-2">Source</span>
|
||||
<span class="nb-badge truncate max-w-xs" title="{{ chunk.source_label }}">{{ chunk.source_label }}</span>
|
||||
</span>
|
||||
|
||||
<span class="inline-flex items-center">
|
||||
<span class="uppercase tracking-wide opacity-60 mr-2">Score</span>
|
||||
<span class="nb-badge">{{ result.score | round(3) }}</span>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
{% elif result.result_type == "knowledge_entity" %}
|
||||
{% set entity = result.knowledge_entity %}
|
||||
<div
|
||||
@@ -84,15 +122,12 @@
|
||||
<div class="flex-1 min-w-0">
|
||||
<h3 class="text-lg font-extrabold mb-1 leading-snug">
|
||||
<a hx-get="/knowledge-entity/{{ entity.id }}" hx-target="#modal" hx-swap="innerHTML" class="nb-link">
|
||||
{% set entity_title = entity.highlighted_name | default(entity.name, true) %}
|
||||
{{ entity_title }}
|
||||
{{ entity.name }}
|
||||
</a>
|
||||
</h3>
|
||||
|
||||
<div class="prose prose-tufte-compact text-base-content/80 mb-3 overflow-hidden line-clamp-6">
|
||||
{% if entity.highlighted_description %}
|
||||
{{ entity.highlighted_description }}
|
||||
{% elif entity.description %}
|
||||
{% if entity.description %}
|
||||
{{ entity.description | escape }}
|
||||
{% else %}
|
||||
<span class="italic opacity-60">No description available.</span>
|
||||
@@ -107,14 +142,14 @@
|
||||
|
||||
{% if entity.source_id %}
|
||||
<span class="inline-flex items-center min-w-0">
|
||||
<span class="uppercase tracking-wide opacity-60 mr-2">Source ID</span>
|
||||
<span class="nb-badge truncate max-w-xs" title="{{ entity.source_id }}">{{ entity.source_id }}</span>
|
||||
<span class="uppercase tracking-wide opacity-60 mr-2">Source</span>
|
||||
<span class="nb-badge truncate max-w-xs" title="{{ entity.source_label }}">{{ entity.source_label }}</span>
|
||||
</span>
|
||||
{% endif %}
|
||||
|
||||
<span class="inline-flex items-center">
|
||||
<span class="uppercase tracking-wide opacity-60 mr-2">Score</span>
|
||||
<span class="nb-badge">{{ result.score }}</span>
|
||||
<span class="nb-badge">{{ result.score | round(3) }}</span>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
@@ -133,4 +168,4 @@
|
||||
<p class="text-lg font-semibold">Enter a term above to search your knowledge base.</p>
|
||||
<p class="text-sm opacity-70">Results will appear here.</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% endif %}
|
||||
|
||||
@@ -189,7 +189,7 @@ impl PipelineServices for DefaultPipelineServices {
|
||||
match retrieval_pipeline::retrieve_entities(
|
||||
&self.db,
|
||||
&self.openai_client,
|
||||
// embedding_provider_ref,
|
||||
Some(&*self.embedding_provider),
|
||||
&input_text,
|
||||
&content.user_id,
|
||||
config,
|
||||
|
||||
@@ -2,18 +2,19 @@ use api_router::{api_routes_v1, api_state::ApiState};
|
||||
use axum::{extract::FromRef, Router};
|
||||
use common::{
|
||||
storage::{
|
||||
db::SurrealDbClient,
|
||||
indexes::ensure_runtime_indexes,
|
||||
store::StorageManager,
|
||||
types::system_settings::SystemSettings,
|
||||
db::SurrealDbClient, indexes::ensure_runtime_indexes, store::StorageManager,
|
||||
types::{
|
||||
knowledge_entity::KnowledgeEntity, system_settings::SystemSettings,
|
||||
text_chunk::TextChunk,
|
||||
},
|
||||
},
|
||||
utils::config::get_config,
|
||||
utils::{config::get_config, embedding::EmbeddingProvider},
|
||||
};
|
||||
use html_router::{html_routes, html_state::HtmlState};
|
||||
use ingestion_pipeline::{pipeline::IngestionPipeline, run_worker_loop};
|
||||
use retrieval_pipeline::reranking::RerankerPool;
|
||||
use std::sync::Arc;
|
||||
use tracing::{error, info};
|
||||
use tracing::{error, info, warn};
|
||||
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
|
||||
|
||||
use tokio::task::LocalSet;
|
||||
@@ -44,8 +45,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
// Ensure db is initialized
|
||||
db.apply_migrations().await?;
|
||||
let settings = SystemSettings::get_current(&db).await?;
|
||||
ensure_runtime_indexes(&db, settings.embedding_dimensions as usize).await?;
|
||||
|
||||
let session_store = Arc::new(db.create_session_store().await?);
|
||||
let openai_client = Arc::new(async_openai::Client::with_config(
|
||||
@@ -54,6 +53,57 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
.with_api_base(&config.openai_base_url),
|
||||
));
|
||||
|
||||
// Create embedding provider based on config before syncing settings.
|
||||
let embedding_provider =
|
||||
Arc::new(EmbeddingProvider::from_config(&config, Some(openai_client.clone())).await?);
|
||||
info!(
|
||||
embedding_backend = ?config.embedding_backend,
|
||||
embedding_dimension = embedding_provider.dimension(),
|
||||
"Embedding provider initialized"
|
||||
);
|
||||
|
||||
// Sync SystemSettings with provider's dimensions/model/backend
|
||||
let (settings, dimensions_changed) =
|
||||
SystemSettings::sync_from_embedding_provider(&db, &embedding_provider).await?;
|
||||
|
||||
// Now ensure runtime indexes with the correct (synced) dimensions
|
||||
ensure_runtime_indexes(&db, settings.embedding_dimensions as usize).await?;
|
||||
|
||||
// If dimensions changed, re-embed existing data to keep queries working.
|
||||
if dimensions_changed {
|
||||
warn!(
|
||||
new_dimensions = settings.embedding_dimensions,
|
||||
"Embedding configuration changed; re-embedding existing data"
|
||||
);
|
||||
|
||||
// Re-embed text chunks
|
||||
info!("Re-embedding TextChunks");
|
||||
if let Err(e) = TextChunk::update_all_embeddings_with_provider(
|
||||
&db,
|
||||
&embedding_provider,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("Failed to re-embed TextChunks: {}. Search results may be stale.", e);
|
||||
}
|
||||
|
||||
// Re-embed knowledge entities
|
||||
info!("Re-embedding KnowledgeEntities");
|
||||
if let Err(e) = KnowledgeEntity::update_all_embeddings_with_provider(
|
||||
&db,
|
||||
&embedding_provider,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!(
|
||||
"Failed to re-embed KnowledgeEntities: {}. Search results may be stale.",
|
||||
e
|
||||
);
|
||||
}
|
||||
|
||||
info!("Re-embedding complete.");
|
||||
}
|
||||
|
||||
let reranker_pool = RerankerPool::maybe_from_config(&config)?;
|
||||
|
||||
// Create global storage manager
|
||||
@@ -66,6 +116,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
storage.clone(),
|
||||
config.clone(),
|
||||
reranker_pool.clone(),
|
||||
embedding_provider.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -114,9 +165,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let settings = SystemSettings::get_current(&worker_db)
|
||||
.await
|
||||
.expect("failed to load system settings");
|
||||
|
||||
// Initialize worker components
|
||||
let openai_client = Arc::new(async_openai::Client::with_config(
|
||||
@@ -125,14 +173,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
.with_api_base(&config.openai_base_url),
|
||||
));
|
||||
|
||||
// Create embedding provider for ingestion
|
||||
// Create embedding provider based on config
|
||||
let embedding_provider = Arc::new(
|
||||
common::utils::embedding::EmbeddingProvider::new_openai(
|
||||
openai_client.clone(),
|
||||
settings.embedding_model,
|
||||
settings.embedding_dimensions,
|
||||
)
|
||||
.expect("failed to create embedding provider"),
|
||||
EmbeddingProvider::from_config(&config, Some(openai_client.clone()))
|
||||
.await
|
||||
.expect("failed to create embedding provider"),
|
||||
);
|
||||
let ingestion_pipeline = Arc::new(
|
||||
IngestionPipeline::new(
|
||||
@@ -226,6 +271,12 @@ mod tests {
|
||||
.await
|
||||
.expect("failed to build storage manager");
|
||||
|
||||
// Use hashed embeddings for tests to avoid external dependencies
|
||||
let embedding_provider = Arc::new(
|
||||
common::utils::embedding::EmbeddingProvider::new_hashed(384)
|
||||
.expect("failed to create hashed embedding provider"),
|
||||
);
|
||||
|
||||
let html_state = HtmlState::new_with_resources(
|
||||
db.clone(),
|
||||
openai_client,
|
||||
@@ -233,6 +284,7 @@ mod tests {
|
||||
storage.clone(),
|
||||
config.clone(),
|
||||
None,
|
||||
embedding_provider,
|
||||
)
|
||||
.await
|
||||
.expect("failed to build html state");
|
||||
|
||||
@@ -3,7 +3,8 @@ use std::sync::Arc;
|
||||
use api_router::{api_routes_v1, api_state::ApiState};
|
||||
use axum::{extract::FromRef, Router};
|
||||
use common::{
|
||||
storage::db::SurrealDbClient, storage::store::StorageManager, utils::config::get_config,
|
||||
storage::{db::SurrealDbClient, store::StorageManager, types::system_settings::SystemSettings},
|
||||
utils::{config::get_config, embedding::EmbeddingProvider},
|
||||
};
|
||||
use html_router::{html_routes, html_state::HtmlState};
|
||||
use retrieval_pipeline::reranking::RerankerPool;
|
||||
@@ -49,6 +50,20 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Create global storage manager
|
||||
let storage = StorageManager::new(&config).await?;
|
||||
|
||||
// Create embedding provider based on config
|
||||
let embedding_provider = Arc::new(
|
||||
EmbeddingProvider::from_config(&config, Some(openai_client.clone())).await?,
|
||||
);
|
||||
info!(
|
||||
embedding_backend = ?config.embedding_backend,
|
||||
embedding_dimension = embedding_provider.dimension(),
|
||||
"Embedding provider initialized"
|
||||
);
|
||||
|
||||
// Sync SystemSettings with provider's dimensions/backend for visibility
|
||||
let (_settings, _dimensions_changed) =
|
||||
SystemSettings::sync_from_embedding_provider(&db, &embedding_provider).await?;
|
||||
|
||||
let html_state = HtmlState::new_with_resources(
|
||||
db,
|
||||
openai_client,
|
||||
@@ -56,6 +71,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
storage.clone(),
|
||||
config.clone(),
|
||||
reranker_pool,
|
||||
embedding_provider,
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::{
|
||||
storage::db::SurrealDbClient, storage::store::StorageManager, utils::config::get_config,
|
||||
storage::db::SurrealDbClient, storage::store::StorageManager,
|
||||
utils::{config::get_config, embedding::EmbeddingProvider},
|
||||
};
|
||||
use ingestion_pipeline::{pipeline::IngestionPipeline, run_worker_loop};
|
||||
use retrieval_pipeline::reranking::RerankerPool;
|
||||
use tracing::info;
|
||||
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
|
||||
|
||||
#[tokio::main]
|
||||
@@ -37,9 +39,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
let reranker_pool = RerankerPool::maybe_from_config(&config)?;
|
||||
|
||||
// Create embedding provider for ingestion
|
||||
let embedding_provider =
|
||||
Arc::new(common::utils::embedding::EmbeddingProvider::new_fastembed(None).await?);
|
||||
// Create embedding provider based on config
|
||||
let embedding_provider = Arc::new(
|
||||
EmbeddingProvider::from_config(&config, Some(openai_client.clone())).await?,
|
||||
);
|
||||
info!(
|
||||
embedding_backend = ?config.embedding_backend,
|
||||
"Embedding provider initialized for worker"
|
||||
);
|
||||
|
||||
// Create global storage manager
|
||||
let storage = StorageManager::new(&config).await?;
|
||||
|
||||
@@ -21,11 +21,29 @@ use tracing::instrument;
|
||||
pub enum StrategyOutput {
|
||||
Entities(Vec<RetrievedEntity>),
|
||||
Chunks(Vec<RetrievedChunk>),
|
||||
Search(SearchResult),
|
||||
}
|
||||
|
||||
/// Unified search result containing both chunks and entities
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchResult {
|
||||
pub chunks: Vec<RetrievedChunk>,
|
||||
pub entities: Vec<RetrievedEntity>,
|
||||
}
|
||||
|
||||
impl SearchResult {
|
||||
pub fn new(chunks: Vec<RetrievedChunk>, entities: Vec<RetrievedEntity>) -> Self {
|
||||
Self { chunks, entities }
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.chunks.is_empty() && self.entities.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
pub use pipeline::{
|
||||
retrieved_entities_to_json, PipelineDiagnostics, PipelineStageTimings, RetrievalConfig,
|
||||
RetrievalStrategy, RetrievalTuning,
|
||||
RetrievalStrategy, RetrievalTuning, SearchTarget,
|
||||
};
|
||||
|
||||
// Captures a supporting chunk plus its fused retrieval score for downstream prompts.
|
||||
@@ -48,6 +66,7 @@ pub struct RetrievedEntity {
|
||||
pub async fn retrieve_entities(
|
||||
db_client: &SurrealDbClient,
|
||||
openai_client: &async_openai::Client<async_openai::config::OpenAIConfig>,
|
||||
embedding_provider: Option<&common::utils::embedding::EmbeddingProvider>,
|
||||
input_text: &str,
|
||||
user_id: &str,
|
||||
config: RetrievalConfig,
|
||||
@@ -56,7 +75,7 @@ pub async fn retrieve_entities(
|
||||
pipeline::run_pipeline(
|
||||
db_client,
|
||||
openai_client,
|
||||
None,
|
||||
embedding_provider,
|
||||
input_text,
|
||||
user_id,
|
||||
config,
|
||||
@@ -252,4 +271,49 @@ mod tests {
|
||||
"Chunk results should contain relevant snippets"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_search_strategy_returns_search_result() {
|
||||
let db = setup_test_db().await;
|
||||
let user_id = "search_user";
|
||||
let chunk = TextChunk::new(
|
||||
"search_src".into(),
|
||||
"Async Rust programming uses Tokio runtime for concurrent tasks.".into(),
|
||||
user_id.into(),
|
||||
);
|
||||
|
||||
TextChunk::store_with_embedding(chunk.clone(), chunk_embedding_primary(), &db)
|
||||
.await
|
||||
.expect("Failed to store chunk");
|
||||
|
||||
let config = RetrievalConfig::for_search(pipeline::SearchTarget::Both);
|
||||
let openai_client = Client::new();
|
||||
let results = pipeline::run_pipeline_with_embedding(
|
||||
&db,
|
||||
&openai_client,
|
||||
None,
|
||||
test_embedding(),
|
||||
"async rust programming",
|
||||
user_id,
|
||||
config,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("Search strategy retrieval failed");
|
||||
|
||||
let search_result = match results {
|
||||
StrategyOutput::Search(sr) => sr,
|
||||
other => panic!("expected Search output, got {:?}", other),
|
||||
};
|
||||
|
||||
// Should return chunks (entities may be empty if none stored)
|
||||
assert!(
|
||||
!search_result.chunks.is_empty(),
|
||||
"Search strategy should return chunks"
|
||||
);
|
||||
assert!(
|
||||
search_result.chunks.iter().any(|c| c.chunk.chunk.contains("Tokio")),
|
||||
"Search results should contain relevant chunks"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,21 @@ pub enum RetrievalStrategy {
|
||||
RelationshipSuggestion,
|
||||
/// Entity retrieval for context during content ingestion
|
||||
Ingestion,
|
||||
/// Unified search returning both chunks and entities
|
||||
Search,
|
||||
}
|
||||
|
||||
/// Configures which result types to include in Search strategy
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SearchTarget {
|
||||
/// Return only text chunks
|
||||
ChunksOnly,
|
||||
/// Return only knowledge entities
|
||||
EntitiesOnly,
|
||||
/// Return both chunks and entities (default)
|
||||
#[default]
|
||||
Both,
|
||||
}
|
||||
|
||||
impl Default for RetrievalStrategy {
|
||||
@@ -37,6 +52,7 @@ impl std::str::FromStr for RetrievalStrategy {
|
||||
}
|
||||
"relationship_suggestion" => Ok(Self::RelationshipSuggestion),
|
||||
"ingestion" => Ok(Self::Ingestion),
|
||||
"search" => Ok(Self::Search),
|
||||
other => Err(format!("unknown retrieval strategy '{other}'")),
|
||||
}
|
||||
}
|
||||
@@ -48,6 +64,7 @@ impl fmt::Display for RetrievalStrategy {
|
||||
RetrievalStrategy::Default => "default",
|
||||
RetrievalStrategy::RelationshipSuggestion => "relationship_suggestion",
|
||||
RetrievalStrategy::Ingestion => "ingestion",
|
||||
RetrievalStrategy::Search => "search",
|
||||
};
|
||||
f.write_str(label)
|
||||
}
|
||||
@@ -140,6 +157,8 @@ impl Default for RetrievalTuning {
|
||||
pub struct RetrievalConfig {
|
||||
pub strategy: RetrievalStrategy,
|
||||
pub tuning: RetrievalTuning,
|
||||
/// Target for Search strategy (chunks, entities, or both)
|
||||
pub search_target: SearchTarget,
|
||||
}
|
||||
|
||||
impl RetrievalConfig {
|
||||
@@ -147,6 +166,7 @@ impl RetrievalConfig {
|
||||
Self {
|
||||
strategy: RetrievalStrategy::Default,
|
||||
tuning,
|
||||
search_target: SearchTarget::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -154,11 +174,16 @@ impl RetrievalConfig {
|
||||
Self {
|
||||
strategy,
|
||||
tuning: RetrievalTuning::default(),
|
||||
search_target: SearchTarget::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_tuning(strategy: RetrievalStrategy, tuning: RetrievalTuning) -> Self {
|
||||
Self { strategy, tuning }
|
||||
Self {
|
||||
strategy,
|
||||
tuning,
|
||||
search_target: SearchTarget::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create config for chat retrieval with strategy selection support
|
||||
@@ -175,6 +200,15 @@ impl RetrievalConfig {
|
||||
pub fn for_ingestion() -> Self {
|
||||
Self::with_strategy(RetrievalStrategy::Ingestion)
|
||||
}
|
||||
|
||||
/// Create config for unified search (chunks and/or entities)
|
||||
pub fn for_search(target: SearchTarget) -> Self {
|
||||
Self {
|
||||
strategy: RetrievalStrategy::Search,
|
||||
tuning: RetrievalTuning::default(),
|
||||
search_target: target,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RetrievalConfig {
|
||||
@@ -182,6 +216,7 @@ impl Default for RetrievalConfig {
|
||||
Self {
|
||||
strategy: RetrievalStrategy::default(),
|
||||
tuning: RetrievalTuning::default(),
|
||||
search_target: SearchTarget::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ mod diagnostics;
|
||||
mod stages;
|
||||
mod strategies;
|
||||
|
||||
pub use config::{RetrievalConfig, RetrievalStrategy, RetrievalTuning};
|
||||
pub use config::{RetrievalConfig, RetrievalStrategy, RetrievalTuning, SearchTarget};
|
||||
pub use diagnostics::{
|
||||
AssembleStats, ChunkEnrichmentStats, CollectCandidatesStats, EntityAssemblyTrace,
|
||||
PipelineDiagnostics,
|
||||
@@ -17,7 +17,7 @@ use std::time::{Duration, Instant};
|
||||
use tracing::info;
|
||||
|
||||
use stages::PipelineContext;
|
||||
use strategies::{DefaultStrategyDriver, IngestionDriver, RelationshipSuggestionDriver};
|
||||
use strategies::{DefaultStrategyDriver, IngestionDriver, RelationshipSuggestionDriver, SearchStrategyDriver};
|
||||
|
||||
// Export StrategyOutput publicly from this module
|
||||
// (it's defined in lib.rs but we re-export it here)
|
||||
@@ -181,6 +181,24 @@ pub async fn run_pipeline(
|
||||
.await?;
|
||||
Ok(StrategyOutput::Entities(run.results))
|
||||
}
|
||||
RetrievalStrategy::Search => {
|
||||
let search_target = config.search_target;
|
||||
let driver = SearchStrategyDriver::new(search_target);
|
||||
let run = execute_strategy(
|
||||
driver,
|
||||
db_client,
|
||||
openai_client,
|
||||
embedding_provider,
|
||||
None,
|
||||
input_text,
|
||||
user_id,
|
||||
config,
|
||||
reranker,
|
||||
false,
|
||||
)
|
||||
.await?;
|
||||
Ok(StrategyOutput::Search(run.results))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -246,6 +264,24 @@ pub async fn run_pipeline_with_embedding(
|
||||
.await?;
|
||||
Ok(StrategyOutput::Entities(run.results))
|
||||
}
|
||||
RetrievalStrategy::Search => {
|
||||
let search_target = config.search_target;
|
||||
let driver = SearchStrategyDriver::new(search_target);
|
||||
let run = execute_strategy(
|
||||
driver,
|
||||
db_client,
|
||||
openai_client,
|
||||
embedding_provider,
|
||||
Some(query_embedding),
|
||||
input_text,
|
||||
user_id,
|
||||
config,
|
||||
reranker,
|
||||
false,
|
||||
)
|
||||
.await?;
|
||||
Ok(StrategyOutput::Search(run.results))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -88,3 +88,63 @@ impl StrategyDriver for IngestionDriver {
|
||||
Ok(ctx.take_entity_results())
|
||||
}
|
||||
}
|
||||
|
||||
use crate::SearchResult;
|
||||
use super::config::SearchTarget;
|
||||
|
||||
/// Search strategy driver that retrieves both chunks and entities
|
||||
pub struct SearchStrategyDriver {
|
||||
target: SearchTarget,
|
||||
}
|
||||
|
||||
impl SearchStrategyDriver {
|
||||
pub fn new(target: SearchTarget) -> Self {
|
||||
Self { target }
|
||||
}
|
||||
}
|
||||
|
||||
impl StrategyDriver for SearchStrategyDriver {
|
||||
type Output = SearchResult;
|
||||
|
||||
fn stages(&self) -> Vec<BoxedStage> {
|
||||
match self.target {
|
||||
SearchTarget::ChunksOnly => vec![
|
||||
Box::new(EmbedStage),
|
||||
Box::new(ChunkVectorStage),
|
||||
Box::new(ChunkRerankStage),
|
||||
Box::new(ChunkAssembleStage),
|
||||
],
|
||||
SearchTarget::EntitiesOnly => vec![
|
||||
Box::new(EmbedStage),
|
||||
Box::new(CollectCandidatesStage),
|
||||
Box::new(GraphExpansionStage),
|
||||
Box::new(RerankStage),
|
||||
Box::new(AssembleEntitiesStage),
|
||||
],
|
||||
SearchTarget::Both => vec![
|
||||
Box::new(EmbedStage),
|
||||
// Chunk retrieval path
|
||||
Box::new(ChunkVectorStage),
|
||||
Box::new(ChunkRerankStage),
|
||||
Box::new(ChunkAssembleStage),
|
||||
// Entity retrieval path (runs after chunk stages)
|
||||
Box::new(CollectCandidatesStage),
|
||||
Box::new(GraphExpansionStage),
|
||||
Box::new(RerankStage),
|
||||
Box::new(AssembleEntitiesStage),
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
fn finalize(&self, ctx: &mut PipelineContext<'_>) -> Result<Self::Output, AppError> {
|
||||
let chunks = match self.target {
|
||||
SearchTarget::EntitiesOnly => Vec::new(),
|
||||
_ => ctx.take_chunk_results(),
|
||||
};
|
||||
let entities = match self.target {
|
||||
SearchTarget::ChunksOnly => Vec::new(),
|
||||
_ => ctx.take_entity_results(),
|
||||
};
|
||||
Ok(SearchResult::new(chunks, entities))
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user