Skip to main content

perspt_sdk/prompt/
compose.rs

1//! Stage composition (PSP-10 system 24).
2//!
3//! A stage composer is ordinary typed Rust: inclusion predicates are code
4//! (`section_if`), variables are typed struct fields, and every section
5//! renders eagerly, so object safety never arises. Compilation fits the
6//! dialect's token budget deterministically and lays the surviving
7//! sections out per the dialect's system-slot policy.
8
9use 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/// An in-progress composition for one stage.
21#[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    /// Use the stage's declared section separator when the dialect
40    /// concatenates system sections into one slot (the migrated platform
41    /// literals' byte identity depends on it). Without this, the dialect's
42    /// boundary marker joins.
43    pub fn system_separator(mut self, separator: &str) -> Self {
44        self.system_separator = Some(separator.to_string());
45        self
46    }
47
48    /// Append a typed base section, rendering it eagerly.
49    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    /// Typed inclusion predicate: append only when `include` holds.
61    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    /// Append an already-resolved section (the template-resolution path;
70    /// carries its override origin).
71    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    /// Compile: deterministic budget fit, dialect layout, digests.
79    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(&section.id.0)
142            .cloned()
143            .unwrap_or(OverrideOrigin::Base),
144    }
145}
146
147/// Lay the surviving sections out per the dialect: separate system messages
148/// under [`SystemSlotPolicy::Many`], or one deterministic concatenation with
149/// the boundary marker under [`SystemSlotPolicy::SingleConcatenated`].
150/// User-role sections are always separate messages.
151fn 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}