orbok_models/
lib.rs

1//! # orbok-models
2//!
3//! Local AI model vocabulary (RFC-012). Milestone M1–M6 only needs the
4//! shared types and the "what is available" summary the UI shows; the
5//! install/locate/validate workflow lands in M12.
6//!
7//! Privacy rule carried from the requirements: model *download* is the
8//! only network operation orbok may ever perform, it is explicit, and
9//! it never involves document contents.
10
11use serde::{Deserialize, Serialize};
12
13/// Model roles (catalog `models.role`).
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(rename_all = "snake_case")]
16pub enum ModelRole {
17    Embedding,
18    Reranker,
19}
20
21impl ModelRole {
22    pub fn as_str(&self) -> &'static str {
23        match self {
24            ModelRole::Embedding => "embedding",
25            ModelRole::Reranker => "reranker",
26        }
27    }
28}
29
30/// Model availability (catalog `models.status`).
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
32#[serde(rename_all = "snake_case")]
33pub enum ModelStatus {
34    Available,
35    Missing,
36    Invalid,
37    Installing,
38    Disabled,
39}
40
41impl ModelStatus {
42    pub fn as_str(&self) -> &'static str {
43        match self {
44            ModelStatus::Available => "available",
45            ModelStatus::Missing => "missing",
46            ModelStatus::Invalid => "invalid",
47            ModelStatus::Installing => "installing",
48            ModelStatus::Disabled => "disabled",
49        }
50    }
51}
52
53/// Search capability derived from model availability. Keyword search
54/// never depends on models (RFC-007: works with zero models installed).
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub enum SearchCapability {
57    /// Keyword only: no embedding model available.
58    KeywordOnly,
59    /// Keyword + semantic: embedding model available.
60    Hybrid,
61    /// Keyword + semantic + rerank refinement.
62    HybridWithRerank,
63}
64
65/// Derive the capability shown in the UI from model statuses.
66pub fn search_capability(
67    embedding: Option<ModelStatus>,
68    reranker: Option<ModelStatus>,
69) -> SearchCapability {
70    match (embedding, reranker) {
71        (Some(ModelStatus::Available), Some(ModelStatus::Available)) => {
72            SearchCapability::HybridWithRerank
73        }
74        (Some(ModelStatus::Available), _) => SearchCapability::Hybrid,
75        _ => SearchCapability::KeywordOnly,
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    // RFC-007/RFC-010: search degrades gracefully without models.
84    #[test]
85    fn capability_degrades_gracefully() {
86        assert_eq!(search_capability(None, None), SearchCapability::KeywordOnly);
87        assert_eq!(
88            search_capability(Some(ModelStatus::Missing), None),
89            SearchCapability::KeywordOnly
90        );
91        assert_eq!(
92            search_capability(Some(ModelStatus::Available), None),
93            SearchCapability::Hybrid
94        );
95        assert_eq!(
96            search_capability(Some(ModelStatus::Available), Some(ModelStatus::Missing)),
97            SearchCapability::Hybrid
98        );
99        assert_eq!(
100            search_capability(Some(ModelStatus::Available), Some(ModelStatus::Available)),
101            SearchCapability::HybridWithRerank
102        );
103    }
104}
105
106/// A vector search candidate (RFC-008 §13).
107#[derive(Debug, Clone)]
108pub struct VectorCandidate {
109    pub chunk_id: orbok_core::ChunkId,
110    pub file_id: orbok_core::FileId,
111    pub rank: u32,
112    pub score: f32,
113}
114
115/// Local embedding model abstraction (RFC-008 §6).
116///
117/// Implementations must not transmit text externally (NFR-001).
118pub trait EmbeddingModel: Send + Sync {
119    /// Stable name stored in `models.model_name`.
120    fn name(&self) -> &str;
121    /// Version string stored in `models.model_version`.
122    fn version(&self) -> &str;
123    /// Output dimension — must match stored embeddings (RFC-008 §11).
124    fn dimension(&self) -> u32;
125    /// Embed a batch of normalized texts. Returns one vector per input,
126    /// each L2-normalized.
127    fn embed_batch(&self, texts: &[&str]) -> orbok_core::OrbokResult<Vec<Vec<f32>>>;
128}
129
130/// Compute cosine similarity between two L2-normalized vectors.
131pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
132    a.iter().zip(b).map(|(x, y)| x * y).sum()
133}
134
135/// L2-normalize a vector in-place. No-op for the zero vector.
136pub fn l2_normalize(v: &mut Vec<f32>) {
137    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
138    if norm > 1e-10 {
139        for x in v.iter_mut() {
140            *x /= norm;
141        }
142    }
143}
144
145/// Serialize a vector to little-endian bytes for BLOB storage (RFC-008
146/// §12.1 "sqlite_blob with FP32").
147pub fn vec_to_blob(v: &[f32]) -> Vec<u8> {
148    v.iter().flat_map(|x| x.to_le_bytes()).collect()
149}
150
151/// Deserialize from BLOB bytes; returns `None` on length mismatch.
152pub fn blob_to_vec(blob: &[u8], expected_dim: u32) -> Option<Vec<f32>> {
153    let dim = expected_dim as usize;
154    if blob.len() != dim * 4 {
155        return None;
156    }
157    Some(
158        blob.chunks_exact(4)
159            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
160            .collect(),
161    )
162}
163
164// ── Mock model ──────────────────────────────────────────────────────
165
166/// Deterministic 8-dimensional mock embedding model.
167///
168/// Uses the SHA-256 of the input text as a pseudo-random source for 8
169/// f32 components, then L2-normalizes the result.  **Never use for
170/// semantic search** — the outputs are semantically meaningless.
171/// Suitable for pipeline correctness tests (RFC-008 §24 tests 1–10).
172pub struct MockEmbeddingModel;
173
174impl EmbeddingModel for MockEmbeddingModel {
175    fn name(&self) -> &str {
176        "mock"
177    }
178    fn version(&self) -> &str {
179        "v1"
180    }
181    fn dimension(&self) -> u32 {
182        8
183    }
184    fn embed_batch(&self, texts: &[&str]) -> orbok_core::OrbokResult<Vec<Vec<f32>>> {
185        use sha2::{Digest, Sha256};
186        texts
187            .iter()
188            .map(|text| {
189                let digest = Sha256::digest(text.as_bytes());
190                let mut v: Vec<f32> = digest[..8]
191                    .iter()
192                    .map(|&b| b as f32 / 255.0)
193                    .collect();
194                l2_normalize(&mut v);
195                Ok(v)
196            })
197            .collect()
198    }
199}
200
201#[cfg(test)]
202mod embedding_tests {
203    use super::*;
204
205    // RFC-008 §24 test 2: embedding generation succeeds for sample chunks.
206    #[test]
207    fn mock_embed_batch() {
208        let model = MockEmbeddingModel;
209        let vecs = model.embed_batch(&["hello world", "foo bar"]).unwrap();
210        assert_eq!(vecs.len(), 2);
211        for v in &vecs {
212            assert_eq!(v.len(), model.dimension() as usize);
213        }
214    }
215
216    // RFC-008 §24 test 3: dimension mismatch can be detected by caller.
217    #[test]
218    fn blob_roundtrip_and_dim_mismatch() {
219        let v = vec![0.1_f32, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
220        let blob = vec_to_blob(&v);
221        assert_eq!(blob.len(), 32);
222        let back = blob_to_vec(&blob, 8).unwrap();
223        for (a, b) in v.iter().zip(&back) {
224            assert!((a - b).abs() < 1e-6);
225        }
226        assert!(blob_to_vec(&blob, 16).is_none(), "dim mismatch must return None");
227    }
228
229    // L2 normalization: unit-length vectors.
230    #[test]
231    fn normalize_produces_unit_vector() {
232        let mut v = vec![3.0_f32, 4.0];
233        l2_normalize(&mut v);
234        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
235        assert!((norm - 1.0).abs() < 1e-6);
236    }
237
238    // RFC-008 §24 test 9: cosine sim of identical vectors = 1.0.
239    #[test]
240    fn cosine_sim_identical_vectors() {
241        let mut v = vec![1.0_f32, 2.0, 3.0];
242        l2_normalize(&mut v);
243        let sim = cosine_similarity(&v, &v);
244        assert!((sim - 1.0).abs() < 1e-6);
245    }
246}
orbok_models/lib.rs

orbok_models/
lib.rs
