cortex_sdk/
context_pack.rs

1use std::collections::BTreeSet;
2
3pub use crate::generated::context_pack_v1::{
4    AnswerGroundingOptionsV1, AnswerGroundingReportV1, AnswerGroundingSpanV1,
5    ContextPackAccessDecisionV1, ContextPackAnomalyV1, ContextPackCellV1, ContextPackExplainV1,
6    ContextPackProvenanceV1, ContextPackSourceRefV1, ContextPackV1,
7    CONTEXT_PACK_V1_REQUIRED_FIELDS, CONTEXT_PACK_V1_SCHEMA_VERSION,
8};
9
10use crate::types::{
11    AnswerGroundingOptionsResponse, AnswerGroundingReportResponse, AnswerGroundingSpanResponse,
12    ContextPackCellResponse, ContextPackResponse,
13};
14
15impl Default for AnswerGroundingOptionsResponse {
16    fn default() -> Self {
17        Self {
18            min_span_support_q16: u16::MAX,
19            require_citations: false,
20            reject_unsupported: false,
21        }
22    }
23}
24
25impl ContextPackResponse {
26    pub const SCHEMA_VERSION_V1: &'static str = CONTEXT_PACK_V1_SCHEMA_VERSION;
27
28    pub fn is_v1(&self) -> bool {
29        self.schema_version == Self::SCHEMA_VERSION_V1
30    }
31
32    pub fn cell_ids(&self) -> impl Iterator<Item = u64> + '_ {
33        self.cells.iter().map(|cell| cell.cell_id)
34    }
35
36    pub fn citation_count(&self) -> usize {
37        self.cells
38            .iter()
39            .filter(|cell| {
40                cell.citation
41                    .as_deref()
42                    .is_some_and(|citation| !citation.is_empty())
43            })
44            .count()
45    }
46
47    pub fn anomaly_count(&self, code: &str) -> usize {
48        self.anomalies
49            .iter()
50            .filter(|anomaly| anomaly.code == code)
51            .count()
52    }
53
54    pub fn is_over_budget(&self) -> bool {
55        self.estimated_tokens > self.token_budget_tokens
56    }
57
58    pub fn ground_answer(&self, answer: &str) -> AnswerGroundingReportResponse {
59        self.ground_answer_with_options(answer, AnswerGroundingOptionsResponse::default())
60    }
61
62    pub fn ground_answer_with_options(
63        &self,
64        answer: &str,
65        options: AnswerGroundingOptionsResponse,
66    ) -> AnswerGroundingReportResponse {
67        let spans = split_answer_spans(answer)
68            .into_iter()
69            .map(|span| ground_span(&self.cells, span, options))
70            .collect::<Vec<_>>();
71        let supported_span_count = spans.iter().filter(|span| span.supported).count() as u32;
72        let unsupported_span_count = spans.iter().filter(|span| !span.supported).count() as u32;
73        let answer_supported = unsupported_span_count == 0;
74        AnswerGroundingReportResponse {
75            answer_supported,
76            rejected: !answer_supported && options.reject_unsupported,
77            support_q16: average_support_q16(&spans),
78            supported_span_count,
79            unsupported_span_count,
80            spans,
81        }
82    }
83}
84
85#[derive(Clone, Copy)]
86struct AnswerSpan<'a> {
87    text: &'a str,
88    start_byte: usize,
89    end_byte: usize,
90}
91
92fn ground_span(
93    cells: &[ContextPackCellResponse],
94    span: AnswerSpan<'_>,
95    options: AnswerGroundingOptionsResponse,
96) -> AnswerGroundingSpanResponse {
97    let span_terms = normalize_terms(&tokenize(span.text));
98    if span_terms.is_empty() {
99        return AnswerGroundingSpanResponse {
100            text: span.text.trim().to_owned(),
101            start_byte: span.start_byte,
102            end_byte: span.end_byte,
103            support_q16: u16::MAX,
104            supported: true,
105            covered_terms: Vec::new(),
106            missing_terms: Vec::new(),
107            supported_by_cell_ids: Vec::new(),
108            citations: Vec::new(),
109        };
110    }
111
112    let mut covered = BTreeSet::new();
113    let mut supporting_cells = BTreeSet::new();
114    let mut citations = BTreeSet::new();
115    for cell in cells {
116        let cell_terms = tokenize(&cell.payload_text)
117            .into_iter()
118            .collect::<BTreeSet<_>>();
119        let mut cell_matched = false;
120        for term in &span_terms {
121            if cell_terms.contains(term) {
122                covered.insert(term.clone());
123                cell_matched = true;
124            }
125        }
126        if cell_matched {
127            supporting_cells.insert(cell.cell_id);
128            if let Some(citation) = &cell.citation {
129                if !citation.trim().is_empty() {
130                    citations.insert(citation.clone());
131                }
132            }
133        }
134    }
135
136    let missing_terms = span_terms
137        .iter()
138        .filter(|term| !covered.contains(*term))
139        .cloned()
140        .collect::<Vec<_>>();
141    let support_q16 = q16_ratio(covered.len(), span_terms.len());
142    let has_required_citation = !options.require_citations || !citations.is_empty();
143    let supported = support_q16 >= options.min_span_support_q16 && has_required_citation;
144
145    AnswerGroundingSpanResponse {
146        text: span.text.trim().to_owned(),
147        start_byte: span.start_byte,
148        end_byte: span.end_byte,
149        support_q16,
150        supported,
151        covered_terms: covered.into_iter().collect(),
152        missing_terms,
153        supported_by_cell_ids: supporting_cells.into_iter().collect(),
154        citations: citations.into_iter().collect(),
155    }
156}
157
158fn split_answer_spans(answer: &str) -> Vec<AnswerSpan<'_>> {
159    let mut spans = Vec::new();
160    let mut start = 0usize;
161    for (index, ch) in answer.char_indices() {
162        if is_span_boundary(answer, index, ch) {
163            push_span(answer, start, index + ch.len_utf8(), &mut spans);
164            start = index + ch.len_utf8();
165        }
166    }
167    push_span(answer, start, answer.len(), &mut spans);
168    spans
169}
170
171fn is_span_boundary(answer: &str, index: usize, ch: char) -> bool {
172    if matches!(ch, '!' | '?' | '\n') {
173        return true;
174    }
175    if ch != '.' {
176        return false;
177    }
178    let previous = answer[..index].chars().next_back();
179    let next = answer[index + ch.len_utf8()..].chars().next();
180    !matches!((previous, next), (Some(prev), Some(next)) if prev.is_ascii_digit() && next.is_ascii_digit())
181}
182
183fn push_span<'a>(answer: &'a str, start: usize, end: usize, spans: &mut Vec<AnswerSpan<'a>>) {
184    let text = &answer[start..end];
185    let trimmed = text.trim();
186    if trimmed.is_empty() {
187        return;
188    }
189    let leading = text.len() - text.trim_start().len();
190    let trailing = text.len() - text.trim_end().len();
191    spans.push(AnswerSpan {
192        text: trimmed,
193        start_byte: start + leading,
194        end_byte: end - trailing,
195    });
196}
197
198fn tokenize(text: &str) -> Vec<String> {
199    text.split(|value: char| !value.is_alphanumeric())
200        .filter_map(normalize_term)
201        .filter(|term| !is_stopword(term))
202        .collect()
203}
204
205fn normalize_term(term: &str) -> Option<String> {
206    let normalized = term
207        .chars()
208        .flat_map(char::to_lowercase)
209        .collect::<String>();
210    (!normalized.is_empty()).then_some(normalized)
211}
212
213fn normalize_terms(terms: &[String]) -> Vec<String> {
214    terms
215        .iter()
216        .cloned()
217        .collect::<BTreeSet<_>>()
218        .into_iter()
219        .collect()
220}
221
222fn is_stopword(term: &str) -> bool {
223    matches!(
224        term,
225        "a" | "an"
226            | "and"
227            | "the"
228            | "or"
229            | "of"
230            | "to"
231            | "in"
232            | "и"
233            | "в"
234            | "на"
235            | "для"
236            | "мен"
237            | "және"
238            | "с"
239            | "со"
240            | "за"
241            | "от"
242            | "до"
243            | "по"
244            | "о"
245            | "об"
246            | "у"
247            | "да"
248            | "де"
249            | "та"
250            | "те"
251            | "үшін"
252    )
253}
254
255fn average_support_q16(spans: &[AnswerGroundingSpanResponse]) -> u16 {
256    if spans.is_empty() {
257        return u16::MAX;
258    }
259    let total = spans
260        .iter()
261        .map(|span| u64::from(span.support_q16))
262        .sum::<u64>();
263    u16::try_from(total / spans.len() as u64).unwrap_or(u16::MAX)
264}
265
266fn q16_ratio(numerator: usize, denominator: usize) -> u16 {
267    if denominator == 0 {
268        return u16::MAX;
269    }
270    ((numerator as u64 * u16::MAX as u64) / denominator as u64) as u16
271}
cortex_sdk/context_pack.rs

cortex_sdk/
context_pack.rs
