vectus 0.1.24

A vector database implemented in Rust for learning purposes.
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

#![allow(unused)]

pub mod document;
pub(crate) mod hnsw;
pub mod model;
pub use hnsw::metric;

use document::{DocBuilder, Document};
use hnsw::{metric::Metric, HNSWInitializer, HNSW};
use model::{Model, ModelType};
use ndarray::{Array1, Array2};
use std::env;
use std::sync::Arc;
use tokio::sync::Mutex;

#[derive(Debug)]
pub enum StorageType {
    InMemory,
    Persistent,
}

pub struct Vectus {
    pub model: Arc<Model>,
    pub embeddings: Arc<Mutex<Array2<f64>>>,
    pub documents: Arc<Mutex<Vec<Document>>>,
    storage_type: StorageType,
    hnsw: Arc<Mutex<HNSW>>,
}

impl Vectus {
    pub fn new(model_name: ModelType, storage_type: StorageType, metric: Metric) -> Vectus {
        let model = Model::new(
            ModelType::OpenAI,
            env::var("OPENAI_API_KEY").expect("Please set the OPENAI_API_KEY environment variable"),
        );

        let initializer = HNSWInitializer {
            max_level: 12,
            ef_construction: 350,
            m: 32,
            m_max: 64,
            norm: 3.0,
            entry: None,
            metric,
        };

        Vectus {
            model: model.into(),
            embeddings: Mutex::new(Array2::zeros((0, 0))).into(),
            documents: Mutex::new(Vec::new()).into(),
            storage_type,
            hnsw: Mutex::new(HNSW::new(initializer)).into(),
        }
    }

    pub async fn get_k_relevant_documents(&self, query: &String, k: usize) -> Vec<Document> {
        let query_embedding = match self.model.get_embedding(query).await {
            Ok(embedding) => embedding,
            Err(e) => panic!("Error getting embedding: {}", e),
        };

        let query_emb = Array1::from_vec(query_embedding.clone());
        let hnsw_guard = self.hnsw.lock().await;
        let result = hnsw_guard.search(query_emb.clone(), hnsw_guard.len(), k);
        drop(hnsw_guard);

        let docs_guard = self.documents.lock().await;
        let mut relevant_docs: Vec<Document> = Vec::new();
        for i in 0..k {
            relevant_docs.push(docs_guard[result[i]].clone());
        }

        relevant_docs
    }

    pub async fn add_documents(&mut self, docs: &Vec<Document>) -> Result<(), String> {
        if docs.is_empty() {
            return Err("No documents to add!".to_string());
        }

        let mut embeddings: Vec<Vec<f64>> = Vec::new();

        for doc in docs {
            let embedding: Vec<f64> = match self.model.get_embedding(&doc.page_content).await {
                Ok(embedding) => embedding,
                Err(e) => panic!("Error getting embedding: {}", e),
            };

            let nembd = Array1::from_vec(embedding.clone());
            self.store_emb_db(&nembd).await;
            embeddings.push(embedding);
        }

        let mut docs_guard = self.documents.lock().await;
        for doc in docs {
            docs_guard.push(doc.clone());
        }
        drop(docs_guard);

        let mut embeddings_guard = self.embeddings.lock().await;
        *embeddings_guard =
            Array2::from_shape_vec((docs.len(), embeddings[0].len()), embeddings.concat()).unwrap();

        Ok(())
    }

    async fn store_emb_db(&self, embedding: &Array1<f64>) {
        match self.storage_type {
            StorageType::InMemory => {
                let mut hnsw_guard = self.hnsw.lock().await;
                let len = hnsw_guard.len();
                hnsw_guard.insert(&embedding, len);
                drop(hnsw_guard);
            }
            StorageType::Persistent => {
                panic!("{:?} Not implemented yet!", self.storage_type);
            }
        }
    }
}