context_core/selection/
budgeting.rs

1use crate::types::context_bundle::{ScoredDocument, SelectedDocument, SelectionWhy};
2
3pub struct BudgetResult {
4    pub selected: Vec<SelectedDocument>,
5    pub tokens_used: usize,
6    pub documents_selected: usize,
7    pub documents_excluded_by_budget: usize,
8}
9
10pub fn apply_budget(scored_docs: Vec<ScoredDocument>, budget: usize) -> BudgetResult {
11    let mut selected = Vec::new();
12    let mut tokens_used = 0;
13    let mut documents_selected = 0;
14    let mut documents_excluded_by_budget = 0;
15
16    for sdoc in scored_docs {
17        // Spec: "Documents with score 0.0 MAY be selected if budget allows."
18        if tokens_used + sdoc.token_count <= budget {
19            selected.push(SelectedDocument {
20                id: sdoc.document.id.as_str().to_string(),
21                version: sdoc.document.version.as_str().to_string(),
22                content: sdoc.document.content.clone(),
23                score: sdoc.score,
24                tokens: sdoc.token_count,
25                why: SelectionWhy {
26                    query_terms: sdoc.score_details.query_terms,
27                    term_matches: sdoc.score_details.term_matches,
28                    total_words: sdoc.score_details.total_words,
29                },
30            });
31            tokens_used += sdoc.token_count;
32            documents_selected += 1;
33        } else {
34            documents_excluded_by_budget += 1;
35        }
36    }
37
38    BudgetResult {
39        selected,
40        tokens_used,
41        documents_selected,
42        documents_excluded_by_budget,
43    }
44}
context_core/selection/budgeting.rs

context_core/selection/
budgeting.rs
