chroma 0.14.0

//! Comprehensive Search API Example
//!
//! This example demonstrates the full functionality of the Chroma Rust client,
//! including schema configuration, data ingestion with sparse vectors, and
//! comprehensive search operations using the Search API.
//!
//! # What This Example Demonstrates
//!
//! 1. **Schema Configuration**
//!    - Creating a collection with a sparse vector index for hybrid search
//!    - Configuring the sparse index with BM25 mode enabled
//!
//! 2. **Data Ingestion**
//!    - Adding documents with dense embeddings, documents, and metadata
//!    - Generating sparse vectors using the built-in BM25 encoder
//!    - Batched insertion for efficient data loading
//!
//! 3. **Search Operations**
//!    - Basic KNN search with dense vectors
//!    - Sparse vector search (keyword-based)
//!    - Hybrid search combining dense and sparse (linear combination)
//!    - Hybrid search with RRF (Reciprocal Rank Fusion)
//!    - Filtered search with metadata conditions
//!    - Complex filters with AND/OR combinations
//!    - Group by with aggregation
//!    - Field selection
//!    - Pagination
//!    - Batch search (multiple queries in one request)
//!
//! # Sparse Vector Generation
//!
//! Sparse vectors can be generated by the user using any method they prefer.
//! This example uses the built-in BM25 sparse embedding function, which tokenizes
//! document text and computes BM25 scores for each token. The sparse vectors are
//! stored in metadata and indexed for efficient retrieval.
//!
//! # Running
//!
//! ```bash
//! cargo run --example collection_search
//! ```

use chroma::client::ChromaHttpClientOptions;
use chroma::embed::bm25::BM25SparseEmbeddingFunction;
use chroma::types::{
    rrf, Aggregate, GroupBy, Key, Metadata, MetadataValue, QueryVector, RankExpr, Schema,
    SearchPayload, SparseVectorIndexConfig,
};
use chroma::ChromaHttpClient;

const TOTAL_DOCS: usize = 256;
const BATCH_SIZE: usize = 128;
const KNN_LIMIT: u32 = 64;
const COLLECTION_NAME: &str = "comprehensive_search_example";

const CATEGORIES: [&str; 3] = ["machine_learning", "quantum_computing", "bioinformatics"];
const AUTHORS: [&str; 5] = [
    "Alice Chen",
    "Bob Smith",
    "Carol Johnson",
    "David Lee",
    "Emma Wilson",
];

/// Generate a realistic paper abstract for the given index.
fn generate_abstract(i: usize) -> String {
    let topics = [
        "deep neural networks for image classification",
        "quantum entanglement in superconducting qubits",
        "protein folding prediction using transformers",
        "reinforcement learning in autonomous systems",
        "quantum error correction codes",
        "genomic sequence alignment algorithms",
        "attention mechanisms in language models",
        "topological quantum computing",
        "single-cell RNA sequencing analysis",
    ];
    let methods = [
        "novel optimization techniques",
        "experimental validation",
        "theoretical framework",
        "large-scale benchmarks",
        "ablation studies",
    ];
    let results = [
        "significant improvements over baselines",
        "state-of-the-art performance",
        "promising preliminary results",
        "robust generalization",
    ];

    format!(
        "This paper investigates {} using {}. Our experiments demonstrate {}, \
         opening new directions for future research in this domain.",
        topics[i % topics.len()],
        methods[i % methods.len()],
        results[i % results.len()]
    )
}

