urna_format/writer/mod.rs
1//! Builder for .urna files.
2//!
3//! The builder owns structured chunk inputs and emits all six required
4//! sections plus the manifest. Output is fully deterministic given the
5//! same inputs and encoding choices.
6//!
7//! Encoding choices:
8//!
9//! - `SectionEncoding::Raw` (default) - sections stored verbatim.
10//! - `SectionEncoding::Zstd` - text-heavy sections (canonical/spans/
11//! provenance/contract) are zstd-compressed on disk; the reader
12//! decompresses transparently.
13//! - `EmbeddingDType::Float32 | Float16 | Int8` - controls the on-disk
14//! representation of the embeddings section. The runtime always
15//! accumulates dot products in f32 regardless of dtype.
16
17mod build;
18mod encoding_choice;
19mod payload;
20#[cfg(test)]
21mod tests;
22mod text_codec;
23
24pub use encoding_choice::{EmbeddingDType, REPRODUCIBLE_CREATED, SectionEncoding};
25
26use crate::chunk::ChunkInput;
27use crate::manifest::Manifest;
28use std::path::Path;
29
30/// High-level builder. Accepts canonical chunks plus an optional provenance
31/// blob. Computes `chunk_id`s, lays out the file, writes deterministic bytes.
32pub struct UrnaFileBuilder {
33 pub(super) manifest: Manifest,
34 pub(super) chunks: Vec<ChunkInput>,
35 pub(super) provenance: serde_json::Value,
36 pub(super) reproducible: bool,
37 pub(super) text_encoding: SectionEncoding,
38 pub(super) dtype: EmbeddingDType,
39 /// Optional HNSW index payload, fully encoded by the caller. The
40 /// builder doesn't know how to build an HNSW graph itself - that's
41 /// the runtime's job.
42 pub(super) hnsw_index: Option<Vec<u8>>,
43 pub(super) bm25_index: Option<Vec<u8>>,
44 /// Optional graph_adjacency (0x0C) csr payload, fully encoded by the
45 /// caller (chunk-to-chunk edges). additive, excluded from content_hash.
46 pub(super) graph_adjacency: Option<Vec<u8>>,
47 /// Optional blob_refs (0x14) payload, fully encoded by the caller
48 /// (`encode_blob_refs`). additive, excluded from content_hash.
49 pub(super) blob_refs: Option<Vec<u8>>,
50 /// Optional blob_data (0x17) payload, fully encoded by the caller
51 /// (`encode_blob_data`). additive, excluded from content_hash.
52 pub(super) blob_data: Option<Vec<u8>>,
53 /// Optional blob_span_overlay (0x16) payload, fully encoded by the
54 /// caller (`encode_blob_span_overlay`). additive, excluded from
55 /// content_hash; the runtime prefers it over 0x03 spans for cite.
56 pub(super) blob_span_overlay: Option<Vec<u8>>,
57 /// Optional space_table (0x15) payload, fully encoded by the caller
58 /// (`encode_space_table`). additive, excluded from content_hash.
59 pub(super) space_table: Option<Vec<u8>>,
60 /// Per-space vector band payloads: (band section id, encoding,
61 /// payload). section id is 0x20 + space_index; encoding is one of the
62 /// dtype encodings (raw/f16/int8/int4), never zstd.
63 pub(super) space_bands: Vec<(u32, u32, Vec<u8>)>,
64}
65
66impl UrnaFileBuilder {
67 pub fn new(manifest: Manifest) -> Self {
68 Self {
69 manifest,
70 chunks: Vec::new(),
71 provenance: serde_json::json!({}),
72 reproducible: false,
73 text_encoding: SectionEncoding::Raw,
74 dtype: EmbeddingDType::Float32,
75 hnsw_index: None,
76 bm25_index: None,
77 graph_adjacency: None,
78 blob_refs: None,
79 blob_data: None,
80 blob_span_overlay: None,
81 space_table: None,
82 space_bands: Vec::new(),
83 }
84 }
85
86 pub fn add_chunk(mut self, c: ChunkInput) -> Self {
87 self.chunks.push(c);
88 self
89 }
90
91 pub fn add_chunks<I: IntoIterator<Item = ChunkInput>>(mut self, chunks: I) -> Self {
92 self.chunks.extend(chunks);
93 self
94 }
95
96 pub fn with_provenance(mut self, v: serde_json::Value) -> Self {
97 self.provenance = v;
98 self
99 }
100
101 /// Reproducible build mode. When enabled, the writer overrides the
102 /// manifest's `created` timestamp to `REPRODUCIBLE_CREATED` so that
103 /// two builds with identical inputs produce byte-identical output.
104 /// Provenance JSON is not rewritten - callers are responsible for
105 /// keeping provenance deterministic if they want bit-for-bit equality.
106 pub fn reproducible(mut self, on: bool) -> Self {
107 self.reproducible = on;
108 self
109 }
110
111 /// Encoding for text-heavy sections (chunks_canonical, original_spans,
112 /// provenance, search_contract). `Zstd` shrinks PT-BR text by ~3-5×
113 /// in practice. chunk_ids stays raw because it is high-entropy and
114 /// almost incompressible.
115 pub fn text_encoding(mut self, enc: SectionEncoding) -> Self {
116 self.text_encoding = enc;
117 self
118 }
119
120 /// Embedding dtype + on-disk encoding. Mutates the manifest's dtype
121 /// to match. Quantized variants (`Float16`, `Int8`) are lossy; the
122 /// runtime always accumulates dot products in f32.
123 pub fn embedding_dtype(mut self, dt: EmbeddingDType) -> Self {
124 self.dtype = dt;
125 self.manifest.dtype = dt.manifest_str().to_string();
126 self
127 }
128
129 /// Attach an HNSW index payload (already encoded by `urna-engine`).
130 /// Sets `index_type=hnsw`, `rerank_policy=exact`, `supports_ann=true`.
131 pub fn hnsw_index(mut self, payload: Vec<u8>) -> Self {
132 self.hnsw_index = Some(payload);
133 self.manifest.index_type = "hnsw".into();
134 self.manifest.rerank_policy = "exact".into();
135 self.manifest.capabilities.supports_ann = true;
136 self
137 }
138
139 /// Attach a BM25 index payload.
140 pub fn bm25_index(mut self, payload: Vec<u8>) -> Self {
141 self.bm25_index = Some(payload);
142 self.manifest.capabilities.supports_bm25 = true;
143 self
144 }
145
146 /// Attach a graph_adjacency (0x0C) csr payload (chunk-to-chunk edges,
147 /// already encoded by `encode_graph_adjacency`). Sets the additive
148 /// `capabilities_ext.graph_present = Some(true)` flag so the runtime opens
149 /// the section behind a capability, exactly like hnsw/bm25. The section is
150 /// EXCLUDED from content_hash, so adding a graph never invalidates a
151 /// urna:// citation. Emitted raw, like hnsw.
152 pub fn graph_adjacency(mut self, payload: Vec<u8>) -> Self {
153 self.graph_adjacency = Some(payload);
154 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
155 ext.graph_present = Some(true);
156 self.manifest.capabilities_ext = Some(ext);
157 self
158 }
159
160 /// Attach a blob_refs (0x14) payload (already encoded by
161 /// `encode_blob_refs`): the content-hash reference table for source
162 /// media blobs. Sets the additive `capabilities_ext.blobs_present`
163 /// flag so the runtime opens the section behind a capability, exactly
164 /// like hnsw/bm25/graph. The section is EXCLUDED from content_hash, so
165 /// a self-contained media corpus keeps the content_hash (and the
166 /// urna:// citations) of its text-only twin. Emitted raw, like hnsw.
167 pub fn blob_refs(mut self, payload: Vec<u8>) -> Self {
168 self.blob_refs = Some(payload);
169 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
170 ext.blobs_present = Some(true);
171 self.manifest.capabilities_ext = Some(ext);
172 self
173 }
174
175 /// Attach a blob_data (0x17) payload (already encoded by
176 /// `encode_blob_data`): the inlined media bytes whose offset table
177 /// parallels the 0x14 record order, making the file self-contained
178 /// with no media sidecar. additive and EXCLUDED from content_hash;
179 /// implies `blobs_present` (the table indexes the 0x14 records).
180 pub fn blob_data(mut self, payload: Vec<u8>) -> Self {
181 self.blob_data = Some(payload);
182 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
183 ext.blobs_present = Some(true);
184 self.manifest.capabilities_ext = Some(ext);
185 self
186 }
187
188 /// Attach a blob_span_overlay (0x16) payload (already encoded by
189 /// `encode_blob_span_overlay`): per-chunk blob-relative spans that the
190 /// runtime prefers over chunks_original_spans (0x03) for cite/retrieve,
191 /// so 0x03 never has to carry an ordinal disguised as a byte range.
192 /// additive and EXCLUDED from content_hash; also implies
193 /// `blobs_present` (the overlay indexes the 0x14 table).
194 pub fn blob_span_overlay(mut self, payload: Vec<u8>) -> Self {
195 self.blob_span_overlay = Some(payload);
196 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
197 ext.blobs_present = Some(true);
198 self.manifest.capabilities_ext = Some(ext);
199 self
200 }
201
202 /// Attach a space_table (0x15) payload (already encoded by
203 /// `encode_space_table`): the multimodal per-space directory. Sets the
204 /// additive `capabilities_ext.supports_multimodal` flag so the runtime
205 /// opens the space bands behind a capability, exactly like the other
206 /// optional sections. EXCLUDED from content_hash, so a multimodal
207 /// corpus keeps the citations of its text-only twin. Emitted raw.
208 pub fn space_table(mut self, payload: Vec<u8>) -> Self {
209 self.space_table = Some(payload);
210 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
211 ext.supports_multimodal = Some(true);
212 self.manifest.capabilities_ext = Some(ext);
213 self
214 }
215
216 /// Attach one per-space vector band: the fixed-stride slab at section
217 /// 0x20 + `space_index`. `encoding` is the dtype encoding
218 /// (`SECTION_ENCODING_RAW` for f32, or float16/int8/int4), NEVER zstd
219 /// (the reader rejects it for band ids). EXCLUDED from content_hash;
220 /// also implies `supports_multimodal` (a band without a table is still
221 /// a multimodal file, though the runtime only opens listed spaces).
222 pub fn space_band(mut self, space_index: u8, encoding: u32, payload: Vec<u8>) -> Self {
223 let band_id = crate::layout::SECTION_SPACE_EMBEDDINGS_BASE + space_index as u32;
224 self.space_bands.push((band_id, encoding, payload));
225 let mut ext = self.manifest.capabilities_ext.take().unwrap_or_default();
226 ext.supports_multimodal = Some(true);
227 self.manifest.capabilities_ext = Some(ext);
228 self
229 }
230
231 /// Mark the search path as hybrid (BM25 + cosine). Requires both an
232 /// HNSW or exact path and a BM25 index. Caller is responsible for
233 /// declaring `score_type=hybrid_rrf` if they want that, otherwise
234 /// score_type stays "cosine".
235 pub fn hybrid(mut self) -> Self {
236 self.manifest.index_type = "hybrid".into();
237 self.manifest.rerank_policy = "exact".into();
238 self.manifest.score_type = "hybrid_rrf".into();
239 self.manifest.capabilities.supports_bm25 = true;
240 self
241 }
242
243 pub fn write_to_path(self, path: impl AsRef<Path>) -> crate::Result<()> {
244 let buf = self.build_bytes()?;
245 std::fs::write(path, buf)?;
246 Ok(())
247 }
248}