kreuzberg 4.0.0.pre.rc.6

data/vendor/kreuzberg/tests/test_fastembed.rs

@@ -0,0 +1,609 @@
+//! Integration tests for fastembed embeddings
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_fastembed_initialization() {
+    use fastembed::TextEmbedding;
+
+    let model = TextEmbedding::try_new(Default::default());
+
+    assert!(model.is_ok(), "Failed to initialize fastembed model: {:?}", model.err());
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_fastembed_embedding_generation() {
+    use fastembed::TextEmbedding;
+
+    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");
+
+    let texts = vec![
+        "Hello world, this is a test.",
+        "Fastembed is a Rust embedding library.",
+        "Testing embedding generation stability.",
+    ];
+
+    let result = model.embed(texts.clone(), None);
+    assert!(result.is_ok(), "Failed to generate embeddings: {:?}", result.err());
+
+    let embeddings = result.unwrap();
+    assert_eq!(embeddings.len(), 3, "Expected 3 embeddings");
+
+    for (i, embedding) in embeddings.iter().enumerate() {
+        assert_eq!(embedding.len(), 384, "Embedding {} has wrong dimensions", i);
+
+        let sum: f32 = embedding.iter().sum();
+        assert!(sum.abs() > 0.0001, "Embedding {} appears to be all zeros", i);
+    }
+
+    println!(
+        "✓ Successfully generated {} embeddings with 384 dimensions each",
+        embeddings.len()
+    );
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_fastembed_batch_processing() {
+    use fastembed::TextEmbedding;
+
+    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");
+
+    let texts: Vec<String> = (0..50)
+        .map(|i| {
+            format!(
+                "This is test sentence number {}. It contains some text for embedding.",
+                i
+            )
+        })
+        .collect();
+
+    let start = std::time::Instant::now();
+    let result = model.embed(texts.clone(), Some(32));
+    let duration = start.elapsed();
+
+    assert!(result.is_ok(), "Batch embedding failed: {:?}", result.err());
+
+    let embeddings = result.unwrap();
+    assert_eq!(embeddings.len(), 50, "Expected 50 embeddings");
+
+    println!(
+        "✓ Batch processed 50 texts in {:?} ({:.2} ms per text)",
+        duration,
+        duration.as_millis() as f64 / 50.0
+    );
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_fastembed_different_models() {
+    use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
+
+    let models_to_test = vec![
+        (EmbeddingModel::AllMiniLML6V2Q, 384, "AllMiniLML6V2Q (fast, quantized)"),
+        (EmbeddingModel::BGEBaseENV15, 768, "BGEBaseENV15 (balanced)"),
+    ];
+
+    let test_text = vec!["Hello world"];
+
+    for (model_type, expected_dims, description) in models_to_test {
+        println!("Testing {}", description);
+
+        let model = TextEmbedding::try_new(InitOptions::new(model_type));
+
+        match model {
+            Ok(mut m) => {
+                let result = m.embed(test_text.clone(), None);
+                assert!(result.is_ok(), "Failed to generate embedding for {}", description);
+
+                let embeddings = result.unwrap();
+                assert_eq!(embeddings.len(), 1);
+                assert_eq!(
+                    embeddings[0].len(),
+                    expected_dims,
+                    "Wrong dimensions for {}",
+                    description
+                );
+
+                println!("  ✓ {} produces {}-dim embeddings", description, expected_dims);
+            }
+            Err(e) => {
+                println!("  ⚠ Failed to initialize {}: {:?}", description, e);
+            }
+        }
+    }
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_fastembed_error_handling() {
+    use fastembed::TextEmbedding;
+
+    let mut model = TextEmbedding::try_new(Default::default()).expect("Failed to initialize model");
+
+    let empty_texts: Vec<String> = vec![];
+    let result = model.embed(empty_texts, None);
+
+    match result {
+        Ok(embeddings) => assert_eq!(embeddings.len(), 0, "Empty input should produce empty output"),
+        Err(_) => {
+            println!("  ℹ fastembed returns error for empty input (acceptable)");
+        }
+    }
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_basic() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks = vec![
+        Chunk {
+            content: "Hello world, this is the first chunk.".to_string(),
+            embedding: None,
+            metadata: ChunkMetadata {
+                char_start: 0,
+                char_end: 38,
+                chunk_index: 0,
+                total_chunks: 1,
+                token_count: None,
+            },
+        },
+        Chunk {
+            content: "This is the second chunk with different content.".to_string(),
+            embedding: None,
+            metadata: ChunkMetadata {
+                char_start: 39,
+                char_end: 87,
+                chunk_index: 1,
+                total_chunks: 1,
+                token_count: None,
+            },
+        },
+        Chunk {
+            content: "And this is the third and final chunk.".to_string(),
+            embedding: None,
+            metadata: ChunkMetadata {
+                char_start: 88,
+                char_end: 126,
+                chunk_index: 2,
+                total_chunks: 1,
+                token_count: None,
+            },
+        },
+    ];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunks, &config);
+    assert!(result.is_ok(), "Failed to generate embeddings: {:?}", result.err());
+
+    for (i, chunk) in chunks.iter().enumerate() {
+        assert!(chunk.embedding.is_some(), "Chunk {} missing embedding", i);
+
+        let embedding = chunk.embedding.as_ref().unwrap();
+        assert_eq!(embedding.len(), 384, "Chunk {} has wrong embedding dimensions", i);
+
+        let sum: f32 = embedding.iter().sum();
+        assert!(sum.abs() > 0.0001, "Chunk {} embedding appears to be all zeros", i);
+    }
+
+    println!("✓ Generated embeddings for {} chunks", chunks.len());
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_normalization() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let test_text = "This is a test sentence for normalization testing.";
+
+    let mut chunks_no_norm = vec![Chunk {
+        content: test_text.to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: test_text.len(),
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config_no_norm = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    generate_embeddings_for_chunks(&mut chunks_no_norm, &config_no_norm)
+        .expect("Failed to generate non-normalized embeddings");
+
+    let mut chunks_norm = vec![Chunk {
+        content: test_text.to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: test_text.len(),
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config_norm = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: true,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    generate_embeddings_for_chunks(&mut chunks_norm, &config_norm).expect("Failed to generate normalized embeddings");
+
+    let embedding_no_norm = chunks_no_norm[0].embedding.as_ref().unwrap();
+    let embedding_norm = chunks_norm[0].embedding.as_ref().unwrap();
+
+    let magnitude_no_norm: f32 = embedding_no_norm.iter().map(|x| x * x).sum::<f32>().sqrt();
+    let magnitude_norm: f32 = embedding_norm.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+    assert!(
+        magnitude_no_norm > 0.9 && magnitude_no_norm < 1.1,
+        "Embedding magnitude should be reasonable (got {})",
+        magnitude_no_norm
+    );
+
+    assert!(
+        (magnitude_norm - 1.0).abs() < 0.001,
+        "Normalized embedding should have unit magnitude (got {})",
+        magnitude_norm
+    );
+
+    println!(
+        "✓ Normalization works: magnitude before={:.4}, after={:.4}",
+        magnitude_no_norm, magnitude_norm
+    );
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_empty_input() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::Chunk;
+
+    let mut empty_chunks: Vec<Chunk> = vec![];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut empty_chunks, &config);
+    assert!(result.is_ok(), "Empty input should be handled gracefully");
+
+    println!("✓ Empty input handled correctly");
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_model_caching() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks1 = vec![Chunk {
+        content: "First batch of text.".to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: 20,
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let start1 = std::time::Instant::now();
+    generate_embeddings_for_chunks(&mut chunks1, &config).expect("First call failed");
+    let duration1 = start1.elapsed();
+
+    let mut chunks2 = vec![Chunk {
+        content: "Second batch of text.".to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: 21,
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let start2 = std::time::Instant::now();
+    generate_embeddings_for_chunks(&mut chunks2, &config).expect("Second call failed");
+    let duration2 = start2.elapsed();
+
+    println!(
+        "✓ Model caching works: first call={:?}, second call={:?}",
+        duration1, duration2
+    );
+
+    assert!(chunks1[0].embedding.is_some());
+    assert!(chunks2[0].embedding.is_some());
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_invalid_preset() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks = vec![Chunk {
+        content: "Test content".to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: 12,
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "nonexistent_preset".to_string(),
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunks, &config);
+    assert!(result.is_err(), "Should return error for unknown preset");
+
+    let error = result.unwrap_err();
+    let error_msg = format!("{}", error);
+    assert!(
+        error_msg.contains("Unknown embedding preset"),
+        "Error should mention unknown preset, got: {}",
+        error_msg
+    );
+
+    println!("✓ Invalid preset error handling works");
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_unknown_model() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks = vec![Chunk {
+        content: "Test content".to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: 12,
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::FastEmbed {
+            model: "UnknownModelXYZ123".to_string(),
+            dimensions: 384,
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunks, &config);
+    assert!(result.is_err(), "Should return error for unknown model");
+
+    let error = result.unwrap_err();
+    let error_msg = format!("{}", error);
+    assert!(
+        error_msg.contains("Unknown fastembed model"),
+        "Error should mention unknown model, got: {}",
+        error_msg
+    );
+
+    println!("✓ Unknown model error handling works");
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_custom_model_not_supported() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks = vec![Chunk {
+        content: "Test content".to_string(),
+        embedding: None,
+        metadata: ChunkMetadata {
+            char_start: 0,
+            char_end: 12,
+            chunk_index: 0,
+            total_chunks: 1,
+            token_count: None,
+        },
+    }];
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Custom {
+            model_id: "hf://custom/model".to_string(),
+            dimensions: 768,
+        },
+        batch_size: 32,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunks, &config);
+    assert!(result.is_err(), "Should return error for custom models");
+
+    let error = result.unwrap_err();
+    let error_msg = format!("{}", error);
+    assert!(
+        error_msg.contains("Custom ONNX models are not yet supported"),
+        "Error should mention custom models not supported, got: {}",
+        error_msg
+    );
+
+    println!("✓ Custom model error handling works");
+}
+
+#[cfg(feature = "embeddings")]
+#[tokio::test]
+async fn test_generate_embeddings_for_chunks_batch_size() {
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+    use kreuzberg::types::{Chunk, ChunkMetadata};
+
+    let mut chunks: Vec<Chunk> = (0..10)
+        .map(|i| Chunk {
+            content: format!("This is test chunk number {}.", i),
+            embedding: None,
+            metadata: ChunkMetadata {
+                char_start: i * 30,
+                char_end: (i + 1) * 30,
+                chunk_index: i,
+                total_chunks: 10,
+                token_count: None,
+            },
+        })
+        .collect();
+
+    let config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 10,
+        normalize: false,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunks, &config);
+    assert!(result.is_ok(), "Processing failed: {:?}", result.err());
+
+    for (i, chunk) in chunks.iter().enumerate() {
+        assert!(
+            chunk.embedding.is_some(),
+            "Chunk {} missing embedding after batch processing",
+            i
+        );
+        assert_eq!(
+            chunk.embedding.as_ref().unwrap().len(),
+            384,
+            "Chunk {} has wrong dimensions",
+            i
+        );
+    }
+
+    println!("✓ Processing with batch_size=10 works for 10 chunks");
+}
+
+#[cfg(all(feature = "embeddings", feature = "chunking"))]
+#[tokio::test]
+async fn test_generate_embeddings_chunking_integration() {
+    use kreuzberg::chunking::{ChunkingConfig, chunk_text};
+    use kreuzberg::core::config::{EmbeddingConfig, EmbeddingModelType};
+    use kreuzberg::embeddings::generate_embeddings_for_chunks;
+
+    let text = "This is a test document. It has multiple sentences. \
+                Each sentence should be chunked appropriately. \
+                The chunking system should create overlapping chunks. \
+                Finally, we will generate embeddings for each chunk.";
+
+    let chunking_config = ChunkingConfig {
+        max_characters: 50,
+        overlap: 10,
+        ..Default::default()
+    };
+
+    let mut chunking_result = chunk_text(text, &chunking_config).expect("Chunking failed");
+
+    assert!(
+        chunking_result.chunks.len() > 1,
+        "Should create multiple chunks from text"
+    );
+    println!("✓ Created {} chunks from text", chunking_result.chunks.len());
+
+    let embedding_config = EmbeddingConfig {
+        model: EmbeddingModelType::Preset {
+            name: "fast".to_string(),
+        },
+        batch_size: 32,
+        normalize: true,
+        show_download_progress: false,
+        cache_dir: None,
+    };
+
+    let result = generate_embeddings_for_chunks(&mut chunking_result.chunks, &embedding_config);
+    assert!(result.is_ok(), "Embedding generation failed: {:?}", result.err());
+
+    for (i, chunk) in chunking_result.chunks.iter().enumerate() {
+        assert!(chunk.embedding.is_some(), "Chunk {} missing embedding", i);
+
+        let embedding = chunk.embedding.as_ref().unwrap();
+        assert_eq!(embedding.len(), 384, "Chunk {} has wrong embedding dimensions", i);
+
+        let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
+        assert!(
+            (magnitude - 1.0).abs() < 0.001,
+            "Chunk {} embedding not normalized (magnitude={})",
+            i,
+            magnitude
+        );
+    }
+
+    println!(
+        "✓ Chunking + Embedding integration works: {} chunks with normalized embeddings",
+        chunking_result.chunks.len()
+    );
+}