use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
pub fn documents_schema(dim: u16) -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("scope", DataType::Utf8, false),
Field::new("path", DataType::Utf8, false),
Field::new("chunk_idx", DataType::UInt32, false),
Field::new("mime_type", DataType::Utf8, false),
Field::new("text", DataType::Utf8, false),
Field::new("byte_start", DataType::UInt32, false),
Field::new("byte_end", DataType::UInt32, false),
Field::new(
"embedding",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
i32::from(dim),
),
false,
),
]))
}
pub fn memory_schema(dim: u16) -> SchemaRef {
let tags_inner = Arc::new(Field::new("item", DataType::Utf8, true));
Arc::new(Schema::new(vec![
Field::new("scope", DataType::Utf8, false),
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Utf8, false),
Field::new("tags", DataType::List(tags_inner), true),
Field::new(
"embedding",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
i32::from(dim),
),
false,
),
Field::new(
"created_at",
DataType::Timestamp(TimeUnit::Microsecond, None),
false,
),
Field::new(
"updated_at",
DataType::Timestamp(TimeUnit::Microsecond, None),
false,
),
]))
}
pub const DOCUMENTS_TABLE: &str = "documents";
pub const MEMORY_TABLE: &str = "memory";