brainwires-storage 0.9.0

Backend-agnostic storage, tiered memory, and document management for the Brainwires Agent Framework
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

//! Unit tests for [`LanceDatabase`].

use super::database::LanceDatabase;
use crate::databases::traits::{StorageBackend, VectorDatabase};
use crate::databases::types::{FieldDef, FieldValue, Filter};
use tempfile::TempDir;

#[tokio::test]
async fn test_lance_database_new() {
    let temp = TempDir::new().unwrap();
    let db_path = temp.path().join("test.lance");
    let db = LanceDatabase::new(db_path.to_str().unwrap()).await.unwrap();
    assert_eq!(db.db_path(), db_path.to_str().unwrap());
}

#[tokio::test]
async fn test_lance_storage_backend_crud() {
    let temp = TempDir::new().unwrap();
    let db_path = temp.path().join("test.lance");
    let db = LanceDatabase::new(db_path.to_str().unwrap()).await.unwrap();

    let schema = vec![
        FieldDef::required("id", crate::databases::types::FieldType::Utf8),
        FieldDef::required("value", crate::databases::types::FieldType::Int64),
    ];
    db.ensure_table("test_table", &schema).await.unwrap();

    let records = vec![vec![
        ("id".to_string(), FieldValue::Utf8(Some("row1".to_string()))),
        ("value".to_string(), FieldValue::Int64(Some(42))),
    ]];
    db.insert("test_table", records).await.unwrap();

    let results = db.query("test_table", None, None).await.unwrap();
    assert_eq!(results.len(), 1);

    let count = db.count("test_table", None).await.unwrap();
    assert_eq!(count, 1);

    db.delete(
        "test_table",
        &Filter::Eq("id".into(), FieldValue::Utf8(Some("row1".into()))),
    )
    .await
    .unwrap();

    let count = db.count("test_table", None).await.unwrap();
    assert_eq!(count, 0);
}

#[tokio::test]
async fn test_lance_vector_search() {
    use crate::databases::types::FieldType;

    let temp = TempDir::new().unwrap();
    let db_path = temp.path().join("vec_search.lance");
    let db = LanceDatabase::new(db_path.to_str().unwrap()).await.unwrap();

    let dim = 4;
    let schema = vec![
        FieldDef::required("id", FieldType::Utf8),
        FieldDef::required("embedding", FieldType::Vector(dim)),
    ];
    db.ensure_table("vectors", &schema).await.unwrap();

    // Insert three records with different vectors.
    let records = vec![
        vec![
            ("id".to_string(), FieldValue::Utf8(Some("a".to_string()))),
            (
                "embedding".to_string(),
                FieldValue::Vector(vec![1.0, 0.0, 0.0, 0.0]),
            ),
        ],
        vec![
            ("id".to_string(), FieldValue::Utf8(Some("b".to_string()))),
            (
                "embedding".to_string(),
                FieldValue::Vector(vec![0.0, 1.0, 0.0, 0.0]),
            ),
        ],
        vec![
            ("id".to_string(), FieldValue::Utf8(Some("c".to_string()))),
            (
                "embedding".to_string(),
                FieldValue::Vector(vec![0.9, 0.1, 0.0, 0.0]),
            ),
        ],
    ];
    db.insert("vectors", records).await.unwrap();

    // Search for a vector closest to [1, 0, 0, 0] — should rank "a" first.
    let results = db
        .vector_search("vectors", "embedding", vec![1.0, 0.0, 0.0, 0.0], 3, None)
        .await
        .unwrap();

    assert!(!results.is_empty(), "vector_search should return results");
    // The first result should be "a" (exact match → distance 0 → highest score).
    let first_id = results[0]
        .record
        .iter()
        .find(|(n, _)| n == "id")
        .and_then(|(_, v)| v.as_str())
        .unwrap();
    assert_eq!(first_id, "a");

    // Scores should be in descending order.
    for w in results.windows(2) {
        assert!(
            w[0].score >= w[1].score,
            "scores should be descending: {} >= {}",
            w[0].score,
            w[1].score
        );
    }
}

#[tokio::test]
async fn test_lance_capabilities() {
    let temp = TempDir::new().unwrap();
    let db_path = temp.path().join("caps.lance");
    let db = LanceDatabase::new(db_path.to_str().unwrap()).await.unwrap();

    let caps = db.capabilities();
    assert!(
        caps.vector_search,
        "LanceDatabase should support vector search"
    );
}

#[tokio::test]
async fn test_lance_shared_connection() {
    use crate::databases::types::FieldType;

    let temp = TempDir::new().unwrap();
    let db_path = temp.path().join("shared.lance");
    let db = LanceDatabase::new(db_path.to_str().unwrap()).await.unwrap();

    // Use StorageBackend trait
    let schema = vec![FieldDef::required("name", FieldType::Utf8)];
    db.ensure_table("store_table", &schema).await.unwrap();
    let records = vec![vec![(
        "name".to_string(),
        FieldValue::Utf8(Some("test".to_string())),
    )]];
    db.insert("store_table", records).await.unwrap();

    // Use VectorDatabase trait on same instance
    db.initialize(4).await.unwrap();

    // Both should work on the same connection
    let store_count = db.count("store_table", None).await.unwrap();
    assert_eq!(store_count, 1);

    let stats = db.get_statistics().await.unwrap();
    assert_eq!(stats.total_vectors, 0);
}