1use crate::error::{Result, SdkError};
10
11use super::budget::fit_budget;
12use super::dialect::{ModelDialect, SystemSlotPolicy};
13use super::digest::{CompiledPromptMessage, CompiledPromptProgram};
14use super::route::PromptRoute;
15use super::section::{
16 OverrideOrigin, PromptMessageRole, PromptSection, RenderedSection, SectionProvenance,
17};
18use super::stage::PromptStage;
19
20#[derive(Debug)]
22pub struct StageComposition {
23 stage: PromptStage,
24 sections: Vec<(RenderedSection, OverrideOrigin)>,
25 error: Option<SdkError>,
26 system_separator: Option<String>,
27}
28
29impl StageComposition {
30 pub fn new(stage: PromptStage) -> Self {
31 Self {
32 stage,
33 sections: Vec::new(),
34 error: None,
35 system_separator: None,
36 }
37 }
38
39 pub fn system_separator(mut self, separator: &str) -> Self {
44 self.system_separator = Some(separator.to_string());
45 self
46 }
47
48 pub fn section(mut self, section: impl PromptSection) -> Self {
50 if self.error.is_some() {
51 return self;
52 }
53 match section.render() {
54 Ok(rendered) => self.sections.push((rendered, OverrideOrigin::Base)),
55 Err(error) => self.error = Some(error),
56 }
57 self
58 }
59
60 pub fn section_if(self, include: bool, section: impl PromptSection) -> Self {
62 if include {
63 self.section(section)
64 } else {
65 self
66 }
67 }
68
69 pub fn resolved(mut self, rendered: RenderedSection, origin: OverrideOrigin) -> Self {
72 if self.error.is_none() {
73 self.sections.push((rendered, origin));
74 }
75 self
76 }
77
78 pub fn compile(
80 self,
81 route: &PromptRoute,
82 dialect: &ModelDialect,
83 tool_spec_hash: &str,
84 ) -> Result<CompiledPromptProgram> {
85 if let Some(error) = self.error {
86 return Err(error);
87 }
88 let token_budget = dialect.context_window_tokens;
89 let mut origins: std::collections::BTreeMap<String, OverrideOrigin> =
90 std::collections::BTreeMap::new();
91 let mut rendered = Vec::with_capacity(self.sections.len());
92 for (section, origin) in self.sections {
93 origins.insert(section.id.0.clone(), origin);
94 rendered.push(section);
95 }
96 let fit = fit_budget(rendered, &dialect.token_accountant, token_budget)?;
97 let separator = self
98 .system_separator
99 .as_deref()
100 .unwrap_or(&dialect.boundary_marker);
101 let messages = layout_messages(&fit.kept, dialect, separator, &origins);
102 let total_bytes: usize = messages.iter().map(|message| message.content.len()).sum();
103 if total_bytes as u64 > dialect.max_prompt_bytes {
104 return Err(SdkError::Domain(format!(
105 "compiled program of {total_bytes} bytes exceeds the dialect's \
106 {}-byte limit",
107 dialect.max_prompt_bytes
108 )));
109 }
110 let program_digest = CompiledPromptProgram::compute_digest(
111 route,
112 self.stage,
113 &dialect.dialect_ref(),
114 &dialect.token_accountant.id,
115 token_budget,
116 &messages,
117 &fit.dropped,
118 tool_spec_hash,
119 );
120 Ok(CompiledPromptProgram {
121 route: route.clone(),
122 stage: self.stage,
123 dialect: dialect.dialect_ref(),
124 messages,
125 dropped_sections: fit.dropped,
126 tool_spec_hash: tool_spec_hash.to_string(),
127 program_digest,
128 })
129 }
130}
131
132fn provenance(
133 section: &RenderedSection,
134 origins: &std::collections::BTreeMap<String, OverrideOrigin>,
135) -> SectionProvenance {
136 SectionProvenance {
137 id: section.id.clone(),
138 version: section.version,
139 content_hash: section.content_hash.clone(),
140 origin: origins
141 .get(§ion.id.0)
142 .cloned()
143 .unwrap_or(OverrideOrigin::Base),
144 }
145}
146
147fn layout_messages(
152 kept: &[RenderedSection],
153 dialect: &ModelDialect,
154 separator: &str,
155 origins: &std::collections::BTreeMap<String, OverrideOrigin>,
156) -> Vec<CompiledPromptMessage> {
157 let mut messages = Vec::new();
158 let system: Vec<&RenderedSection> = kept
159 .iter()
160 .filter(|section| section.role == PromptMessageRole::System)
161 .collect();
162 match dialect.system_slots {
163 SystemSlotPolicy::Many => {
164 for section in &system {
165 messages.push(message(
166 PromptMessageRole::System,
167 section.content.clone(),
168 vec![provenance(section, origins)],
169 ));
170 }
171 }
172 SystemSlotPolicy::SingleConcatenated => {
173 if !system.is_empty() {
174 let content = system
175 .iter()
176 .map(|section| section.content.as_str())
177 .collect::<Vec<_>>()
178 .join(separator);
179 let sections = system
180 .iter()
181 .map(|section| provenance(section, origins))
182 .collect();
183 messages.push(message(PromptMessageRole::System, content, sections));
184 }
185 }
186 }
187 for section in kept
188 .iter()
189 .filter(|section| section.role == PromptMessageRole::User)
190 {
191 messages.push(message(
192 PromptMessageRole::User,
193 section.content.clone(),
194 vec![provenance(section, origins)],
195 ));
196 }
197 messages
198}
199
200fn message(
201 role: PromptMessageRole,
202 content: String,
203 sections: Vec<SectionProvenance>,
204) -> CompiledPromptMessage {
205 let rendered_hash = CompiledPromptMessage::hash_content(&content);
206 CompiledPromptMessage {
207 role,
208 content,
209 sections,
210 rendered_hash,
211 }
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 use crate::model::ModelFamily;
218 use crate::prompt::section::{PromptSectionId, PromptSectionVersion};
219
220 fn rendered(id: &str, role: PromptMessageRole, body: &str) -> RenderedSection {
221 RenderedSection {
222 id: PromptSectionId(id.into()),
223 version: PromptSectionVersion(1),
224 role,
225 required: true,
226 priority: 0,
227 content_hash: format!("sha256:{id}"),
228 content: body.into(),
229 }
230 }
231
232 fn route() -> PromptRoute {
233 PromptRoute {
234 adapter: "scripted".into(),
235 family: ModelFamily::Other("scripted".into()),
236 exact_model: None,
237 }
238 }
239
240 #[test]
241 fn the_two_dialects_lay_out_the_same_program_differently() {
242 let build = || {
243 StageComposition::new(PromptStage::BranchCorrect)
244 .resolved(
245 rendered("branch_correct/role", PromptMessageRole::System, "You fix."),
246 OverrideOrigin::Base,
247 )
248 .resolved(
249 rendered(
250 "branch_correct/output_contract",
251 PromptMessageRole::System,
252 "Reply with JSON.",
253 ),
254 OverrideOrigin::Base,
255 )
256 };
257 let many = build()
258 .compile(&route(), &ModelDialect::native_tools_v1(), "sha256:tools")
259 .unwrap();
260 let single = build()
261 .compile(
262 &route(),
263 &ModelDialect::json_single_slot_v1(),
264 "sha256:tools",
265 )
266 .unwrap();
267 assert_eq!(many.messages.len(), 2);
268 assert_eq!(single.messages.len(), 1);
269 assert!(single.messages[0].content.contains("perspt:boundary"));
270 assert_ne!(many.program_digest, single.program_digest);
271 }
272
273 #[test]
274 fn identical_inputs_yield_identical_digests() {
275 let build = || {
276 StageComposition::new(PromptStage::BranchCorrect).resolved(
277 rendered("branch_correct/role", PromptMessageRole::System, "You fix."),
278 OverrideOrigin::Base,
279 )
280 };
281 let dialect = ModelDialect::native_tools_v1();
282 let first = build().compile(&route(), &dialect, "sha256:tools").unwrap();
283 let second = build().compile(&route(), &dialect, "sha256:tools").unwrap();
284 assert_eq!(first.program_digest, second.program_digest);
285 assert_eq!(first.messages, second.messages);
286 assert_eq!(first.dropped_sections, second.dropped_sections);
287 }
288
289 #[test]
290 fn oversize_programs_fail_the_dialect_byte_limit() {
291 let mut dialect = ModelDialect::native_tools_v1();
292 dialect.max_prompt_bytes = 4;
293 let composition = StageComposition::new(PromptStage::BranchCorrect).resolved(
294 rendered("branch_correct/role", PromptMessageRole::System, "too long"),
295 OverrideOrigin::Base,
296 );
297 assert!(composition
298 .compile(&route(), &dialect, "sha256:tools")
299 .is_err());
300 }
301}