Enhance JSON ingestion: enforce file type validation, split chunks into sub-chunks, and improve error handling and logging for embedding failures

Author: BryanChasko

rust_ingest/src/ingest.rs

@@ -354,18 +354,71 @@ async fn process_single_file_for_embedding(
     config: &IngestConfig,
     client: &reqwest::Client,
 ) -> Result<Vec<f32>> {
-    // Read and truncate file content
+    // Only process .json or .jsonl files
+    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
+    if ext != "json" && ext != "jsonl" {
+        return Err(anyhow::anyhow!("Only .json and .jsonl files are supported for ingestion: {}", path.display()));
+    }
+
     let content = std::fs::read_to_string(path)
         .with_context(|| format!("Failed to read file: {}", path.display()))?;
 
-    let truncated_content = truncate_content(&content, config.max_chars);
-
-    // Generate embedding vector
-    let embedding = embed::embed(truncated_content, config.max_tokens, client)
-        .await
-        .with_context(|| format!("Failed to generate embedding for file: {}", path.display()))?;
+    // Parse as JSON array (or linewise for JSONL)
+    let objects: Vec<serde_json::Value> = if ext == "jsonl" {
+        content
+            .lines()
+            .filter_map(|line| serde_json::from_str(line).ok())
+            .collect()
+    } else {
+        match serde_json::from_str(&content) {
+            Ok(serde_json::Value::Array(arr)) => arr,
+            Ok(obj) => vec![obj],
+            Err(e) => return Err(anyhow::anyhow!("Failed to parse JSON: {}", e)),
+        }
+    };
 
-    Ok(embedding)
+    // For each object, iterate fields and chunk
+    let mut indexed_embeddings = Vec::new();
+    let mut total_sub_chunks = 0;
+    for (obj_idx, obj) in objects.iter().enumerate() {
+        if let Some(map) = obj.as_object() {
+            for (field, value) in map.iter() {
+                let field_str = match value {
+                    serde_json::Value::String(s) => s.clone(),
+                    _ => value.to_string(),
+                };
+                // Split into sub-chunks ≤200 chars
+                let mut start = 0;
+                let chunk_len = field_str.chars().count();
+                let sub_chunk_size = 200;
+                let mut field_sub_chunks = 0;
+                while start < chunk_len {
+                    let end = (start + sub_chunk_size).min(chunk_len);
+                    let sub_chunk: String = field_str.chars().skip(start).take(end - start).collect();
+                    // Embed each sub-chunk
+                    let embedding = embed::embed(&sub_chunk, config.max_tokens, client)
+                        .await
+                        .with_context(|| format!("Failed to embed field '{}', chunk {}-{} in file: {}", field, start, end, path.display()))?;
+                    // Store embedding with metadata
+                    indexed_embeddings.push((obj_idx, field.clone(), start, end, embedding));
+                    start = end;
+                    field_sub_chunks += 1;
+                }
+                if field_sub_chunks > 0 {
+                    println!("[DEBUG] Field '{}' produced {} sub-chunks for file: {}", field, field_sub_chunks, path.display());
+                }
+                total_sub_chunks += field_sub_chunks;
+            }
+        }
+    }
+    println!("[DEBUG] Total sub-chunks embedded for file {}: {}", path.display(), total_sub_chunks);
+    println!("[DEBUG] Total embeddings indexed: {}", indexed_embeddings.len());
+    // Instead of returning a single embedding, return an error if none found
+    if indexed_embeddings.is_empty() {
+        return Err(anyhow::anyhow!("No embeddings generated for file: {}", path.display()));
+    }
+    // For compatibility, return the first embedding (could be refactored to index all)
+    Ok(indexed_embeddings[0].4.clone())
 }
 
 /// Processes a single file and adds it to the index.

src/ingest/marvelai_ingest.rs

