tea_context/
compiler.rs

1use std::collections::BTreeMap;
2
3use crate::budget::{PROMPT_SEPARATOR, effective_remaining_bytes, estimate_tokens, truncate};
4use crate::{
5    BudgetBehavior, ConflictKey, ConflictMode, ContextError, ContextErrorCode, PromptBudget,
6    PromptDiagnostic, PromptDiagnosticCode, PromptInspectionEntry, PromptModule, PromptModuleId,
7    PromptSegment, PromptSegmentId, SegmentDisposition,
8};
9
10/// Maximum modules accepted by one compilation.
11pub const MAX_COMPILE_MODULES: usize = 1024;
12
13/// Byte-identical compiled system prompt and explainability data.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct CompiledPrompt {
16    text: String,
17    estimated_tokens: usize,
18    diagnostics: Vec<PromptDiagnostic>,
19    inspection: Vec<PromptInspectionEntry>,
20}
21
22impl CompiledPrompt {
23    /// Returns exact rendered prompt text without trailing newline.
24    #[must_use]
25    pub fn text(&self) -> &str {
26        &self.text
27    }
28    /// Returns exact UTF-8 output bytes.
29    #[must_use]
30    pub fn bytes(&self) -> usize {
31        self.text.len()
32    }
33    /// Returns conservative token estimate for the complete output.
34    #[must_use]
35    pub const fn estimated_tokens(&self) -> usize {
36        self.estimated_tokens
37    }
38    /// Returns stable ordered diagnostics.
39    #[must_use]
40    pub fn diagnostics(&self) -> &[PromptDiagnostic] {
41        &self.diagnostics
42    }
43    /// Returns one explainability row per unique input segment.
44    #[must_use]
45    pub fn inspection(&self) -> &[PromptInspectionEntry] {
46        &self.inspection
47    }
48}
49
50#[derive(Clone)]
51struct Candidate {
52    module_id: PromptModuleId,
53    authority: crate::PromptAuthority,
54    priority: crate::PromptPriority,
55    segment_order: usize,
56    segment: PromptSegment,
57}
58
59impl Candidate {
60    fn same_precedence(&self, other: &Self) -> bool {
61        self.authority == other.authority && self.priority == other.priority
62    }
63}
64
65type Selection = (
66    Vec<Candidate>,
67    Vec<PromptDiagnostic>,
68    Vec<PromptInspectionEntry>,
69);
70
71/// Pure deterministic prompt compiler.
72#[derive(Debug, Clone, Copy, Default)]
73pub struct PromptCompiler;
74
75impl PromptCompiler {
76    /// Selects, budgets, renders, and inspects prompt modules.
77    ///
78    /// # Errors
79    ///
80    /// Returns an error for oversized input, divergent duplicate identities,
81    /// ambiguous equal-precedence conflicts, or required budget overflow.
82    pub fn compile(
83        &self,
84        modules: impl IntoIterator<Item = PromptModule>,
85        budget: PromptBudget,
86    ) -> Result<CompiledPrompt, ContextError> {
87        let modules = modules.into_iter().collect::<Vec<_>>();
88        if modules.len() > MAX_COMPILE_MODULES {
89            return Err(ContextError::new(
90                ContextErrorCode::BoundsExceeded,
91                "prompt compilation contains too many modules",
92            ));
93        }
94        let mut candidates = flatten(modules);
95        candidates.sort_by(|left, right| {
96            left.authority
97                .cmp(&right.authority)
98                .then_with(|| right.priority.cmp(&left.priority))
99                .then_with(|| left.module_id.cmp(&right.module_id))
100                .then_with(|| left.segment_order.cmp(&right.segment_order))
101        });
102        let (selected, mut diagnostics, mut inspection) = select(candidates)?;
103        render(selected, budget, &mut diagnostics, &mut inspection)
104    }
105}
106
107fn flatten(modules: Vec<PromptModule>) -> Vec<Candidate> {
108    modules
109        .into_iter()
110        .flat_map(|module| {
111            let module_id = module.id().clone();
112            let authority = module.authority();
113            let priority = module.priority();
114            let segments = module.segments().to_vec();
115            segments
116                .into_iter()
117                .enumerate()
118                .map(move |(segment_order, segment)| Candidate {
119                    module_id: module_id.clone(),
120                    authority,
121                    priority,
122                    segment_order,
123                    segment,
124                })
125        })
126        .collect()
127}
128
129fn select(candidates: Vec<Candidate>) -> Result<Selection, ContextError> {
130    let mut identities: BTreeMap<PromptSegmentId, Candidate> = BTreeMap::new();
131    let mut conflicts: BTreeMap<ConflictKey, Candidate> = BTreeMap::new();
132    let mut selected = Vec::new();
133    let mut diagnostics = Vec::new();
134    let mut inspection = Vec::new();
135    for candidate in candidates {
136        if let Some(existing) = identities.get(candidate.segment.id()) {
137            if existing.segment != candidate.segment {
138                return Err(ContextError::new(
139                    ContextErrorCode::DuplicateIdentity,
140                    "prompt segment identity has divergent definitions",
141                ));
142            }
143            diagnostics.push(PromptDiagnostic::new(
144                PromptDiagnosticCode::ExactDuplicate,
145                candidate.segment.id().clone(),
146                Some(existing.segment.id().clone()),
147                None,
148            ));
149            inspection.push(nonrendered(&candidate, SegmentDisposition::Duplicate));
150            continue;
151        }
152        identities.insert(candidate.segment.id().clone(), candidate.clone());
153        if let Some(claim) = candidate.segment.conflict() {
154            if let Some(winner) = conflicts.get(claim.key()) {
155                if candidate.same_precedence(winner)
156                    && candidate.segment.content() != winner.segment.content()
157                {
158                    return Err(ContextError::new(
159                        ContextErrorCode::AmbiguousConflict,
160                        "equal-precedence prompt conflict is ambiguous",
161                    ));
162                }
163                let code = if winner
164                    .segment
165                    .conflict()
166                    .is_some_and(|value| value.mode() == ConflictMode::Protected)
167                {
168                    PromptDiagnosticCode::ProtectedConflict
169                } else {
170                    PromptDiagnosticCode::ConflictShadowed
171                };
172                diagnostics.push(PromptDiagnostic::new(
173                    code,
174                    candidate.segment.id().clone(),
175                    Some(winner.segment.id().clone()),
176                    Some(claim.key().clone()),
177                ));
178                inspection.push(nonrendered(
179                    &candidate,
180                    SegmentDisposition::ConflictShadowed,
181                ));
182                continue;
183            }
184            conflicts.insert(claim.key().clone(), candidate.clone());
185        }
186        selected.push(candidate);
187    }
188    Ok((selected, diagnostics, inspection))
189}
190
191fn render(
192    selected: Vec<Candidate>,
193    budget: PromptBudget,
194    diagnostics: &mut Vec<PromptDiagnostic>,
195    inspection: &mut Vec<PromptInspectionEntry>,
196) -> Result<CompiledPrompt, ContextError> {
197    let mut text = String::new();
198    for candidate in selected {
199        let separator_bytes = if text.is_empty() {
200            0
201        } else {
202            PROMPT_SEPARATOR.len()
203        };
204        let used_with_separator = text.len().saturating_add(separator_bytes);
205        let remaining = effective_remaining_bytes(budget, used_with_separator);
206        let content = candidate.segment.content();
207        let (rendered, disposition) = if content.len() <= remaining {
208            (Some(content.to_owned()), SegmentDisposition::Included)
209        } else {
210            match candidate.segment.budget_behavior() {
211                BudgetBehavior::Required => {
212                    return Err(ContextError::new(
213                        ContextErrorCode::BudgetExceeded,
214                        "required prompt segment exceeds compilation budget",
215                    ));
216                }
217                BudgetBehavior::Omit => (None, SegmentDisposition::OmittedForBudget),
218                BudgetBehavior::Truncate => truncate(content, remaining)
219                    .map_or((None, SegmentDisposition::OmittedForBudget), |value| {
220                        (Some(value), SegmentDisposition::Truncated)
221                    }),
222            }
223        };
224        let Some(rendered) = rendered else {
225            diagnostics.push(PromptDiagnostic::new(
226                PromptDiagnosticCode::OmittedForBudget,
227                candidate.segment.id().clone(),
228                None,
229                None,
230            ));
231            inspection.push(nonrendered(
232                &candidate,
233                SegmentDisposition::OmittedForBudget,
234            ));
235            continue;
236        };
237        if !text.is_empty() {
238            text.push_str(PROMPT_SEPARATOR);
239        }
240        let start = text.len();
241        text.push_str(&rendered);
242        let end = text.len();
243        if disposition == SegmentDisposition::Truncated {
244            diagnostics.push(PromptDiagnostic::new(
245                PromptDiagnosticCode::TruncatedForBudget,
246                candidate.segment.id().clone(),
247                None,
248                None,
249            ));
250        }
251        inspection.push(PromptInspectionEntry::new(
252            candidate.module_id,
253            candidate.segment.id().clone(),
254            candidate.segment.provenance().clone(),
255            candidate.segment.trust(),
256            candidate.segment.cache_scope(),
257            disposition,
258            Some(start..end),
259            rendered.len(),
260            estimate_tokens(rendered.len()),
261        ));
262    }
263    inspection.sort_by(|left, right| {
264        left.byte_range()
265            .map_or(usize::MAX, |range| range.start)
266            .cmp(&right.byte_range().map_or(usize::MAX, |range| range.start))
267            .then_with(|| left.module_id().cmp(right.module_id()))
268            .then_with(|| left.segment_id().cmp(right.segment_id()))
269    });
270    diagnostics.sort_by(|left, right| {
271        left.segment_id()
272            .cmp(right.segment_id())
273            .then_with(|| left.code().cmp(&right.code()))
274    });
275    let estimated_tokens = estimate_tokens(text.len());
276    if text.len() > budget.max_bytes() || estimated_tokens > budget.max_estimated_tokens() {
277        return Err(ContextError::new(
278            ContextErrorCode::BudgetExceeded,
279            "compiled prompt is empty or exceeds final budget",
280        ));
281    }
282    Ok(CompiledPrompt {
283        text,
284        estimated_tokens,
285        diagnostics: diagnostics.clone(),
286        inspection: inspection.clone(),
287    })
288}
289
290fn nonrendered(candidate: &Candidate, disposition: SegmentDisposition) -> PromptInspectionEntry {
291    PromptInspectionEntry::new(
292        candidate.module_id.clone(),
293        candidate.segment.id().clone(),
294        candidate.segment.provenance().clone(),
295        candidate.segment.trust(),
296        candidate.segment.cache_scope(),
297        disposition,
298        None,
299        0,
300        0,
301    )
302}
tea_context/compiler.rs

tea_context/
compiler.rs
