1use std::collections::BTreeMap;
5
6use serde::{Deserialize, Serialize};
7
8use crate::ids::ArtifactId;
9use crate::limits::BACKEND_OUTPUT_V1;
10use crate::schema;
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
13#[serde(deny_unknown_fields)]
14pub struct BackendOutputDoc {
15 pub schema: String,
16 pub step: String,
17 pub instance: Vec<String>,
18 pub artifact: ArtifactId,
20 pub output: RawOutput,
21}
22
23crate::document::document!(BackendOutputDoc, schema::BACKEND_OUTPUT, BACKEND_OUTPUT_V1);
24
25#[derive(Debug, Clone, PartialEq, Serialize, schemars::JsonSchema)]
27#[serde(rename_all = "snake_case", deny_unknown_fields)]
28pub enum RawOutput {
29 Ranking {
30 #[schemars(length(max = 1024))]
31 ids: Vec<String>,
32 scores: RankingScores,
33 },
34 Label(String),
35 Distribution(BTreeMap<String, f64>),
36 Logits(BTreeMap<String, f64>),
37 Scores(BTreeMap<String, f64>),
38 Screen {
39 blocked: bool,
40 categories: Vec<String>,
41 units: ScreenUnits,
42 },
43 TextRecord(BTreeMap<String, TextValue>),
44 TypedRecord {
45 fields: BTreeMap<String, serde_json::Value>,
46 logprobs: Vec<FieldLogprobs>,
47 #[serde(deserialize_with = "provider_reported")]
48 provider_reported: BTreeMap<String, String>,
49 },
50}
51
52#[derive(Deserialize)]
53#[serde(rename_all = "snake_case", deny_unknown_fields)]
54enum RawOutputWire {
55 Ranking {
56 ids: Vec<String>,
57 scores: RankingScores,
58 },
59 Label(String),
60 Distribution(BTreeMap<String, f64>),
61 Logits(BTreeMap<String, f64>),
62 Scores(BTreeMap<String, f64>),
63 Screen {
64 blocked: bool,
65 categories: Vec<String>,
66 units: ScreenUnits,
67 },
68 TextRecord(BTreeMap<String, TextValue>),
69 TypedRecord {
70 fields: BTreeMap<String, serde_json::Value>,
71 logprobs: Vec<FieldLogprobs>,
72 #[serde(deserialize_with = "provider_reported")]
73 provider_reported: BTreeMap<String, String>,
74 },
75}
76
77impl<'de> Deserialize<'de> for RawOutput {
78 fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
79 Ok(match RawOutputWire::deserialize(d)? {
80 RawOutputWire::Ranking { ids, scores } => {
81 if ids.len() > 1024 || matches!(&scores, RankingScores::List(s) if s.len() > 1024) {
82 return Err(serde::de::Error::custom("ranking output bound exceeded"));
83 }
84 if let RankingScores::List(s) = &scores
85 && (s.is_empty() || s.len() != ids.len() || s.iter().any(|n| !n.is_finite()))
86 {
87 return Err(serde::de::Error::custom(
88 "ranking scores must be nonempty, finite and match ids length",
89 ));
90 }
91 Self::Ranking { ids, scores }
92 }
93 RawOutputWire::Label(v) => Self::Label(v),
94 RawOutputWire::Distribution(v) => Self::Distribution(v),
95 RawOutputWire::Logits(v) => Self::Logits(v),
96 RawOutputWire::Scores(v) => Self::Scores(v),
97 RawOutputWire::Screen {
98 blocked,
99 categories,
100 units,
101 } => Self::Screen {
102 blocked,
103 categories,
104 units,
105 },
106 RawOutputWire::TextRecord(v) => Self::TextRecord(v),
107 RawOutputWire::TypedRecord {
108 fields,
109 logprobs,
110 provider_reported,
111 } => Self::TypedRecord {
112 fields,
113 logprobs,
114 provider_reported,
115 },
116 })
117 }
118}
119
120#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
122#[serde(rename_all = "snake_case", deny_unknown_fields)]
123pub enum RankingScores {
124 None,
125 #[serde(untagged)]
129 List(#[schemars(length(min = 1, max = 1024))] Vec<f64>),
130}
131
132#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
134#[serde(rename_all = "snake_case", deny_unknown_fields)]
135pub enum ScreenUnits {
136 Reported { n: u64 },
137 NotReported,
138}
139
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
142#[serde(rename_all = "snake_case", deny_unknown_fields)]
143pub enum TextValue {
144 Text(String),
145 List(Vec<String>),
146}
147
148#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
150#[serde(deny_unknown_fields)]
151pub struct FieldLogprobs {
152 pub field: String,
153 pub item: LogprobItem,
154 #[serde(deserialize_with = "finite_logprobs")]
155 pub values: BTreeMap<String, f64>,
156}
157
158#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
159#[serde(rename_all = "snake_case")]
160pub enum LogprobItem {
161 None,
162 #[serde(untagged)]
164 Index(u64),
165}
166
167fn finite_logprobs<'de, D: serde::Deserializer<'de>>(
168 d: D,
169) -> Result<BTreeMap<String, f64>, D::Error> {
170 let v = BTreeMap::<String, f64>::deserialize(d)?;
171 if v.values().any(|n| !n.is_finite()) {
172 return Err(serde::de::Error::custom("non-finite extraction logprob"));
173 }
174 Ok(v)
175}
176
177fn provider_reported<'de, D: serde::Deserializer<'de>>(
178 d: D,
179) -> Result<BTreeMap<String, String>, D::Error> {
180 let v = BTreeMap::<String, String>::deserialize(d)?;
181 if v.values().any(|s| s.len() > 256) {
182 return Err(serde::de::Error::custom(
183 "provider_reported exceeds 256 bytes",
184 ));
185 }
186 Ok(v)
187}