docling_rag/
model.rs

1//! Core data types shared across the pipeline: documents, chunks, scored hits.
2
3use serde::{Deserialize, Serialize};
4use std::str::FromStr;
5
6/// A source document, stored once with its metadata. Its text lives in [`Chunk`]s.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Document {
9    /// Stable id (UUID v4, hyphenated) — primary key.
10    pub id: String,
11    /// Where the document came from (file path, ftp/sftp URI, …).
12    pub source_uri: String,
13    /// Human-readable title (first heading, or file stem).
14    pub title: String,
15    /// Content hash (sha256 hex) used to skip re-ingesting unchanged documents.
16    pub hash: String,
17    /// Arbitrary metadata (format, byte length, custom fields).
18    #[serde(default)]
19    pub metadata: serde_json::Value,
20    /// RFC-3339 creation timestamp (stored as TEXT so it is DB-portable).
21    pub created_at: String,
22}
23
24impl Document {
25    /// Build a document with a fresh UUID and the given content hash.
26    pub fn new(
27        source_uri: impl Into<String>,
28        title: impl Into<String>,
29        hash: impl Into<String>,
30    ) -> Self {
31        Document {
32            id: new_id(),
33            source_uri: source_uri.into(),
34            title: title.into(),
35            hash: hash.into(),
36            metadata: serde_json::Value::Null,
37            created_at: now_rfc3339(),
38        }
39    }
40
41    /// Attach a metadata object, replacing any existing one.
42    pub fn with_metadata(mut self, meta: serde_json::Value) -> Self {
43        self.metadata = meta;
44        self
45    }
46}
47
48/// One retrievable chunk of a document.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct Chunk {
51    /// Stable id (UUID v4, hyphenated) — primary key.
52    pub id: String,
53    /// Owning [`Document::id`].
54    pub doc_id: String,
55    /// 0-based position of this chunk within its document.
56    pub ordinal: i64,
57    /// The chunk text (already includes any prepended heading context).
58    pub text: String,
59    /// Number of units (words/tokens) counted at chunk time — for eval/telemetry.
60    pub token_count: i64,
61    /// Arbitrary per-chunk metadata (heading path, page, …).
62    #[serde(default)]
63    pub metadata: serde_json::Value,
64    /// The embedding vector; `None` until embedded / when loaded without it.
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub embedding: Option<Vec<f32>>,
67}
68
69impl Chunk {
70    /// Construct a chunk (without embedding) for a document.
71    pub fn new(
72        doc_id: impl Into<String>,
73        ordinal: i64,
74        text: impl Into<String>,
75        token_count: i64,
76    ) -> Self {
77        Chunk {
78            id: new_id(),
79            doc_id: doc_id.into(),
80            ordinal,
81            text: text.into(),
82            token_count,
83            metadata: serde_json::Value::Null,
84            embedding: None,
85        }
86    }
87}
88
89/// A retrieval hit: a chunk plus the score that ranked it.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct Scored {
92    /// The retrieved chunk.
93    pub chunk: Chunk,
94    /// Higher is better. Cosine similarity for vector search, BM25 score for
95    /// keyword search, fused rank score for hybrid/multi-query.
96    pub score: f32,
97}
98
99impl Scored {
100    /// Pair a chunk with a score.
101    pub fn new(chunk: Chunk, score: f32) -> Self {
102        Scored { chunk, score }
103    }
104}
105
106/// The retrieval strategy to run for a query.
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
108#[serde(rename_all = "kebab-case")]
109pub enum RetrievalMode {
110    /// Dense vector (semantic) search only.
111    Vector,
112    /// Sparse keyword search only (Okapi BM25).
113    Bm25,
114    /// Vector + BM25, fused with Reciprocal Rank Fusion.
115    Hybrid,
116    /// LLM rewrites the query into N variants; results are fused (RRF).
117    MultiQuery,
118    /// LLM writes a hypothetical answer; its embedding drives vector search.
119    Hyde,
120}
121
122impl RetrievalMode {
123    /// Whether this mode needs a [`crate::llm::ChatModel`] to run.
124    pub fn needs_llm(self) -> bool {
125        matches!(self, RetrievalMode::MultiQuery | RetrievalMode::Hyde)
126    }
127
128    /// All modes, for the evaluation matrix.
129    pub const ALL: [RetrievalMode; 5] = [
130        RetrievalMode::Vector,
131        RetrievalMode::Bm25,
132        RetrievalMode::Hybrid,
133        RetrievalMode::MultiQuery,
134        RetrievalMode::Hyde,
135    ];
136
137    /// Modes that run without any network LLM — used by offline eval.
138    pub const OFFLINE: [RetrievalMode; 3] = [
139        RetrievalMode::Vector,
140        RetrievalMode::Bm25,
141        RetrievalMode::Hybrid,
142    ];
143}
144
145impl std::fmt::Display for RetrievalMode {
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        let s = match self {
148            RetrievalMode::Vector => "vector",
149            RetrievalMode::Bm25 => "bm25",
150            RetrievalMode::Hybrid => "hybrid",
151            RetrievalMode::MultiQuery => "multi-query",
152            RetrievalMode::Hyde => "hyde",
153        };
154        f.write_str(s)
155    }
156}
157
158impl FromStr for RetrievalMode {
159    type Err = crate::RagError;
160    fn from_str(s: &str) -> Result<Self, Self::Err> {
161        match s.trim().to_ascii_lowercase().replace('_', "-").as_str() {
162            "vector" | "dense" | "semantic" => Ok(RetrievalMode::Vector),
163            "bm25" | "keyword" | "sparse" => Ok(RetrievalMode::Bm25),
164            "hybrid" => Ok(RetrievalMode::Hybrid),
165            "multi-query" | "multiquery" | "fusion" => Ok(RetrievalMode::MultiQuery),
166            "hyde" => Ok(RetrievalMode::Hyde),
167            other => Err(crate::RagError::config(format!(
168                "unknown retrieval mode '{other}'"
169            ))),
170        }
171    }
172}
173
174/// Generate a fresh hyphenated UUID v4.
175pub fn new_id() -> String {
176    uuid::Uuid::new_v4().to_string()
177}
178
179/// sha256 hex digest of the given bytes — used for document dedup.
180pub fn content_hash(bytes: &[u8]) -> String {
181    use sha2::{Digest, Sha256};
182    let mut h = Sha256::new();
183    h.update(bytes);
184    let digest = h.finalize();
185    let mut out = String::with_capacity(64);
186    for b in digest {
187        out.push_str(&format!("{b:02x}"));
188    }
189    out
190}
191
192/// Current time as an RFC-3339 string, without pulling in `chrono`.
193///
194/// Uses the system clock; formatted as UTC epoch seconds when the platform
195/// clock is available, falling back to `"1970-01-01T00:00:00Z"`.
196pub fn now_rfc3339() -> String {
197    use std::time::{SystemTime, UNIX_EPOCH};
198    match SystemTime::now().duration_since(UNIX_EPOCH) {
199        Ok(d) => format_epoch_utc(d.as_secs()),
200        Err(_) => "1970-01-01T00:00:00Z".to_string(),
201    }
202}
203
204/// Minimal civil-date formatter (UTC) so we avoid a chrono dependency for a
205/// single timestamp field. Correct for all dates after the Unix epoch.
206fn format_epoch_utc(secs: u64) -> String {
207    let days = (secs / 86_400) as i64;
208    let rem = secs % 86_400;
209    let (hh, mm, ss) = (rem / 3600, (rem % 3600) / 60, rem % 60);
210    // Howard Hinnant's civil_from_days algorithm.
211    let z = days + 719_468;
212    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
213    let doe = z - era * 146_097;
214    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
215    let y = yoe + era * 400;
216    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
217    let mp = (5 * doy + 2) / 153;
218    let d = doy - (153 * mp + 2) / 5 + 1;
219    let m = if mp < 10 { mp + 3 } else { mp - 9 };
220    let y = if m <= 2 { y + 1 } else { y };
221    format!("{y:04}-{m:02}-{d:02}T{hh:02}:{mm:02}:{ss:02}Z")
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[test]
229    fn retrieval_mode_roundtrip() {
230        for m in RetrievalMode::ALL {
231            let s = m.to_string();
232            assert_eq!(RetrievalMode::from_str(&s).unwrap(), m, "roundtrip {s}");
233        }
234        assert_eq!(
235            RetrievalMode::from_str("KEYWORD").unwrap(),
236            RetrievalMode::Bm25
237        );
238        assert!(RetrievalMode::from_str("nope").is_err());
239    }
240
241    #[test]
242    fn hash_is_stable_and_hex() {
243        let h = content_hash(b"hello world");
244        assert_eq!(h.len(), 64);
245        assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
246        assert_eq!(h, content_hash(b"hello world"));
247        assert_ne!(h, content_hash(b"hello worlds"));
248    }
249
250    #[test]
251    fn epoch_formats_known_dates() {
252        assert_eq!(format_epoch_utc(0), "1970-01-01T00:00:00Z");
253        // 2021-01-01T00:00:00Z = 1609459200
254        assert_eq!(format_epoch_utc(1_609_459_200), "2021-01-01T00:00:00Z");
255    }
256}
docling_rag/model.rs

docling_rag/
model.rs
