mod budget;
pub mod chunk;
mod classify;
mod dedup;
pub mod estimator;
pub mod insights;
pub mod model;
pub(crate) mod provenance;
mod relevance;
pub use chunk::{chunk_text, ChunkBoundary, ChunkPolicy, TextChunk};
pub use estimator::{DynTokenEstimator, HeuristicEstimator, TokenEstimator};
pub use insights::{CompilationInsights, ModelPricing, PricingTable};
pub use model::{
CompilePolicy, CompileRequest, CompiledContext, CompiledSection, ContextAction,
ContextDecision, ContextDecisionRef, ContextFact, ContextFragment, ContextSavings,
FidelityRisk, MemoryScope, RetrievalHandle, SectionKind, SourceReference, WorkingContext,
};
pub use relevance::DeterministicReranker;
use std::collections::BTreeMap;
use crate::error::MemoryError;
use crate::id::stable_id;
use crate::limits;
use budget::PackItem;
use classify::RuleMatch;
use dedup::{DupKind, Duplicate};
#[must_use]
pub fn fragment_id(content: &str) -> u64 {
stable_id(content)
}
pub struct ContextCompiler {
policy: CompilePolicy,
estimator: DynTokenEstimator,
pricing: Option<PricingTable>,
}
impl ContextCompiler {
#[must_use]
pub fn new(policy: CompilePolicy) -> Self {
Self {
policy,
estimator: Box::new(HeuristicEstimator),
pricing: None,
}
}
#[must_use]
pub fn with_estimator(mut self, estimator: DynTokenEstimator) -> Self {
self.estimator = estimator;
self
}
#[must_use]
pub fn with_pricing(mut self, pricing: PricingTable) -> Self {
self.pricing = Some(pricing);
self
}
pub(crate) fn effective_policy<'a>(&'a self, request: &'a CompileRequest) -> &'a CompilePolicy {
request.policy.as_ref().unwrap_or(&self.policy)
}
pub fn compile(&self, request: &CompileRequest) -> Result<CompiledContext, MemoryError> {
let policy = self.effective_policy(request);
let usable = validate(request, policy)?;
let analyses = analyze(request, policy, self.estimator.as_ref());
let items = pack_items(&analyses, policy, usable, self.estimator.as_ref());
let taken = budget::pack(&items, usable, &self.estimator);
let emissions = emissions(&items, &taken);
Ok(self.finish(request, &analyses, &emissions))
}
fn finish(
&self,
request: &CompileRequest,
analyses: &[Analysis],
emissions: &BTreeMap<usize, Emission>,
) -> CompiledContext {
let sections = sections(analyses, emissions);
let content = sections
.iter()
.map(|section| section.content.as_str())
.collect::<Vec<_>>()
.join(budget::JOINER);
let decisions: Vec<ContextDecision> = analyses
.iter()
.map(|analysis| decision(analysis, analyses, emissions))
.collect();
let insights = self.insights(request, analyses, &decisions, emissions, &content);
CompiledContext {
retrieval_handles: retrieval_handles(analyses, &decisions),
sources: analyses
.iter()
.filter(|analysis| analysis.dup.is_none())
.map(|analysis| provenance::source_for(analysis.fragment_id, analysis.content_hash))
.collect(),
risk: decisions
.iter()
.map(|decision| decision.risk)
.max()
.unwrap_or_default(),
content,
sections,
decisions,
insights,
}
}
fn insights(
&self,
request: &CompileRequest,
analyses: &[Analysis],
decisions: &[ContextDecision],
emissions: &BTreeMap<usize, Emission>,
content: &str,
) -> CompilationInsights {
let estimator = self.estimator.as_ref();
let tokens_in: u64 = analyses
.iter()
.map(|analysis| analysis.tokens)
.fold(0, u64::saturating_add);
let tokens_out = estimator.estimate(content);
let tokens_saved = tokens_in.saturating_sub(tokens_out);
let mut insights = CompilationInsights {
tokens_in,
tokens_out,
tokens_saved,
tokens_saved_by_rule: saved_by_rule(analyses, decisions, emissions, estimator),
..CompilationInsights::default()
};
let cost = request.target_model.as_deref().and_then(|model| {
let pricing = self
.effective_policy(request)
.pricing
.as_ref()
.or(self.pricing.as_ref())?;
let micros = pricing.cost_micros(model, tokens_saved)?;
Some((micros, pricing.currency.clone(), pricing.version.clone()))
});
if let Some((micros, currency, version)) = cost {
insights.estimated_cost_saved_micros = Some(micros);
insights.currency = Some(currency);
insights.pricing_version = Some(version);
}
insights
}
}
struct Analysis<'a> {
seq: usize,
fragment_id: u64,
content_hash: u64,
original: &'a str,
tokens: u64,
rule: RuleMatch,
relevance: f32,
priority: u8,
dup: Option<Duplicate>,
}
struct Emission {
text: String,
taken: usize,
total: usize,
}
impl Emission {
fn is_full(&self) -> bool {
self.taken == self.total
}
}
fn validate(request: &CompileRequest, policy: &CompilePolicy) -> Result<u64, MemoryError> {
if request.fragments.len() > limits::MAX_FRAGMENTS {
return Err(MemoryError::ContextOverLimit(format!(
"{} fragments exceed the cap of {}",
request.fragments.len(),
limits::MAX_FRAGMENTS
)));
}
if let Some(oversized) = request
.fragments
.iter()
.find(|fragment| fragment.content.len() > limits::MAX_FRAGMENT_BYTES)
{
return Err(MemoryError::ContextOverLimit(format!(
"a fragment of {} bytes exceeds the cap of {} bytes",
oversized.content.len(),
limits::MAX_FRAGMENT_BYTES
)));
}
let budget = limits::clamp_token_budget(request.token_budget);
let usable = budget.saturating_sub(policy.response_reserve_tokens);
if usable == 0 {
return Err(MemoryError::ContextBudget {
budget,
reserve: policy.response_reserve_tokens,
});
}
Ok(usable)
}
fn analyze<'a>(
request: &'a CompileRequest,
policy: &CompilePolicy,
estimator: &dyn TokenEstimator,
) -> Vec<Analysis<'a>> {
let contents: Vec<&str> = request
.fragments
.iter()
.map(|fragment| fragment.content.as_str())
.collect();
let duplicates = dedup::find_duplicates(&contents, policy.near_dup_dedup);
let query_terms = relevance::terms(&request.query);
let mut analyses: Vec<Analysis<'a>> = request
.fragments
.iter()
.zip(duplicates)
.enumerate()
.map(|(seq, (fragment, dup))| {
let content_hash = stable_id(&fragment.content);
Analysis {
seq,
fragment_id: fragment.id.unwrap_or(content_hash),
content_hash,
original: &fragment.content,
tokens: estimator.estimate(&fragment.content),
rule: classify::classify(fragment, policy),
relevance: relevance::lexical_relevance(&query_terms, &fragment.content),
priority: fragment.priority.unwrap_or(0),
dup,
}
})
.collect();
retain_safe_duplicates(&mut analyses);
analyses
}
fn retain_safe_duplicates(analyses: &mut [Analysis<'_>]) {
for index in 0..analyses.len() {
let Some(dup) = analyses[index].dup else {
continue;
};
let twin_verbatim = matches!(
analyses[dup.kept_seq].rule.action,
ContextAction::Preserve | ContextAction::Cache
);
let critical_near = dup.kind == DupKind::Near && analyses[index].rule.critical;
if !twin_verbatim || critical_near {
analyses[index].dup = None;
}
}
}
fn pack_items(
analyses: &[Analysis],
policy: &CompilePolicy,
usable: u64,
estimator: &dyn TokenEstimator,
) -> Vec<PackItem> {
let chunk_policy = effective_chunk_policy(policy, usable, estimator);
analyses
.iter()
.filter(|analysis| analysis.dup.is_none())
.map(|analysis| PackItem {
seq: analysis.seq,
critical: analysis.rule.critical,
priority: analysis.priority,
relevance: analysis.relevance,
pieces: pieces(analysis, &chunk_policy),
})
.collect()
}
fn pieces(analysis: &Analysis, chunk_policy: &ChunkPolicy) -> Vec<String> {
if analysis.rule.action == ContextAction::Abstract {
return vec![classify::collapse_repeated_lines(analysis.original)];
}
chunk_text(analysis.original, chunk_policy)
.into_iter()
.map(|chunk| chunk.text)
.collect()
}
const MIN_CHUNK_BYTES: usize = 256;
fn effective_chunk_policy(
policy: &CompilePolicy,
usable: u64,
estimator: &dyn TokenEstimator,
) -> ChunkPolicy {
let budget_bytes = usize::try_from(usable.saturating_mul(estimator.bytes_per_token_hint()))
.unwrap_or(usize::MAX);
ChunkPolicy {
max_chunk_bytes: policy
.chunk
.max_chunk_bytes
.min(budget_bytes)
.max(MIN_CHUNK_BYTES),
overlap_bytes: 0,
boundary: policy.chunk.boundary,
}
}
fn emissions(items: &[PackItem], taken: &[usize]) -> BTreeMap<usize, Emission> {
items
.iter()
.zip(taken.iter().copied())
.filter(|&(item, count)| count > 0 || item.pieces.is_empty())
.map(|(item, count)| {
(
item.seq,
Emission {
text: item.pieces[..count].concat(),
taken: count,
total: item.pieces.len(),
},
)
})
.collect()
}
fn sections(analyses: &[Analysis], emissions: &BTreeMap<usize, Emission>) -> Vec<CompiledSection> {
let mut result = Vec::new();
for kind in [SectionKind::Cache, SectionKind::Body] {
let mut blocks: Vec<&str> = Vec::new();
let mut ids: Vec<u64> = Vec::new();
for analysis in analyses {
let cache = analysis.rule.action == ContextAction::Cache;
let wanted = (kind == SectionKind::Cache) == cache;
if let Some(emission) = emissions
.get(&analysis.seq)
.filter(|emission| wanted && !emission.text.is_empty())
{
blocks.push(&emission.text);
ids.push(analysis.fragment_id);
}
}
if !blocks.is_empty() {
result.push(CompiledSection {
kind,
content: blocks.join(budget::JOINER),
fragment_ids: ids,
});
}
}
result
}
fn decision(
analysis: &Analysis,
all: &[Analysis],
emissions: &BTreeMap<usize, Emission>,
) -> ContextDecision {
let emission = emissions.get(&analysis.seq);
let (action, rule_id, risk, reason, handle) = match (&analysis.dup, emission) {
(Some(dup), _) => dup_verdict(analysis, *dup, &all[dup.kept_seq], emissions),
(None, Some(emission)) if emission.is_full() => full_verdict(analysis),
(None, Some(emission)) => partial_verdict(analysis, emission),
(None, None) => externalized_verdict(analysis),
};
ContextDecision {
fragment_id: analysis.fragment_id,
content_hash: analysis.content_hash,
action,
rule_id,
relevance: analysis.relevance,
risk,
reason,
memory_id: None,
handle,
}
}
type Verdict = (ContextAction, String, FidelityRisk, String, Option<String>);
fn critical_risk(critical: bool) -> FidelityRisk {
if critical {
FidelityRisk::High
} else {
FidelityRisk::Medium
}
}
fn dup_verdict(
analysis: &Analysis,
dup: Duplicate,
twin: &Analysis,
emissions: &BTreeMap<usize, Emission>,
) -> Verdict {
let (rule_id, variant) = match dup.kind {
DupKind::Exact => ("drop.duplicate", "exact duplicate"),
DupKind::Near => ("drop.near_duplicate", "near-duplicate"),
};
let twin_full = emissions.get(&twin.seq).is_some_and(Emission::is_full);
let (risk, fate) = if twin_full {
(FidelityRisk::Low, "content survives through it")
} else {
(
critical_risk(analysis.rule.critical),
"but that twin was not fully emitted — recover via the handle",
)
};
(
ContextAction::Drop,
rule_id.to_owned(),
risk,
format!("{variant} of fragment #{} — {fate}", dup.kept_seq),
Some(provenance::handle_for(analysis.content_hash)),
)
}
fn full_verdict(analysis: &Analysis) -> Verdict {
let risk = if analysis.rule.action == ContextAction::Abstract {
FidelityRisk::Medium
} else {
FidelityRisk::Low
};
(
analysis.rule.action,
analysis.rule.id.to_owned(),
risk,
analysis.rule.reason.to_owned(),
None,
)
}
fn partial_verdict(analysis: &Analysis, emission: &Emission) -> Verdict {
(
analysis.rule.action,
analysis.rule.id.to_owned(),
critical_risk(analysis.rule.critical),
format!(
"{} — packed {}/{} chunks, the rest stays retrievable",
analysis.rule.reason, emission.taken, emission.total
),
Some(provenance::handle_for(analysis.content_hash)),
)
}
fn externalized_verdict(analysis: &Analysis) -> Verdict {
(
ContextAction::Retrieve,
"budget.externalize".to_owned(),
critical_risk(analysis.rule.critical),
format!(
"did not fit the budget ({}); retrievable via its handle",
analysis.rule.reason
),
Some(provenance::handle_for(analysis.content_hash)),
)
}
fn retrieval_handles(analyses: &[Analysis], decisions: &[ContextDecision]) -> Vec<RetrievalHandle> {
analyses
.iter()
.zip(decisions)
.filter(|(_, decision)| decision.action == ContextAction::Retrieve)
.map(|(analysis, _)| RetrievalHandle {
handle: provenance::handle_for(analysis.content_hash),
fragment_id: analysis.fragment_id,
estimated_tokens: analysis.tokens,
})
.collect()
}
fn saved_by_rule(
analyses: &[Analysis],
decisions: &[ContextDecision],
emissions: &BTreeMap<usize, Emission>,
estimator: &dyn TokenEstimator,
) -> BTreeMap<String, u64> {
let mut by_rule = BTreeMap::new();
for (analysis, decision) in analyses.iter().zip(decisions) {
let emitted = emissions
.get(&analysis.seq)
.map_or(0, |emission| estimator.estimate(&emission.text));
let saved = analysis.tokens.saturating_sub(emitted);
if saved > 0 {
*by_rule.entry(decision.rule_id.clone()).or_insert(0) += saved;
}
}
by_rule
}
#[cfg(test)]
mod chunk_policy_tests {
use super::{effective_chunk_policy, MIN_CHUNK_BYTES};
use crate::context::estimator::HeuristicEstimator;
use crate::context::model::CompilePolicy;
#[test]
fn test_effective_chunk_policy_floors_chunk_size_under_a_tiny_budget() {
let policy = CompilePolicy::default();
let effective = effective_chunk_policy(&policy, 1, &HeuristicEstimator);
assert!(
effective.max_chunk_bytes >= MIN_CHUNK_BYTES,
"tiny budget drove chunk size to {} bytes, below the {MIN_CHUNK_BYTES}-byte floor",
effective.max_chunk_bytes
);
}
#[test]
fn test_effective_chunk_policy_floors_a_caller_supplied_tiny_chunk_size() {
let mut policy = CompilePolicy::default();
policy.chunk.max_chunk_bytes = 1;
let effective = effective_chunk_policy(&policy, 1_000, &HeuristicEstimator);
assert!(
effective.max_chunk_bytes >= MIN_CHUNK_BYTES,
"caller max_chunk_bytes=1 bypassed the {MIN_CHUNK_BYTES}-byte floor, got {}",
effective.max_chunk_bytes
);
}
}