/// Generate test data: IDs, embeddings, documents, and metadata.
#[allow(clippy::type_complexity)]
fn generate_test_data() -> (
    Vec<String>,
    Vec<Vec<f32>>,
    Vec<Option<String>>,
    Vec<Option<Metadata>>,
) {
    // BM25 encoder for generating sparse vectors from document text.
    // Note: Users can generate sparse vectors using any method they prefer.
    // The BM25 encoder is provided as a convenience for keyword-based search.
    let bm25 = BM25SparseEmbeddingFunction::default_murmur3_abs();

    let mut ids = Vec::with_capacity(TOTAL_DOCS);
    let mut embeddings = Vec::with_capacity(TOTAL_DOCS);
    let mut documents = Vec::with_capacity(TOTAL_DOCS);
    let mut metadatas = Vec::with_capacity(TOTAL_DOCS);

    for i in 0..TOTAL_DOCS {
        // ID
        ids.push(format!("paper_{:03}", i + 1));

        // Dense embedding: interpolate from [1,0,0] to [0,0,1] for predictable distances
        let t = i as f32 / (TOTAL_DOCS - 1) as f32;
        embeddings.push(vec![1.0 - t, 0.0, t]);

        // Document text
        let doc_text = generate_abstract(i);
        documents.push(Some(doc_text.clone()));

        // Generate sparse vector from document text using BM25
        let sparse_vector = bm25
            .encode(&doc_text)
            .expect("BM25 encoding should not fail");

        // Metadata
        let mut metadata = Metadata::new();
        metadata.insert("category".into(), CATEGORIES[i % 3].into());
        metadata.insert("year".into(), MetadataValue::Int(2020 + (i % 5) as i64));
        metadata.insert("author".into(), AUTHORS[i % 5].into());
        metadata.insert("citations".into(), MetadataValue::Int((i * 7 % 500) as i64));
        metadata.insert(
            "quality_score".into(),
            MetadataValue::Float(0.5 + (i % 50) as f64 / 100.0),
        );
        metadata.insert("peer_reviewed".into(), MetadataValue::Bool(i % 2 == 0));
        // Store sparse vector in metadata for hybrid search
        metadata.insert("sparse_embedding".into(), sparse_vector.into());

        metadatas.push(Some(metadata));
    }

    (ids, embeddings, documents, metadatas)
}

