1use std::collections::{BTreeMap, BTreeSet};
3
4use crate::{generation_eval, kinds::Distribution, value::Value};
5use rustev_contract::{
6 Decimal,
7 extraction::{AbsenceReason as Reason, Distributions, ExtractionSchema, ExtractionType},
8 extraction_results::{
9 DistributionAbsence, ExtractionCalibration, ExtractionFieldEvidence,
10 ExtractionItemEvidence, ExtractionOrigin, FieldDistribution, RejectedExtractionItem,
11 },
12 generation::{FieldInvalid, GeneratedField, GeneratedType, Lines},
13 ids::ContentDigest,
14 judgment::Derivation,
15 output::{FieldLogprobs, LogprobItem, RawOutput, TextValue},
16 time::Timestamp,
17};
18use serde_json::Value as Json;
19
20type Vocabularies = BTreeMap<String, Vec<String>>;
21
22#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct InvalidOutput(pub &'static str);
26impl std::fmt::Display for InvalidOutput {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 f.write_str(self.0)
29 }
30}
31impl std::error::Error for InvalidOutput {}
32
33#[derive(Debug, Clone, PartialEq)]
35pub struct ValidatedField {
36 pub value: Option<Value>,
37 pub evidence: ExtractionFieldEvidence,
38}
39
40pub fn fallback_derivation(sources: impl IntoIterator<Item = Derivation>) -> Derivation {
42 if sources.into_iter().all(|d| d == Derivation::ExactDerived) {
43 Derivation::ExactDerived
44 } else {
45 Derivation::MixedDerived
46 }
47}
48
49fn digest(value: &Json) -> Result<ContentDigest, InvalidOutput> {
50 let s = rustev_contract::grade::value_digest(value)
51 .map_err(|_| InvalidOutput("value has no canonical encoding"))?;
52 ContentDigest::parse(&s).map_err(|_| InvalidOutput("invalid value digest"))
53}
54
55fn options(ty: &ExtractionType, vocab: &Vocabularies) -> Option<Vec<String>> {
56 match ty {
57 ExtractionType::Bool {} => Some(vec!["false".into(), "true".into()]),
58 ExtractionType::Enum { vocabulary } => vocab.get(vocabulary).cloned(),
59 _ => None,
60 }
61}
62
63pub fn validate(
66 schema: &ExtractionSchema,
67 vocab: &Vocabularies,
68 raw: &RawOutput,
69 distributions: Distributions,
70 option_logprobs: bool,
71) -> Result<Vec<ValidatedField>, InvalidOutput> {
72 let RawOutput::TypedRecord {
73 fields,
74 logprobs,
75 provider_reported,
76 } = raw
77 else {
78 return Err(InvalidOutput("expected typed_record"));
79 };
80 if fields
81 .keys()
82 .chain(provider_reported.keys())
83 .any(|n| !schema.fields.iter().any(|f| &f.name == n))
84 {
85 return Err(InvalidOutput("undeclared member"));
86 }
87 if provider_reported.values().any(|s| s.len() > 256) {
88 return Err(InvalidOutput("provider_reported exceeds 256 bytes"));
89 }
90 if !option_logprobs && !logprobs.is_empty() {
91 return Err(InvalidOutput("binding does not support option_logprobs"));
92 }
93 let mut seen = BTreeSet::new();
94 for lp in logprobs {
95 let field = schema
96 .fields
97 .iter()
98 .find(|f| f.name == lp.field)
99 .ok_or(InvalidOutput("undeclared logprob field"))?;
100 let index = match lp.item {
101 LogprobItem::None => None,
102 LogprobItem::Index(i) => Some(i),
103 };
104 if !seen.insert((&lp.field, index)) {
105 return Err(InvalidOutput("duplicate logprob target"));
106 }
107 let ty = match (&field.ty, &lp.item) {
108 (ExtractionType::List { item, .. }, LogprobItem::Index(i))
109 if fields
110 .get(&field.name)
111 .and_then(Json::as_array)
112 .is_some_and(|v| (*i as u128) < v.len() as u128) =>
113 {
114 item.as_ref()
115 }
116 (ty, LogprobItem::None) => ty,
117 _ => return Err(InvalidOutput("invalid logprob item target")),
118 };
119 let declared =
120 options(ty, vocab).ok_or(InvalidOutput("logprob target has no resolved options"))?;
121 if lp
122 .values
123 .iter()
124 .any(|(k, v)| !declared.contains(k) || !v.is_finite())
125 {
126 return Err(InvalidOutput("undeclared option or non-finite logprob"));
127 }
128 }
129 schema
130 .fields
131 .iter()
132 .map(|field| {
133 let raw_value = fields.get(&field.name);
134 let mut rejected = vec![];
135 let mut indices = vec![];
136 let result = field_value(&field.ty, raw_value, vocab, &mut rejected, &mut indices)?;
137 let (value, reason, subfield) = match result {
138 Ok(v) => (Some(v), None, None),
139 Err((r, sub)) => (None, Some(r), sub),
140 };
141 let origin = reason.map_or(ExtractionOrigin::Extracted, |reason| {
142 ExtractionOrigin::Absent { reason }
143 });
144 let derivation = value.as_ref().map(|_| Derivation::ModelDerived);
145 let value_digest = value.as_ref().map(|v| digest(&v.to_json())).transpose()?;
146 let mut items = vec![];
147 if let (Some(Value::List(values)), ExtractionType::List { item, .. }) =
148 (&value, &field.ty)
149 {
150 for (v, index) in values.iter().zip(indices) {
151 items.push(ExtractionItemEvidence {
152 index,
153 value_digest: digest(&v.to_json())?,
154 value: Some(v.to_json()),
156 distribution: distribution(
157 item,
158 &field.name,
159 LogprobItem::Index(index),
160 vocab,
161 logprobs,
162 &distributions,
163 option_logprobs,
164 ),
165 calibration: ExtractionCalibration::None,
166 });
167 }
168 }
169 let retained = value.as_ref().map(|v| {
171 if matches!(field.ty, ExtractionType::Date {}) {
172 raw_value.cloned().unwrap_or(Json::Null)
173 } else {
174 v.to_json()
175 }
176 });
177 let mut evidence = ExtractionFieldEvidence::new(
178 field.name.clone(),
179 &field.ty,
180 origin,
181 derivation,
182 value_digest,
183 retained,
184 items,
185 );
186 evidence.subfield = subfield;
187 evidence.rejected = rejected;
188 evidence.provider_reported = provider_reported.get(&field.name).cloned();
189 if value.is_some() {
190 evidence.distribution = distribution(
191 &field.ty,
192 &field.name,
193 LogprobItem::None,
194 vocab,
195 logprobs,
196 &distributions,
197 option_logprobs,
198 );
199 }
200 Ok(ValidatedField { value, evidence })
201 })
202 .collect()
203}
204
205fn distribution(
206 ty: &ExtractionType,
207 field: &str,
208 item: LogprobItem,
209 vocab: &Vocabularies,
210 logprobs: &[FieldLogprobs],
211 mode: &Distributions,
212 capable: bool,
213) -> FieldDistribution {
214 let absent = |reason| FieldDistribution::Absent { reason };
215 if !capable || *mode != Distributions::WhenAvailable {
216 return absent(DistributionAbsence::NotReturned);
217 }
218 let Some(options) = options(ty, vocab) else {
219 return absent(DistributionAbsence::NotReturned);
220 };
221 let Some(lp) = logprobs
222 .iter()
223 .find(|lp| lp.field == field && lp.item == item)
224 else {
225 return absent(DistributionAbsence::NotReturned);
226 };
227 match Distribution::from_logits(&options, &lp.values, Decimal::from_units(1)) {
229 Ok(d) => FieldDistribution::Masses(d.masses().to_vec()),
230 Err(_) => absent(DistributionAbsence::IncompleteLogprobs),
231 }
232}
233
234type FieldResult = Result<Value, (Reason, Option<String>)>;
235fn field_value(
236 ty: &ExtractionType,
237 raw: Option<&Json>,
238 vocab: &Vocabularies,
239 rejected: &mut Vec<RejectedExtractionItem>,
240 indices: &mut Vec<u64>,
241) -> Result<FieldResult, InvalidOutput> {
242 let Some(raw) = raw.filter(|v| !v.is_null()) else {
243 return Ok(Err((Reason::Missing, None)));
244 };
245 let result = match ty {
246 ExtractionType::List { item, max_items } => {
247 let Some(array) = raw.as_array() else {
248 return Ok(Err((Reason::WrongType, None)));
249 };
250 if array.len() as u64 > *max_items {
251 return Ok(Err((Reason::TooManyItems, None)));
252 }
253 let mut values = vec![];
254 for (i, raw) in array.iter().enumerate() {
255 match scalar(item, raw, vocab) {
256 Ok(v) => {
257 values.push(v);
258 indices.push(i as u64);
259 }
260 Err(reason) => rejected.push(RejectedExtractionItem {
261 index: i as u64,
262 reason,
263 value_digest: digest(raw)?,
264 }),
265 }
266 }
267 if values
268 .iter()
269 .enumerate()
270 .any(|(i, v)| values[..i].contains(v))
271 {
272 Err(Reason::DuplicateItem)
273 } else if !array.is_empty() && values.is_empty() {
274 Err(rejected
275 .first()
276 .map(|r| r.reason)
277 .unwrap_or(Reason::WrongType))
278 } else {
279 Ok(Value::List(values))
280 }
281 }
282 ExtractionType::Record { fields } => {
283 let Some(object) = raw.as_object() else {
284 return Ok(Err((Reason::WrongType, None)));
285 };
286 if object.keys().any(|n| !fields.iter().any(|f| &f.name == n)) {
287 return Err(InvalidOutput("undeclared record subfield"));
288 }
289 let mut values = BTreeMap::new();
290 for f in fields {
291 let result = match object.get(&f.name).filter(|v| !v.is_null()) {
292 Some(v) => scalar(&f.ty, v, vocab),
293 None => Err(Reason::Missing),
294 };
295 match result {
296 Ok(v) => {
297 values.insert(f.name.clone(), v);
298 }
299 Err(r) => return Ok(Err((r, Some(f.name.clone())))),
300 }
301 }
302 Ok(Value::Record(values))
303 }
304 _ => scalar(ty, raw, vocab),
305 };
306 Ok(result.map_err(|r| (r, None)))
307}
308
309fn shared(ty: GeneratedType, value: TextValue, vocab: &Vocabularies) -> Result<(), Reason> {
312 let field = GeneratedField {
313 name: String::new(),
314 ty,
315 required: false,
316 };
317 generation_eval::validate_field(&field, Some(&value), vocab)
318 .map(|_| ())
319 .map_err(|r| match r {
320 FieldInvalid::Missing => Reason::Missing,
321 FieldInvalid::WrongType | FieldInvalid::Empty => Reason::WrongType,
322 FieldInvalid::TooLong => Reason::TooLong,
323 FieldInvalid::ControlCharacter => Reason::ControlCharacter,
324 FieldInvalid::NotInVocabulary => Reason::NotInVocabulary,
325 FieldInvalid::DuplicateItem => Reason::DuplicateItem,
326 FieldInvalid::TooManyItems => Reason::TooManyItems,
327 })
328}
329
330fn scalar(ty: &ExtractionType, raw: &Json, vocab: &Vocabularies) -> Result<Value, Reason> {
331 match ty {
332 ExtractionType::Bool {} => raw.as_bool().map(Value::Bool).ok_or(Reason::WrongType),
333 ExtractionType::Integer { bounds: [lo, hi] } => {
334 let v = raw.as_i64().ok_or_else(|| {
335 if raw.is_u64() {
336 Reason::OutOfBounds
337 } else {
338 Reason::WrongType
339 }
340 })?;
341 if v < *lo || v > *hi {
342 Err(Reason::OutOfBounds)
343 } else {
344 Ok(Value::Integer(v))
345 }
346 }
347 ExtractionType::Decimal(bounds) => {
348 let s = raw.as_str().ok_or(Reason::WrongType)?;
349 let unsigned = s.strip_prefix('-').unwrap_or(s);
351 let (whole, frac) = unsigned
352 .split_once('.')
353 .map_or((unsigned, None), |(a, b)| (a, Some(b)));
354 if whole.is_empty()
355 || !whole.bytes().all(|b| b.is_ascii_digit())
356 || (whole.len() > 1 && whole.starts_with('0'))
357 || frac.is_some_and(|f| f.is_empty() || !f.bytes().all(|b| b.is_ascii_digit()))
358 {
359 return Err(Reason::WrongType);
360 }
361 if frac.is_some_and(|f| f.len() > bounds.scale as usize) {
362 return Err(Reason::OutOfBounds);
363 }
364 let v = Decimal::parse(s).map_err(|_| Reason::OutOfBounds)?;
365 let lo = Decimal::parse(&bounds.bounds[0]).map_err(|_| Reason::OutOfBounds)?;
366 let hi = Decimal::parse(&bounds.bounds[1]).map_err(|_| Reason::OutOfBounds)?;
367 if v < lo || v > hi {
368 Err(Reason::OutOfBounds)
369 } else {
370 Ok(Value::Decimal(v))
371 }
372 }
373 ExtractionType::Text { max_bytes } => {
374 let s = raw.as_str().ok_or(Reason::WrongType)?;
375 if !s.is_empty() {
376 shared(
377 GeneratedType::Text {
378 max_bytes: *max_bytes,
379 lines: Lines::Single,
380 },
381 TextValue::Text(s.into()),
382 vocab,
383 )?;
384 }
385 Ok(Value::Text(s.into()))
386 }
387 ExtractionType::Enum { vocabulary } => {
388 let s = raw.as_str().ok_or(Reason::WrongType)?;
389 if !vocab.contains_key(vocabulary) {
390 return Err(Reason::VocabularyUnavailable);
391 }
392 shared(
393 GeneratedType::List {
394 vocabulary: vocabulary.clone(),
395 max_items: 1,
396 },
397 TextValue::List(vec![s.into()]),
398 vocab,
399 )?;
400 Ok(Value::Enum(s.into()))
401 }
402 ExtractionType::Date {} => {
403 date(raw.as_str().ok_or(Reason::WrongType)?).map(Value::Timestamp)
404 }
405 _ => Err(Reason::WrongType),
406 }
407}
408
409fn date(s: &str) -> Result<Timestamp, Reason> {
410 let b = s.as_bytes();
411 if b.len() != 10
412 || b[4] != b'-'
413 || b[7] != b'-'
414 || b.iter()
415 .enumerate()
416 .any(|(i, b)| i != 4 && i != 7 && !b.is_ascii_digit())
417 {
418 return Err(Reason::InvalidDate);
419 }
420 let number = |a: usize, z: usize| {
421 b[a..z]
422 .iter()
423 .fold(0i64, |n, b| n * 10 + i64::from(b - b'0'))
424 };
425 let (year, month, day) = (number(0, 4), number(5, 7), number(8, 10));
426 let leap = |y: i64| y % 4 == 0 && (y % 100 != 0 || y % 400 == 0);
427 let days = |m| match m {
428 2 => {
429 if leap(year) {
430 29
431 } else {
432 28
433 }
434 }
435 4 | 6 | 9 | 11 => 30,
436 _ => 31,
437 };
438 if !(1..=12).contains(&month) || day < 1 || day > days(month) {
439 return Err(Reason::InvalidDate);
440 }
441 if year < 1970 {
444 return Err(Reason::OutOfBounds);
445 }
446 let prior = year - 1;
447 let leap_days = prior / 4 - prior / 100 + prior / 400 - (1969 / 4 - 1969 / 100 + 1969 / 400);
448 let total = (year - 1970) * 365 + leap_days + (1..month).map(days).sum::<i64>() + day - 1;
449 Timestamp::from_ms(total * 86_400_000).map_err(|_| Reason::OutOfBounds)
450}