Skip to main content

uqa_scoring/
vector_calibration.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Persistable vector-calibration models and their compatibility contract.
8//!
9//! A bare [`crate::VectorProbabilityTransform`] is only a numeric mapping.
10//! Reusing one safely requires knowing the corpus, physical index, embedding
11//! model, candidate-pool size, and versions it was fitted against. This module
12//! keeps that provenance inseparable from a reusable model and rejects a
13//! runtime target that does not match it exactly.
14
15use serde::{Deserialize, Serialize};
16
17use crate::error::{invalid_input, require_finite};
18use crate::{ScoringError, ScoringResult, VectorProbabilityTransform};
19
20/// JSON schema version for [`VectorCalibrationModel`].
21pub const VECTOR_CALIBRATION_MODEL_SCHEMA_VERSION: u32 = 1;
22
23/// Identity of the retrieval surface to which a calibration model applies. Versions are opaque to Scoring; Execution binds `DiskANN` corpus/index versions to verified Storage metadata. Other integrations and embedding-model versions retain their explicit caller identity contract.
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25pub struct VectorCalibrationTarget {
26    pub corpus_id: String,
27    pub corpus_version: String,
28    pub index_id: String,
29    pub index_version: String,
30    pub index_kind: String,
31    pub embedding_model_id: String,
32    pub embedding_model_version: String,
33    pub candidate_k: usize,
34    pub dimensions: u32,
35}
36
37impl VectorCalibrationTarget {
38    pub fn validate(&self) -> ScoringResult<()> {
39        for (name, value) in [
40            ("corpus_id", self.corpus_id.as_str()),
41            ("corpus_version", self.corpus_version.as_str()),
42            ("index_id", self.index_id.as_str()),
43            ("index_version", self.index_version.as_str()),
44            ("index_kind", self.index_kind.as_str()),
45            ("embedding_model_id", self.embedding_model_id.as_str()),
46            (
47                "embedding_model_version",
48                self.embedding_model_version.as_str(),
49            ),
50        ] {
51            if value.trim().is_empty() {
52                return Err(invalid_input(format!("{name} must not be empty")));
53            }
54        }
55        if self.candidate_k == 0 {
56            return Err(invalid_input("candidate_k must be greater than zero"));
57        }
58        if self.dimensions == 0 {
59            return Err(invalid_input("dimensions must be greater than zero"));
60        }
61        Ok(())
62    }
63}
64
65/// Provenance stored with a fitted vector-calibration model.
66#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
67pub struct VectorCalibrationProvenance {
68    /// Version of the fitted parameters or training recipe.
69    pub model_version: String,
70    /// Exact retrieval target used for fitting.
71    pub target: VectorCalibrationTarget,
72    /// Number of labeled or background samples used to fit the model.
73    pub fit_sample_count: usize,
74}
75
76impl VectorCalibrationProvenance {
77    pub fn validate(&self) -> ScoringResult<()> {
78        if self.model_version.trim().is_empty() {
79            return Err(invalid_input("model_version must not be empty"));
80        }
81        if self.fit_sample_count < 2 {
82            return Err(invalid_input(format!(
83                "fit_sample_count must be at least 2, got {}",
84                self.fit_sample_count
85            )));
86        }
87        self.target.validate()
88    }
89}
90
91/// A reusable vector-calibration transform with mandatory provenance.
92#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
93pub struct VectorCalibrationModel {
94    schema_version: u32,
95    transform: VectorProbabilityTransform,
96    provenance: VectorCalibrationProvenance,
97}
98
99impl VectorCalibrationModel {
100    /// Apply the fixed transform to canonical cosine postings and retain the existing document ranking. No query-pool parameters are fitted.
101    pub fn calibrate_postings(
102        &self,
103        postings: &uqa_core::PostingList,
104        target: &VectorCalibrationTarget,
105    ) -> ScoringResult<Vec<uqa_core::ScoredEntry>> {
106        self.validate_for(target)?;
107        let mut calibrated = Vec::with_capacity(postings.len());
108        for entry in postings {
109            let score = entry.payload.score;
110            if !score.is_finite() || !(-1.0..=1.0).contains(&score) {
111                return Err(invalid_input(format!(
112                    "calibrated-vector KNN returned invalid cosine score {} for document {}",
113                    score, entry.doc_id
114                )));
115            }
116            calibrated.push(uqa_core::ScoredEntry {
117                doc_id: entry.doc_id,
118                score: self.transform.calibrate_one(1.0 - score)?,
119            });
120        }
121        Ok(crate::rank_scored_entries_top_k(
122            calibrated,
123            target.candidate_k,
124        ))
125    }
126    pub fn new(
127        transform: VectorProbabilityTransform,
128        provenance: VectorCalibrationProvenance,
129    ) -> ScoringResult<Self> {
130        let model = Self {
131            schema_version: VECTOR_CALIBRATION_MODEL_SCHEMA_VERSION,
132            transform,
133            provenance,
134        };
135        model.validate()?;
136        Ok(model)
137    }
138
139    pub fn transform(&self) -> VectorProbabilityTransform {
140        self.transform
141    }
142
143    pub fn provenance(&self) -> &VectorCalibrationProvenance {
144        &self.provenance
145    }
146
147    pub fn validate_for(&self, target: &VectorCalibrationTarget) -> ScoringResult<()> {
148        self.validate()?;
149        target.validate()?;
150        if self.provenance.target != *target {
151            return Err(invalid_input(format!(
152                "vector calibration target mismatch: model={:?}, runtime={target:?}",
153                self.provenance.target
154            )));
155        }
156        Ok(())
157    }
158
159    pub fn calibrate_one(
160        &self,
161        distance: f64,
162        target: &VectorCalibrationTarget,
163    ) -> ScoringResult<f64> {
164        self.validate_for(target)?;
165        self.transform.calibrate_one(distance)
166    }
167
168    pub fn calibrate(
169        &self,
170        distances: &[f64],
171        target: &VectorCalibrationTarget,
172    ) -> ScoringResult<Vec<f64>> {
173        self.validate_for(target)?;
174        self.transform.calibrate(distances, None)
175    }
176
177    pub fn to_json(&self) -> ScoringResult<String> {
178        self.validate()?;
179        serde_json::to_string(self)
180            .map_err(|error| invalid_input(format!("serialize vector calibration model: {error}")))
181    }
182
183    pub fn from_json(json: &str) -> ScoringResult<Self> {
184        let model: Self = serde_json::from_str(json).map_err(|error| {
185            invalid_input(format!("deserialize vector calibration model: {error}"))
186        })?;
187        model.validate()?;
188        Ok(model)
189    }
190
191    fn validate(&self) -> ScoringResult<()> {
192        if self.schema_version != VECTOR_CALIBRATION_MODEL_SCHEMA_VERSION {
193            return Err(invalid_input(format!(
194                "unsupported vector calibration schema version {}, expected {}",
195                self.schema_version, VECTOR_CALIBRATION_MODEL_SCHEMA_VERSION
196            )));
197        }
198        VectorProbabilityTransform::new(
199            self.transform.mu_match,
200            self.transform.mu_random,
201            self.transform.sigma,
202            self.transform.base_rate,
203        )?;
204        self.provenance.validate()
205    }
206}
207
208/// Probability drift observed when two calibration models score the same
209/// distance probes. This makes candidate-`K` sensitivity a measured contract
210/// instead of an informal observation.
211#[derive(Debug, Clone, Copy, PartialEq)]
212pub struct VectorCalibrationStabilityReport {
213    pub reference_k: usize,
214    pub candidate_k: usize,
215    pub probe_count: usize,
216    pub mean_absolute_drift: f64,
217    pub max_absolute_drift: f64,
218}
219
220impl VectorCalibrationStabilityReport {
221    pub fn compare(
222        reference: &VectorCalibrationModel,
223        candidate: &VectorCalibrationModel,
224        probe_distances: &[f64],
225    ) -> ScoringResult<Self> {
226        if probe_distances.is_empty() {
227            return Err(invalid_input(
228                "calibration stability probes must not be empty",
229            ));
230        }
231        let reference_target = &reference.provenance.target;
232        let candidate_target = &candidate.provenance.target;
233        let reference_probabilities = reference.calibrate(probe_distances, reference_target)?;
234        let candidate_probabilities = candidate.calibrate(probe_distances, candidate_target)?;
235        let mut sum = 0.0;
236        let mut max = 0.0_f64;
237        for (left, right) in reference_probabilities
238            .iter()
239            .zip(candidate_probabilities.iter())
240        {
241            let drift = (left - right).abs();
242            require_finite(drift, "vector calibration probability drift")?;
243            sum += drift;
244            max = max.max(drift);
245            if !sum.is_finite() {
246                return Err(ScoringError::ArithmeticOverflow(
247                    "vector calibration drift accumulation is not finite".into(),
248                ));
249            }
250        }
251        Ok(Self {
252            reference_k: reference_target.candidate_k,
253            candidate_k: candidate_target.candidate_k,
254            probe_count: probe_distances.len(),
255            mean_absolute_drift: sum / probe_distances.len() as f64,
256            max_absolute_drift: max,
257        })
258    }
259}