1use crate::{
4 Decimal,
5 definition::{BackendReq, ItemSource, Literal},
6 ids::{ArtifactId, DescriptorId, PromptTemplateId},
7};
8use serde::{Deserialize, Deserializer, Serialize};
9use std::collections::{BTreeMap, BTreeSet};
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
12#[serde(deny_unknown_fields)]
13pub struct Generate {
14 pub when: GenerateWhen,
15 pub for_each: ModelForEach,
16 pub template: TemplateRef,
17 pub language: Language,
18 pub projection: Vec<String>,
23 #[serde(deserialize_with = "carry")]
24 pub carry: Vec<String>,
25 pub vocabularies: BTreeMap<String, Vocabulary>,
26 pub output: OutputSchema,
27 pub backstop: Backstop,
28 pub repair: Repair,
29 pub screen: GenerateScreen,
30 pub limits: ModelLimits,
31 pub backend: BackendReq,
32}
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
34#[serde(rename_all = "snake_case", deny_unknown_fields)]
35pub enum GenerateWhen {
36 AnyPropose,
37 Actions(Vec<String>),
38}
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
40#[serde(rename_all = "snake_case", deny_unknown_fields)]
41pub enum ModelForEach {
42 None,
43 #[serde(untagged)]
44 List(#[serde(deserialize_with = "object")] ModelList),
45}
46#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
47#[serde(deny_unknown_fields)]
48pub struct ModelList {
49 pub over: String,
50 pub bind: String,
51 pub items: ItemSource,
52 pub max_items: ModelBound<256>,
53}
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
55#[serde(rename_all = "snake_case", deny_unknown_fields)]
56pub enum ModelBound<const MAX: u64> {
57 Literal(#[serde(deserialize_with = "bound::<_, MAX>")] u64),
58 Param(String),
59}
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
61#[serde(untagged, deny_unknown_fields)]
62pub enum Language {
63 Literal(String),
64 Parameter(#[serde(deserialize_with = "object")] crate::parameter::ParamRef),
65}
66#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
67#[serde(untagged, deny_unknown_fields)]
68pub enum Vocabulary {
69 Literal(#[serde(deserialize_with = "options")] Vec<String>),
70 Parameter(#[serde(deserialize_with = "object")] crate::parameter::ParamRef),
71}
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
73#[serde(deny_unknown_fields)]
74pub struct TemplateRef {
75 #[serde(deserialize_with = "template_name")]
76 pub id: String,
77 pub digest: PromptTemplateId,
78 pub params: BTreeMap<String, TemplateBinding>,
79}
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
81#[serde(untagged, deny_unknown_fields)]
82pub enum TemplateBinding {
83 Literal { literal: Literal },
84 Parameter(#[serde(deserialize_with = "object")] crate::parameter::ParamRef),
85 Reference(String),
86}
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
88#[serde(deny_unknown_fields)]
89pub struct OutputSchema {
90 #[serde(deserialize_with = "fields")]
91 pub fields: Vec<GeneratedField>,
92}
93#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
94#[serde(deny_unknown_fields)]
95pub struct GeneratedField {
96 pub name: String,
97 #[serde(rename = "type")]
98 pub ty: GeneratedType,
99 pub required: bool,
100}
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
102#[serde(rename_all = "snake_case", deny_unknown_fields)]
103pub enum GeneratedType {
104 Text {
105 #[serde(deserialize_with = "bound::<_, 16384>")]
106 max_bytes: u64,
107 lines: Lines,
108 },
109 List {
110 vocabulary: String,
111 #[serde(deserialize_with = "bound::<_, 64>")]
112 max_items: u64,
113 },
114}
115#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
116#[serde(rename_all = "snake_case", deny_unknown_fields)]
117pub enum Lines {
118 Single,
119 Multi,
120}
121#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
122#[serde(deny_unknown_fields)]
123pub struct Backstop {
124 pub locale: String,
125 pub escape: Escape,
126 pub overflow: Overflow,
127 pub fields: BTreeMap<String, BackstopField>,
128}
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
130#[serde(rename_all = "snake_case", deny_unknown_fields)]
131pub enum Escape {
132 Text,
133 Html,
134}
135#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
136#[serde(rename_all = "snake_case", deny_unknown_fields)]
137pub enum Overflow {
138 Refuse,
139 Truncate,
140}
141#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
142#[serde(rename_all = "snake_case", deny_unknown_fields)]
143pub enum BackstopField {
144 Text(String),
145 List(Vec<String>),
146}
147#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
148#[serde(deny_unknown_fields)]
149pub struct Repair {
150 #[serde(deserialize_with = "repair_attempts")]
151 pub max_attempts: u64,
152 pub only_missing_fields: bool,
153}
154#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
155#[serde(deny_unknown_fields)]
156pub struct GenerateScreen {
157 pub input: ScreenInput,
158 pub output: ScreenOutput,
159}
160#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
161#[serde(rename_all = "snake_case", deny_unknown_fields)]
162pub enum ScreenInput {
163 UntrustedFields,
164}
165#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
166#[serde(rename_all = "snake_case", deny_unknown_fields)]
167pub enum ScreenOutput {
168 None,
169 Screener { backend: BackendReq },
170}
171#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
172#[serde(deny_unknown_fields)]
173pub struct ModelLimits {
174 #[serde(deserialize_with = "deadline_share")]
175 pub deadline_share: Decimal,
176 pub per_item_deadline_ms: ModelBound<600000>,
177 #[serde(deserialize_with = "bound::<_, 64>")]
178 pub concurrency: u64,
179 pub max_cost_units: CostLimit,
180}
181#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
182#[serde(rename_all = "snake_case", deny_unknown_fields)]
183pub enum CostLimit {
184 None,
185 Units {
186 #[serde(deserialize_with = "bound::<_, {u64::MAX}>")]
187 n: u64,
188 },
189}
190#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
191#[serde(deny_unknown_fields)]
192pub struct PromptTemplate {
193 pub schema: String,
194 #[serde(deserialize_with = "template_name")]
195 pub id: String,
196 #[serde(deserialize_with = "version")]
197 pub version: String,
198 #[serde(deserialize_with = "instructions")]
199 pub instructions: String,
200 #[serde(deserialize_with = "params")]
201 pub params: Vec<PromptParam>,
202}
203#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
204#[serde(deny_unknown_fields)]
205pub struct PromptParam {
206 pub name: String,
207 #[serde(rename = "type")]
208 pub ty: PromptParamType,
209}
210#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
211#[serde(rename_all = "snake_case", deny_unknown_fields)]
212pub enum PromptParamType {
213 Bool,
214 Integer,
215 Decimal,
216 Text {
217 #[serde(deserialize_with = "bound::<_, 1024>")]
218 max_bytes: u64,
219 },
220 Enum {
221 #[serde(deserialize_with = "options")]
222 options: Vec<String>,
223 },
224}
225#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
226#[serde(deny_unknown_fields)]
227pub struct GenerateBinding {
228 pub template_id: PromptTemplateId,
229 pub template: PromptTemplate,
230 pub helpers_version: String,
231 pub backend_id: String,
232 pub artifact: ArtifactId,
233 pub descriptor: DescriptorId,
234 pub output: OutputSchema,
235 pub vocabularies: BTreeMap<String, Vec<String>>,
236 pub max_items: u64,
237 pub worst_case_requests: u64,
238 pub worst_case_projection_bytes: u64,
239}
240
241crate::document::document!(
242 PromptTemplate,
243 crate::schema::PROMPT_TEMPLATE,
244 crate::limits::DEFINITION_V1,
245 PromptTemplateId
246);
247
248pub(crate) fn bound<'de, D: Deserializer<'de>, const MAX: u64>(d: D) -> Result<u64, D::Error> {
249 let n = u64::deserialize(d)?;
250 if !(1..=MAX).contains(&n) {
251 return Err(serde::de::Error::custom("generation bound out of range"));
252 }
253 Ok(n)
254}
255fn repair_attempts<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
256 let n = u64::deserialize(d)?;
257 if n > 2 {
258 return Err(serde::de::Error::custom("repair.max_attempts exceeds 2"));
259 }
260 Ok(n)
261}
262pub(crate) fn deadline_share<'de, D: Deserializer<'de>>(d: D) -> Result<Decimal, D::Error> {
263 let text = String::deserialize(d)?;
264 let n: Decimal = Decimal::parse(&text).map_err(serde::de::Error::custom)?;
265 if n <= Decimal::ZERO
266 || n > Decimal::from_i64(1)
267 || text.split_once('.').is_some_and(|(_, f)| f.len() > 4)
268 {
269 return Err(serde::de::Error::custom(
270 "deadline_share outside (0,1] or over four fractional digits",
271 ));
272 }
273 Ok(n)
274}
275fn carry<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<String>, D::Error> {
276 let v = Vec::deserialize(d)?;
277 if v.len() > 16 {
278 return Err(serde::de::Error::custom("carry exceeds 16 references"));
279 }
280 Ok(v)
281}
282fn options<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<String>, D::Error> {
283 let v = Vec::<String>::deserialize(d)?;
284 if !(1..=256).contains(&v.len()) || v.iter().any(|x| !(1..=64).contains(&x.len())) {
285 return Err(serde::de::Error::custom(
286 "vocabulary outside generation bounds",
287 ));
288 }
289 Ok(v)
290}
291fn fields<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<GeneratedField>, D::Error> {
292 let v = Vec::<GeneratedField>::deserialize(d)?;
293 let mut names = BTreeSet::new();
294 if !(1..=32).contains(&v.len())
295 || v.iter()
296 .any(|f| !crate::parameter::is_parameter_name(&f.name) || !names.insert(&f.name))
297 {
298 return Err(serde::de::Error::custom(
299 "invalid or duplicate generated field name or field count",
300 ));
301 }
302 Ok(v)
303}
304fn template_name<'de, D: Deserializer<'de>>(d: D) -> Result<String, D::Error> {
305 let s = String::deserialize(d)?;
306 if !(1..=128).contains(&s.len())
307 || !s
308 .bytes()
309 .next()
310 .is_some_and(|b| b.is_ascii_lowercase() || b.is_ascii_digit())
311 || !s
312 .bytes()
313 .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b"._-".contains(&b))
314 {
315 return Err(serde::de::Error::custom("invalid prompt template id"));
316 }
317 Ok(s)
318}
319fn version<'de, D: Deserializer<'de>>(d: D) -> Result<String, D::Error> {
320 let s = String::deserialize(d)?;
321 if s.split('.').count() != 3
322 || s.split('.').any(|p| {
323 p.is_empty()
324 || (p.len() > 1 && p.starts_with('0'))
325 || !p.bytes().all(|b| b.is_ascii_digit())
326 })
327 {
328 return Err(serde::de::Error::custom("invalid prompt template version"));
329 }
330 Ok(s)
331}
332fn instructions<'de, D: Deserializer<'de>>(d: D) -> Result<String, D::Error> {
333 let s = String::deserialize(d)?;
334 if s.len() > 16384 {
336 return Err(serde::de::Error::custom(
337 "prompt instructions exceed 16 KiB",
338 ));
339 }
340 Ok(s)
341}
342fn params<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<PromptParam>, D::Error> {
343 let v = Vec::<PromptParam>::deserialize(d)?;
344 let mut names = BTreeSet::new();
345 if v.len() > 32
346 || v.iter()
347 .any(|p| !crate::parameter::is_parameter_name(&p.name) || !names.insert(&p.name))
348 {
349 return Err(serde::de::Error::custom("invalid prompt parameters"));
350 }
351 Ok(v)
352}
353
354pub(crate) fn steps<'de, D: Deserializer<'de>>(
355 d: D,
356) -> Result<Vec<crate::definition::StepDecl>, D::Error> {
357 let v = Vec::<crate::definition::StepDecl>::deserialize(d)?;
358 if v.iter()
359 .filter(|s| matches!(s.body, crate::definition::StepBody::Generate(_)))
360 .count()
361 > 8
362 {
363 return Err(serde::de::Error::custom("more than eight generate steps"));
364 }
365 if v.iter()
366 .filter(|s| matches!(s.body, crate::definition::StepBody::Extract(_)))
367 .count()
368 > 8
369 {
370 return Err(serde::de::Error::custom("more than eight extract steps"));
371 }
372 Ok(v)
373}
374
375pub(crate) fn object<'de, D: Deserializer<'de>, T: serde::de::DeserializeOwned>(
377 d: D,
378) -> Result<T, D::Error> {
379 let v = serde_json::Value::deserialize(d)?;
380 if !v.is_object() {
381 return Err(serde::de::Error::custom("expected an object"));
382 }
383 serde_json::from_value(v).map_err(serde::de::Error::custom)
384}
385
386#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
388#[serde(deny_unknown_fields)]
389pub struct GenerateRequest {
390 pub schema: String,
391 pub step: String,
392 pub instance: Vec<String>,
393 pub template: String,
394 pub template_id: PromptTemplateId,
395 pub instructions: String,
396 pub language: String,
397 pub projection: BTreeMap<String, serde_json::Value>,
398 pub output: OutputSchema,
399 pub vocabularies: BTreeMap<String, Vec<String>>,
400 pub repair: RepairRequest,
401}
402#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
403#[serde(rename_all = "snake_case", deny_unknown_fields)]
404pub enum RepairRequest {
405 None,
406 #[serde(untagged)]
412 Round {
413 round: u64,
414 fields: Vec<RepairField>,
415 },
416}
417#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
418#[serde(deny_unknown_fields)]
419pub struct RepairField {
420 pub name: String,
421 pub reason: FieldInvalid,
422}
423#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
424#[serde(rename_all = "snake_case")]
425pub enum FieldInvalid {
426 Missing,
427 WrongType,
428 Empty,
429 TooLong,
430 ControlCharacter,
431 NotInVocabulary,
432 DuplicateItem,
433 TooManyItems,
434}
435crate::document::document!(
436 GenerateRequest,
437 crate::schema::GENERATE_REQUEST,
438 crate::limits::RECORD_V1
439);
440
441#[cfg(test)]
442mod tests {
443 use super::RepairRequest;
444 use serde_json::json;
445
446 #[test]
447 fn repair_round_rejects_unknown_fields() {
448 let round = json!({"round": 1, "fields": []});
449 let parsed: RepairRequest = serde_json::from_value(round.clone()).unwrap();
450 assert!(matches!(parsed, RepairRequest::Round { round: 1, .. }));
451 assert_eq!(serde_json::to_value(parsed).unwrap(), round);
452
453 let mut extra = round;
454 extra["unexpected"] = json!(true);
455 assert!(serde_json::from_str::<RepairRequest>(&extra.to_string()).is_err());
456 assert!(serde_json::from_value::<RepairRequest>(extra).is_err());
457 }
458}