MindWorkAI/AI-Studio
1
0 Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-10-05 01:09:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added local RAG (#756)

Browse Source 
Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>
This commit is contained in:
Paul KoudelkaandThorsten Sommer authored and GitHub committed 2026-09-09 18:43:37 +02:00
1 parent d043fbc8f0
commit c7b42bee96
179 files changed
+276977 -1421
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

runtime/Cargo.lock
Generated
+196
View File
@@ -731,6 +731,12 @@ dependencies = [
"windows-link 0.2.1",
]
[[package]]
name = "base64"
version = "0.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
[[package]]
name = "base64"
version = "0.21.7"
@@ -1130,6 +1136,15 @@ dependencies = [
"toml 0.9.12+spec-1.1.0",
]
[[package]]
name = "castaway"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
dependencies = [
"rustversion",
]
[[package]]
name = "cbc"
version = "0.1.2"
@@ -1398,6 +1413,21 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "compact_str"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab"
dependencies = [
"castaway",
"cfg-if",
"itoa",
"rustversion",
"ryu",
"serde",
"static_assertions",
]
[[package]]
name = "concurrent-queue"
version = "2.5.0"
@@ -1770,6 +1800,9 @@ name = "dary_heap"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe"
dependencies = [
"serde",
]
[[package]]
name = "data-encoding"
@@ -1877,6 +1910,37 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "derive_builder"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
dependencies = [
"derive_builder_macro",
]
[[package]]
name = "derive_builder_core"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
dependencies = [
"darling 0.20.10",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "derive_builder_macro"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
dependencies = [
"derive_builder_core",
"syn 2.0.117",
]
[[package]]
name = "derive_more"
version = "2.1.1"
@@ -2295,6 +2359,15 @@ version = "3.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a0474425d51df81997e2f90a21591180b38eccf27292d755f3e30750225c175b"
[[package]]
name = "esaxx-rs"
version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
dependencies = [
"cc",
]
[[package]]
name = "event-listener"
version = "5.4.1"
@@ -4320,6 +4393,7 @@ dependencies = [
"tauri-plugin-updater",
"tauri-plugin-window-state",
"tempfile",
"tokenizers",
"tokio",
"tokio-stream",
"toml 1.1.4+spec-1.1.0",
@@ -4375,6 +4449,28 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "monostate"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
dependencies = [
"monostate-impl",
"serde",
"serde_core",
]
[[package]]
name = "monostate-impl"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "moxcms"
version = "0.8.1"
@@ -4996,6 +5092,28 @@ version = "1.70.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
[[package]]
name = "onig"
version = "6.5.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2"
dependencies = [
"bitflags 2.11.1",
"libc",
"once_cell",
"onig_sys",
]
[[package]]
name = "onig_sys"
version = "69.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7"
dependencies = [
"cc",
"pkg-config",
]
[[package]]
name = "open"
version = "5.3.4"
@@ -6006,6 +6124,17 @@ dependencies = [
"rayon-core",
]
[[package]]
name = "rayon-cond"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
dependencies = [
"either",
"itertools",
"rayon",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
@@ -7120,12 +7249,30 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "spm_precompiled"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
dependencies = [
"base64 0.13.1",
"nom 7.1.3",
"serde",
"unicode-segmentation",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
[[package]]
name = "static_assertions"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
[[package]]
name = "strength_reduce"
version = "0.2.4"
@@ -8099,6 +8246,40 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokenizers"
version = "0.23.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112"
dependencies = [
"ahash",
"compact_str",
"daachorse",
"dary_heap",
"derive_builder",
"esaxx-rs",
"getrandom 0.3.1",
"indicatif",
"itertools",
"log",
"macro_rules_attribute",
"monostate",
"onig",
"paste",
"rand 0.9.4",
"rayon",
"rayon-cond",
"regex",
"regex-syntax",
"serde",
"serde_json",
"spm_precompiled",
"thiserror 2.0.18",
"unicode-normalization-alignments",
"unicode-segmentation",
"unicode_categories",
]
[[package]]
name = "tokio"
version = "1.52.3"
@@ -8536,6 +8717,15 @@ dependencies = [
"tinyvec",
]
[[package]]
name = "unicode-normalization-alignments"
version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
dependencies = [
"smallvec",
]
[[package]]
name = "unicode-segmentation"
version = "1.11.0"
@@ -8554,6 +8744,12 @@ version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "229730647fbc343e3a80e463c1db7f78f3855d3f3739bee0dda773c9a037c90a"
[[package]]
name = "unicode_categories"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
[[package]]
name = "unit-prefix"
version = "0.5.2"
runtime/Cargo.toml
+1
View File
@@ -76,6 +76,7 @@ regex = "1.13.1"
aho-corasick = "1.1.5"
toml = "1.1.4"
image = { version = "0.25.10", default-features = false, features = ["jpeg", "png", "webp"] }
tokenizers = "0.23.1"
[patch.crates-io]
# Issue: It was not possible to build qdrant-edge for macOS. See PR 9312: https://github.com/qdrant/qdrant/pull/9312
runtime/resources/tokenizers/tokenizer.json
+263174
View File
File diff suppressed because it is too large. Load diff
runtime/src/app_window.rs
+3
View File
@@ -33,6 +33,7 @@ use crate::global_shortcuts::{RegisterShortcutRequest, ShortcutResponse};
#[cfg(debug_assertions)]
use crate::dotnet::create_startup_env_file;
use crate::tokenizer::set_default_tokenizer_path;
#[cfg(target_os = "linux")]
use webkit2gtk::glib::Cast;
@@ -173,6 +174,8 @@ pub fn start_tauri(tauri_context: tauri::Context<tauri::Wry>) {
start_qdrant_edge_database(app.handle().clone());
set_default_tokenizer_path(app.handle().clone());
info!(Source = "Bootloader Tauri"; "Reconfigure the file logger to use the app data directory {data_path:?}");
switch_to_file_logging(data_path).map_err(|e| error!("Failed to switch logging to file: {e}")).unwrap();
set_pdfium_path(app.path());
runtime/src/file_data.rs
+345 -43
View File
@@ -12,12 +12,12 @@ use axum::response::sse::{Event, Sse};
use base64::{engine::general_purpose, Engine as _};
use calamine::{open_workbook_auto, Error as CalamineError, Reader};
use chardetng::{EncodingDetector, Iso2022JpDetection, Utf8Detection};
use docx_to_md::{DocumentContainer, ImageHandlingMode as DocumentImageHandlingMode, Metadata as DocumentMetadata, ParserConfig as DocumentParserConfig};
use docx_to_md::{DocumentContainer, Error as DocumentError, ImageHandlingMode as DocumentImageHandlingMode, Metadata as DocumentMetadata, ParserConfig as DocumentParserConfig};
use encoding_rs::Encoding;
use file_format::{FileFormat, Kind};
use futures::{Stream, StreamExt};
use pdfium_render::prelude::{Pdfium, PdfiumError, PdfiumInternalError};
use pptx_to_md::{DiagnosticSeverity, ImageHandlingMode, MarkdownOptions, ParserConfig, PresentationContainer, PresentationFormat, PresentationMetadata, ReadingOrder};
use pptx_to_md::{DiagnosticSeverity, Error as PresentationError, ImageHandlingMode, MarkdownOptions, ParserConfig, PresentationContainer, PresentationFormat, PresentationMetadata, ReadingOrder};
use serde::{Deserialize, Deserializer, Serialize};
use serde::de::{Error as SerdeError, Visitor};
use std::path::Path;
@@ -27,17 +27,20 @@ use log::{debug, error, warn};
use tokio::io::AsyncReadExt;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokenizers::tokenizer::Tokenizer;
#[derive(Debug, Serialize)]
pub struct Chunk {
pub content: String,
pub stream_id: String,
pub metadata: Metadata,
#[serde(skip_serializing_if = "Option::is_none")]
pub token_count: Option<usize>,
}
impl Chunk {
pub fn new(content: String, metadata: Metadata) -> Self {
Chunk { content, stream_id: String::new(), metadata }
Chunk { content, stream_id: String::new(), metadata, token_count: None }
}
/// Creates a chunk which reports a failed extraction. Errors travel through the same
@@ -53,11 +56,17 @@ impl Chunk {
page_number: error.page_number,
detected_format: error.detected_format.clone(),
},
token_count: None,
}
}
pub fn set_stream_id(&mut self, stream_id: &str) { self.stream_id = stream_id.to_string(); }
pub fn set_token_count(&mut self, tokenizer: &Tokenizer) -> std::result::Result<(), String> {
self.token_count = Some(crate::tokenizer::get_segment_token_count(tokenizer, &self.content)?);
Ok(())
}
/// Whether this chunk's content is prose a prompt injection could hide in.
///
/// Image chunks carry base64 data, which must never reach the filter: it is not text, and
@@ -73,9 +82,34 @@ impl Chunk {
| Metadata::Presentation { image: Some(_), .. }
)
}
/// Splits an oversized chunk into segments the embedding side can still handle.
///
/// Only prose is split. Everything the filter leaves untouched -- image data, error notices --
/// is passed on whole: cutting base64 in half would corrupt it, which is exactly the set
/// `carries_filterable_text` describes.
fn into_bounded_text_segments(self) -> Vec<Self> {
if !self.carries_filterable_text() {
return vec![self];
}
let ranges = bounded_text_segment_ranges(&self.content);
if ranges.len() == 1 {
return vec![self];
}
ranges
.into_iter()
.map(|(start, end)| {
let mut segment = Chunk::new(self.content[start..end].to_string(), self.metadata.clone());
segment.stream_id = self.stream_id.clone();
segment
})
.collect()
}
}
#[derive(Debug, Serialize)]
#[derive(Clone, Debug, Serialize)]
pub enum Metadata {
Text {
line_number: usize
@@ -141,6 +175,11 @@ pub enum ExtractionErrorCode {
FormatDetectionFailed,
NotAValidPdf,
NotAValidSpreadsheet,
/// The package of a document or presentation is broken, e.g. a damaged ZIP or a missing part.
/// The counterpart of `NotAValidPdf` and `NotAValidSpreadsheet` for the OOXML and ODF formats.
NotAValidDocument,
PdfiumUnavailable,
PdfEncrypted,
PageExtractionFailed,
@@ -246,7 +285,7 @@ fn classify_io_error(error: &std::io::Error) -> ExtractionErrorCode {
}
}
#[derive(Debug, Serialize)]
#[derive(Clone, Debug, Serialize)]
pub struct Base64Image {
pub id: String,
pub content: String,
@@ -271,13 +310,27 @@ const DOCX: &str = "docx";
const ODT: &str = "odt";
const HTML: &str = "html";
const IMAGE_SEGMENT_SIZE_IN_CHARS: usize = 8_192; // equivalent to ~ 5500 token
const MAX_TEXT_SEGMENT_LENGTH_IN_CHARS: usize = 100_000;
/// Every PDF file starts with this signature.
/// The signature which identifies a PDF file.
///
/// It does not have to sit at the very beginning, which is why we search for it instead of
/// comparing against it.
const PDF_MAGIC: &[u8] = b"%PDF-";
/// How many bytes we probe to verify the PDF signature. The few extra bytes beyond the
/// signature itself make the diagnostics useful when the signature does not match.
const PDF_HEADER_PROBE_SIZE: u64 = 8;
/// How far into the file we look for the PDF signature.
///
/// ISO 32000-1 (7.5.2) allows the header anywhere within the first 1024 bytes, and PDFium
/// searches exactly that far. Files carrying something in front of their header do occur:
/// a raw HTTP response saved with a `.pdf` extension has its response headers there. PDFium
/// treats the offset it finds as the origin of the file and shifts every cross-reference
/// offset by it, so such a file reads just fine. The signature itself is added on top of the
/// window, so a header at its very end is still found completely.
const PDF_HEADER_SEARCH_SIZE: u64 = 1024 + PDF_MAGIC.len() as u64;
/// How many of the leading bytes we name when we refuse a file. Enough to recognize what was
/// really saved there, short enough to keep the log line readable.
const PDF_HEADER_DIAGNOSTIC_SIZE: usize = 8;
/// Last-resort payload used when even an error event cannot be serialized. It keeps the
/// chunk schema intact, so the .NET app never has to parse a bare string.
@@ -292,6 +345,10 @@ pub struct ExtractDataQuery {
stream_id: String,
#[serde(deserialize_with = "deserialize_bool_case_insensitive")]
extract_images: bool,
#[serde(default, deserialize_with = "deserialize_bool_case_insensitive")]
include_token_count: bool,
#[serde(default)]
tokenizer_path: String,
}
fn deserialize_bool_case_insensitive<'de, D>(deserializer: D) -> std::result::Result<bool, D::Error>
@@ -354,7 +411,11 @@ fn content_event(chunk: &Chunk, stream_id: &str, path: &str) -> Event {
/// The sanitizer holds chunks back until it has seen enough text to scan across their
/// boundaries, and releases them in order. Their metadata waited here in the meantime,
/// which is what keeps a page's text under its own page number.
fn take_released(held: &mut VecDeque<(u64, Chunk)>, released: Vec<(u64, String)>) -> Vec<Chunk> {
///
/// Splitting oversized chunks and counting their tokens happens here as well, and for the same
/// reason the filter sits where it does: this is where the text the app actually receives comes
/// into being. Counting earlier would report numbers for text the filter had not finished with.
fn take_released(held: &mut VecDeque<(u64, Chunk)>, released: Vec<(u64, String)>, tokenizer: Option<&Tokenizer>) -> Vec<Chunk> {
let mut chunks = Vec::with_capacity(released.len());
for (id, text) in released {
@@ -365,7 +426,21 @@ fn take_released(held: &mut VecDeque<(u64, Chunk)>, released: Vec<(u64, String)>
debug_assert_eq!(held_id, id, "chunks must be released in the order they arrived");
chunk.content = text;
chunks.push(chunk);
for mut segment in chunk.into_bounded_text_segments() {
//
// A count we cannot produce is left out instead of failing the extraction: the app
// treats a missing count as "not counted yet" and counts that segment itself, so the
// document still arrives complete.
//
if let Some(tokenizer) = tokenizer
&& let Err(e) = segment.set_token_count(tokenizer)
{
warn!("Failed to count the tokens of a released chunk: {e}");
}
chunks.push(segment);
}
}
chunks
@@ -443,7 +518,25 @@ pub async fn extract_data(
let stream = stream! {
match query {
Ok(query) => {
Ok(query) => 'request: {
//
// The tokenizer is loaded once, before any chunk is read: it is the same for the
// whole file, and a failure here means we cannot answer the request at all.
//
let tokenizer = if query.include_token_count {
match crate::tokenizer::get_tokenizer(&query.tokenizer_path) {
Ok(tokenizer) => Some(tokenizer),
Err(e) => {
let error = ExtractionError::new(ExtractionErrorCode::InvalidRequest, format!("The tokenizer could not be loaded: {e}"));
warn!("{}", error.message);
yield Ok(error_event(&error, Some(&query.stream_id)));
break 'request;
},
}
} else {
None
};
let stream_result = stream_data(&query.path, query.extract_images, &query.stream_id).await;
let id_ref = &query.stream_id;
let path_ref = &query.path;
@@ -475,7 +568,7 @@ pub async fn extract_data(
break;
};
for released in take_released(&mut held, released_chunks) {
for released in take_released(&mut held, released_chunks, tokenizer.as_deref()) {
yield Ok(content_event(&released, id_ref, path_ref));
}
@@ -494,7 +587,7 @@ pub async fn extract_data(
break;
};
for released in take_released(&mut held, released_chunks) {
for released in take_released(&mut held, released_chunks, tokenizer.as_deref()) {
yield Ok(content_event(&released, id_ref, path_ref));
}
},
@@ -508,7 +601,7 @@ pub async fn extract_data(
// that failed on top of that releases nothing; the extraction
// error below is reported either way.
if let Some(released_chunks) = scan_off_worker(&mut sanitizer, Sanitizer::flush).await {
for released in take_released(&mut held, released_chunks) {
for released in take_released(&mut held, released_chunks, tokenizer.as_deref()) {
yield Ok(content_event(&released, id_ref, path_ref));
}
}
@@ -529,7 +622,7 @@ pub async fn extract_data(
return;
};
for released in take_released(&mut held, released_chunks) {
for released in take_released(&mut held, released_chunks, tokenizer.as_deref()) {
yield Ok(content_event(&released, id_ref, path_ref));
}
}
@@ -578,6 +671,79 @@ pub async fn extract_data(
Sse::new(stream)
}
/// Counts the characters of a text which carry meaning for an embedding.
///
/// Whitespace says nothing, so the readers ask for this count instead of the length of their
/// output: a page which is left with nothing but blank lines is a page without text. Counting
/// the raw length would report a document of scanned images as readable, hand the whitespace to
/// the embedding provider, and leave the user without any hint that the file was never read.
fn readable_character_count(text: &str) -> usize {
text.chars().filter(|character| !character.is_whitespace()).count()
}
/// Counts the readable characters of content whose reader marks the structure it found with HTML
/// comments.
///
/// The presentation reader notes every slide number that way, which means a deck of scanned
/// slides consists of nothing but those markers. They are ours, not the author's, so they must
/// not make such a file look readable.
///
/// Only the readers which add markers of their own use this. In a file the user wrote, a comment
/// is their own text and counts like every other character.
fn readable_character_count_outside_comments(text: &str) -> usize {
const COMMENT_START: &str = "<!--";
const COMMENT_END: &str = "-->";
let mut count = 0;
let mut remaining = text;
loop {
let (before_comment, rest) = match remaining.find(COMMENT_START) {
Some(index) => (&remaining[..index], &remaining[index + COMMENT_START.len()..]),
None => return count + readable_character_count(remaining),
};
count += readable_character_count(before_comment);
//
// An unterminated comment swallows the rest of the text, exactly as a Markdown reader
// would render it: everything behind it is comment and therefore says nothing.
//
remaining = match rest.find(COMMENT_END) {
Some(index) => &rest[index + COMMENT_END.len()..],
None => return count,
};
}
}
/// Splits content into ranges no longer than the segment limit, cutting on character boundaries.
fn bounded_text_segment_ranges(content: &str) -> Vec<(usize, usize)> {
let mut ranges = Vec::new();
let mut start = 0;
while start < content.len() {
let remaining = &content[start..];
let Some(maximum_end_offset) = remaining
.char_indices()
.nth(MAX_TEXT_SEGMENT_LENGTH_IN_CHARS)
.map(|(index, _)| index)
else {
ranges.push((start, content.len()));
break;
};
let end = start + maximum_end_offset;
ranges.push((start, end));
start = end;
}
if ranges.is_empty() {
ranges.push((0, 0));
}
ranges
}
/// How a file is read.
///
/// Deriving the route from the extension and from the content separately is what lets us notice
@@ -835,6 +1001,22 @@ async fn read_text_file(file_path: &str) -> Result<String> {
async fn stream_text_file(file_path: &str, use_md_fences: bool, fence_language: Option<String>) -> Result<ChunkStream> {
let text = read_text_file(file_path).await?;
//
// An empty file, or one holding nothing but blank lines, decodes without a complaint. Reported
// as content, it would arrive as a document the AI is asked to work with, and the Markdown
// fences below would even make it look like one. The whole file is in hand here and nothing was
// sent yet, so this refuses the extraction instead of marking it afterwards.
//
if readable_character_count(&text) == 0 {
warn!("No readable text could be extracted from '{file_path}': the file holds {length} character(s), none of which carry text.", length = text.chars().count());
return Err(ExtractionError::new(
ExtractionErrorCode::NoTextExtracted,
"The file holds no readable text.",
).into());
}
let mut line_number = 0;
let stream = stream! {
@@ -874,6 +1056,10 @@ async fn stream_text_file(file_path: &str, use_md_fences: bool, fence_language:
/// Verifies the file really is a PDF before handing it to PDFium. Without this check, a file
/// which only carries the `.pdf` extension, or whose bytes are not available, would end up in
/// the text branch and silently produce empty content.
///
/// The signature is searched for rather than expected at the beginning, because the format
/// allows it anywhere within the first bytes of the file. Insisting on offset zero would turn
/// documents which every PDF viewer opens into unreadable ones.
async fn ensure_pdf_header(file_path: &str) -> Result<()> {
let file = tokio::fs::File::open(file_path).await.map_err(|error| ExtractionError::new(
classify_io_error(&error),
@@ -885,23 +1071,35 @@ async fn ensure_pdf_header(file_path: &str) -> Result<()> {
format!("The file size could not be read: {error}"),
))?.len();
let mut header = Vec::with_capacity(PDF_HEADER_PROBE_SIZE as usize);
file.take(PDF_HEADER_PROBE_SIZE).read_to_end(&mut header).await.map_err(|error| ExtractionError::new(
let mut header = Vec::with_capacity(PDF_HEADER_SEARCH_SIZE as usize);
file.take(PDF_HEADER_SEARCH_SIZE).read_to_end(&mut header).await.map_err(|error| ExtractionError::new(
classify_io_error(&error),
format!("The first bytes of the file could not be read: {error}"),
))?;
if header.starts_with(PDF_MAGIC) {
return Ok(());
match header.windows(PDF_MAGIC.len()).position(|window| window == PDF_MAGIC) {
Some(0) => Ok(()),
//
// Something sits in front of the header, e.g. the response headers of a raw HTTP
// response which was saved with a `.pdf` extension. PDFium finds the very same offset
// and reads the document from there, so this is worth a note instead of a refusal.
//
Some(offset) => {
warn!("The PDF signature of '{file_path}' begins at offset {offset} instead of at the start of the file; size: {file_size} bytes. PDFium reads the document from that offset.");
Ok(())
},
None => {
let header_hex = header.iter().take(PDF_HEADER_DIAGNOSTIC_SIZE).map(|byte| format!("{byte:02x}")).collect::<Vec<_>>().join(" ");
error!("The file '{file_path}' carries no PDF signature within its first {PDF_HEADER_SEARCH_SIZE} bytes; size: {file_size} bytes, first bytes: [{header_hex}].");
Err(ExtractionError::new(
ExtractionErrorCode::NotAValidPdf,
format!("The file carries no PDF signature within its first {PDF_HEADER_SEARCH_SIZE} bytes. Size: {file_size} bytes, first bytes: [{header_hex}]."),
).into())
},
}
let header_hex = header.iter().map(|byte| format!("{byte:02x}")).collect::<Vec<_>>().join(" ");
error!("The file '{file_path}' does not start with the PDF signature; size: {file_size} bytes, first bytes: [{header_hex}].");
Err(ExtractionError::new(
ExtractionErrorCode::NotAValidPdf,
format!("The file does not start with the PDF signature. Size: {file_size} bytes, first bytes: [{header_hex}]."),
).into())
}
/// Classifies why PDFium refused to open a document, so the cause reaches the user instead of
@@ -982,7 +1180,7 @@ async fn stream_pdf(file_path: &str) -> Result<ChunkStream> {
}
};
number_of_characters += content.chars().count();
number_of_characters += readable_character_count(&content);
if tx.blocking_send(Ok(Chunk::new(
content,
@@ -998,7 +1196,7 @@ async fn stream_pdf(file_path: &str) -> Result<ChunkStream> {
return;
}
debug!("Extracted {number_of_characters} character(s) from {number_of_pages} page(s) of '{path}'; failed pages: {number_of_failed_pages}.");
debug!("Extracted {number_of_characters} readable character(s) from {number_of_pages} page(s) of '{path}'; failed pages: {number_of_failed_pages}.");
//
// Without this marker, a PDF without a text layer and a broken extraction both arrive as
@@ -1026,6 +1224,48 @@ fn classify_spreadsheet_error_code(error: &CalamineError) -> ExtractionErrorCode
}
}
/// Classifies a failure of the document reader, so a broken package is told apart from a file
/// which is merely out of reach right now.
///
/// The distinction decides how long a file stays out of the index: a damaged ZIP or a missing
/// `content.xml` is a property of the file and will fail the same way on every run, while a
/// network share which went away is worth another attempt. Without this, both arrived as an
/// unclassified failure and every run read the broken file again.
///
/// What remains uncoded are the failures of our own image handling. They say nothing about the
/// document, so they keep the generic code.
fn classify_document_error(error: &DocumentError) -> ExtractionErrorCode {
match error {
DocumentError::Io(io_error) => classify_io_error(io_error),
DocumentError::Zip(_)
| DocumentError::Xml { .. }
| DocumentError::Utf8 { .. }
| DocumentError::UnknownFormat
| DocumentError::FormatMismatch { .. }
| DocumentError::MissingPart(_)
| DocumentError::InvalidRelationship { .. } => ExtractionErrorCode::NotAValidDocument,
_ => ExtractionErrorCode::Internal,
}
}
/// Classifies a failure of the presentation reader. Same reasoning as for documents above.
fn classify_presentation_error(error: &PresentationError) -> ExtractionErrorCode {
match error {
PresentationError::Io(io_error) => classify_io_error(io_error),
PresentationError::Zip(_)
| PresentationError::Xml { .. }
| PresentationError::Utf8(_)
| PresentationError::ParseError(_)
| PresentationError::SlideNotFound
| PresentationError::RelationshipNotFound => ExtractionErrorCode::NotAValidDocument,
_ => ExtractionErrorCode::Internal,
}
}
async fn stream_spreadsheet_as_csv(file_path: &str) -> Result<ChunkStream> {
let path = file_path.to_owned();
let (tx, rx) = mpsc::channel(10);
@@ -1042,6 +1282,9 @@ async fn stream_spreadsheet_as_csv(file_path: &str) -> Result<ChunkStream> {
}
};
let mut number_of_sheets = 0;
let mut number_of_characters = 0;
for sheet_name in workbook.sheet_names() {
let range = match workbook.worksheet_range(&sheet_name) {
Ok(r) => r,
@@ -1064,6 +1307,7 @@ async fn stream_spreadsheet_as_csv(file_path: &str) -> Result<ChunkStream> {
}
};
number_of_sheets += 1;
let mut row_idx = 0;
tx.blocking_send(Ok(Chunk::new(
"```csv".to_string(),
@@ -1075,10 +1319,15 @@ async fn stream_spreadsheet_as_csv(file_path: &str) -> Result<ChunkStream> {
for row in range.rows() {
row_idx += 1;
let content = row.iter()
.map(|cell| cell.to_string())
.collect::<Vec<_>>()
.join(",");
let cells = row.iter().map(|cell| cell.to_string()).collect::<Vec<_>>();
//
// The cells are counted one by one, before they are joined: a row of empty cells
// joins into a line of commas, and those would pass for content although the row
// holds nothing. The fences around each sheet are left out for the same reason.
//
number_of_characters += cells.iter().map(|cell| readable_character_count(cell)).sum::<usize>();
let content = cells.join(",");
if tx.blocking_send(Ok(Chunk::new(
content,
@@ -1099,6 +1348,21 @@ async fn stream_spreadsheet_as_csv(file_path: &str) -> Result<ChunkStream> {
}
))).ok();
}
debug!("Extracted {number_of_characters} readable character(s) from {number_of_sheets} sheet(s) of '{path}'.");
//
// Without this marker, an empty workbook arrives as a handful of Markdown fences with
// nothing between them, and the AI would answer as if that were the content of the file.
//
if number_of_characters == 0 {
warn!("No text could be extracted from '{path}': {number_of_sheets} sheet(s), all of them without any cell content.");
let _ = tx.blocking_send(Ok(Chunk::from_error(&ExtractionError::new(
ExtractionErrorCode::NoTextExtracted,
format!("No text could be extracted from {number_of_sheets} sheet(s) of the spreadsheet."),
))));
}
});
Ok(Box::pin(ReceiverStream::new(rx)))
@@ -1248,7 +1512,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
Ok(document) => document,
Err(e) => {
let _ = tx.blocking_send(Err(ExtractionError::new(
ExtractionErrorCode::FileNotReadable,
classify_document_error(&e),
format!("The document could not be read: {e}"),
).into()));
return;
@@ -1259,7 +1523,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
Ok(pages) => pages,
Err(e) => {
let _ = tx.blocking_send(Err(ExtractionError::new(
ExtractionErrorCode::FileNotReadable,
classify_document_error(&e),
format!("The pages of the document could not be read: {e}"),
).into()));
return;
@@ -1281,7 +1545,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
Ok(page) => page,
Err(e) => {
let _ = tx.blocking_send(Err(ExtractionError::new(
ExtractionErrorCode::Internal,
classify_document_error(&e),
format!("A page of the document could not be read: {e}"),
).into()));
return;
@@ -1291,7 +1555,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
Ok(content) => content,
Err(e) => {
let _ = tx.blocking_send(Err(ExtractionError::new(
ExtractionErrorCode::Internal,
classify_document_error(&e),
format!("Page {page_number} of the document could not be converted: {e}", page_number = page.page_number),
).into()));
return;
@@ -1299,7 +1563,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
};
number_of_pages = page.page_number;
number_of_characters += content.chars().count();
number_of_characters += readable_character_count(&content);
if let Some(metadata) = metadata_md.take() {
content = format!("{metadata}\n\n{content}");
@@ -1331,7 +1595,7 @@ async fn stream_document(file_path: &str, extract_images: bool, stream_id: &str)
}
}
debug!("Extracted {number_of_characters} character(s) from {number_of_pages} page(s) of '{path}'.", path = path.display());
debug!("Extracted {number_of_characters} readable character(s) from {number_of_pages} page(s) of '{path}'.", path = path.display());
//
// Without this marker, a document without any text and a broken extraction both arrive as
@@ -1363,6 +1627,10 @@ async fn stream_presentation(file_path: &str, extract_images: bool, format: Pres
let path = Path::new(file_path).to_owned();
let stream_id = stream_id.to_owned();
// The path itself is moved into the task which opens the presentation, so the diagnostics of
// the worker below keep their own copy:
let log_path = file_path.to_owned();
let parser_config = ParserConfig::builder()
.extract_images(extract_images)
.compress_images(true)
@@ -1380,7 +1648,10 @@ async fn stream_presentation(file_path: &str, extract_images: bool, format: Pres
};
let mut streamer = tokio::task::spawn_blocking(move || {
PresentationContainer::open_as(&path, parser_config, format).map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
PresentationContainer::open_as(&path, parser_config, format).map_err(|e| Box::new(ExtractionError::new(
classify_presentation_error(&e),
format!("The presentation could not be read: {e}"),
)) as Box<dyn std::error::Error + Send + Sync>)
}).await??;
let (tx, rx) = mpsc::channel(32);
@@ -1390,12 +1661,17 @@ async fn stream_presentation(file_path: &str, extract_images: bool, format: Pres
// so the complete producer must stay outside Tokio's asynchronous workers.
let worker = tokio::task::spawn_blocking(move || {
let mut metadata_md = presentation_metadata_to_markdown(streamer.metadata());
let mut number_of_slides = 0;
let mut number_of_characters = 0;
for slide_result in streamer.iter_slides() {
let slide = match slide_result {
Ok(slide) => slide,
Err(e) => {
let _ = tx.blocking_send(Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>));
let _ = tx.blocking_send(Err(ExtractionError::new(
classify_presentation_error(&e),
format!("A slide of the presentation could not be read: {e}"),
).into()));
return;
},
};
@@ -1421,11 +1697,22 @@ async fn stream_presentation(file_path: &str, extract_images: bool, format: Pres
let mut content = match slide.to_markdown(&markdown_options) {
Ok(content) => content,
Err(e) => {
let _ = tx.blocking_send(Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>));
let _ = tx.blocking_send(Err(ExtractionError::new(
classify_presentation_error(&e),
format!("Slide {slide_number} of the presentation could not be converted: {e}", slide_number = slide.slide_number),
).into()));
return;
},
};
//
// Counted here, before the metadata of the presentation is put in front of the first
// slide: its title and author belong to the file, not to the slides, and a deck of
// scanned images would look readable through them alone.
//
number_of_slides += 1;
number_of_characters += readable_character_count_outside_comments(&content);
if let Some(metadata) = metadata_md.take() {
content = format!("{metadata}\n\n{content}");
}
@@ -1486,6 +1773,21 @@ async fn stream_presentation(file_path: &str, extract_images: bool, format: Pres
}
}
}
debug!("Extracted {number_of_characters} readable character(s) from {number_of_slides} slide(s) of '{log_path}'.");
//
// Without this marker, a presentation of nothing but pictures arrives as a row of slide
// number comments, and the AI would answer as if that were the content of the file.
//
if number_of_characters == 0 {
warn!("No text could be extracted from '{log_path}': {number_of_slides} slide(s). The presentation may consist of images only.");
let _ = tx.blocking_send(Ok(Chunk::from_error(&ExtractionError::new(
ExtractionErrorCode::NoTextExtracted,
format!("No text could be extracted from {number_of_slides} slide(s). The presentation may consist of images only."),
))));
}
});
tokio::spawn(async move {
runtime/src/lib.rs
+1
View File
@@ -22,5 +22,6 @@ pub mod runtime_api_token;
pub mod stale_process_cleanup;
pub mod share_sheet;
mod sidecar_types;
pub mod tokenizer;
mod file_actions;
pub mod global_shortcuts;
runtime/src/main.rs
+1 -1
View File
@@ -62,4 +62,4 @@ fn main() {
start_runtime_api();
start_tauri(tauri_context);
}
}
runtime/src/qdrant_edge_database.rs
+397 -56
View File
@@ -6,12 +6,13 @@ 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::serde_json::{json, Value};
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,
FieldCondition, Filter, HnswIndexConfig, Match, MatchValue, NamedQuery, Payload, PointId,
PointInsertOperations, PointOperations, PointStruct, QueryEnum, ScoredPoint, SearchRequest,
UpdateOperation, ValueVariants, VectorInternal, Vectors, WithPayloadInterface, WithVector,
};
use serde::{Deserialize, Serialize};
use tauri::Manager;
@@ -26,6 +27,10 @@ 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";
type QdrantEdgeResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
@@ -67,19 +72,28 @@ pub struct QdrantEdgeStoragePoint {
pub data_source_id: String,
pub data_source_name: 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,
pub confidence_level: String,
pub confidence_level_rank: i32,
}
#[derive(Deserialize)]
pub struct EnsureQdrantEdgeStoreRequest {
pub store_name: String,
pub data_source_name: String,
pub vector_size: usize,
}
@@ -89,21 +103,64 @@ pub struct InsertQdrantEdgeEmbeddingRequest {
pub points: Vec<QdrantEdgeStoragePoint>,
}
#[derive(Deserialize)]
pub struct SearchQdrantEdgeEmbeddingRequest {
pub store_name: String,
pub vector: Vec<f32>,
pub max_matches: usize,
}
#[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 QdrantEdgeOperationResponse {
pub struct QdrantEdgeResponse<T> {
pub success: bool,
pub issue: String,
pub data: Option<T>,
}
#[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_name: 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,
pub confidence_level: String,
pub confidence_level_rank: i32,
}
#[derive(Clone, Serialize)]
@@ -129,49 +186,88 @@ impl QdrantEdgeDatabase {
fn store_path(&self, store_name: &str) -> QdrantEdgeResult<PathBuf> {
validate_store_name(store_name)?;
Ok(self.base_path.join("stores").join(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> {
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());
return Ok((self.shards.get(store_name).unwrap(), false));
}
let path = self.store_path(store_name)?;
let shard = if has_existing_store(&path) {
EdgeShard::load(&path, None)?
let shard = if is_initialized {
EdgeShard::load(&path, None).map_err(|error| {
format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display())
})?
} else {
fs::create_dir_all(&path)?;
EdgeShard::new(&path, edge_config(vector_size))?
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.shards.insert(store_name.to_string(), shard);
Ok(self.shards.get(store_name).unwrap())
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));
}
let path = self.store_path(store_name)?;
if !has_existing_store(&path) {
if !is_initialized {
return Ok(None);
}
let shard = EdgeShard::load(&path, None)?;
let shard = EdgeShard::load(&path, None).map_err(|error| {
format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display())
})?;
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().is_dir())
.filter(|entry| entry.path().join(STORE_INITIALIZATION_MARKER).is_file())
.count()
} else {
0
@@ -185,10 +281,15 @@ impl QdrantEdgeDatabase {
})
}
fn ensure_store_exists(&mut self, store_name: &str, vector_size: usize) -> QdrantEdgeResult<()> {
fn ensure_store_exists(&mut self, store_name: &str, data_source_name: &str, vector_size: usize) -> QdrantEdgeResult<QdrantEdgeEnsureStoreResult> {
validate_vector_size(vector_size)?;
self.get_or_create_store(store_name, vector_size)?;
Ok(())
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<()> {
@@ -202,11 +303,19 @@ impl QdrantEdgeDatabase {
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 data_source_name = first_point.data_source_name.clone();
validate_data_source_name(&data_source_name)?;
if points.iter().any(|point| point.data_source_name != data_source_name) {
return Err("All points in one insert request must belong to the same data source name.".into());
}
let store_path = self.store_path(store_name)?;
let (shard, _) = self.get_or_create_store(store_name, vector_size)?;
write_store_display_name(&store_path, &data_source_name)?;
let points = points
.into_iter()
.map(to_qdrant_edge_point)
.collect::<Vec<_>>();
.collect::<QdrantEdgeResult<Vec<_>>>()?;
shard.update(UpdateOperation::PointOperation(
PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)),
@@ -215,6 +324,36 @@ impl QdrantEdgeDatabase {
Ok(())
}
fn search_embedding(&mut self, store_name: &str, vector: Vec<f32>, max_matches: usize) -> QdrantEdgeResult<Vec<QdrantEdgeSearchResult>> {
if max_matches == 0 {
return Ok(vec![]);
}
validate_vector_size(vector.len())?;
let Some(shard) = self.get_existing_store(store_name)? else {
return Ok(vec![]);
};
let search_results = shard.search(SearchRequest {
query: QueryEnum::Nearest(NamedQuery::new(
VectorInternal::Dense(vector),
VECTOR_NAME,
)),
filter: None,
params: None,
limit: max_matches,
offset: 0,
with_payload: Some(WithPayloadInterface::Bool(true)),
with_vector: Some(WithVector::Bool(false)),
score_threshold: None,
})?;
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(());
@@ -227,6 +366,19 @@ impl QdrantEdgeDatabase {
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();
Ok(())
}
fn delete_store(&mut self, store_name: &str) -> QdrantEdgeResult<()> {
self.shards.remove(store_name);
@@ -243,6 +395,11 @@ impl QdrantEdgeDatabase {
}
}
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()
@@ -276,26 +433,38 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
})
}
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 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<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
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 delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
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)
})
}
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 delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
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)
})
}
@@ -338,29 +507,33 @@ pub fn stop_qdrant_edge_database() {
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
}
fn execute_qdrant_edge_operation<F>(operation: F) -> Json<QdrantEdgeOperationResponse>
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
where
F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<()>,
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(QdrantEdgeOperationResponse {
return Json(QdrantEdgeResponse {
success: false,
issue: "Qdrant Edge is not available.".to_string(),
data: None,
});
};
match operation(database) {
Ok(_) => Json(QdrantEdgeOperationResponse {
Ok(data) => Json(QdrantEdgeResponse {
success: true,
issue: String::new(),
data: Some(data),
}),
Err(e) => {
let issue = e.to_string();
error!(Source = "Qdrant Edge"; "Qdrant Edge operation failed: {issue}");
Json(QdrantEdgeOperationResponse {
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
Json(QdrantEdgeResponse {
success: false,
issue,
data: None,
})
},
}
@@ -477,8 +650,53 @@ fn edge_optimizers_config() -> EdgeOptimizersConfig {
}
}
fn has_existing_store(path: &Path) -> bool {
path.join("edge_config.json").exists() || path.join("segments").exists()
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}"),
}
}
fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
@@ -489,6 +707,24 @@ fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
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()
@@ -500,41 +736,90 @@ fn vector_store_version() -> QdrantEdgeResult<String> {
Ok(metadata.vector_store_version.clone())
}
fn to_qdrant_edge_point(point: QdrantEdgeStoragePoint) -> qdrant_edge::PointStructPersisted {
PointStruct::new(
to_point_id(&point.point_id),
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_name": point.data_source_name,
"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,
"confidence_level": point.confidence_level,
"confidence_level_rank": point.confidence_level_rank,
}),
)
.into()
.into())
}
fn to_point_id(point_id: &str) -> PointId {
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_name: payload_string(&payload, "data_source_name"),
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"),
confidence_level: payload_string(&payload, "confidence_level"),
confidence_level_rank: payload_i32(&payload, "confidence_level_rank").unwrap_or_default(),
}
}
fn to_point_id(point_id: &str) -> QdrantEdgeResult<PointId> {
Uuid::parse_str(point_id)
.map(PointId::Uuid)
.unwrap_or_else(|_| PointId::NumId(stable_u64(point_id)))
.map_err(|_| "Vector point ID must be a valid UUID.".into())
}
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);
fn point_id_to_string(point_id: PointId) -> String {
match point_id {
PointId::NumId(id) => id.to_string(),
PointId::Uuid(uuid) => uuid.to_string(),
}
}
hash
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> {
@@ -554,17 +839,19 @@ fn match_keyword_filter(field_name: &str, value: &str) -> QdrantEdgeResult<Filte
}
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 matches!(store_name, "." | "..") {
return Err(format!("Vector store name '{store_name}' is not supported.").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 == '-' || c == '.')
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
{
return Ok(());
}
@@ -578,12 +865,66 @@ mod tests {
#[test]
fn validate_store_name_allows_safe_store_names() {
assert!(validate_store_name("rag_1234-abcd.ef").is_ok());
assert!(validate_store_name("rag_1234-abcd").is_ok());
}
#[test]
fn validate_store_name_rejects_path_traversal_names() {
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 = std::env::temp_dir().join(format!(
"ai-studio-qdrant-ensure-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
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 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());
}
}
runtime/src/runtime_api.rs
+6
View File
@@ -40,7 +40,9 @@ pub fn start_runtime_api() {
.route("/system/qdrant-edge/info", get(crate::qdrant_edge_database::qdrant_edge_info))
.route("/system/qdrant-edge/ensure", post(crate::qdrant_edge_database::ensure_qdrant_edge_store))
.route("/system/qdrant-edge/insert", post(crate::qdrant_edge_database::insert_qdrant_edge_embedding))
.route("/system/qdrant-edge/search", post(crate::qdrant_edge_database::search_qdrant_edge_embeddings))
.route("/system/qdrant-edge/delete-file", post(crate::qdrant_edge_database::delete_qdrant_edge_embedding_by_file))
.route("/system/qdrant-edge/optimize", post(crate::qdrant_edge_database::optimize_qdrant_edge_store))
.route("/system/qdrant-edge/delete-store", post(crate::qdrant_edge_database::delete_qdrant_edge_store))
.route("/clipboard/set", post(crate::clipboard::set_clipboard))
.route("/share/file", post(crate::share_sheet::share_file))
@@ -81,6 +83,10 @@ pub fn start_runtime_api() {
.route("/image/prepare", post(crate::image::prepare_image))
.route("/log/paths", get(crate::log::get_log_paths))
.route("/log/event", post(crate::log::log_event))
.route("/tokenizer/count", post(crate::tokenizer::token_count))
.route("/tokenizer/validate", post(crate::tokenizer::validate_tokenizer))
.route("/tokenizer/store", post(crate::tokenizer::store_tokenizer))
.route("/tokenizer/delete", post(crate::tokenizer::delete_tokenizer))
.route("/shortcuts/register", post(crate::app_window::register_shortcut))
.route("/shortcuts/validate", post(crate::app_window::validate_shortcut))
.route("/shortcuts/suspend", post(crate::app_window::suspend_shortcuts))
runtime/src/tokenizer.rs
+305
View File
@@ -0,0 +1,305 @@
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, OnceLock, RwLock};
use axum::Json;
use log::{error, warn};
use serde::{Deserialize, Serialize};
use tauri::path::BaseDirectory;
use tauri::Manager;
use tokenizers::tokenizer::Tokenizer;
use crate::api_token::APIToken;
use crate::environment::DATA_DIRECTORY;
const DEFAULT_TOKENIZER_RESOURCE_PATH: &str = "resources/tokenizers/tokenizer.json";
static TOKENIZERS: OnceLock<RwLock<HashMap<PathBuf, Arc<Tokenizer>>>> = OnceLock::new();
static DEFAULT_TOKENIZER_PATH: OnceLock<PathBuf> = OnceLock::new();
static TOKENIZER_STORAGE_LOCK: Mutex<()> = Mutex::new(());
#[derive(Deserialize)]
pub struct SetTokenText {
text: String,
tokenizer_path: String,
}
#[derive(Clone, Deserialize)]
pub struct TokenizerStorage {
model_id: String,
file_path: String,
}
#[derive(Clone, Deserialize)]
pub struct TokenizerDelete {
model_id: String,
}
#[derive(Clone, Deserialize)]
pub struct TokenizerPath {
file_path: String,
}
#[derive(Serialize)]
pub struct TokenizerResponse {
success: bool,
token_count: usize,
message: String,
stored_path: String,
}
impl TokenizerResponse {
fn available(token_count: usize) -> Self {
TokenizerResponse {
success: true,
token_count,
message: String::new(),
stored_path: String::new(),
}
}
fn stored(stored_path: String) -> Self {
TokenizerResponse {
success: true,
token_count: 0,
message: String::new(),
stored_path,
}
}
fn unavailable(reason: String) -> Self {
TokenizerResponse {
success: false,
token_count: 0,
message: reason,
stored_path: String::new(),
}
}
}
pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
let tokenizer_path = match app_handle
.path()
.resolve(DEFAULT_TOKENIZER_RESOURCE_PATH, BaseDirectory::Resource)
{
Ok(path) => path,
Err(e) => {
let reason = format!("The default tokenizer file '{DEFAULT_TOKENIZER_RESOURCE_PATH}' could not be resolved: {e}");
error!(Source = "Tokenizer"; "{reason}");
return;
}
};
if !tokenizer_path.is_file() {
let reason = format!("The default tokenizer file was not found: {}", tokenizer_path.display());
error!(Source = "Tokenizer"; "{reason}");
return;
}
match DEFAULT_TOKENIZER_PATH.set(tokenizer_path) {
Ok(_) => (),
Err(e) => warn!(Source = "Tokenizer"; "Could not set the default tokenizer path: {:?}", e),
}
}
pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
match get_token_count(&req.tokenizer_path, &req.text) {
Ok(count) => Json(TokenizerResponse::available(count)),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
pub async fn validate_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
match handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())) {
Ok(count) => Json(TokenizerResponse::available(count)),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
pub async fn store_tokenizer(_token: APIToken, payload: Json<TokenizerStorage>) -> Json<TokenizerResponse> {
match handle_tokenizer_store(&payload) {
Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
}
}
pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>) -> Json<TokenizerResponse> {
match handle_tokenizer_delete(&payload) {
Ok(_) => Json(TokenizerResponse::stored(String::new())),
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
}
}
fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
validate_tokenizer_file(path)
}
pub fn get_token_count(path: &str, text: &str) -> Result<usize, String> {
let tokenizer = get_tokenizer(path)?;
get_token_count_internal(&tokenizer, text, true)
}
pub fn get_segment_token_count(tokenizer: &Tokenizer, text: &str) -> Result<usize, String> {
// Special tokens belong to the final encoding and would inflate sums across many segments.
get_token_count_internal(tokenizer, text, false)
}
fn get_token_count_internal(tokenizer: &Tokenizer, text: &str, add_special_tokens: bool) -> Result<usize, String> {
if text.trim().is_empty() {
return Ok(0);
}
tokenizer
.encode(text, add_special_tokens)
.map(|encoding| encoding.len())
.map_err(|e| format!("Failed to tokenize text: {e}"))
}
fn validate_tokenizer_file(path: &PathBuf) -> Result<usize, String> {
let tokenizer = load_tokenizer_from_file(path)?;
let test_string = "Hello, world! This is a test string for tokenizer validation.";
let encoding = tokenizer
.encode(test_string, true)
.map_err(|e| format!("Tokenizer failed to encode validation string: {e}"))?;
let token_count = encoding.len();
if token_count == 0 {
return Err("Tokenizer produced 0 tokens for test string. The tokenizer is likely invalid or misconfigured.".to_string());
}
if encoding.get_tokens().iter().any(|t| t.is_empty()) {
return Err("Tokenizer produced empty tokens. The tokenizer is invalid.".to_string());
}
Ok(token_count)
}
fn handle_tokenizer_store(payload: &TokenizerStorage) -> Result<String, std::io::Error> {
let data_dir = DATA_DIRECTORY
.get()
.ok_or_else(|| std::io::Error::other("DATA_DIRECTORY not initialized"))?;
let base_path = PathBuf::from(data_dir).join("tokenizers");
let source_path = PathBuf::from(&payload.file_path);
let source_name = source_path
.file_name()
.and_then(|n| n.to_str())
.ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidInput, "Invalid tokenizer file path"))?;
let model_path = base_path.join(&payload.model_id);
let destination_path = model_path.join(source_name);
if source_path.eq(&destination_path) {
return Ok(destination_path.to_string_lossy().to_string());
}
let _storage_guard = TOKENIZER_STORAGE_LOCK
.lock()
.map_err(|_| std::io::Error::other("Tokenizer storage lock is poisoned."))?;
if model_path.try_exists()? {
invalidate_tokenizers_under(&model_path);
fs::remove_dir_all(&model_path)?;
}
if payload.file_path.trim().is_empty() {
return Ok(String::new());
}
fs::create_dir_all(&model_path)?;
fs::copy(&source_path, &destination_path)?;
Ok(destination_path.to_string_lossy().to_string())
}
fn handle_tokenizer_delete(payload: &TokenizerDelete) -> Result<(), std::io::Error> {
if payload.model_id.trim().is_empty() {
return Ok(());
}
let data_dir = DATA_DIRECTORY
.get()
.ok_or_else(|| std::io::Error::other("DATA_DIRECTORY not initialized"))?;
let tokenizer_path = PathBuf::from(data_dir)
.join("tokenizers")
.join(&payload.model_id);
let _storage_guard = TOKENIZER_STORAGE_LOCK
.lock()
.map_err(|_| std::io::Error::other("Tokenizer storage lock is poisoned."))?;
if tokenizer_path.exists() {
invalidate_tokenizers_under(&tokenizer_path);
fs::remove_dir_all(tokenizer_path)?;
}
Ok(())
}
fn tokenizer_cache() -> &'static RwLock<HashMap<PathBuf, Arc<Tokenizer>>> {
TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
}
pub fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
let resolved_path = resolve_tokenizer_path(path)?;
let tokenizer_path = fs::canonicalize(&resolved_path)
.map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
if let Some(tokenizer) = tokenizer_cache()
.read()
.map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?
.get(&tokenizer_path)
.cloned()
{
return Ok(tokenizer);
}
let _storage_guard = TOKENIZER_STORAGE_LOCK
.lock()
.map_err(|_| "Tokenizer storage lock is poisoned.".to_string())?;
if let Some(tokenizer) = tokenizer_cache()
.read()
.map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?
.get(&tokenizer_path)
.cloned()
{
return Ok(tokenizer);
}
let loaded_tokenizer = Arc::new(load_tokenizer_from_file(&tokenizer_path)?);
let mut cache = tokenizer_cache()
.write()
.map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?;
Ok(cache
.entry(tokenizer_path)
.or_insert_with(|| loaded_tokenizer)
.clone())
}
fn invalidate_tokenizers_under(path: &PathBuf) {
let cache_path = fs::canonicalize(path).unwrap_or_else(|_| path.clone());
match tokenizer_cache().write() {
Ok(mut cache) => cache.retain(|tokenizer_path, _| !tokenizer_path.starts_with(&cache_path)),
Err(_) => warn!(Source = "Tokenizer"; "Could not invalidate tokenizer cache because its lock is poisoned."),
}
}
fn resolve_tokenizer_path(path: &str) -> Result<PathBuf, String> {
if !path.trim().is_empty() {
return Ok(PathBuf::from(path));
}
DEFAULT_TOKENIZER_PATH
.get()
.cloned()
.ok_or_else(|| "Default tokenizer path is not initialized.".to_string())
}
fn load_tokenizer_from_file(path: &PathBuf) -> Result<Tokenizer, String> {
if !path.is_file() {
return Err(format!("Tokenizer file was not found: {}", path.display()));
}
Tokenizer::from_file(path)
.map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
}
runtime/tauri.conf.json
+2 -1
View File
@@ -28,7 +28,8 @@
],
"resources": [
"resources/libraries/*",
"resources/notices/*"
"resources/notices/*",
"resources/tokenizers/*"
],
"macOS": {
"exceptionDomain": "localhost"