abu-rag 0.2.0

Retrieval-Augmented Generation system
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

mod flat;
pub use flat::*;
mod lsh;
pub use lsh::*;
mod ivf;
pub use ivf::*;

use super::{Float, VectorId};

#[derive(Clone, Debug)]
pub struct ScoredId<F: Float> {
    pub id: VectorId,
    pub score: F,
}

#[allow(async_fn_in_trait)]
pub trait VectorIndex: Send + Sync {
    type F: Float;
    type Error: std::error::Error + 'static + Send + Sync;

    async fn add(&mut self, id: VectorId, vector: Vec<Self::F>) -> Result<(), Self::Error>;
    async fn add_batch<IS, VS>(&mut self, ids: IS, vectors: VS) -> Result<(), Self::Error>
    where 
        IS: IntoIterator<Item = VectorId>,
        VS: IntoIterator<Item = Vec<Self::F>>,
    {
        for (id, vector) in ids.into_iter().zip(vectors.into_iter()) {
            self.add(id, vector).await?;
        }
        Ok(())
    }

    async fn delete(&mut self, id: VectorId) -> Result<(), Self::Error>;
    async fn delete_batch<IS>(&mut self, ids: IS) -> Result<(), Self::Error> 
    where 
        IS: IntoIterator<Item = VectorId>,
    {
        for id in ids {
            self.delete(id).await?;
        }
        Ok(())
    }

    async fn search(&self, query: &[Self::F], top_k: usize) -> Result<Vec<ScoredId<Self::F>>, Self::Error>;

    async fn clear(&mut self) -> Result<(), Self::Error>;

    fn len(&self) -> usize;
}

impl<F: Float> ScoredId<F> {
    pub fn new(id: VectorId, score: F) -> Self {
        Self { id, score }
    }
}


// #[derive(Clone, Debug)]
// pub struct ScoredId {
//     pub id: VectorId,
//     pub score: f32,
// }

// #[allow(async_fn_in_trait)]
// pub trait VectorIndex: Send + Sync {
//     type Error: std::error::Error + 'static + Send + Sync;

//     async fn add(&mut self, id: VectorId, vector: Vec<f32>) -> Result<(), Self::Error>;
//     async fn add_batch<IS, VS>(&mut self, ids: IS, vectors: VS) -> Result<(), Self::Error>
//     where 
//         IS: IntoIterator<Item = VectorId>,
//         VS: IntoIterator<Item = Vec<f32>>,
//     {
//         for (id, vector) in ids.into_iter().zip(vectors.into_iter()) {
//             self.add(id, vector).await?;
//         }
//         Ok(())
//     }

//     async fn delete(&mut self, id: VectorId) -> Result<(), Self::Error>;
//     async fn delete_batch<IS>(&mut self, ids: IS) -> Result<(), Self::Error> 
//     where 
//         IS: IntoIterator<Item = VectorId>,
//     {
//         for id in ids {
//             self.delete(id).await?;
//         }
//         Ok(())
//     }

//     async fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<ScoredId>, Self::Error>;

//     async fn clear(&mut self) -> Result<(), Self::Error>;

//     fn len(&self) -> usize;
// }

// impl ScoredId {
//     pub fn new(id: VectorId, score: f32) -> Self {
//         Self { id, score }
//     }
// }