vectus 0.1.2

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

#![allow(unused)]
use ndarray::{Array1, Array2};
use std::env;

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};

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

pub struct Vectus {
    pub model: Model,
    pub embeddings: Array2<f64>,
    pub documents: Vec<Document>,
    storage_type: StorageType,
    hnsw: 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,
            embeddings: Array2::zeros((0, 0)),
            documents: Vec::new(),
            storage_type,
            hnsw: HNSW::new(initializer),
        }
    }

    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 mut result = self.hnsw.search(query_emb.clone(), self.hnsw.len(), k);

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

        relevant_docs
    }

    pub async fn add_documents(&mut self, docs: &Vec<Document>) -> Result<(), String> {
        let mut embeddings: Vec<Vec<f64>> = Vec::new();

        if docs.len() == 0 {
            return Err("No documents to add!".to_string());
        }

        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);
            self.documents.push(doc.clone());
            embeddings.push(embedding);
        }

        self.embeddings =
            Array2::from_shape_vec((docs.len(), embeddings[0].len()), embeddings.concat()).unwrap();

        Ok(())
    }

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