tidying stuff up, dto for search

This commit is contained in:
Per Stark
2025-12-20 22:30:31 +01:00
parent a5bc72aedf
commit 79ea007b0a
23 changed files with 936 additions and 73 deletions

2
Cargo.lock generated
View File

@@ -2194,7 +2194,7 @@ dependencies = [
]
[[package]]
name = "eval"
name = "evaluations"
version = "0.1.0"
dependencies = [
"anyhow",

View File

@@ -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"

View File

@@ -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;

View File

@@ -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)]

View File

@@ -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)]

View File

@@ -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)]

View File

@@ -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(),
}
}
}

View File

@@ -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

View File

@@ -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,
})
}

View File

@@ -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()
};

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,
}),
});

View File

@@ -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 %}

View File

@@ -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 %}

View File

@@ -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,

View File

@@ -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");

View File

@@ -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?;

View File

@@ -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?;

View File

@@ -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"
);
}
}

View File

@@ -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(),
}
}
}

View File

@@ -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))
}
}
}

View File

@@ -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))
}
}