src/core.rs - master - nicktretyakov1/vectorDBML

src/core.rs
55 строк2 KB
nicktretyakov
update4
19 янв 2026, 20:08
19 янв 2026, 20:083f2d861
// src/core.rs (исправлен импорт NearestNeighborSearcher)
use burn::{prelude::*, tensor::Tensor};

use crate::{
    config::VectorDBConfig,
    metrics,
    processors::{IdentityProcessor, L2Normalizer, VectorProcessor},
    storage::traits::NearestNeighborSearcher,  // Правильный путь через traits
};

pub struct VectorDB<B: Backend> {
    repository: Box<dyn NearestNeighborSearcher<B>>,
    processor: Box<dyn VectorProcessor<B>>,
    default_computer: Box<dyn metrics::DistanceComputer<B>>,
    dim: usize,
}

impl<B: Backend> VectorDB<B> {
    pub(crate) fn new(config: VectorDBConfig, device: &Device<B>) -> Self {
        let repository: Box<dyn NearestNeighborSearcher<B>> =
            Box::new(storage::InMemoryRepository::new(device, config.dim));

        let processor: Box<dyn VectorProcessor<B>> = if config.use_normalization {
            Box::new(L2Normalizer { epsilon: config.normalization_epsilon })
        } else {
            Box::new(IdentityProcessor)
        };

        let default_computer = metrics::create_computer(config.default_metric);

        Self {
            repository,
            processor,
            default_computer,
            dim: config.dim,
        }
    }

    pub fn add_vector(&mut self, vector: Tensor<B, 1>) -> usize {
        let processed = self.processor.process(vector);
        self.repository.add(processed)
    }

    pub fn get_by_index(&self, id: usize) -> Option<Tensor<B, 1>> {
        self.repository.get(id)
    }

    pub fn search_nearest(&self, query: Tensor<B, 1>, k: usize) -> Vec<(usize, f64)> {
        self.repository.search_nearest(query, k, self.default_computer.as_ref())
    }

    pub fn set_metric(&mut self, computer: Box<dyn metrics::DistanceComputer<B>>) {
        self.default_computer = computer;
    }
}