/// Print search results in a readable format.
fn print_results(
    title: &str,
    response: &chroma::types::SearchResponse,
    fields: &[&str],
    max_display: usize,
) {
    println!("\n{}", title);
    println!("{}", "=".repeat(title.len()));

    let total = response.ids[0].len();
    let display_count = total.min(max_display);

    for (i, id) in response.ids[0].iter().take(display_count).enumerate() {
        let score = response.scores[0]
            .as_ref()
            .and_then(|s| s.get(i))
            .and_then(|s| *s)
            .map(|s| format!("{:.4}", s))
            .unwrap_or_else(|| "N/A".to_string());

        let metadata = response.metadatas[0]
            .as_ref()
            .and_then(|m| m.get(i))
            .and_then(|m| m.as_ref());

        let field_values: Vec<String> = fields
            .iter()
            .filter_map(|f| {
                metadata
                    .and_then(|m| m.get(*f))
                    .map(|v| format!("{}={:?}", f, v))
            })
            .collect();

        let fields_str = if field_values.is_empty() {
            String::new()
        } else {
            format!(", {}", field_values.join(", "))
        };

        println!("  {}. {} (score={}{})", i + 1, id, score, fields_str);
    }

    if total > display_count {
        println!("  ... and {} more results", total - display_count);
    }
    println!("  Total: {} results", total);
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Comprehensive Chroma Search API Example ===\n");

    // 1. Connect to Chroma Cloud
    // Configure `ChromaHttpClientOptions` for different setups, such as
    // - Load config from environment
    // - Local connection
    let client = ChromaHttpClient::new(ChromaHttpClientOptions::cloud(
        "<chroma-api-key>",
        "<chroma-database-name>",
    )?);
    println!("Connected to Chroma");

    // 2. Delete existing collection if present
    println!("Deleting existing collection if present...");
    let _ = client.delete_collection(COLLECTION_NAME).await;

    // 3. Create schema with sparse vector index
    // The sparse vector index enables efficient keyword-based search.
    // Setting bm25=true indicates this index will store BM25 sparse embeddings.
    // For other sparse embedding, such as splade, this should be set to false.
    let schema = Schema::default().create_index(
        Some("sparse_embedding"),
        SparseVectorIndexConfig {
            embedding_function: None,
            source_key: None,
            bm25: Some(true), // Enable BM25 mode for this sparse index
        }
        .into(),
    )?;

    // 4. Create collection with schema
    let collection = client
        .get_or_create_collection(COLLECTION_NAME, Some(schema), None)
        .await?;
    let initial_version = collection.version();
    println!(
        "Created collection '{}' (version: {})",
        collection.name(),
        initial_version
    );

    // 5. Generate test data with BM25 sparse vectors
    println!("\nGenerating {} test documents...", TOTAL_DOCS);
    let (ids, embeddings, documents, metadatas) = generate_test_data();

    // 6. Insert data in batches
    println!("Inserting documents in batches of {}...", BATCH_SIZE);
    let num_batches = TOTAL_DOCS / BATCH_SIZE;
    for batch_idx in 0..num_batches {
        let start = batch_idx * BATCH_SIZE;
        let end = start + BATCH_SIZE;

        collection
            .add(
                ids[start..end].to_vec(),
                embeddings[start..end].to_vec(),
                Some(documents[start..end].to_vec()),
                None, // URIs not used in this example
                Some(metadatas[start..end].to_vec()),
            )
            .await?;

        println!("  Batch {}/{} inserted", batch_idx + 1, num_batches);
    }

    let count = collection.count().await?;
    println!("Collection ready with {} documents\n", count);

    // =========================================================================
    // SEARCH DEMONSTRATIONS
    // =========================================================================

    println!("\n>>> Starting Search Demonstrations <<<\n");

    // Query vectors for searches
    let dense_query = vec![1.0, 0.0, 0.0]; // Closest to paper_001
    let bm25 = BM25SparseEmbeddingFunction::default_murmur3_abs();
    let sparse_query = bm25
        .encode("deep learning neural networks image classification")
        .expect("BM25 encoding should not fail");

    // -------------------------------------------------------------------------
    // 1. Basic KNN Search with Dense Vectors
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([Key::Score, Key::field("category"), Key::field("year")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "1. Basic KNN Search (Dense Vectors)",
        &response,
        &["category", "year"],
        5,
    );

    // -------------------------------------------------------------------------
    // 2. Sparse Vector Search (Keyword-based)
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .rank(RankExpr::Knn {
            query: QueryVector::Sparse(sparse_query.clone()),
            key: Key::field("sparse_embedding"),
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([Key::Score, Key::field("category"), Key::Document]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "2. Sparse Vector Search (Keyword-based with BM25)",
        &response,
        &["category"],
        5,
    );

    // -------------------------------------------------------------------------
    // 3. Hybrid Search - Linear Combination
    // -------------------------------------------------------------------------
    let dense_knn = RankExpr::Knn {
        query: QueryVector::Dense(dense_query.clone()),
        key: Key::Embedding,
        limit: KNN_LIMIT,
        default: Some(10.0), // Default score for docs not in dense results
        return_rank: false,
    };

    let sparse_knn = RankExpr::Knn {
        query: QueryVector::Sparse(sparse_query.clone()),
        key: Key::field("sparse_embedding"),
        limit: KNN_LIMIT,
        default: Some(10.0), // Default score for docs not in sparse results
        return_rank: false,
    };

    // Weighted combination: 70% dense + 30% sparse
    let hybrid_rank = dense_knn * 0.7 + sparse_knn * 0.3;

    let search = SearchPayload::default()
        .rank(hybrid_rank)
        .limit(Some(5), 0)
        .select([Key::Score, Key::field("category"), Key::field("author")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "3. Hybrid Search (Linear Combination: 70% dense + 30% sparse)",
        &response,
        &["category", "author"],
        5,
    );

    // -------------------------------------------------------------------------
    // 4. Hybrid Search with RRF (Reciprocal Rank Fusion)
    // -------------------------------------------------------------------------
    let dense_knn_rank = RankExpr::Knn {
        query: QueryVector::Dense(dense_query.clone()),
        key: Key::Embedding,
        limit: KNN_LIMIT,
        default: None,
        return_rank: true, // Required for RRF
    };

    let sparse_knn_rank = RankExpr::Knn {
        query: QueryVector::Sparse(sparse_query.clone()),
        key: Key::field("sparse_embedding"),
        limit: KNN_LIMIT,
        default: None,
        return_rank: true, // Required for RRF
    };

    // RRF fusion with custom weights
    let rrf_rank = rrf(
        vec![dense_knn_rank, sparse_knn_rank],
        Some(60),             // k parameter
        Some(vec![0.7, 0.3]), // weights
        false,                // normalize
    )?;

    let search = SearchPayload::default()
        .rank(rrf_rank)
        .limit(Some(5), 0)
        .select([Key::Score, Key::field("category"), Key::field("citations")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "4. Hybrid Search (RRF Fusion)",
        &response,
        &["category", "citations"],
        5,
    );

    // -------------------------------------------------------------------------
    // 5. Metadata Filter
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .r#where(Key::field("category").eq("machine_learning"))
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([Key::Score, Key::field("category"), Key::field("year")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "5. Metadata Filter (category = 'machine_learning')",
        &response,
        &["category", "year"],
        5,
    );

    // -------------------------------------------------------------------------
    // 6. Full-Text Search (Document Contains)
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .r#where(Key::Document.contains("neural networks"))
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([Key::Score, Key::Document, Key::field("category")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "6. Full-Text Search (document contains 'neural networks')",
        &response,
        &["category"],
        5,
    );

    // -------------------------------------------------------------------------
    // 7. Regex Filter on Document
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .r#where(Key::Document.regex(r"quantum\s+\w+"))
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([Key::Score, Key::Document, Key::field("category")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        r"7. Regex Filter (document matches 'quantum\s+\w+')",
        &response,
        &["category"],
        5,
    );

    // -------------------------------------------------------------------------
    // 8. Complex Filter (Metadata + Full-Text Search)
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .r#where(
            (Key::field("year").gte(2022))
                & (Key::field("peer_reviewed").eq(true))
                & (Key::Document.contains("learning")),
        )
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 0)
        .select([
            Key::Score,
            Key::Document,
            Key::field("year"),
            Key::field("peer_reviewed"),
        ]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "8. Complex Filter (year >= 2022 AND peer_reviewed AND document contains 'learning')",
        &response,
        &["year", "peer_reviewed"],
        5,
    );

    // -------------------------------------------------------------------------
    // 9. Group By with Aggregation
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .group_by(GroupBy {
            keys: vec![Key::field("category")],
            aggregate: Some(Aggregate::MinK {
                keys: vec![Key::Score],
                k: 2, // Top 2 per category
            }),
        })
        .limit(Some(10), 0)
        .select([Key::Score, Key::field("category"), Key::field("author")]);

    let response = collection.search(vec![search]).await?;
    print_results(
        "9. Group By Category (Top 2 per category by score)",
        &response,
        &["category", "author"],
        10,
    );

    // -------------------------------------------------------------------------
    // 10. Field Selection
    // -------------------------------------------------------------------------
    let search = SearchPayload::default()
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(3), 0)
        .select([
            Key::Document,
            Key::Score,
            Key::field("author"),
            Key::field("quality_score"),
        ]);

    let response = collection.search(vec![search]).await?;
    println!("\n10. Field Selection (Document, Score, author, quality_score)");
    println!("=============================================================");
    for (i, id) in response.ids[0].iter().take(3).enumerate() {
        let doc = response.documents[0]
            .as_ref()
            .and_then(|d| d.get(i))
            .and_then(|d| d.as_ref())
            .map(|d| {
                if d.len() > 80 {
                    format!("{}...", &d[..80])
                } else {
                    d.clone()
                }
            })
            .unwrap_or_default();

        let score = response.scores[0]
            .as_ref()
            .and_then(|s| s.get(i))
            .and_then(|s| *s)
            .map(|s| format!("{:.4}", s))
            .unwrap_or_else(|| "N/A".to_string());

        let metadata = response.metadatas[0]
            .as_ref()
            .and_then(|m| m.get(i))
            .and_then(|m| m.as_ref());

        let author = metadata
            .and_then(|m| m.get("author"))
            .map(|v| format!("{:?}", v))
            .unwrap_or_default();

        let quality = metadata
            .and_then(|m| m.get("quality_score"))
            .map(|v| format!("{:?}", v))
            .unwrap_or_default();

        println!("  {}. {}", i + 1, id);
        println!("     Score: {}", score);
        println!("     Author: {}", author);
        println!("     Quality: {}", quality);
        println!("     Doc: {}", doc);
    }

    // -------------------------------------------------------------------------
    // 11. Pagination
    // -------------------------------------------------------------------------
    println!("\n11. Pagination (limit=5, offset=10)");
    println!("====================================");

    let search = SearchPayload::default()
        .rank(RankExpr::Knn {
            query: QueryVector::Dense(dense_query.clone()),
            key: Key::Embedding,
            limit: KNN_LIMIT,
            default: None,
            return_rank: false,
        })
        .limit(Some(5), 10) // Skip first 10, take 5
        .select([Key::Score, Key::field("category")]);

    let response = collection.search(vec![search]).await?;
    print_results("Results (offset=10)", &response, &["category"], 5);

    // -------------------------------------------------------------------------
    // 12. Batch Search (Multiple queries in one request)
    // -------------------------------------------------------------------------
    println!("\n12. Batch Search (3 queries in one request)");
    println!("============================================");

    let searches = vec![
        // Query 1: Machine learning papers
        SearchPayload::default()
            .r#where(Key::field("category").eq("machine_learning"))
            .rank(RankExpr::Knn {
                query: QueryVector::Dense(dense_query.clone()),
                key: Key::Embedding,
                limit: KNN_LIMIT,
                default: None,
                return_rank: false,
            })
            .limit(Some(3), 0)
            .select([Key::Score, Key::field("category")]),
        // Query 2: Quantum computing papers
        SearchPayload::default()
            .r#where(Key::field("category").eq("quantum_computing"))
            .rank(RankExpr::Knn {
                query: QueryVector::Dense(dense_query.clone()),
                key: Key::Embedding,
                limit: KNN_LIMIT,
                default: None,
                return_rank: false,
            })
            .limit(Some(3), 0)
            .select([Key::Score, Key::field("category")]),
        // Query 3: Bioinformatics papers
        SearchPayload::default()
            .r#where(Key::field("category").eq("bioinformatics"))
            .rank(RankExpr::Knn {
                query: QueryVector::Dense(dense_query.clone()),
                key: Key::Embedding,
                limit: KNN_LIMIT,
                default: None,
                return_rank: false,
            })
            .limit(Some(3), 0)
            .select([Key::Score, Key::field("category")]),
    ];

    let response = collection.search(searches).await?;

    for (query_idx, category) in ["machine_learning", "quantum_computing", "bioinformatics"]
        .iter()
        .enumerate()
    {
        println!("\n  Query {}: {} papers", query_idx + 1, category);
        for (i, id) in response.ids[query_idx].iter().enumerate() {
            let score = response.scores[query_idx]
                .as_ref()
                .and_then(|s| s.get(i))
                .and_then(|s| *s)
                .map(|s| format!("{:.4}", s))
                .unwrap_or_else(|| "N/A".to_string());
            println!("    {}. {} (score={})", i + 1, id, score);
        }
    }

    // =========================================================================
    // CLEANUP
    // =========================================================================

    println!("\n\n>>> Cleanup <<<");
    client.delete_collection(COLLECTION_NAME).await?;
    println!("Deleted collection '{}'", COLLECTION_NAME);

    println!("\n=== Example Complete ===");

    Ok(())
}