1mod budget;
33pub mod chunk;
34mod classify;
35mod dedup;
36pub mod estimator;
37pub mod insights;
38pub mod model;
39pub(crate) mod provenance;
40mod relevance;
41
42pub use chunk::{chunk_text, ChunkBoundary, ChunkPolicy, TextChunk};
43pub use estimator::{DynTokenEstimator, HeuristicEstimator, TokenEstimator};
44pub use insights::{CompilationInsights, ModelPricing, PricingTable};
45pub use model::{
46 CompilePolicy, CompileRequest, CompiledContext, CompiledSection, ContextAction,
47 ContextDecision, ContextDecisionRef, ContextFact, ContextFragment, ContextSavings,
48 FidelityRisk, MemoryScope, RetrievalHandle, SectionKind, SourceReference, WorkingContext,
49};
50pub use relevance::DeterministicReranker;
51
52use std::collections::BTreeMap;
53
54use crate::error::MemoryError;
55use crate::id::stable_id;
56use crate::limits;
57
58use budget::PackItem;
59use classify::RuleMatch;
60use dedup::{DupKind, Duplicate};
61
62#[must_use]
66pub fn fragment_id(content: &str) -> u64 {
67 stable_id(content)
68}
69
70pub struct ContextCompiler {
99 policy: CompilePolicy,
100 estimator: DynTokenEstimator,
101 pricing: Option<PricingTable>,
102}
103
104impl ContextCompiler {
105 #[must_use]
108 pub fn new(policy: CompilePolicy) -> Self {
109 Self {
110 policy,
111 estimator: Box::new(HeuristicEstimator),
112 pricing: None,
113 }
114 }
115
116 #[must_use]
118 pub fn with_estimator(mut self, estimator: DynTokenEstimator) -> Self {
119 self.estimator = estimator;
120 self
121 }
122
123 #[must_use]
126 pub fn with_pricing(mut self, pricing: PricingTable) -> Self {
127 self.pricing = Some(pricing);
128 self
129 }
130
131 pub(crate) fn effective_policy<'a>(&'a self, request: &'a CompileRequest) -> &'a CompilePolicy {
135 request.policy.as_ref().unwrap_or(&self.policy)
136 }
137
138 pub fn compile(&self, request: &CompileRequest) -> Result<CompiledContext, MemoryError> {
147 let policy = self.effective_policy(request);
148 let usable = validate(request, policy)?;
149 let analyses = analyze(request, policy, self.estimator.as_ref());
150 let items = pack_items(&analyses, policy, usable, self.estimator.as_ref());
151 let taken = budget::pack(&items, usable, &self.estimator);
152 let emissions = emissions(&items, &taken);
153 Ok(self.finish(request, &analyses, &emissions))
154 }
155
156 fn finish(
158 &self,
159 request: &CompileRequest,
160 analyses: &[Analysis],
161 emissions: &BTreeMap<usize, Emission>,
162 ) -> CompiledContext {
163 let sections = sections(analyses, emissions);
164 let content = sections
165 .iter()
166 .map(|section| section.content.as_str())
167 .collect::<Vec<_>>()
168 .join(budget::JOINER);
169 let decisions: Vec<ContextDecision> = analyses
170 .iter()
171 .map(|analysis| decision(analysis, analyses, emissions))
172 .collect();
173 let insights = self.insights(request, analyses, &decisions, emissions, &content);
174 CompiledContext {
175 retrieval_handles: retrieval_handles(analyses, &decisions),
176 sources: analyses
177 .iter()
178 .filter(|analysis| analysis.dup.is_none())
179 .map(|analysis| provenance::source_for(analysis.fragment_id, analysis.content_hash))
180 .collect(),
181 risk: decisions
182 .iter()
183 .map(|decision| decision.risk)
184 .max()
185 .unwrap_or_default(),
186 content,
187 sections,
188 decisions,
189 insights,
190 }
191 }
192
193 fn insights(
195 &self,
196 request: &CompileRequest,
197 analyses: &[Analysis],
198 decisions: &[ContextDecision],
199 emissions: &BTreeMap<usize, Emission>,
200 content: &str,
201 ) -> CompilationInsights {
202 let estimator = self.estimator.as_ref();
203 let tokens_in: u64 = analyses
204 .iter()
205 .map(|analysis| analysis.tokens)
206 .fold(0, u64::saturating_add);
207 let tokens_out = estimator.estimate(content);
208 let tokens_saved = tokens_in.saturating_sub(tokens_out);
209 let mut insights = CompilationInsights {
210 tokens_in,
211 tokens_out,
212 tokens_saved,
213 tokens_saved_by_rule: saved_by_rule(analyses, decisions, emissions, estimator),
214 ..CompilationInsights::default()
215 };
216 let cost = request.target_model.as_deref().and_then(|model| {
217 let pricing = self
221 .effective_policy(request)
222 .pricing
223 .as_ref()
224 .or(self.pricing.as_ref())?;
225 let micros = pricing.cost_micros(model, tokens_saved)?;
226 Some((micros, pricing.currency.clone(), pricing.version.clone()))
227 });
228 if let Some((micros, currency, version)) = cost {
229 insights.estimated_cost_saved_micros = Some(micros);
230 insights.currency = Some(currency);
231 insights.pricing_version = Some(version);
232 }
233 insights
234 }
235}
236
237struct Analysis<'a> {
241 seq: usize,
243 fragment_id: u64,
245 content_hash: u64,
248 original: &'a str,
250 tokens: u64,
253 rule: RuleMatch,
255 relevance: f32,
257 priority: u8,
259 dup: Option<Duplicate>,
262}
263
264struct Emission {
266 text: String,
268 taken: usize,
270 total: usize,
272}
273
274impl Emission {
275 fn is_full(&self) -> bool {
277 self.taken == self.total
278 }
279}
280
281fn validate(request: &CompileRequest, policy: &CompilePolicy) -> Result<u64, MemoryError> {
284 if request.fragments.len() > limits::MAX_FRAGMENTS {
285 return Err(MemoryError::ContextOverLimit(format!(
286 "{} fragments exceed the cap of {}",
287 request.fragments.len(),
288 limits::MAX_FRAGMENTS
289 )));
290 }
291 if let Some(oversized) = request
292 .fragments
293 .iter()
294 .find(|fragment| fragment.content.len() > limits::MAX_FRAGMENT_BYTES)
295 {
296 return Err(MemoryError::ContextOverLimit(format!(
297 "a fragment of {} bytes exceeds the cap of {} bytes",
298 oversized.content.len(),
299 limits::MAX_FRAGMENT_BYTES
300 )));
301 }
302 let budget = limits::clamp_token_budget(request.token_budget);
303 let usable = budget.saturating_sub(policy.response_reserve_tokens);
304 if usable == 0 {
305 return Err(MemoryError::ContextBudget {
306 budget,
307 reserve: policy.response_reserve_tokens,
308 });
309 }
310 Ok(usable)
311}
312
313fn analyze<'a>(
316 request: &'a CompileRequest,
317 policy: &CompilePolicy,
318 estimator: &dyn TokenEstimator,
319) -> Vec<Analysis<'a>> {
320 let contents: Vec<&str> = request
321 .fragments
322 .iter()
323 .map(|fragment| fragment.content.as_str())
324 .collect();
325 let duplicates = dedup::find_duplicates(&contents, policy.near_dup_dedup);
326 let query_terms = relevance::terms(&request.query);
327 let mut analyses: Vec<Analysis<'a>> = request
328 .fragments
329 .iter()
330 .zip(duplicates)
331 .enumerate()
332 .map(|(seq, (fragment, dup))| {
333 let content_hash = stable_id(&fragment.content);
334 Analysis {
335 seq,
336 fragment_id: fragment.id.unwrap_or(content_hash),
337 content_hash,
338 original: &fragment.content,
339 tokens: estimator.estimate(&fragment.content),
340 rule: classify::classify(fragment, policy),
341 relevance: relevance::lexical_relevance(&query_terms, &fragment.content),
342 priority: fragment.priority.unwrap_or(0),
343 dup,
344 }
345 })
346 .collect();
347 retain_safe_duplicates(&mut analyses);
348 analyses
349}
350
351fn retain_safe_duplicates(analyses: &mut [Analysis<'_>]) {
357 for index in 0..analyses.len() {
358 let Some(dup) = analyses[index].dup else {
359 continue;
360 };
361 let twin_verbatim = matches!(
362 analyses[dup.kept_seq].rule.action,
363 ContextAction::Preserve | ContextAction::Cache
364 );
365 let critical_near = dup.kind == DupKind::Near && analyses[index].rule.critical;
366 if !twin_verbatim || critical_near {
367 analyses[index].dup = None;
368 }
369 }
370}
371
372fn pack_items(
376 analyses: &[Analysis],
377 policy: &CompilePolicy,
378 usable: u64,
379 estimator: &dyn TokenEstimator,
380) -> Vec<PackItem> {
381 let chunk_policy = effective_chunk_policy(policy, usable, estimator);
382 analyses
383 .iter()
384 .filter(|analysis| analysis.dup.is_none())
385 .map(|analysis| PackItem {
386 seq: analysis.seq,
387 critical: analysis.rule.critical,
388 priority: analysis.priority,
389 relevance: analysis.relevance,
390 pieces: pieces(analysis, &chunk_policy),
391 })
392 .collect()
393}
394
395fn pieces(analysis: &Analysis, chunk_policy: &ChunkPolicy) -> Vec<String> {
397 if analysis.rule.action == ContextAction::Abstract {
398 return vec![classify::collapse_repeated_lines(analysis.original)];
399 }
400 chunk_text(analysis.original, chunk_policy)
401 .into_iter()
402 .map(|chunk| chunk.text)
403 .collect()
404}
405
406const MIN_CHUNK_BYTES: usize = 256;
414
415fn effective_chunk_policy(
427 policy: &CompilePolicy,
428 usable: u64,
429 estimator: &dyn TokenEstimator,
430) -> ChunkPolicy {
431 let budget_bytes = usize::try_from(usable.saturating_mul(estimator.bytes_per_token_hint()))
432 .unwrap_or(usize::MAX);
433 ChunkPolicy {
434 max_chunk_bytes: policy
435 .chunk
436 .max_chunk_bytes
437 .min(budget_bytes)
438 .max(MIN_CHUNK_BYTES),
439 overlap_bytes: 0,
440 boundary: policy.chunk.boundary,
441 }
442}
443
444fn emissions(items: &[PackItem], taken: &[usize]) -> BTreeMap<usize, Emission> {
450 items
451 .iter()
452 .zip(taken.iter().copied())
453 .filter(|&(item, count)| count > 0 || item.pieces.is_empty())
454 .map(|(item, count)| {
455 (
456 item.seq,
457 Emission {
458 text: item.pieces[..count].concat(),
459 taken: count,
460 total: item.pieces.len(),
461 },
462 )
463 })
464 .collect()
465}
466
467fn sections(analyses: &[Analysis], emissions: &BTreeMap<usize, Emission>) -> Vec<CompiledSection> {
470 let mut result = Vec::new();
471 for kind in [SectionKind::Cache, SectionKind::Body] {
472 let mut blocks: Vec<&str> = Vec::new();
473 let mut ids: Vec<u64> = Vec::new();
474 for analysis in analyses {
475 let cache = analysis.rule.action == ContextAction::Cache;
476 let wanted = (kind == SectionKind::Cache) == cache;
477 if let Some(emission) = emissions
483 .get(&analysis.seq)
484 .filter(|emission| wanted && !emission.text.is_empty())
485 {
486 blocks.push(&emission.text);
487 ids.push(analysis.fragment_id);
488 }
489 }
490 if !blocks.is_empty() {
491 result.push(CompiledSection {
492 kind,
493 content: blocks.join(budget::JOINER),
494 fragment_ids: ids,
495 });
496 }
497 }
498 result
499}
500
501fn decision(
503 analysis: &Analysis,
504 all: &[Analysis],
505 emissions: &BTreeMap<usize, Emission>,
506) -> ContextDecision {
507 let emission = emissions.get(&analysis.seq);
508 let (action, rule_id, risk, reason, handle) = match (&analysis.dup, emission) {
509 (Some(dup), _) => dup_verdict(analysis, *dup, &all[dup.kept_seq], emissions),
510 (None, Some(emission)) if emission.is_full() => full_verdict(analysis),
511 (None, Some(emission)) => partial_verdict(analysis, emission),
512 (None, None) => externalized_verdict(analysis),
513 };
514 ContextDecision {
515 fragment_id: analysis.fragment_id,
516 content_hash: analysis.content_hash,
517 action,
518 rule_id,
519 relevance: analysis.relevance,
520 risk,
521 reason,
522 memory_id: None,
523 handle,
524 }
525}
526
527type Verdict = (ContextAction, String, FidelityRisk, String, Option<String>);
529
530fn critical_risk(critical: bool) -> FidelityRisk {
535 if critical {
536 FidelityRisk::High
537 } else {
538 FidelityRisk::Medium
539 }
540}
541
542fn dup_verdict(
548 analysis: &Analysis,
549 dup: Duplicate,
550 twin: &Analysis,
551 emissions: &BTreeMap<usize, Emission>,
552) -> Verdict {
553 let (rule_id, variant) = match dup.kind {
554 DupKind::Exact => ("drop.duplicate", "exact duplicate"),
555 DupKind::Near => ("drop.near_duplicate", "near-duplicate"),
556 };
557 let twin_full = emissions.get(&twin.seq).is_some_and(Emission::is_full);
558 let (risk, fate) = if twin_full {
559 (FidelityRisk::Low, "content survives through it")
560 } else {
561 (
562 critical_risk(analysis.rule.critical),
563 "but that twin was not fully emitted — recover via the handle",
564 )
565 };
566 (
567 ContextAction::Drop,
568 rule_id.to_owned(),
569 risk,
570 format!("{variant} of fragment #{} — {fate}", dup.kept_seq),
571 Some(provenance::handle_for(analysis.content_hash)),
572 )
573}
574
575fn full_verdict(analysis: &Analysis) -> Verdict {
577 let risk = if analysis.rule.action == ContextAction::Abstract {
578 FidelityRisk::Medium
579 } else {
580 FidelityRisk::Low
581 };
582 (
583 analysis.rule.action,
584 analysis.rule.id.to_owned(),
585 risk,
586 analysis.rule.reason.to_owned(),
587 None,
588 )
589}
590
591fn partial_verdict(analysis: &Analysis, emission: &Emission) -> Verdict {
593 (
594 analysis.rule.action,
595 analysis.rule.id.to_owned(),
596 critical_risk(analysis.rule.critical),
597 format!(
598 "{} — packed {}/{} chunks, the rest stays retrievable",
599 analysis.rule.reason, emission.taken, emission.total
600 ),
601 Some(provenance::handle_for(analysis.content_hash)),
602 )
603}
604
605fn externalized_verdict(analysis: &Analysis) -> Verdict {
607 (
608 ContextAction::Retrieve,
609 "budget.externalize".to_owned(),
610 critical_risk(analysis.rule.critical),
611 format!(
612 "did not fit the budget ({}); retrievable via its handle",
613 analysis.rule.reason
614 ),
615 Some(provenance::handle_for(analysis.content_hash)),
616 )
617}
618
619fn retrieval_handles(analyses: &[Analysis], decisions: &[ContextDecision]) -> Vec<RetrievalHandle> {
621 analyses
622 .iter()
623 .zip(decisions)
624 .filter(|(_, decision)| decision.action == ContextAction::Retrieve)
625 .map(|(analysis, _)| RetrievalHandle {
626 handle: provenance::handle_for(analysis.content_hash),
627 fragment_id: analysis.fragment_id,
628 estimated_tokens: analysis.tokens,
629 })
630 .collect()
631}
632
633fn saved_by_rule(
638 analyses: &[Analysis],
639 decisions: &[ContextDecision],
640 emissions: &BTreeMap<usize, Emission>,
641 estimator: &dyn TokenEstimator,
642) -> BTreeMap<String, u64> {
643 let mut by_rule = BTreeMap::new();
644 for (analysis, decision) in analyses.iter().zip(decisions) {
645 let emitted = emissions
646 .get(&analysis.seq)
647 .map_or(0, |emission| estimator.estimate(&emission.text));
648 let saved = analysis.tokens.saturating_sub(emitted);
649 if saved > 0 {
650 *by_rule.entry(decision.rule_id.clone()).or_insert(0) += saved;
651 }
652 }
653 by_rule
654}
655
656#[cfg(test)]
657mod chunk_policy_tests {
658 use super::{effective_chunk_policy, MIN_CHUNK_BYTES};
659 use crate::context::estimator::HeuristicEstimator;
660 use crate::context::model::CompilePolicy;
661
662 #[test]
663 fn test_effective_chunk_policy_floors_chunk_size_under_a_tiny_budget() {
664 let policy = CompilePolicy::default();
668 let effective = effective_chunk_policy(&policy, 1, &HeuristicEstimator);
669 assert!(
670 effective.max_chunk_bytes >= MIN_CHUNK_BYTES,
671 "tiny budget drove chunk size to {} bytes, below the {MIN_CHUNK_BYTES}-byte floor",
672 effective.max_chunk_bytes
673 );
674 }
675
676 #[test]
677 fn test_effective_chunk_policy_floors_a_caller_supplied_tiny_chunk_size() {
678 let mut policy = CompilePolicy::default();
681 policy.chunk.max_chunk_bytes = 1;
682 let effective = effective_chunk_policy(&policy, 1_000, &HeuristicEstimator);
683 assert!(
684 effective.max_chunk_bytes >= MIN_CHUNK_BYTES,
685 "caller max_chunk_bytes=1 bypassed the {MIN_CHUNK_BYTES}-byte floor, got {}",
686 effective.max_chunk_bytes
687 );
688 }
689}