mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-10 13:13:47 +00:00
1365 lines
50 KiB
Rust
1365 lines
50 KiB
Rust
use std::collections::{HashMap, HashSet};
|
|
use std::fs;
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::Mutex;
|
|
|
|
use axum::Json;
|
|
use log::{error, info, warn};
|
|
use once_cell::sync::Lazy;
|
|
use qdrant_edge::external::serde_json::{json, Value};
|
|
use qdrant_edge::external::uuid::Uuid;
|
|
use qdrant_edge::{
|
|
Condition, Distance, EdgeConfig, EdgeOptimizersConfig, EdgeShard, EdgeVectorParams,
|
|
FieldCondition, Filter, HasIdCondition, HnswIndexConfig, Match, MatchValue, NamedQuery,
|
|
Payload, PointId, PointInsertOperations, PointOperations, PointStruct, QueryEnum,
|
|
QueryRequest, ScoredPoint, ScoringQuery, UpdateOperation, ValueVariants, VectorInternal,
|
|
Vectors, WithPayloadInterface, WithVector,
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use tauri::Manager;
|
|
|
|
use crate::api_token::APIToken;
|
|
use crate::environment::DATA_DIRECTORY;
|
|
use crate::metadata::META_DATA;
|
|
|
|
const VECTOR_NAME: &str = "embedding";
|
|
const HNSW_M: usize = 16;
|
|
const HNSW_EF_CONSTRUCT: usize = 100;
|
|
const HNSW_FULL_SCAN_THRESHOLD_KB: usize = 10_000;
|
|
const HNSW_MAX_INDEXING_THREADS: usize = 0;
|
|
const VECTOR_INDEXING_THRESHOLD_KB: usize = 10_000;
|
|
const STORE_INITIALIZATION_MARKER: &str = "store_name.txt";
|
|
const STORE_INITIALIZATION_MARKER_TEMP: &str = "store_name.tmp";
|
|
const STORE_DISPLAY_NAME_MARKER: &str = "data_source_name.txt";
|
|
const STORE_DISPLAY_NAME_MARKER_TEMP: &str = "data_source_name.tmp";
|
|
const WAL_DIRECTORY: &str = "wal";
|
|
const CLOSED_WAL_SEGMENT_PREFIX: &str = "closed-";
|
|
|
|
/// Marks a response whose store exists on disk but cannot be opened. The .NET side keys its repair
|
|
/// offer off this value instead of parsing `issue`, so rewording the message stays harmless.
|
|
const ISSUE_CODE_STORE_UNREADABLE: &str = "store-unreadable";
|
|
|
|
type QdrantEdgeResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
|
|
|
|
static QDRANT_EDGE_DATABASE: Lazy<Mutex<Option<QdrantEdgeDatabase>>> =
|
|
Lazy::new(|| Mutex::new(None));
|
|
|
|
static QDRANT_EDGE_STATUS: Lazy<Mutex<QdrantEdgeStatusInfo>> =
|
|
Lazy::new(|| Mutex::new(QdrantEdgeStatusInfo::default()));
|
|
|
|
#[derive(Default)]
|
|
struct QdrantEdgeStatusInfo {
|
|
status: QdrantEdgeStatus,
|
|
unavailable_reason: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Default, Serialize, PartialEq, Eq)]
|
|
pub enum QdrantEdgeStatus {
|
|
#[default]
|
|
Starting,
|
|
Available,
|
|
Unavailable,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
pub struct QdrantEdgeServiceInfo {
|
|
pub status: QdrantEdgeStatus,
|
|
pub name: String,
|
|
pub version: String,
|
|
pub path: String,
|
|
pub stores_count: usize,
|
|
pub is_available: bool,
|
|
pub unavailable_reason: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, Deserialize)]
|
|
pub struct QdrantEdgeStoragePoint {
|
|
pub point_id: String,
|
|
pub vector: Vec<f32>,
|
|
pub data_source_id: String,
|
|
pub data_source_type: String,
|
|
pub chunk_id: String,
|
|
pub parent_file_id: String,
|
|
pub file_path: String,
|
|
pub absolute_path: String,
|
|
pub file_name: String,
|
|
pub relative_path: String,
|
|
pub file_type: String,
|
|
pub page_number: Option<i32>,
|
|
pub chunk_index: i32,
|
|
pub text: String,
|
|
pub fingerprint: String,
|
|
pub creation_utc: String,
|
|
pub last_write_utc: String,
|
|
pub embedded_at_utc: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct EnsureQdrantEdgeStoreRequest {
|
|
pub store_name: String,
|
|
pub data_source_name: String,
|
|
pub vector_size: usize,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct InsertQdrantEdgeEmbeddingRequest {
|
|
pub store_name: String,
|
|
pub points: Vec<QdrantEdgeStoragePoint>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct SearchQdrantEdgeEmbeddingRequest {
|
|
pub store_name: String,
|
|
pub vector: Vec<f32>,
|
|
pub max_matches: usize,
|
|
|
|
/// The only points the search may return. Missing means the whole store. An empty list means
|
|
/// no point at all, so the search finds nothing -- it never falls back to the whole store.
|
|
pub point_ids: Option<Vec<String>>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct DeleteQdrantEdgeEmbeddingByFileRequest {
|
|
pub store_name: String,
|
|
pub file_path: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct OptimizeQdrantEdgeStoreRequest {
|
|
pub store_name: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct DeleteQdrantEdgeStoreRequest {
|
|
pub store_name: String,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
pub struct QdrantEdgeResponse<T> {
|
|
pub success: bool,
|
|
pub issue: String,
|
|
pub issue_code: &'static str,
|
|
pub data: Option<T>,
|
|
}
|
|
|
|
/// A vector store which is initialized on disk but which Qdrant Edge refuses to open.
|
|
///
|
|
/// This is deliberately its own error type rather than one more formatted string: a broken store
|
|
/// is the one failure the user can act on, and the request layer has to recognize it to label the
|
|
/// response. Nothing here deletes the store -- rebuilding the embeddings costs the user time and,
|
|
/// with a cloud embedding provider, money, so that stays their decision.
|
|
#[derive(Debug)]
|
|
struct StoreUnreadableError {
|
|
store_name: String,
|
|
message: String,
|
|
}
|
|
|
|
impl StoreUnreadableError {
|
|
fn new(store_name: &str, path: &Path, source: impl std::fmt::Display) -> Self {
|
|
Self {
|
|
store_name: store_name.to_string(),
|
|
message: format!("Failed to load vector store '{store_name}' from '{}': {source}", path.display()),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Display for StoreUnreadableError {
|
|
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
formatter.write_str(&self.message)
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for StoreUnreadableError {}
|
|
|
|
#[derive(Serialize)]
|
|
pub struct QdrantEdgeEnsureStoreResult {
|
|
pub created: bool,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
pub struct QdrantEdgeSearchResult {
|
|
pub point_id: String,
|
|
pub score: f32,
|
|
pub data_source_id: String,
|
|
pub data_source_type: String,
|
|
pub chunk_id: String,
|
|
pub parent_file_id: String,
|
|
pub file_path: String,
|
|
pub absolute_path: String,
|
|
pub file_name: String,
|
|
pub relative_path: String,
|
|
pub file_type: String,
|
|
pub page_number: Option<i32>,
|
|
pub chunk_index: i32,
|
|
pub text: String,
|
|
pub fingerprint: String,
|
|
pub creation_utc: String,
|
|
pub last_write_utc: String,
|
|
pub embedded_at_utc: String,
|
|
}
|
|
|
|
#[derive(Clone, Serialize)]
|
|
pub struct QdrantEdgeInfo {
|
|
pub name: String,
|
|
pub version: String,
|
|
pub path: String,
|
|
pub stores_count: usize,
|
|
}
|
|
|
|
pub struct QdrantEdgeDatabase {
|
|
base_path: PathBuf,
|
|
shards: HashMap<String, EdgeShard>,
|
|
|
|
/// Stores whose unreadability has already been logged. A broken store is hit by every single
|
|
/// request against it, and one log line per request would bury everything else.
|
|
reported_unreadable_stores: HashSet<String>,
|
|
}
|
|
|
|
impl QdrantEdgeDatabase {
|
|
pub fn new(base_path: PathBuf) -> Self {
|
|
Self {
|
|
base_path,
|
|
shards: HashMap::new(),
|
|
reported_unreadable_stores: HashSet::new(),
|
|
}
|
|
}
|
|
|
|
/// Whether this store's defect still has to be written to the log. True exactly once per store,
|
|
/// until the store loads again.
|
|
fn report_unreadable_store(&mut self, store_name: &str) -> bool {
|
|
self.reported_unreadable_stores.insert(store_name.to_string())
|
|
}
|
|
|
|
fn store_path(&self, store_name: &str) -> QdrantEdgeResult<PathBuf> {
|
|
validate_store_name(store_name)?;
|
|
Ok(self.base_path.join("stores").join(store_directory_name(store_name)))
|
|
}
|
|
|
|
// To ensure a shard exists and that you can insert a vector
|
|
fn get_or_create_store(&mut self, store_name: &str, vector_size: usize) -> QdrantEdgeResult<(&EdgeShard, bool)> {
|
|
let (path, is_initialized) = self.reconcile_store_state(store_name)?;
|
|
if self.shards.contains_key(store_name) {
|
|
return Ok((self.shards.get(store_name).unwrap(), false));
|
|
}
|
|
|
|
let shard = if is_initialized {
|
|
compact_wal(&path, store_name);
|
|
match EdgeShard::load(&path, None) {
|
|
Ok(shard) => shard,
|
|
Err(error) => return Err(StoreUnreadableError::new(store_name, &path, error).into()),
|
|
}
|
|
} else {
|
|
fs::create_dir_all(&path).map_err(|error| {
|
|
format!("Failed to create directory for vector store '{store_name}' at '{}': {error}", path.display())
|
|
})?;
|
|
let shard = match EdgeShard::new(&path, edge_config(vector_size)) {
|
|
Ok(shard) => shard,
|
|
Err(error) => {
|
|
let cleanup_issue = remove_partial_store(&path);
|
|
return Err(format!("Failed to create vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into());
|
|
},
|
|
};
|
|
|
|
if let Err(error) = write_store_initialization_marker(&path, store_name) {
|
|
drop(shard);
|
|
let cleanup_issue = remove_partial_store(&path);
|
|
return Err(format!("Failed to finalize vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into());
|
|
}
|
|
|
|
shard
|
|
};
|
|
|
|
self.reported_unreadable_stores.remove(store_name);
|
|
self.shards.insert(store_name.to_string(), shard);
|
|
Ok((self.shards.get(store_name).unwrap(), !is_initialized))
|
|
}
|
|
|
|
// To check whether a shard exists so you can delete a file from it
|
|
fn get_existing_store(&mut self, store_name: &str) -> QdrantEdgeResult<Option<&EdgeShard>> {
|
|
let (path, is_initialized) = self.reconcile_store_state(store_name)?;
|
|
if self.shards.contains_key(store_name) {
|
|
return Ok(self.shards.get(store_name));
|
|
}
|
|
|
|
if !is_initialized {
|
|
return Ok(None);
|
|
}
|
|
|
|
compact_wal(&path, store_name);
|
|
let shard = match EdgeShard::load(&path, None) {
|
|
Ok(shard) => shard,
|
|
Err(error) => return Err(StoreUnreadableError::new(store_name, &path, error).into()),
|
|
};
|
|
|
|
self.reported_unreadable_stores.remove(store_name);
|
|
self.shards.insert(store_name.to_string(), shard);
|
|
Ok(self.shards.get(store_name))
|
|
}
|
|
|
|
fn reconcile_store_state(&mut self, store_name: &str) -> QdrantEdgeResult<(PathBuf, bool)> {
|
|
let path = self.store_path(store_name)?;
|
|
let is_initialized = store_is_initialized(&path, store_name)?;
|
|
|
|
if self.shards.contains_key(store_name) && !is_initialized {
|
|
warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name);
|
|
self.shards.remove(store_name);
|
|
}
|
|
|
|
if path.exists() && !is_initialized {
|
|
warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before continuing.", store_name);
|
|
fs::remove_dir_all(&path).map_err(|error| {
|
|
format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display())
|
|
})?;
|
|
}
|
|
|
|
Ok((path, is_initialized))
|
|
}
|
|
|
|
fn info(&self) -> QdrantEdgeResult<QdrantEdgeInfo> {
|
|
let stores_path = self.base_path.join("stores");
|
|
let stores_count = if stores_path.exists() {
|
|
fs::read_dir(stores_path)?
|
|
.filter_map(Result::ok)
|
|
.filter(|entry| entry.path().join(STORE_INITIALIZATION_MARKER).is_file())
|
|
.count()
|
|
} else {
|
|
0
|
|
};
|
|
|
|
Ok(QdrantEdgeInfo {
|
|
name: "Qdrant Edge".to_string(),
|
|
version: vector_store_version()?,
|
|
path: self.base_path.to_string_lossy().to_string(),
|
|
stores_count,
|
|
})
|
|
}
|
|
|
|
fn ensure_store_exists(&mut self, store_name: &str, data_source_name: &str, vector_size: usize) -> QdrantEdgeResult<QdrantEdgeEnsureStoreResult> {
|
|
validate_vector_size(vector_size)?;
|
|
validate_data_source_name(data_source_name)?;
|
|
let store_path = self.store_path(store_name)?;
|
|
let (_, created) = self.get_or_create_store(store_name, vector_size)?;
|
|
write_store_display_name(&store_path, data_source_name)?;
|
|
Ok(QdrantEdgeEnsureStoreResult {
|
|
created,
|
|
})
|
|
}
|
|
|
|
fn insert_embedding(&mut self, store_name: &str, points: Vec<QdrantEdgeStoragePoint>) -> QdrantEdgeResult<()> {
|
|
let Some(first_point) = points.first() else {
|
|
return Ok(());
|
|
};
|
|
|
|
let vector_size = first_point.vector.len();
|
|
validate_vector_size(vector_size)?;
|
|
if points.iter().any(|point| point.vector.len() != vector_size) {
|
|
return Err("All vectors in one insert request must have the same size.".into());
|
|
}
|
|
|
|
let (shard, _) = self.get_or_create_store(store_name, vector_size)?;
|
|
let points = points
|
|
.into_iter()
|
|
.map(to_qdrant_edge_point)
|
|
.collect::<QdrantEdgeResult<Vec<_>>>()?;
|
|
|
|
shard.update(UpdateOperation::PointOperation(
|
|
PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)),
|
|
))?;
|
|
shard.flush()?;
|
|
Ok(())
|
|
}
|
|
|
|
fn search_embedding(&mut self, store_name: &str, vector: Vec<f32>, max_matches: usize, point_ids: Option<Vec<String>>) -> QdrantEdgeResult<Vec<QdrantEdgeSearchResult>> {
|
|
if max_matches == 0 {
|
|
return Ok(vec![]);
|
|
}
|
|
|
|
let filter = match point_ids {
|
|
None => None,
|
|
Some(point_ids) if point_ids.is_empty() => return Ok(vec![]),
|
|
Some(point_ids) => Some(point_id_filter(&point_ids)?),
|
|
};
|
|
|
|
validate_vector_size(vector.len())?;
|
|
let Some(shard) = self.get_existing_store(store_name)? else {
|
|
return Ok(vec![]);
|
|
};
|
|
|
|
let search_results = shard.query(QueryRequest {
|
|
prefetches: Vec::new(),
|
|
query: Some(ScoringQuery::Vector(QueryEnum::Nearest(NamedQuery::new(
|
|
VectorInternal::Dense(vector),
|
|
VECTOR_NAME,
|
|
)))),
|
|
filter,
|
|
score_threshold: None,
|
|
limit: max_matches,
|
|
offset: 0,
|
|
params: None,
|
|
with_vector: WithVector::Bool(false),
|
|
with_payload: WithPayloadInterface::Bool(true),
|
|
})?;
|
|
|
|
Ok(search_results
|
|
.into_iter()
|
|
.map(to_qdrant_edge_search_result)
|
|
.collect())
|
|
}
|
|
|
|
fn delete_embedding_by_file(&mut self, store_name: &str, file_path: &str) -> QdrantEdgeResult<()> {
|
|
let Some(shard) = self.get_existing_store(store_name)? else {
|
|
return Ok(());
|
|
};
|
|
|
|
shard.update(UpdateOperation::PointOperation(
|
|
PointOperations::DeletePointsByFilter(match_keyword_filter("file_path", file_path)?),
|
|
))?;
|
|
shard.flush()?;
|
|
Ok(())
|
|
}
|
|
|
|
fn optimize_store(&mut self, store_name: &str) -> QdrantEdgeResult<()> {
|
|
let Some(shard) = self.get_existing_store(store_name)? else {
|
|
return Ok(());
|
|
};
|
|
|
|
let optimized = shard.optimize()?;
|
|
if optimized {
|
|
info!(Source = "Qdrant Edge"; "Optimized vector store '{}'.", store_name);
|
|
}
|
|
shard.flush()?;
|
|
|
|
// Temporary compatibility shim, see compact_wal: without this, the WAL of a long session
|
|
// keeps growing until the next start. The store has to be unloaded first, because Windows
|
|
// cannot delete memory-mapped files. The next request loads it again.
|
|
let path = self.store_path(store_name)?;
|
|
if has_obsolete_wal_segments(&path) {
|
|
drop(self.shards.remove(store_name));
|
|
compact_wal(&path, store_name);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn delete_store(&mut self, store_name: &str) -> QdrantEdgeResult<()> {
|
|
self.shards.remove(store_name);
|
|
|
|
let path = self.store_path(store_name)?;
|
|
if path.exists() {
|
|
fs::remove_dir_all(path)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn base_path(&self) -> PathBuf {
|
|
self.base_path.clone()
|
|
}
|
|
}
|
|
|
|
fn store_directory_name(store_name: &str) -> String {
|
|
let stable_id = store_name.strip_prefix("rag_").unwrap_or(store_name);
|
|
format!("store_{stable_id}")
|
|
}
|
|
|
|
fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
|
|
let data_directory = DATA_DIRECTORY
|
|
.get()
|
|
.ok_or("The data directory has not been initialized.")?;
|
|
|
|
Ok(Path::new(data_directory)
|
|
.join("databases")
|
|
.join("vector_database"))
|
|
}
|
|
|
|
pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
|
let status = QDRANT_EDGE_STATUS.lock().unwrap();
|
|
let current_status = status.status;
|
|
let unavailable_reason = status.unavailable_reason.clone();
|
|
drop(status);
|
|
|
|
let database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
|
|
let database_info = database_guard
|
|
.as_ref()
|
|
.and_then(|database| database.info().ok());
|
|
|
|
let is_available = current_status == QdrantEdgeStatus::Available && database_info.is_some();
|
|
Json(QdrantEdgeServiceInfo {
|
|
status: current_status,
|
|
name: database_info.as_ref().map(|info| info.name.clone()).unwrap_or_default(),
|
|
version: database_info.as_ref().map(|info| info.version.clone()).unwrap_or_default(),
|
|
path: database_info.as_ref().map(|info| info.path.clone()).unwrap_or_default(),
|
|
stores_count: database_info.as_ref().map(|info| info.stores_count).unwrap_or_default(),
|
|
is_available,
|
|
unavailable_reason,
|
|
})
|
|
}
|
|
|
|
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<QdrantEdgeEnsureStoreResult>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
|
|
})
|
|
}
|
|
|
|
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<()>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.insert_embedding(&request.store_name, request.points)
|
|
})
|
|
}
|
|
|
|
pub async fn search_qdrant_edge_embeddings(_token: APIToken, Json(request): Json<SearchQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<Vec<QdrantEdgeSearchResult>>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.search_embedding(&request.store_name, request.vector, request.max_matches, request.point_ids)
|
|
})
|
|
}
|
|
|
|
pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeResponse<()>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.delete_embedding_by_file(&request.store_name, &request.file_path)
|
|
})
|
|
}
|
|
|
|
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.optimize_store(&request.store_name)
|
|
})
|
|
}
|
|
|
|
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
|
execute_qdrant_edge_request(|database| {
|
|
database.delete_store(&request.store_name)
|
|
})
|
|
}
|
|
|
|
pub fn start_qdrant_edge_database<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {
|
|
set_qdrant_edge_starting();
|
|
remove_obsolete_qdrant_sidecar_files(&app_handle);
|
|
|
|
let path = match qdrant_edge_base_path() {
|
|
Ok(path) => path,
|
|
Err(e) => {
|
|
let reason = format!("Qdrant Edge cannot be started: {e}");
|
|
error!(Source = "Qdrant Edge"; "{reason}");
|
|
set_qdrant_edge_unavailable(reason);
|
|
return;
|
|
},
|
|
};
|
|
|
|
match fs::create_dir_all(&path) {
|
|
Ok(_) => {
|
|
let database = QdrantEdgeDatabase::new(path.clone());
|
|
*QDRANT_EDGE_DATABASE.lock().unwrap() = Some(database);
|
|
set_qdrant_edge_available();
|
|
info!(Source = "Qdrant Edge"; "Qdrant Edge is available at '{}'.", path.display());
|
|
},
|
|
Err(e) => {
|
|
let reason = format!("The Qdrant Edge data directory could not be created: {e}");
|
|
error!(Source = "Qdrant Edge"; "{reason}");
|
|
set_qdrant_edge_unavailable(reason);
|
|
},
|
|
}
|
|
}
|
|
|
|
pub fn stop_qdrant_edge_database() {
|
|
if let Some(database) = QDRANT_EDGE_DATABASE.lock().unwrap().take() {
|
|
info!(Source = "Qdrant Edge"; "Stopping Qdrant Edge at '{}'.", database.base_path().display());
|
|
drop(database);
|
|
}
|
|
|
|
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
|
|
}
|
|
|
|
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
|
where
|
|
T: Serialize,
|
|
F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<T>,
|
|
{
|
|
let mut database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
|
|
let Some(database) = database_guard.as_mut() else {
|
|
return Json(QdrantEdgeResponse {
|
|
success: false,
|
|
issue: "Qdrant Edge is not available.".to_string(),
|
|
issue_code: "",
|
|
data: None,
|
|
});
|
|
};
|
|
|
|
match operation(database) {
|
|
Ok(data) => Json(QdrantEdgeResponse {
|
|
success: true,
|
|
issue: String::new(),
|
|
issue_code: "",
|
|
data: Some(data),
|
|
}),
|
|
Err(e) => {
|
|
let issue = e.to_string();
|
|
|
|
//
|
|
// An unreadable store keeps failing for as long as the user leaves it alone, so it is
|
|
// logged once and then only answered. Every other failure is logged as it happens,
|
|
// because those are one-offs worth seeing each time.
|
|
//
|
|
let issue_code = match e.downcast_ref::<StoreUnreadableError>() {
|
|
Some(unreadable) => {
|
|
if database.report_unreadable_store(&unreadable.store_name) {
|
|
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
|
|
}
|
|
|
|
ISSUE_CODE_STORE_UNREADABLE
|
|
},
|
|
|
|
None => {
|
|
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
|
|
""
|
|
},
|
|
};
|
|
|
|
Json(QdrantEdgeResponse {
|
|
success: false,
|
|
issue,
|
|
issue_code,
|
|
data: None,
|
|
})
|
|
},
|
|
}
|
|
}
|
|
|
|
fn set_qdrant_edge_available() {
|
|
let mut status = QDRANT_EDGE_STATUS.lock().unwrap();
|
|
status.status = QdrantEdgeStatus::Available;
|
|
status.unavailable_reason = None;
|
|
}
|
|
|
|
fn set_qdrant_edge_starting() {
|
|
let mut status = QDRANT_EDGE_STATUS.lock().unwrap();
|
|
status.status = QdrantEdgeStatus::Starting;
|
|
status.unavailable_reason = None;
|
|
}
|
|
|
|
fn set_qdrant_edge_unavailable(reason: String) {
|
|
let mut status = QDRANT_EDGE_STATUS.lock().unwrap();
|
|
status.status = QdrantEdgeStatus::Unavailable;
|
|
status.unavailable_reason = Some(reason);
|
|
}
|
|
|
|
// Temporary compatibility shim until 2026-12-02:
|
|
// documentation/compatibility-shims/2026-06-qdrant-edge-migration.md
|
|
fn remove_obsolete_qdrant_sidecar_files<R: tauri::Runtime>(app_handle: &tauri::AppHandle<R>) {
|
|
let mut paths = Vec::new();
|
|
|
|
if let Some(data_directory) = DATA_DIRECTORY.get() {
|
|
let databases_directory = Path::new(data_directory).join("databases");
|
|
paths.push(databases_directory.join("qdrant"));
|
|
paths.push(databases_directory.join("qdrant_test"));
|
|
}
|
|
|
|
if let Ok(resource_dir) = app_handle.path().resource_dir() {
|
|
paths.push(resource_dir.join("target").join("databases").join("qdrant"));
|
|
paths.push(resource_dir.join("resources").join("databases").join("qdrant"));
|
|
}
|
|
|
|
cfg_if::cfg_if! {
|
|
if #[cfg(any(target_os = "windows", target_os = "macos"))]{
|
|
if let Ok(current_exe) = std::env::current_exe() && let Some(exe_dir) = current_exe.parent()
|
|
&& exe_dir.to_string_lossy().contains("MindWork AI Studio") {
|
|
paths.push(exe_dir.join("target").join("databases").join("qdrant"));
|
|
paths.push(exe_dir.join("qdrant.exe"));
|
|
paths.push(exe_dir.join("qdrant"));
|
|
}
|
|
}
|
|
}
|
|
|
|
for path in paths {
|
|
remove_obsolete_qdrant_path(&path);
|
|
}
|
|
}
|
|
|
|
fn remove_obsolete_qdrant_path(path: &Path) {
|
|
if !path.exists() {
|
|
info!(Source = "Qdrant Edge"; "Obsolete file or directory '{}' was not found.", path.display());
|
|
return;
|
|
}
|
|
|
|
let result = if path.is_dir() {
|
|
fs::remove_dir_all(path)
|
|
} else {
|
|
fs::remove_file(path)
|
|
};
|
|
|
|
match result {
|
|
Ok(_) => warn!(Source = "Qdrant Edge"; "Removed obsolete Qdrant sidecar file or directory '{}'.", path.display()),
|
|
Err(e) => warn!(Source = "Qdrant Edge"; "Could not remove obsolete Qdrant sidecar file or directory '{}': {e}", path.display()),
|
|
}
|
|
}
|
|
|
|
fn edge_config(vector_size: usize) -> EdgeConfig {
|
|
EdgeConfig {
|
|
on_disk_payload: Some(true),
|
|
vectors: HashMap::from([(
|
|
VECTOR_NAME.to_string(),
|
|
EdgeVectorParams {
|
|
size: vector_size,
|
|
distance: Distance::Cosine,
|
|
on_disk: Some(true),
|
|
quantization_config: None,
|
|
multivector_config: None,
|
|
datatype: None,
|
|
hnsw_config: Some(hnsw_config()),
|
|
},
|
|
)]),
|
|
sparse_vectors: HashMap::new(),
|
|
hnsw_config: Some(hnsw_config()),
|
|
quantization_config: None,
|
|
optimizers: Some(edge_optimizers_config()),
|
|
wal_options: None,
|
|
max_search_threads: None,
|
|
search_pool_core: None,
|
|
}
|
|
}
|
|
|
|
// `on_disk` is deprecated in favor of `memory`, but Qdrant Edge does not re-export the `Memory`
|
|
// type, so the new field cannot be named from here. Leaving both unset is not an option either:
|
|
// the effective placement would fall back to cached instead of on-disk, which is a real change
|
|
// and would have the optimizers rebuild the HNSW graph.
|
|
#[allow(deprecated)]
|
|
fn hnsw_config() -> HnswIndexConfig {
|
|
HnswIndexConfig {
|
|
m: HNSW_M,
|
|
ef_construct: HNSW_EF_CONSTRUCT,
|
|
full_scan_threshold: HNSW_FULL_SCAN_THRESHOLD_KB,
|
|
max_indexing_threads: HNSW_MAX_INDEXING_THREADS,
|
|
on_disk: Some(true),
|
|
memory: None,
|
|
payload_m: None,
|
|
inline_storage: None,
|
|
}
|
|
}
|
|
|
|
fn edge_optimizers_config() -> EdgeOptimizersConfig {
|
|
EdgeOptimizersConfig {
|
|
indexing_threshold: Some(VECTOR_INDEXING_THRESHOLD_KB),
|
|
prevent_unoptimized: Some(false),
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn store_is_initialized(path: &Path, store_name: &str) -> QdrantEdgeResult<bool> {
|
|
if !path.join("edge_config.json").is_file() || !path.join("segments").is_dir() {
|
|
return Ok(false);
|
|
}
|
|
|
|
let marker_path = path.join(STORE_INITIALIZATION_MARKER);
|
|
if !marker_path.exists() {
|
|
return Ok(false);
|
|
}
|
|
|
|
let initialized_store_name = fs::read_to_string(&marker_path).map_err(|error| {
|
|
format!("Failed to read vector store initialization marker '{}': {error}", marker_path.display())
|
|
})?;
|
|
if initialized_store_name != store_name {
|
|
return Err(format!("Vector store path collision at '{}': expected store '{}', but the path belongs to '{}'.", path.display(), store_name, initialized_store_name).into());
|
|
}
|
|
|
|
Ok(true)
|
|
}
|
|
|
|
fn write_store_initialization_marker(path: &Path, store_name: &str) -> std::io::Result<()> {
|
|
write_store_marker(path, STORE_INITIALIZATION_MARKER, STORE_INITIALIZATION_MARKER_TEMP, store_name)
|
|
}
|
|
|
|
fn write_store_display_name(path: &Path, data_source_name: &str) -> std::io::Result<()> {
|
|
write_store_marker(path, STORE_DISPLAY_NAME_MARKER, STORE_DISPLAY_NAME_MARKER_TEMP, data_source_name)
|
|
}
|
|
|
|
fn write_store_marker(path: &Path, marker_name: &str, temporary_marker_name: &str, value: &str) -> std::io::Result<()> {
|
|
let marker_path = path.join(marker_name);
|
|
if fs::read_to_string(&marker_path).is_ok_and(|current_value| current_value == value) {
|
|
return Ok(());
|
|
}
|
|
|
|
let temporary_marker_path = path.join(temporary_marker_name);
|
|
fs::write(&temporary_marker_path, value)?;
|
|
if marker_path.exists() {
|
|
fs::remove_file(&marker_path)?;
|
|
}
|
|
fs::rename(temporary_marker_path, marker_path)
|
|
}
|
|
|
|
fn remove_partial_store(path: &Path) -> String {
|
|
match fs::remove_dir_all(path) {
|
|
Ok(()) => String::new(),
|
|
Err(error) => format!(" The incomplete store could not be removed: {error}"),
|
|
}
|
|
}
|
|
|
|
// Temporary compatibility shim until a qdrant-edge release acknowledges its WAL by itself:
|
|
// documentation/compatibility-shims/2026-10-qdrant-edge-wal-compaction.md
|
|
//
|
|
// Qdrant Edge writes every update to its WAL, but never acknowledges the WAL and never replays it,
|
|
// so the WAL grows with every update. We flush after every update, which means the segments already
|
|
// hold everything the WAL contains. The store must not be loaded while this runs.
|
|
fn compact_wal(store_path: &Path, store_name: &str) {
|
|
match remove_obsolete_wal_segments(&store_path.join(WAL_DIRECTORY)) {
|
|
Ok((0, _)) => {},
|
|
Ok((removed_segments, removed_bytes)) => info!(Source = "Qdrant Edge"; "Removed {removed_segments} obsolete WAL segment(s) with {:.1} MiB from vector store '{store_name}'.", removed_bytes as f64 / (1024.0 * 1024.0)),
|
|
Err(error) => warn!(Source = "Qdrant Edge"; "Could not remove the obsolete WAL segments of vector store '{store_name}': {error}"),
|
|
}
|
|
}
|
|
|
|
/// Whether the store's WAL holds segments which `compact_wal` would remove.
|
|
fn has_obsolete_wal_segments(store_path: &Path) -> bool {
|
|
wal_file_names(&store_path.join(WAL_DIRECTORY))
|
|
.is_ok_and(|file_names| !obsolete_wal_segments(&file_names).is_empty())
|
|
}
|
|
|
|
/// Removes the obsolete closed segments, oldest first, and returns their number and size. When a
|
|
/// removal fails, the segments left over are still contiguous, so Qdrant Edge can open the WAL.
|
|
fn remove_obsolete_wal_segments(wal_path: &Path) -> std::io::Result<(usize, u64)> {
|
|
let mut removed_segments = 0;
|
|
let mut removed_bytes = 0;
|
|
for file_name in obsolete_wal_segments(&wal_file_names(wal_path)?) {
|
|
let segment_path = wal_path.join(file_name);
|
|
let segment_size = fs::metadata(&segment_path)?.len();
|
|
fs::remove_file(&segment_path)?;
|
|
removed_segments += 1;
|
|
removed_bytes += segment_size;
|
|
}
|
|
|
|
Ok((removed_segments, removed_bytes))
|
|
}
|
|
|
|
fn wal_file_names(wal_path: &Path) -> std::io::Result<Vec<String>> {
|
|
if !wal_path.is_dir() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
// Every file counts: a closed segment missing from this list would leave a gap behind.
|
|
let mut file_names = Vec::new();
|
|
for entry in fs::read_dir(wal_path)? {
|
|
let entry = entry?;
|
|
if entry.file_type()?.is_file() && let Ok(file_name) = entry.file_name().into_string() {
|
|
file_names.push(file_name);
|
|
}
|
|
}
|
|
|
|
Ok(file_names)
|
|
}
|
|
|
|
/// The closed WAL segments which may go, oldest first. That is all of them except the newest:
|
|
/// Qdrant Edge numbers its operations from the first closed segment and would start again at zero
|
|
/// without one. The segments then ignore every later update of a point they already know.
|
|
fn obsolete_wal_segments(file_names: &[String]) -> Vec<String> {
|
|
let mut closed_segments = file_names
|
|
.iter()
|
|
.filter_map(|file_name| {
|
|
let start_index = file_name.strip_prefix(CLOSED_WAL_SEGMENT_PREFIX)?.parse::<u64>().ok()?;
|
|
Some((start_index, file_name.clone()))
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
closed_segments.sort_unstable_by_key(|(start_index, _)| *start_index);
|
|
closed_segments.pop();
|
|
closed_segments
|
|
.into_iter()
|
|
.map(|(_, file_name)| file_name)
|
|
.collect()
|
|
}
|
|
|
|
fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
|
|
if vector_size == 0 {
|
|
return Err("Vector size must be greater than zero.".into());
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_data_source_name(data_source_name: &str) -> QdrantEdgeResult<()> {
|
|
const MAX_DATA_SOURCE_NAME_LENGTH: usize = 40;
|
|
|
|
if data_source_name.trim().is_empty() {
|
|
return Err("Data source name cannot be empty.".into());
|
|
}
|
|
|
|
if data_source_name.chars().count() > MAX_DATA_SOURCE_NAME_LENGTH {
|
|
return Err(format!("Data source name exceeds the maximum length of {MAX_DATA_SOURCE_NAME_LENGTH} characters.").into());
|
|
}
|
|
|
|
if data_source_name.chars().any(|c| c.is_control()) {
|
|
return Err("Data source name contains unsupported control characters.".into());
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn vector_store_version() -> QdrantEdgeResult<String> {
|
|
let metadata = META_DATA
|
|
.lock()
|
|
.map_err(|_| "Metadata lock was poisoned.")?;
|
|
let Some(metadata) = metadata.as_ref() else {
|
|
return Err("Metadata was not initialized.".into());
|
|
};
|
|
|
|
Ok(metadata.vector_store_version.clone())
|
|
}
|
|
|
|
fn to_qdrant_edge_point(point: QdrantEdgeStoragePoint) -> QdrantEdgeResult<qdrant_edge::PointStructPersisted> {
|
|
Ok(PointStruct::new(
|
|
to_point_id(&point.point_id)?,
|
|
Vectors::new_named([(VECTOR_NAME, point.vector)]),
|
|
json!({
|
|
"data_source_id": point.data_source_id,
|
|
"data_source_type": point.data_source_type,
|
|
"chunk_id": point.chunk_id,
|
|
"parent_file_id": point.parent_file_id,
|
|
"file_path": point.file_path,
|
|
"absolute_path": point.absolute_path,
|
|
"file_name": point.file_name,
|
|
"relative_path": point.relative_path,
|
|
"file_type": point.file_type,
|
|
"page_number": point.page_number,
|
|
"chunk_index": point.chunk_index,
|
|
"text": point.text,
|
|
"fingerprint": point.fingerprint,
|
|
"creation_utc": point.creation_utc,
|
|
"last_write_utc": point.last_write_utc,
|
|
"embedded_at_utc": point.embedded_at_utc,
|
|
}),
|
|
)
|
|
.into())
|
|
}
|
|
|
|
fn to_qdrant_edge_search_result(point: ScoredPoint) -> QdrantEdgeSearchResult {
|
|
let payload = point.payload.unwrap_or_default();
|
|
QdrantEdgeSearchResult {
|
|
point_id: point_id_to_string(point.id),
|
|
score: point.score,
|
|
data_source_id: payload_string(&payload, "data_source_id"),
|
|
data_source_type: payload_string(&payload, "data_source_type"),
|
|
chunk_id: payload_string(&payload, "chunk_id"),
|
|
parent_file_id: payload_string(&payload, "parent_file_id"),
|
|
file_path: payload_string(&payload, "file_path"),
|
|
absolute_path: payload_string(&payload, "absolute_path"),
|
|
file_name: payload_string(&payload, "file_name"),
|
|
relative_path: payload_string(&payload, "relative_path"),
|
|
file_type: payload_string(&payload, "file_type"),
|
|
page_number: payload_i32(&payload, "page_number"),
|
|
chunk_index: payload_i32(&payload, "chunk_index").unwrap_or_default(),
|
|
text: payload_string(&payload, "text"),
|
|
fingerprint: payload_string(&payload, "fingerprint"),
|
|
creation_utc: payload_string(&payload, "creation_utc"),
|
|
last_write_utc: payload_string(&payload, "last_write_utc"),
|
|
embedded_at_utc: payload_string(&payload, "embedded_at_utc"),
|
|
}
|
|
}
|
|
|
|
fn to_point_id(point_id: &str) -> QdrantEdgeResult<PointId> {
|
|
Uuid::parse_str(point_id)
|
|
.map(PointId::Uuid)
|
|
.map_err(|_| "Vector point ID must be a valid UUID.".into())
|
|
}
|
|
|
|
fn point_id_to_string(point_id: PointId) -> String {
|
|
match point_id {
|
|
PointId::NumId(id) => id.to_string(),
|
|
PointId::Uuid(uuid) => uuid.to_string(),
|
|
}
|
|
}
|
|
|
|
fn payload_string(payload: &Payload, key: &str) -> String {
|
|
payload
|
|
.0
|
|
.get(key)
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default()
|
|
.to_string()
|
|
}
|
|
|
|
fn payload_i32(payload: &Payload, key: &str) -> Option<i32> {
|
|
payload
|
|
.0
|
|
.get(key)
|
|
.and_then(Value::as_i64)
|
|
.and_then(|value| i32::try_from(value).ok())
|
|
}
|
|
|
|
fn match_keyword_filter(field_name: &str, value: &str) -> QdrantEdgeResult<Filter> {
|
|
Ok(Filter {
|
|
should: None,
|
|
min_should: None,
|
|
must: Some(vec![Condition::Field(FieldCondition::new_match(
|
|
field_name
|
|
.try_into()
|
|
.map_err(|_| format!("Invalid payload field name '{field_name}'."))?,
|
|
Match::Value(MatchValue {
|
|
value: ValueVariants::String(value.to_string()),
|
|
}),
|
|
))]),
|
|
must_not: None,
|
|
})
|
|
}
|
|
|
|
/// Restricts a search to the given points. Every id has to be a valid UUID, the same as when the
|
|
/// points were inserted, so a malformed id is an error rather than a point which is silently missed.
|
|
fn point_id_filter(point_ids: &[String]) -> QdrantEdgeResult<Filter> {
|
|
let has_id = point_ids
|
|
.iter()
|
|
.map(|point_id| to_point_id(point_id))
|
|
.collect::<QdrantEdgeResult<HasIdCondition>>()?;
|
|
|
|
Ok(Filter::new_must(Condition::HasId(has_id)))
|
|
}
|
|
|
|
fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
|
|
const MAX_STORE_NAME_LENGTH: usize = 128;
|
|
|
|
if store_name.is_empty() {
|
|
return Err("Vector store name cannot be empty.".into());
|
|
}
|
|
|
|
if store_name.len() > MAX_STORE_NAME_LENGTH {
|
|
return Err(format!("Vector store name exceeds the maximum length of {MAX_STORE_NAME_LENGTH} bytes.").into());
|
|
}
|
|
|
|
if store_name
|
|
.chars()
|
|
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
|
{
|
|
return Ok(());
|
|
}
|
|
|
|
Err(format!("Vector store name '{store_name}' contains unsupported characters.").into())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use qdrant_edge::WalOptions;
|
|
|
|
#[test]
|
|
fn validate_store_name_allows_safe_store_names() {
|
|
assert!(validate_store_name("rag_1234-abcd").is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn validate_store_name_rejects_path_syntax() {
|
|
assert!(validate_store_name(".").is_err());
|
|
assert!(validate_store_name("..").is_err());
|
|
assert!(validate_store_name("../store").is_err());
|
|
assert!(validate_store_name("store\\name").is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn validate_store_name_rejects_oversized_names() {
|
|
assert!(validate_store_name(&"a".repeat(129)).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn store_directory_name_contains_the_stable_data_source_id() {
|
|
assert_eq!(
|
|
store_directory_name("rag_6cc665a82b1e4d42bc748015b7b391ec"),
|
|
"store_6cc665a82b1e4d42bc748015b7b391ec"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn validate_data_source_name_allows_display_names_but_rejects_invalid_values() {
|
|
assert!(validate_data_source_name("Mäßig Confidence C#").is_ok());
|
|
assert!(validate_data_source_name(" ").is_err());
|
|
assert!(validate_data_source_name("invalid\nname").is_err());
|
|
assert!(validate_data_source_name(&"a".repeat(41)).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn ensure_store_reports_creation_and_updates_the_display_name() {
|
|
let test_directory = test_directory("ensure");
|
|
let store_name = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.clone());
|
|
|
|
let created = database.ensure_store_exists(store_name, "Original name", 3).unwrap();
|
|
assert!(created.created);
|
|
|
|
let existing = database.ensure_store_exists(store_name, "Renamed source", 3).unwrap();
|
|
assert!(!existing.created);
|
|
let display_name_path = database.store_path(store_name).unwrap().join(STORE_DISPLAY_NAME_MARKER);
|
|
assert_eq!(fs::read_to_string(display_name_path).unwrap(), "Renamed source");
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn an_unreadable_store_is_reported_but_never_deleted() {
|
|
let test_directory = test_directory("unreadable");
|
|
let store_name = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
|
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.clone());
|
|
assert!(database.ensure_store_exists(store_name, "Some source", 3).unwrap().created);
|
|
let store_path = database.store_path(store_name).unwrap();
|
|
|
|
// Release the shard before breaking it, so the files are not held open any more.
|
|
drop(database);
|
|
fs::write(store_path.join("edge_config.json"), "this is not a config").unwrap();
|
|
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.clone());
|
|
let error = database.get_existing_store(store_name).unwrap_err();
|
|
assert!(
|
|
error.downcast_ref::<StoreUnreadableError>().is_some(),
|
|
"a store which cannot be opened has to be recognizable as such, not just a message"
|
|
);
|
|
|
|
// The whole point: the user's embeddings survive a defect until they ask for a rebuild.
|
|
assert!(store_path.join("segments").is_dir());
|
|
assert!(store_path.join(STORE_INITIALIZATION_MARKER).is_file());
|
|
|
|
// And the defect is logged once, not once per request.
|
|
assert!(database.report_unreadable_store(store_name));
|
|
assert!(!database.report_unreadable_store(store_name));
|
|
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn point_ids_must_be_valid_uuids() {
|
|
assert!(to_point_id("6cc665a8-2b1e-4d42-bc74-8015b7b391ec").is_ok());
|
|
assert!(to_point_id("deliberate-collision-input").is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn a_search_restricted_to_no_points_finds_nothing() {
|
|
let test_directory = test_directory("no-points");
|
|
let mut database = store_with_one_point_per_axis(&test_directory);
|
|
|
|
let found = database.search_embedding(SEARCH_STORE, search_vector(), 10, Some(vec![])).unwrap();
|
|
|
|
// An empty restriction comes from conditions which matched nothing. Searching the whole
|
|
// store instead would answer with exactly what those conditions ruled out.
|
|
assert!(found.is_empty());
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn a_search_restricted_to_some_points_finds_only_those() {
|
|
let test_directory = test_directory("some-points");
|
|
let mut database = store_with_one_point_per_axis(&test_directory);
|
|
|
|
let unrestricted = database.search_embedding(SEARCH_STORE, search_vector(), 10, None).unwrap();
|
|
assert_eq!(found_point_ids(&unrestricted), vec![POINT_X, POINT_Y, POINT_Z]);
|
|
|
|
let restricted = database
|
|
.search_embedding(SEARCH_STORE, search_vector(), 10, Some(vec![POINT_Y.to_string(), POINT_Z.to_string()]))
|
|
.unwrap();
|
|
assert_eq!(
|
|
found_point_ids(&restricted),
|
|
vec![POINT_Y, POINT_Z],
|
|
"the closest point lies outside the restriction, so it must not be found"
|
|
);
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn a_restriction_with_a_malformed_point_id_is_rejected() {
|
|
assert!(point_id_filter(&[POINT_X.to_string()]).is_ok());
|
|
assert!(
|
|
point_id_filter(&[POINT_X.to_string(), "not-a-point-id".to_string()]).is_err(),
|
|
"skipping the malformed id would quietly search fewer points than asked for"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn obsolete_wal_segments_keeps_the_newest_closed_segment() {
|
|
let file_names = ["closed-120", "open-7", "closed-0", "tmp-open-8", "closed-40", ".wal", "first-index", "closed-abc"].map(String::from);
|
|
|
|
// Sorted by number, not by name: as text, "closed-40" would come after "closed-120".
|
|
assert_eq!(obsolete_wal_segments(&file_names), vec!["closed-0", "closed-40"]);
|
|
|
|
assert!(obsolete_wal_segments(&["closed-0", "open-1"].map(String::from)).is_empty());
|
|
assert!(obsolete_wal_segments(&["open-1"].map(String::from)).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn a_reloaded_store_drops_obsolete_wal_segments_and_keeps_its_points() {
|
|
let test_directory = test_directory("wal-reload");
|
|
let (database, point_ids) = store_with_several_closed_wal_segments(&test_directory);
|
|
let store_path = database.store_path(SEARCH_STORE).unwrap();
|
|
drop(database);
|
|
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.clone());
|
|
let found = database.search_embedding(SEARCH_STORE, search_vector(), point_ids.len(), None).unwrap();
|
|
|
|
assert_eq!(closed_wal_segments(&store_path).len(), 1);
|
|
assert_eq!(
|
|
found_point_ids(&found).into_iter().collect::<HashSet<_>>(),
|
|
point_ids.iter().map(String::as_str).collect::<HashSet<_>>(),
|
|
"the points live in the segments, so removing WAL segments must not lose any of them"
|
|
);
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn an_update_after_wal_compaction_replaces_the_existing_point() {
|
|
let test_directory = test_directory("wal-update");
|
|
let (database, point_ids) = store_with_several_closed_wal_segments(&test_directory);
|
|
drop(database);
|
|
|
|
// The newest point carries the highest operation number so far. Should the numbering start
|
|
// again after the compaction, the segment would ignore this update as outdated.
|
|
let newest_point_id = point_ids.last().unwrap();
|
|
let mut updated_point = test_point(newest_point_id, vec![0.0, 0.0, 1.0]);
|
|
updated_point.text = "Updated text.".to_string();
|
|
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.clone());
|
|
database.insert_embedding(SEARCH_STORE, vec![updated_point]).unwrap();
|
|
|
|
let found = database
|
|
.search_embedding(SEARCH_STORE, search_vector(), 1, Some(vec![newest_point_id.clone()]))
|
|
.unwrap();
|
|
assert_eq!(found.len(), 1);
|
|
assert_eq!(found[0].text, "Updated text.");
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn optimize_removes_obsolete_wal_segments_of_a_loaded_store() {
|
|
let test_directory = test_directory("wal-optimize");
|
|
let (mut database, point_ids) = store_with_several_closed_wal_segments(&test_directory);
|
|
let store_path = database.store_path(SEARCH_STORE).unwrap();
|
|
|
|
database.optimize_store(SEARCH_STORE).unwrap();
|
|
assert_eq!(closed_wal_segments(&store_path).len(), 1);
|
|
|
|
// The next request loads the store again, with all its points and still taking updates.
|
|
let newest_point_id = point_ids.last().unwrap();
|
|
let mut updated_point = test_point(newest_point_id, vec![0.0, 0.0, 1.0]);
|
|
updated_point.text = "Updated text.".to_string();
|
|
database.insert_embedding(SEARCH_STORE, vec![updated_point]).unwrap();
|
|
|
|
let found = database.search_embedding(SEARCH_STORE, search_vector(), point_ids.len(), None).unwrap();
|
|
assert_eq!(
|
|
found_point_ids(&found).into_iter().collect::<HashSet<_>>(),
|
|
point_ids.iter().map(String::as_str).collect::<HashSet<_>>()
|
|
);
|
|
|
|
let updated = found.iter().find(|result| &result.point_id == newest_point_id).unwrap();
|
|
assert_eq!(updated.text, "Updated text.");
|
|
|
|
drop(database);
|
|
fs::remove_dir_all(test_directory).unwrap();
|
|
}
|
|
|
|
const SEARCH_STORE: &str = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
|
const POINT_X: &str = "0b5f1e8a-3c2d-4e6f-9a1b-7c8d9e0f1a2b";
|
|
const POINT_Y: &str = "1c6a2f9b-4d3e-4f70-8b2c-8d9e0f1a2b3c";
|
|
const POINT_Z: &str = "2d7b3a0c-5e4f-4a81-9c3d-9e0f1a2b3c4d";
|
|
|
|
/// A directory of its own for one test, so tests running in parallel never share a store.
|
|
fn test_directory(label: &str) -> PathBuf {
|
|
std::env::temp_dir().join(format!(
|
|
"ai-studio-qdrant-{label}-{}-{}",
|
|
std::process::id(),
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
))
|
|
}
|
|
|
|
/// A store with one point on each axis.
|
|
fn store_with_one_point_per_axis(test_directory: &Path) -> QdrantEdgeDatabase {
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.to_path_buf());
|
|
database.ensure_store_exists(SEARCH_STORE, "Some source", 3).unwrap();
|
|
database
|
|
.insert_embedding(SEARCH_STORE, vec![
|
|
test_point(POINT_X, vec![1.0, 0.0, 0.0]),
|
|
test_point(POINT_Y, vec![0.0, 1.0, 0.0]),
|
|
test_point(POINT_Z, vec![0.0, 0.0, 1.0]),
|
|
])
|
|
.unwrap();
|
|
|
|
database
|
|
}
|
|
|
|
/// A store with tiny WAL segments, filled one point at a time until the WAL has closed several
|
|
/// of them. Returns the database with the store still loaded, and the IDs in insertion order.
|
|
fn store_with_several_closed_wal_segments(test_directory: &Path) -> (QdrantEdgeDatabase, Vec<String>) {
|
|
const SMALL_WAL_SEGMENT_CAPACITY: usize = 8 * 1024;
|
|
const MAX_POINTS: usize = 1_000;
|
|
|
|
let mut database = QdrantEdgeDatabase::new(test_directory.to_path_buf());
|
|
let store_path = database.store_path(SEARCH_STORE).unwrap();
|
|
fs::create_dir_all(&store_path).unwrap();
|
|
|
|
// The production config, except for the WAL. Qdrant Edge keeps these options in the store,
|
|
// so they stay in effect whenever the store gets loaded again.
|
|
let config = EdgeConfig {
|
|
wal_options: Some(WalOptions {
|
|
segment_capacity: SMALL_WAL_SEGMENT_CAPACITY,
|
|
..WalOptions::default()
|
|
}),
|
|
..edge_config(3)
|
|
};
|
|
drop(EdgeShard::new(&store_path, config).unwrap());
|
|
write_store_initialization_marker(&store_path, SEARCH_STORE).unwrap();
|
|
|
|
let mut point_ids = Vec::new();
|
|
while closed_wal_segments(&store_path).len() < 3 {
|
|
assert!(point_ids.len() < MAX_POINTS, "the WAL never closed enough segments");
|
|
|
|
let point_id = Uuid::from_u128(point_ids.len() as u128 + 1).to_string();
|
|
database
|
|
.insert_embedding(SEARCH_STORE, vec![test_point(&point_id, vec![1.0, point_ids.len() as f32, 1.0])])
|
|
.unwrap();
|
|
point_ids.push(point_id);
|
|
}
|
|
|
|
(database, point_ids)
|
|
}
|
|
|
|
fn closed_wal_segments(store_path: &Path) -> Vec<String> {
|
|
fs::read_dir(store_path.join(WAL_DIRECTORY))
|
|
.unwrap()
|
|
.map(|entry| entry.unwrap().file_name().into_string().unwrap())
|
|
.filter(|file_name| file_name.starts_with(CLOSED_WAL_SEGMENT_PREFIX))
|
|
.collect()
|
|
}
|
|
|
|
/// Closest to the point on the x axis, then the one on the y axis, then the one on the z axis.
|
|
fn search_vector() -> Vec<f32> {
|
|
vec![1.0, 0.5, 0.0]
|
|
}
|
|
|
|
fn found_point_ids(found: &[QdrantEdgeSearchResult]) -> Vec<&str> {
|
|
found.iter().map(|result| result.point_id.as_str()).collect()
|
|
}
|
|
|
|
fn test_point(point_id: &str, vector: Vec<f32>) -> QdrantEdgeStoragePoint {
|
|
QdrantEdgeStoragePoint {
|
|
point_id: point_id.to_string(),
|
|
vector,
|
|
data_source_id: "6cc665a8-2b1e-4d42-bc74-8015b7b391ec".to_string(),
|
|
data_source_type: "LOCAL_DIRECTORY".to_string(),
|
|
chunk_id: point_id.to_string(),
|
|
parent_file_id: String::new(),
|
|
file_path: "/tmp/test-data/notes.md".to_string(),
|
|
absolute_path: "/tmp/test-data/notes.md".to_string(),
|
|
file_name: "notes.md".to_string(),
|
|
relative_path: "notes.md".to_string(),
|
|
file_type: "md".to_string(),
|
|
page_number: None,
|
|
chunk_index: 0,
|
|
text: "Some text.".to_string(),
|
|
fingerprint: String::new(),
|
|
creation_utc: String::new(),
|
|
last_write_utc: String::new(),
|
|
embedded_at_utc: String::new(),
|
|
}
|
|
}
|
|
}
|