1use crate::error::UrnaError;
9use sha2::{Digest, Sha256};
10
11#[derive(Clone, Debug, PartialEq)]
17pub struct ChunkInput {
18 pub canonical_text: String,
19 pub source_uri: String,
20 pub byte_start: u64,
21 pub byte_end: u64,
22 pub embedding: Vec<f32>,
23}
24
25pub fn chunk_id(
43 canonical_text: &str,
44 source_uri: &str,
45 byte_start: u64,
46 byte_end: u64,
47 chunker_version: &str,
48) -> String {
49 const DOMAIN: &[u8] = b"urna:chunk_id:v1\n";
50 let mut h = Sha256::new();
51 h.update(DOMAIN);
52 write_lp(&mut h, canonical_text.as_bytes());
53 write_lp(&mut h, source_uri.as_bytes());
54 h.update(byte_start.to_le_bytes());
55 h.update(byte_end.to_le_bytes());
56 write_lp(&mut h, chunker_version.as_bytes());
57 let digest = h.finalize();
58 format!("sha256:{}", hex::encode(digest))
59}
60
61fn write_lp(h: &mut Sha256, bytes: &[u8]) {
62 let len = bytes.len() as u32;
63 h.update(len.to_le_bytes());
64 h.update(bytes);
65}
66
67pub fn validate_chunk(c: &ChunkInput, embedding_dim: usize) -> crate::Result<()> {
70 if c.embedding.len() != embedding_dim {
71 return Err(UrnaError::DimensionMismatch {
72 expected: embedding_dim,
73 got: c.embedding.len(),
74 });
75 }
76 for v in &c.embedding {
77 if v.is_nan() || v.is_infinite() {
78 return Err(UrnaError::InvalidEmbeddingValue);
79 }
80 }
81 if c.byte_end < c.byte_start {
82 return Err(UrnaError::InvalidInput(format!(
83 "byte_end ({}) < byte_start ({})",
84 c.byte_end, c.byte_start
85 )));
86 }
87 Ok(())
88}
89
90#[cfg(test)]
91mod tests {
92 use super::*;
93
94 #[test]
95 fn chunk_id_is_deterministic() {
96 let a = chunk_id("hello", "doc.txt", 0, 5, "v1");
97 let b = chunk_id("hello", "doc.txt", 0, 5, "v1");
98 assert_eq!(a, b);
99 assert!(a.starts_with("sha256:"));
100 assert_eq!(a.len(), "sha256:".len() + 64);
101 }
102
103 #[test]
104 fn chunk_id_changes_with_inputs() {
105 let base = chunk_id("hello", "doc.txt", 0, 5, "v1");
106 assert_ne!(base, chunk_id("HELLO", "doc.txt", 0, 5, "v1"));
107 assert_ne!(base, chunk_id("hello", "other.txt", 0, 5, "v1"));
108 assert_ne!(base, chunk_id("hello", "doc.txt", 1, 5, "v1"));
109 assert_ne!(base, chunk_id("hello", "doc.txt", 0, 6, "v1"));
110 assert_ne!(base, chunk_id("hello", "doc.txt", 0, 5, "v2"));
111 }
112
113 #[test]
114 fn validate_chunk_dim_mismatch() {
115 let c = ChunkInput {
116 canonical_text: "x".into(),
117 source_uri: "y".into(),
118 byte_start: 0,
119 byte_end: 1,
120 embedding: vec![1.0, 0.0],
121 };
122 assert!(validate_chunk(&c, 4).is_err());
123 }
124
125 #[test]
126 fn validate_chunk_nan() {
127 let c = ChunkInput {
128 canonical_text: "x".into(),
129 source_uri: "y".into(),
130 byte_start: 0,
131 byte_end: 1,
132 embedding: vec![f32::NAN, 0.0, 0.0, 0.0],
133 };
134 assert!(matches!(
135 validate_chunk(&c, 4),
136 Err(UrnaError::InvalidEmbeddingValue)
137 ));
138 }
139
140 #[test]
141 fn validate_chunk_bad_span() {
142 let c = ChunkInput {
143 canonical_text: "x".into(),
144 source_uri: "y".into(),
145 byte_start: 10,
146 byte_end: 5,
147 embedding: vec![1.0, 0.0, 0.0, 0.0],
148 };
149 assert!(matches!(
150 validate_chunk(&c, 4),
151 Err(UrnaError::InvalidInput(_))
152 ));
153 }
154}