vectus 0.1.21-experimental

#![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};
use std::sync::{Arc, Mutex, RwLock};

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

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

unsafe impl Send for Vectus {}
unsafe impl Sync for Vectus {}

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: RwLock::new(Array2::zeros((0, 0))).into(),
            documents: RwLock::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().unwrap();
        let result = hnsw_guard.search(query_emb.clone(), hnsw_guard.len(), k);
        drop(hnsw_guard);

        let docs_guard = self.documents.read().unwrap();
        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);
            embeddings.push(embedding);
        }

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

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

        Ok(())
    }

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