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

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