adk-rag 0.8.0

Retrieval-Augmented Generation for Rust Agent Development Kit (ADK-Rust) agents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

//! Property tests for in-memory vector store search ordering.

use std::collections::HashMap;

use adk_rag::document::Chunk;
use adk_rag::inmemory::InMemoryVectorStore;
use adk_rag::vectorstore::VectorStore;
use proptest::prelude::*;

/// Generate a non-zero L2-normalized embedding of the given dimension.
fn arb_normalized_embedding(dim: usize) -> impl Strategy<Value = Vec<f32>> {
    proptest::collection::vec(-1.0f32..1.0f32, dim).prop_filter_map(
        "non-zero embedding",
        |mut v| {
            let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm < 1e-8 {
                return None;
            }
            for val in &mut v {
                *val /= norm;
            }
            Some(v)
        },
    )
}

/// Generate a chunk with a normalized embedding.
fn arb_chunk(dim: usize) -> impl Strategy<Value = Chunk> {
    ("[a-z]{3,8}", "[a-z ]{5,30}", arb_normalized_embedding(dim)).prop_map(
        |(id, text, embedding)| Chunk {
            id,
            text,
            embedding,
            metadata: HashMap::new(),
            document_id: "doc_1".to_string(),
        },
    )
}

/// **Feature: adk-rag, Property 3: In-memory vector store search ordering**
/// *For any* set of chunks with embeddings stored in an InMemoryVectorStore,
/// searching with a query embedding SHALL return results ordered by descending
/// cosine similarity score, and the number of results SHALL be at most top_k.
/// **Validates: Requirements 2.3**
mod prop_inmemory_search_ordering {
    use super::*;

    const DIM: usize = 16;

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(100))]

        #[test]
        fn results_ordered_descending_and_bounded_by_top_k(
            chunks in proptest::collection::vec(arb_chunk(DIM), 1..20),
            query in arb_normalized_embedding(DIM),
            top_k in 1usize..25,
        ) {
            let rt = tokio::runtime::Runtime::new().unwrap();
            let results = rt.block_on(async {
                let store = InMemoryVectorStore::new();
                store.create_collection("test", DIM).await.unwrap();

                // Deduplicate chunks by id to avoid upsert overwriting
                let mut deduped: HashMap<String, Chunk> = HashMap::new();
                for chunk in &chunks {
                    deduped.entry(chunk.id.clone()).or_insert_with(|| chunk.clone());
                }
                let unique_chunks: Vec<Chunk> = deduped.into_values().collect();
                let count = unique_chunks.len();

                store.upsert("test", &unique_chunks).await.unwrap();
                let results = store.search("test", &query, top_k).await.unwrap();
                (results, count)
            });

            let (results, unique_count) = results;

            // Result count is at most top_k and at most the number of stored chunks
            prop_assert!(results.len() <= top_k);
            prop_assert!(results.len() <= unique_count);

            // Results are ordered by descending score
            for window in results.windows(2) {
                prop_assert!(
                    window[0].score >= window[1].score,
                    "results not in descending order: {} < {}",
                    window[0].score,
                    window[1].score,
                );
            }
        }
    }
}