1use std::collections::BTreeMap;
3
4use serde::{Deserialize, Serialize};
5
6use crate::extraction::ExtractionSchema;
7use crate::ids::{ContentDigest, PromptTemplateId};
8
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
10#[serde(deny_unknown_fields)]
11pub struct ExtractRules {
12 pub schema: String,
13 pub name: String,
14 pub version: String,
15 pub sources: Vec<ExtractionSource>,
16 pub fields: Vec<FieldRules>,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
20#[serde(deny_unknown_fields)]
21pub struct ExtractionSource {
22 pub name: String,
23 #[serde(rename = "ref", deserialize_with = "crate::extraction::input_ref")]
24 pub reference: String,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
28#[serde(deny_unknown_fields)]
29pub struct FieldRules {
30 pub field: String,
31 pub rules: Vec<ExtractionRule>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
35#[serde(deny_unknown_fields)]
36pub struct ExtractionRule {
37 pub name: String,
38 pub when: ExtractionCondition,
39 pub then: ExtractionValue,
40}
41
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
43#[serde(rename_all = "snake_case", deny_unknown_fields)]
44pub enum ExtractionCondition {
45 Always,
46 ContainsAny { source: String, terms: Vec<String> },
47 ContainsAll { source: String, terms: Vec<String> },
48 All(Vec<ExtractionCondition>),
49 Any(Vec<ExtractionCondition>),
50 Not(Box<ExtractionCondition>),
51}
52
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
54#[serde(rename_all = "snake_case", deny_unknown_fields)]
55pub enum ExtractionValue {
56 Literal(serde_json::Value),
57 IntegerAfter {
58 source: String,
59 terms: Vec<String>,
60 max_digits: u8,
61 },
62 DateIso {
63 source: String,
64 occurrence: u8,
65 },
66 Record(BTreeMap<String, ExtractionValue>),
67}
68
69crate::document::versioned_document!(
70 ExtractRules,
71 [crate::schema::EXTRACT_RULES],
72 "rustev.extract-rules/1",
73 crate::limits::DESCRIPTOR_V1,
74 check = ExtractRules::validate
75);
76
77impl crate::Identified for ExtractRules {
78 type Id = ContentDigest;
79 fn wrap(digest: String) -> ContentDigest {
80 ContentDigest::from_digest(digest)
81 }
82}
83
84fn terms(v: &[String]) -> Result<(), String> {
85 if v.len() > 64 || v.iter().any(|s| s.is_empty() || s.len() > 64) {
89 return Err("extraction terms require at most 64 entries of 1 to 64 bytes".into());
90 }
91 Ok(())
92}
93
94impl ExtractionCondition {
95 fn validate(&self, depth: usize) -> Result<(), String> {
96 if depth > 8 {
97 return Err("extraction condition depth exceeds 8".into());
98 }
99 match self {
100 Self::Always => Ok(()),
101 Self::ContainsAny { terms: t, .. } | Self::ContainsAll { terms: t, .. } => terms(t),
102 Self::All(v) | Self::Any(v) => v.iter().try_for_each(|c| c.validate(depth + 1)),
103 Self::Not(c) => c.validate(depth + 1),
104 }
105 }
106}
107
108impl ExtractionValue {
109 fn validate(&self) -> Result<(), String> {
110 match self {
111 Self::Literal(_) => Ok(()),
112 Self::IntegerAfter {
113 terms: t,
114 max_digits,
115 ..
116 } => {
117 terms(t)?;
118 if !(1..=18).contains(max_digits) {
119 return Err("max_digits must be 1 to 18".into());
120 }
121 Ok(())
122 }
123 Self::DateIso { occurrence, .. } => {
124 if !(1..=2).contains(occurrence) {
125 return Err("occurrence must be 1 or 2".into());
126 }
127 Ok(())
128 }
129 Self::Record(v) => v.values().try_for_each(|value| {
130 if matches!(value, Self::Record(_)) {
131 return Err("a record value's subfields are scalar".into());
132 }
133 value.validate()
134 }),
135 }
136 }
137}
138
139impl ExtractRules {
140 fn validate(&self) -> Result<(), String> {
141 if self.fields.len() > 32 {
142 return Err("extraction program exceeds 32 fields".into());
143 }
144 for f in &self.fields {
145 if f.rules.len() > 128 {
146 return Err("extraction field exceeds 128 rules".into());
147 }
148 for r in &f.rules {
149 r.when.validate(1)?;
150 r.then.validate()?;
151 }
152 }
153 Ok(())
154 }
155}
156
157#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
159#[serde(deny_unknown_fields)]
160pub struct ExtractRequest {
161 pub schema: String,
162 pub step: String,
163 pub template: String,
164 pub template_id: PromptTemplateId,
165 pub instructions: String,
166 pub projection: BTreeMap<String, serde_json::Value>,
167 pub output: ExtractionSchema,
168 pub schema_digest: ContentDigest,
169 pub vocabularies: BTreeMap<String, Vec<String>>,
170}
171
172crate::document::versioned_document!(
173 ExtractRequest,
174 [crate::schema::EXTRACT_REQUEST],
175 "rustev.extract-request/1",
176 crate::limits::RECORD_V1,
177 check = ExtractRequest::validate
178);
179
180impl ExtractRequest {
181 fn validate(&self) -> Result<(), String> {
182 if self.schema_digest
183 != extraction_schema_digest(&self.output).map_err(|e| e.to_string())?
184 {
185 return Err("extraction schema digest disagrees with output".into());
186 }
187 Ok(())
188 }
189}
190
191pub fn extraction_schema_digest(
194 schema: &ExtractionSchema,
195) -> Result<ContentDigest, crate::canonical::CanonicalError> {
196 crate::generation_result::generation_value_digest(schema)
197}