use crate::Metadata;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Document {
pub id: String,
pub text: String,
pub metadata: Metadata,
}
impl Document {
pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
Self {
id: id.into(),
text: text.into(),
metadata: Metadata::new(),
}
}
pub fn with_metadata(mut self, metadata: Metadata) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Chunk {
pub id: String,
pub document_id: String,
pub text: String,
pub metadata: Metadata,
pub start_char: usize,
pub end_char: usize,
}
impl Chunk {
pub fn new(
id: impl Into<String>,
document_id: impl Into<String>,
text: impl Into<String>,
start_char: usize,
end_char: usize,
) -> Self {
Self {
id: id.into(),
document_id: document_id.into(),
text: text.into(),
metadata: Metadata::new(),
start_char,
end_char,
}
}
pub fn with_metadata(mut self, metadata: Metadata) -> Self {
self.metadata = metadata;
self
}
}
#[cfg(test)]
mod tests {
use super::{Chunk, Document};
use crate::Metadata;
#[test]
fn document_constructor_sets_id_and_text() {
let doc = Document::new("doc-1", "hello");
assert_eq!(doc.id, "doc-1");
assert_eq!(doc.text, "hello");
}
#[test]
fn chunk_constructor_sets_offsets() {
let chunk = Chunk::new("chunk-1", "doc-1", "hello", 0, 5);
assert_eq!(chunk.document_id, "doc-1");
assert_eq!(chunk.start_char, 0);
assert_eq!(chunk.end_char, 5);
}
#[test]
fn document_with_metadata_replaces_default_metadata() {
let metadata = Metadata::new().with("source", "runbook.md");
let doc = Document::new("doc-1", "hello").with_metadata(metadata.clone());
assert_eq!(doc.metadata, metadata);
}
#[test]
fn chunk_with_metadata_replaces_default_metadata() {
let metadata = Metadata::new().with("score", 0.9);
let chunk = Chunk::new("chunk-1", "doc-1", "hello", 0, 5).with_metadata(metadata.clone());
assert_eq!(chunk.metadata, metadata);
}
}