@@ -286,55 +286,65 @@ async fn main() {
         let chunk_success = true;
         for obj in &chunk_json {
         for (field, chunk) in chunk_entity_fields(obj, max_embed_len) {
-            if debug {
-                println!("    [DEBUG] Field '{}' chunk size: {} chars", field, chunk.chars().count());
-                println!("    [DEBUG] Chunk content for '{}.{}':\n{}", field, char_name, chunk);
+            // Further split any chunk >200 chars into sub-chunks ≤200 chars
+            let mut sub_chunks = Vec::new();
+            let mut start = 0;
+            let chunk_len = chunk.chars().count();
+            let sub_chunk_size = 200;
+            while start < chunk_len {
+                let end = (start + sub_chunk_size).min(chunk_len);
+                let sub_chunk: String = chunk.chars().skip(start).take(end - start).collect();
+                sub_chunks.push(sub_chunk);
+                start = end;
             }
-            let mut success = false;
-            // Use a fresh client for each request to avoid connection reuse issues
-            let embed_client = reqwest::Client::builder()
-                .timeout(std::time::Duration::from_secs(30))
-                .build()
-                .expect("Failed to build reqwest client");
-            for attempt in 1..=max_retries {
-                println!("  Attempt {}/{}: embedding '{}.{}' ({} chars)", attempt, max_retries, field, char_name, chunk.chars().count());
-                let result = embed_with_config(&chunk, 100, &embed_client, embed_config.clone()).await;
-                match result {
-                    Ok(embedding) => {
-                        println!("  ✅ Embedding vectors received successfully (dim: {})", embedding.len());
-                        success = true;
-                        break;
-                    }
-                    Err(e) => {
-                        println!("  ❌ Embedding request failed for '{}.{}': {}", field, char_name, e);
-                        // Print full error chain if available
-                        let mut source = e.source();
-                        while let Some(s) = source {
-                            println!("      [Error source] {}", s);
-                            source = s.source();
+            for (sub_idx, sub_chunk) in sub_chunks.iter().enumerate() {
+                if debug {
+                    println!("    [DEBUG] Field '{}' sub-chunk {} size: {} chars", field, sub_idx + 1, sub_chunk.chars().count());
+                    println!("    [DEBUG] Sub-chunk content for '{}.{}[{}]':\n{}", field, char_name, sub_idx + 1, sub_chunk);
+                }
+                let mut success = false;
+                let embed_client = reqwest::Client::builder()
+                    .timeout(std::time::Duration::from_secs(120))
+                    .build()
+                    .expect("Failed to build reqwest client");
+                for attempt in 1..=max_retries {
+                    println!("  Attempt {}/{}: embedding '{}.{}[{}]' ({} chars)", attempt, max_retries, field, char_name, sub_idx + 1, sub_chunk.chars().count());
+                    let result = embed_with_config(sub_chunk, 100, &embed_client, embed_config.clone()).await;
+                    match result {
+                        Ok(embedding) => {
+                            println!("  ✅ Embedding vectors received successfully (dim: {})", embedding.len());
+                            success = true;
+                            break;
+                        }
+                        Err(e) => {
+                            println!("  ❌ Embedding request failed for '{}.{}[{}]': {}", field, char_name, sub_idx + 1, e);
+                            let mut source = e.source();
+                            while let Some(s) = source {
+                                println!("      [Error source] {}", s);
+                                source = s.source();
+                            }
+                            let log_path = marvel_dir.join("failed_chunks.log");
+                            let log_content = format!(
+                                "Failed embedding: field='{}', char_name='{}', sub_chunk_idx={}, chunk_size={}, attempt={}, error='{}', chunk='{}'\n",
+                                field, char_name, sub_idx + 1, sub_chunk.chars().count(), attempt, e, sub_chunk.replace('\n', " ")
+                            );
+                            fs::OpenOptions::new()
+                                .create(true)
+                                .append(true)
+                                .open(&log_path)
+                                .and_then(|mut f| std::io::Write::write_all(&mut f, log_content.as_bytes()))
+                                .unwrap_or_else(|err| println!("      [Log error] Could not write to failed_chunks.log: {}", err));
                         }
-                        // Log failed chunk metadata for retry
-                        let log_path = marvel_dir.join("failed_chunks.log");
-                        let log_content = format!(
-                            "Failed embedding: field='{}', char_name='{}', chunk_size={}, attempt={}, error='{}', chunk='{}'\n",
-                            field, char_name, chunk.chars().count(), attempt, e, chunk.replace('\n', " ")
-                        );
-                        fs::OpenOptions::new()
-                            .create(true)
-                            .append(true)
-                            .open(&log_path)
-                            .and_then(|mut f| std::io::Write::write_all(&mut f, log_content.as_bytes()))
-                            .unwrap_or_else(|err| println!("      [Log error] Could not write to failed_chunks.log: {}", err));
                     }
+                    let backoff = retry_delay * attempt;
+                    println!("      ⏳ Waiting {}s before next attempt...", backoff);
+                    tokio::time::sleep(std::time::Duration::from_secs(backoff as u64)).await;
+                }
+                if !success {
+                    println!("  ❌ Failed to generate embeddings for '{}.{}[{}]' after {} attempts. Aborting further chunk processing.", field, char_name, sub_idx + 1, max_retries);
+                    failed_count += 1;
+                    break;
                 }
-                let backoff = retry_delay * attempt;
-                println!("      ⏳ Waiting {}s before next attempt...", backoff);
-                tokio::time::sleep(std::time::Duration::from_secs(backoff as u64)).await;
-            }
-            if !success {
-                println!("  ❌ Failed to generate embeddings for '{}.{}' after {} attempts. Aborting further chunk processing.", field, char_name, max_retries);
-                failed_count += 1;
-                break;
             }
         }
         // ...existing code...