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("visibility", DataType::Utf8, false),
Field::new("agent_id", DataType::Utf8, false),
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),
]))
}
#[cfg(feature = "code-search")]
pub fn code_chunks_schema(dim: u16) -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("scope", DataType::Utf8, false),
Field::new("path", DataType::Utf8, false),
Field::new("chunk_id", DataType::Utf8, false),
Field::new("symbol", DataType::Utf8, false),
Field::new("kind", DataType::Utf8, false),
Field::new("lang", DataType::Utf8, false),
Field::new("line_start", DataType::UInt32, false),
Field::new("line_end", DataType::UInt32, false),
Field::new("byte_start", DataType::UInt32, false),
Field::new("byte_end", DataType::UInt32, false),
Field::new("text", DataType::Utf8, false),
Field::new(
"embedding",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), i32::from(dim)),
false,
),
]))
}
pub const DOCUMENTS_TABLE: &str = "documents";
pub const MEMORY_TABLE: &str = "memory";
#[cfg(feature = "code-search")]
pub const CODE_CHUNKS_TABLE: &str = "code_chunks";