Replace Qdrant with Qdrant Edge (#783)

runtime/src/qdrant_edge_database.rs

@@ -0,0 +1,587 @@
+use std::collections::HashMap;
+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;
+use qdrant_edge::external::uuid::Uuid;
+use qdrant_edge::{
+    Condition, Distance, EdgeConfig, EdgeOptimizersConfig, EdgeShard, EdgeVectorParams,
+    FieldCondition, Filter, HnswIndexConfig, Match, MatchValue, PointId, PointInsertOperations,
+    PointOperations, PointStruct, UpdateOperation, ValueVariants, Vectors,
+};
+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;
+
+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_name: String,
+    pub data_source_type: String,
+    pub file_path: String,
+    pub file_name: String,
+    pub relative_path: String,
+    pub chunk_index: i32,
+    pub text: String,
+    pub fingerprint: String,
+    pub last_write_utc: String,
+    pub embedded_at_utc: String,
+}
+
+#[derive(Deserialize)]
+pub struct EnsureQdrantEdgeStoreRequest {
+    pub store_name: String,
+    pub vector_size: usize,
+}
+
+#[derive(Deserialize)]
+pub struct InsertQdrantEdgeEmbeddingRequest {
+    pub store_name: String,
+    pub points: Vec<QdrantEdgeStoragePoint>,
+}
+
+#[derive(Deserialize)]
+pub struct DeleteQdrantEdgeEmbeddingByFileRequest {
+    pub store_name: String,
+    pub file_path: String,
+}
+
+#[derive(Deserialize)]
+pub struct DeleteQdrantEdgeStoreRequest {
+    pub store_name: String,
+}
+
+#[derive(Serialize)]
+pub struct QdrantEdgeOperationResponse {
+    pub success: bool,
+    pub issue: 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>,
+}
+
+impl QdrantEdgeDatabase {
+    pub fn new(base_path: PathBuf) -> Self {
+        Self {
+            base_path,
+            shards: HashMap::new(),
+        }
+    }
+
+    fn store_path(&self, store_name: &str) -> QdrantEdgeResult<PathBuf> {
+        validate_store_name(store_name)?;
+        Ok(self.base_path.join("stores").join(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> {
+        if self.shards.contains_key(store_name) {
+            return Ok(self.shards.get(store_name).unwrap());
+        }
+
+        let path = self.store_path(store_name)?;
+        let shard = if has_existing_store(&path) {
+            EdgeShard::load(&path, None)?
+        } else {
+            fs::create_dir_all(&path)?;
+            EdgeShard::new(&path, edge_config(vector_size))?
+        };
+
+        self.shards.insert(store_name.to_string(), shard);
+        Ok(self.shards.get(store_name).unwrap())
+    }
+
+    // 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>> {
+        if self.shards.contains_key(store_name) {
+            return Ok(self.shards.get(store_name));
+        }
+
+        let path = self.store_path(store_name)?;
+        if !has_existing_store(&path) {
+            return Ok(None);
+        }
+
+        let shard = EdgeShard::load(&path, None)?;
+        self.shards.insert(store_name.to_string(), shard);
+        Ok(self.shards.get(store_name))
+    }
+
+    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().is_dir())
+                .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, vector_size: usize) -> QdrantEdgeResult<()> {
+        validate_vector_size(vector_size)?;
+        self.get_or_create_store(store_name, vector_size)?;
+        Ok(())
+    }
+
+    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::<Vec<_>>();
+
+        shard.update(UpdateOperation::PointOperation(
+            PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)),
+        ))?;
+        shard.flush();
+        Ok(())
+    }
+
+    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 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 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<QdrantEdgeOperationResponse> {
+    execute_qdrant_edge_operation(|database| {
+        database.ensure_store_exists(&request.store_name, request.vector_size)
+    })
+}
+
+pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeOperationResponse> {
+    execute_qdrant_edge_operation(|database| {
+        database.insert_embedding(&request.store_name, request.points)
+    })
+}
+
+pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeOperationResponse> {
+    execute_qdrant_edge_operation(|database| {
+        database.delete_embedding_by_file(&request.store_name, &request.file_path)
+    })
+}
+
+pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
+    execute_qdrant_edge_operation(|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_operation<F>(operation: F) -> Json<QdrantEdgeOperationResponse>
+where
+    F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<()>,
+{
+    let mut database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
+    let Some(database) = database_guard.as_mut() else {
+        return Json(QdrantEdgeOperationResponse {
+            success: false,
+            issue: "Qdrant Edge is not available.".to_string(),
+        });
+    };
+
+    match operation(database) {
+        Ok(_) => Json(QdrantEdgeOperationResponse {
+            success: true,
+            issue: String::new(),
+        }),
+        Err(e) => {
+            let issue = e.to_string();
+            error!(Source = "Qdrant Edge"; "Qdrant Edge operation failed: {issue}");
+            Json(QdrantEdgeOperationResponse {
+                success: false,
+                issue,
+            })
+        },
+    }
+}
+
+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);
+}
+
+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() {
+                if (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: 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: hnsw_config(),
+        quantization_config: None,
+        optimizers: edge_optimizers_config(),
+    }
+}
+
+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),
+        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 has_existing_store(path: &Path) -> bool {
+    path.join("edge_config.json").exists() || path.join("segments").exists()
+}
+
+fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
+    if vector_size == 0 {
+        return Err("Vector size must be greater than zero.".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) -> qdrant_edge::PointStructPersisted {
+    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_name": point.data_source_name,
+            "data_source_type": point.data_source_type,
+            "file_path": point.file_path,
+            "file_name": point.file_name,
+            "relative_path": point.relative_path,
+            "chunk_index": point.chunk_index,
+            "text": point.text,
+            "fingerprint": point.fingerprint,
+            "last_write_utc": point.last_write_utc,
+            "embedded_at_utc": point.embedded_at_utc,
+        }),
+    )
+    .into()
+}
+
+fn to_point_id(point_id: &str) -> PointId {
+    Uuid::parse_str(point_id)
+        .map(PointId::Uuid)
+        .unwrap_or_else(|_| PointId::NumId(stable_u64(point_id)))
+}
+
+fn stable_u64(value: &str) -> u64 {
+    let mut hash = 0xcbf29ce484222325_u64;
+    for byte in value.as_bytes() {
+        hash ^= u64::from(*byte);
+        hash = hash.wrapping_mul(0x100000001b3);
+    }
+
+    hash
+}
+
+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,
+    })
+}
+
+fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
+    if store_name.is_empty() {
+        return Err("Vector store name cannot be empty.".into());
+    }
+
+    if matches!(store_name, "." | "..") {
+        return Err(format!("Vector store name '{store_name}' is not supported.").into());
+    }
+
+    if store_name
+        .chars()
+        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')
+    {
+        return Ok(());
+    }
+
+    Err(format!("Vector store name '{store_name}' contains unsupported characters.").into())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn validate_store_name_allows_safe_store_names() {
+        assert!(validate_store_name("rag_1234-abcd.ef").is_ok());
+    }
+
+    #[test]
+    fn validate_store_name_rejects_path_traversal_names() {
+        assert!(validate_store_name(".").is_err());
+        assert!(validate_store_name("..").is_err());
+    }
+}