mirror of
https://github.com/perstarkse/minne.git
synced 2026-07-04 03:51:43 +02:00
chore: harden common storage bootstrap and slim embedded db assets
Unify embedding config, build providers from system settings, and fail startup when index builds error or time out. Move Surreal assets under common/db so embeds exclude crate source, and read storage via streams.
This commit is contained in:
@@ -1,9 +1,18 @@
|
||||
use config::{Config, ConfigError, Environment, File};
|
||||
use serde::Deserialize;
|
||||
use std::env;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::{env, sync::Once, str::FromStr};
|
||||
use thiserror::Error;
|
||||
|
||||
/// Error returned when parsing an embedding backend name.
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
#[error("unknown embedding backend '{input}': expected 'openai', 'hashed', or 'fastembed'")]
|
||||
pub struct ParseEmbeddingBackendError {
|
||||
/// The unrecognized input string.
|
||||
pub input: String,
|
||||
}
|
||||
|
||||
/// Selects the embedding backend for vector generation.
|
||||
#[derive(Clone, Copy, Deserialize, Debug, Default, PartialEq)]
|
||||
#[derive(Clone, Copy, Deserialize, Serialize, Debug, Default, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum EmbeddingBackend {
|
||||
/// Use OpenAI-compatible API for embeddings.
|
||||
@@ -15,6 +24,32 @@ pub enum EmbeddingBackend {
|
||||
Hashed,
|
||||
}
|
||||
|
||||
impl EmbeddingBackend {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::OpenAI => "openai",
|
||||
Self::FastEmbed => "fastembed",
|
||||
Self::Hashed => "hashed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for EmbeddingBackend {
|
||||
type Err = ParseEmbeddingBackendError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_ascii_lowercase().as_str() {
|
||||
"openai" => Ok(Self::OpenAI),
|
||||
"hashed" => Ok(Self::Hashed),
|
||||
"fastembed" | "fast-embed" | "fast" => Ok(Self::FastEmbed),
|
||||
other => Err(ParseEmbeddingBackendError {
|
||||
input: other.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Deserialize, Debug, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum StorageKind {
|
||||
@@ -133,11 +168,17 @@ fn default_ingest_max_category_bytes() -> usize {
|
||||
128
|
||||
}
|
||||
|
||||
static ORT_PATH_INIT: Once = Once::new();
|
||||
|
||||
/// Sets `ORT_DYLIB_PATH` once per process when a bundled ONNX runtime library is found.
|
||||
pub fn ensure_ort_path() {
|
||||
if env::var_os("ORT_DYLIB_PATH").is_some() {
|
||||
return;
|
||||
}
|
||||
if let Ok(mut exe) = env::current_exe() {
|
||||
ORT_PATH_INIT.call_once(|| {
|
||||
if env::var_os("ORT_DYLIB_PATH").is_some() {
|
||||
return;
|
||||
}
|
||||
let Ok(mut exe) = env::current_exe() else {
|
||||
return;
|
||||
};
|
||||
exe.pop();
|
||||
|
||||
if cfg!(target_os = "windows") {
|
||||
@@ -160,7 +201,7 @@ pub fn ensure_ort_path() {
|
||||
if p.exists() {
|
||||
env::set_var("ORT_DYLIB_PATH", p);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
impl Default for AppConfig {
|
||||
|
||||
@@ -8,59 +8,15 @@ use std::{
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use async_openai::{types::CreateEmbeddingRequestArgs, Client};
|
||||
use fastembed::{EmbeddingModel, ModelTrait, TextEmbedding, TextInitOptions};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror::Error;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::{
|
||||
error::AppError,
|
||||
storage::{db::SurrealDbClient, types::system_settings::SystemSettings},
|
||||
utils::config::AppConfig,
|
||||
};
|
||||
|
||||
/// Error returned when parsing an embedding backend name.
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
#[error("unknown embedding backend '{input}': expected 'openai', 'hashed', or 'fastembed'")]
|
||||
pub struct ParseEmbeddingBackendError {
|
||||
/// The unrecognized input string.
|
||||
pub input: String,
|
||||
}
|
||||
|
||||
/// Supported embedding backends.
|
||||
#[allow(clippy::module_name_repetitions)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum EmbeddingBackend {
|
||||
#[default]
|
||||
OpenAI,
|
||||
FastEmbed,
|
||||
Hashed,
|
||||
}
|
||||
|
||||
impl EmbeddingBackend {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::OpenAI => "openai",
|
||||
Self::FastEmbed => "fastembed",
|
||||
Self::Hashed => "hashed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for EmbeddingBackend {
|
||||
type Err = ParseEmbeddingBackendError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_ascii_lowercase().as_str() {
|
||||
"openai" => Ok(Self::OpenAI),
|
||||
"hashed" => Ok(Self::Hashed),
|
||||
"fastembed" | "fast-embed" | "fast" => Ok(Self::FastEmbed),
|
||||
other => Err(ParseEmbeddingBackendError {
|
||||
input: other.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
pub use crate::utils::config::{EmbeddingBackend, ParseEmbeddingBackendError};
|
||||
|
||||
/// Wrapper around the chosen embedding backend.
|
||||
#[allow(clippy::module_name_repetitions)]
|
||||
@@ -281,30 +237,31 @@ impl EmbeddingProvider {
|
||||
})
|
||||
}
|
||||
|
||||
/// Creates an embedding provider based on application configuration.
|
||||
/// Creates an embedding provider from persisted settings and bootstrap config.
|
||||
///
|
||||
/// 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,
|
||||
/// Model name and dimensions come from [`SystemSettings`]. The active backend is taken
|
||||
/// from `config.embedding_backend` at startup; [`SystemSettings::sync_from_embedding_provider`]
|
||||
/// persists the resolved backend to the database.
|
||||
pub async fn from_system_settings(
|
||||
settings: &SystemSettings,
|
||||
config: &AppConfig,
|
||||
openai_client: Option<Arc<Client<async_openai::config::OpenAIConfig>>>,
|
||||
) -> Result<Self> {
|
||||
use crate::utils::config::EmbeddingBackend;
|
||||
|
||||
let dimensions = settings.embedding_dimensions;
|
||||
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)
|
||||
Self::new_openai(client, settings.embedding_model.clone(), dimensions)
|
||||
}
|
||||
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
|
||||
Self::new_fastembed(Some(settings.embedding_model.clone())).await
|
||||
}
|
||||
EmbeddingBackend::Hashed => {
|
||||
let dimension = usize::try_from(dimensions)
|
||||
.map_err(|_| anyhow!("embedding_dimensions exceeds usize::MAX"))?;
|
||||
Self::new_hashed(dimension)
|
||||
}
|
||||
EmbeddingBackend::Hashed => Self::new_hashed(384),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -460,6 +417,14 @@ mod tests {
|
||||
use crate::storage::types::system_settings::SystemSettings;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn embedding_backend_defaults_to_fastembed() {
|
||||
assert_eq!(
|
||||
EmbeddingBackend::default(),
|
||||
EmbeddingBackend::FastEmbed
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_backend_as_str_matches_serde_names() {
|
||||
assert_eq!(EmbeddingBackend::OpenAI.as_str(), "openai");
|
||||
|
||||
@@ -29,6 +29,14 @@ pub fn validate_ingest_input(
|
||||
category: &str,
|
||||
file_count: usize,
|
||||
) -> Result<(), IngestValidationError> {
|
||||
let text_field_bytes = content.map(str::len).unwrap_or(0) + ctx.len() + category.len();
|
||||
if text_field_bytes > config.ingest_max_body_bytes {
|
||||
return Err(IngestValidationError::PayloadTooLarge(format!(
|
||||
"request text fields exceed maximum allowed body size of {} bytes",
|
||||
config.ingest_max_body_bytes
|
||||
)));
|
||||
}
|
||||
|
||||
if file_count > config.ingest_max_files {
|
||||
return Err(IngestValidationError::BadRequest(format!(
|
||||
"too many files: maximum allowed is {}",
|
||||
@@ -127,4 +135,18 @@ mod tests {
|
||||
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_ingest_input_rejects_oversized_text_fields() {
|
||||
let config = AppConfig {
|
||||
ingest_max_body_bytes: 10,
|
||||
..Default::default()
|
||||
};
|
||||
let result = validate_ingest_input(&config, Some("123456"), "ctx", "cat", 0);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(IngestValidationError::PayloadTooLarge(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user