use anyhow::Result;
use arrow_array::RecordBatch;
use super::{CodeBlock, CommitBlock, DocumentBlock, TextBlock};
pub trait BlockType: Clone + Send + Sync + 'static {
const TABLE_NAME: &'static str;
fn to_batch(blocks: &[Self], embeddings: &[Vec<f32>], vector_dim: usize)
-> Result<RecordBatch>;
fn from_batch(batch: &RecordBatch) -> Result<Vec<Self>>;
fn distance(&self) -> Option<f32>;
fn set_distance(&mut self, distance: f32);
fn get_hash(&self) -> String;
}
impl BlockType for CodeBlock {
const TABLE_NAME: &'static str = "code_blocks";
fn to_batch(
blocks: &[Self],
embeddings: &[Vec<f32>],
vector_dim: usize,
) -> Result<RecordBatch> {
super::batch_converter::BatchConverter::new(vector_dim)
.code_block_to_batch(blocks, embeddings)
}
fn from_batch(batch: &RecordBatch) -> Result<Vec<Self>> {
super::batch_converter::BatchConverter::new(0).batch_to_code_blocks(batch, None)
}
fn distance(&self) -> Option<f32> {
self.distance
}
fn set_distance(&mut self, distance: f32) {
self.distance = Some(distance);
}
fn get_hash(&self) -> String {
self.hash.clone()
}
}
impl BlockType for TextBlock {
const TABLE_NAME: &'static str = "text_blocks";
fn to_batch(
blocks: &[Self],
embeddings: &[Vec<f32>],
vector_dim: usize,
) -> Result<RecordBatch> {
super::batch_converter::BatchConverter::new(vector_dim)
.text_block_to_batch(blocks, embeddings)
}
fn from_batch(batch: &RecordBatch) -> Result<Vec<Self>> {
super::batch_converter::BatchConverter::new(0).batch_to_text_blocks(batch, None)
}
fn distance(&self) -> Option<f32> {
self.distance
}
fn set_distance(&mut self, distance: f32) {
self.distance = Some(distance);
}
fn get_hash(&self) -> String {
self.hash.clone()
}
}
impl BlockType for DocumentBlock {
const TABLE_NAME: &'static str = "document_blocks";
fn to_batch(
blocks: &[Self],
embeddings: &[Vec<f32>],
vector_dim: usize,
) -> Result<RecordBatch> {
super::batch_converter::BatchConverter::new(vector_dim)
.document_block_to_batch(blocks, embeddings)
}
fn from_batch(batch: &RecordBatch) -> Result<Vec<Self>> {
super::batch_converter::BatchConverter::new(0).batch_to_document_blocks(batch, None)
}
fn distance(&self) -> Option<f32> {
self.distance
}
fn set_distance(&mut self, distance: f32) {
self.distance = Some(distance);
}
fn get_hash(&self) -> String {
self.hash.clone()
}
}
impl BlockType for CommitBlock {
const TABLE_NAME: &'static str = "commit_blocks";
fn to_batch(
blocks: &[Self],
embeddings: &[Vec<f32>],
vector_dim: usize,
) -> Result<RecordBatch> {
super::batch_converter::BatchConverter::new(vector_dim)
.commit_block_to_batch(blocks, embeddings)
}
fn from_batch(batch: &RecordBatch) -> Result<Vec<Self>> {
super::batch_converter::BatchConverter::new(0).batch_to_commit_blocks(batch)
}
fn distance(&self) -> Option<f32> {
self.distance
}
fn set_distance(&mut self, distance: f32) {
self.distance = Some(distance);
}
fn get_hash(&self) -> String {
self.hash.clone()
}
}