langchain-rust 5.2.0

LangChain for Rust, the easiest way to write LLM-based programs in Rust
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

//! Offline tests against an in-memory SurrealDB (`kv-mem`, dev-dependency only).

use ::surrealdb::engine::local::{Db, Mem};
use ::surrealdb::Surreal;
use serde_json::{json, Value};

use super::{Store, StoreBuilder};
use crate::{
    schemas::Document,
    test_utils::FakeEmbedder,
    vectorstore::{VecStoreOptions, VectorStore},
};

fn embedder() -> FakeEmbedder {
    FakeEmbedder::new(
        3,
        &[
            ("cats purr", vec![1.0, 0.0, 0.0]),
            ("dogs bark", vec![0.0, 1.0, 0.0]),
            ("kittens purr", vec![0.9, 0.1, 0.0]),
            ("cat", vec![1.0, 0.0, 0.0]),
        ],
    )
}

async fn mem_db() -> Surreal<Db> {
    let db = Surreal::new::<Mem>(()).await.unwrap();
    db.use_ns("test").use_db("test").await.unwrap();
    db
}

fn docs() -> Vec<Document> {
    vec![
        Document::new("cats purr").with_metadata([("k".to_string(), json!("a"))].into()),
        Document::new("dogs bark").with_metadata([("k".to_string(), json!("b"))].into()),
        Document::new("kittens purr"),
    ]
}

async fn add_and_search(store: &Store<Db>) {
    let opts = VecStoreOptions::<Value>::default();
    let ids = store.add_documents(&docs(), &opts).await.unwrap();
    assert_eq!(ids.len(), 3);
    assert!(ids.iter().all(|id| !id.is_empty()));

    let results = store.similarity_search("cat", 2, &opts).await.unwrap();
    assert_eq!(results.len(), 2, "{results:?}");
    assert_eq!(results[0].page_content, "cats purr");
    assert!((results[0].score - 1.0).abs() < 1e-9);
    assert_eq!(results[0].metadata["k"], json!("a"));
    assert_eq!(results[1].page_content, "kittens purr");
    assert!(results[0].score >= results[1].score);

    // "dogs bark" is orthogonal to the query (cosine 0) and is cut by the threshold.
    let thresholded = store
        .similarity_search("cat", 10, &VecStoreOptions::new().with_score_threshold(0.5))
        .await
        .unwrap();
    assert_eq!(thresholded.len(), 2);
}

#[tokio::test]
async fn schemafull_shared_table_round_trip() {
    let store = StoreBuilder::new()
        .db(mem_db().await)
        .vector_dimensions(3)
        .embedder(embedder())
        .build()
        .await
        .unwrap();
    store.initialize().await.unwrap();

    add_and_search(&store).await;
}

#[tokio::test]
async fn shared_table_isolates_collections() {
    let db = mem_db().await;
    let build = |name: &str| {
        StoreBuilder::new()
            .db(db.clone())
            .collection_name(name)
            .vector_dimensions(3)
            .embedder(embedder())
            .build()
    };
    let a = build("a").await.unwrap();
    let b = build("b").await.unwrap();
    a.initialize().await.unwrap();

    let opts = VecStoreOptions::<Value>::default();
    a.add_documents(&[Document::new("cats purr")], &opts)
        .await
        .unwrap();
    b.add_documents(&[Document::new("dogs bark")], &opts)
        .await
        .unwrap();

    let results = a.similarity_search("cat", 10, &opts).await.unwrap();
    assert_eq!(results.len(), 1, "{results:?}");
    assert_eq!(results[0].page_content, "cats purr");
    assert_eq!(results[0].metadata["collection"], json!("a"));
}

#[tokio::test]
async fn compatibility_mode_table_per_collection() {
    let store = StoreBuilder::new_with_compatiblity()
        .db(mem_db().await)
        .embedder(embedder())
        .build()
        .await
        .unwrap();
    store.initialize().await.unwrap();

    add_and_search(&store).await;
}

#[tokio::test]
async fn schemafull_rejects_wrong_dimensions() {
    let store = StoreBuilder::new()
        .db(mem_db().await)
        .vector_dimensions(5)
        .embedder(embedder())
        .build()
        .await
        .unwrap();
    store.initialize().await.unwrap();

    let err = store
        .add_documents(&[Document::new("cats purr")], &VecStoreOptions::default())
        .await;
    assert!(err.is_err());
}

#[tokio::test]
async fn builder_requires_db_and_embedder() {
    let no_embedder = StoreBuilder::<Db>::new().db(mem_db().await).build().await;
    assert!(no_embedder.is_err());
    let no_db = StoreBuilder::<Db>::new().embedder(embedder()).build().await;
    assert!(no_db.is_err());
}