fixed tokenization issue for long files

runtime/src/file_data.rs

@@ -28,17 +28,40 @@ 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 }
     }
     
     pub fn set_stream_id(&mut self, stream_id: &str) { self.stream_id = stream_id.to_string(); }
+
+    pub fn set_token_count(&mut self) -> std::result::Result<(), String> {
+        self.token_count = Some(crate::tokenizer::get_segment_token_count(&self.content)?);
+        Ok(())
+    }
+
+    fn into_bounded_text_segments(self) -> Vec<Self> {
+        if matches!(&self.metadata, Metadata::Image {}) {
+            return vec![self];
+        }
+
+        let ranges = bounded_text_segment_ranges(&self.content);
+        if ranges.len() == 1 {
+            return vec![self];
+        }
+
+        ranges
+            .into_iter()
+            .map(|(start, end)| Chunk::new(self.content[start..end].to_string(), self.metadata.clone()))
+            .collect()
+    }
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Clone, Debug, Serialize)]
 pub enum Metadata {
     Text {
         line_number: usize
@@ -62,7 +85,7 @@ pub enum Metadata {
     },
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Clone, Debug, Serialize)]
 pub struct Base64Image {
     pub id: String,
     pub content: String,
@@ -80,6 +103,7 @@ const TO_MARKDOWN: &str = "markdown";
 const DOCX: &str = "docx";
 const ODT: &str = "odt";
 const IMAGE_SEGMENT_SIZE_IN_CHARS: usize = 8_192; // equivalent to ~ 5500 token
+const MAX_TEXT_SEGMENT_LENGTH_IN_CHARS: usize = 100_000;
 
 type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
 type ChunkStream = Pin<Box<dyn Stream<Item = Result<Chunk>> + Send>>;
@@ -90,6 +114,8 @@ 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,
 }
 
 fn deserialize_bool_case_insensitive<'de, D>(deserializer: D) -> std::result::Result<bool, D::Error>
@@ -145,11 +171,21 @@ pub async fn extract_data(
 
                 match stream_result {
                     Ok(mut stream) => {
-                        while let Some(chunk) = stream.next().await {
+                        'stream_chunks: while let Some(chunk) = stream.next().await {
                             match chunk {
-                                Ok(mut chunk) => {
-                                    chunk.set_stream_id(id_ref);
-                                    yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}"))));
+                                Ok(chunk) => {
+                                    let chunks = chunk.into_bounded_text_segments();
+
+                                    for mut chunk in chunks {
+                                        chunk.set_stream_id(id_ref);
+                                        if query.include_token_count {
+                                            if let Err(e) = chunk.set_token_count() {
+                                                yield Ok(Event::default().json_data(format!("Error counting tokens: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error counting tokens: {e}"))));
+                                                break 'stream_chunks;
+                                            }
+                                        }
+                                        yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}"))));
+                                    }
                                 },
 
                                 Err(e) => {
@@ -175,6 +211,33 @@ pub async fn extract_data(
     Sse::new(stream)
 }
 
+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
+}
+
 async fn stream_data(file_path: &str, extract_images: bool) -> Result<ChunkStream> {
     if !Path::new(file_path).exists() {
         error!("File does not exist: '{file_path}'");

runtime/src/qdrant_edge_database.rs

@@ -15,6 +15,7 @@ use qdrant_edge::{
     UpdateOperation, ValueVariants, VectorInternal, Vectors, WithPayloadInterface, WithVector,
 };
 use serde::{Deserialize, Serialize};
+use sha2::{Digest, Sha256};
 use tauri::Manager;
 
 use crate::api_token::APIToken;
@@ -27,6 +28,8 @@ 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";
 
 type QdrantEdgeResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
 
@@ -182,21 +185,52 @@ 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> {
+        let path = self.store_path(store_name)?;
+        let is_initialized = store_is_initialized(&path, store_name)?;
         if self.shards.contains_key(store_name) {
-            return Ok(self.shards.get(store_name).unwrap());
+            if is_initialized {
+                return Ok(self.shards.get(store_name).unwrap());
+            }
+
+            warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name);
+            self.shards.remove(store_name);
         }
 
-        let path = self.store_path(store_name)?;
-        let shard = if has_existing_store(&path) {
-            EdgeShard::load(&path, None)?
+        if path.exists() && !is_initialized {
+            warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before recreating it.", store_name);
+            fs::remove_dir_all(&path).map_err(|error| {
+                format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display())
+            })?;
+        }
+
+        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);
@@ -205,16 +239,31 @@ impl QdrantEdgeDatabase {
 
     // 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 = self.store_path(store_name)?;
+        let is_initialized = store_is_initialized(&path, store_name)?;
         if self.shards.contains_key(store_name) {
-            return Ok(self.shards.get(store_name));
+            if is_initialized {
+                return Ok(self.shards.get(store_name));
+            }
+
+            warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name);
+            self.shards.remove(store_name);
         }
 
-        let path = self.store_path(store_name)?;
-        if !has_existing_store(&path) {
+        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())
+            })?;
+        }
+
+        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))
     }
@@ -224,7 +273,7 @@ impl QdrantEdgeDatabase {
         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
@@ -339,6 +388,16 @@ impl QdrantEdgeDatabase {
     }
 }
 
+fn store_directory_name(store_name: &str) -> String {
+    // Qdrant creates deeply nested files, so keep the physical path short on Windows.
+    let digest = Sha256::digest(store_name.as_bytes());
+    let short_hash = digest[..12]
+        .iter()
+        .map(|byte| format!("{byte:02x}"))
+        .collect::<String>();
+    format!("store_{short_hash}")
+}
+
 fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
     let data_directory = DATA_DIRECTORY
         .get()
@@ -616,8 +675,38 @@ 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<()> {
+    let marker_path = path.join(STORE_INITIALIZATION_MARKER);
+    let temporary_marker_path = path.join(STORE_INITIALIZATION_MARKER_TEMP);
+    fs::write(&temporary_marker_path, store_name)?;
+    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<()> {

runtime/src/tokenizer.rs

@@ -212,6 +212,15 @@ fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
 }
 
 pub fn get_token_count(text: &str) -> Result<usize, String> {
+    get_token_count_internal(text, true)
+}
+
+pub fn get_segment_token_count(text: &str) -> Result<usize, String> {
+    // Special tokens belong to the final encoding and would inflate sums across many segments.
+    get_token_count_internal(text, false)
+}
+
+fn get_token_count_internal(text: &str, add_special_tokens: bool) -> Result<usize, String> {
     if text.trim().is_empty() {
         return Ok(0);
     }
@@ -234,7 +243,7 @@ pub fn get_token_count(text: &str) -> Result<usize, String> {
             return Err(unavailable_with_status_update("Tokenizer not initialized."));
         }
     };
-    let token_count = match tokenizer.encode(text, true) {
+    let token_count = match tokenizer.encode(text, add_special_tokens) {
         Ok(enc) => enc.len(),
         Err(e) => {
             let reason = format!("Failed to tokenize text: {e}");