use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::bounded::{ParseError, ParseLimits, parse_bounded};
use crate::canonical::{CanonicalError, canonical_bytes, tagged_digest};
use crate::decimal::Decimal;
use crate::definition::ProvenanceClass;
use crate::descriptor::InputExcess;
use crate::execution::FailureClass;
use crate::ids::{ArtifactId, DynamicCapabilityId, DynamicRequestId, TemplateId};
use crate::judgment::Derivation;
use crate::limits::{DESCRIPTOR_V1, DYNAMIC_QUESTION_V1, RECORD_V1};
use crate::remote::{CancelSupport, Disclosure, ServedIdentity, ServedTerms};
use crate::run::{AttemptCost, AttemptEnd, Cancellation, CostModel, CostSummary, RemoteState};
use crate::schema;
use crate::scope::Scope;
use crate::time::Timestamp;
use crate::{Document, DocumentError, Identified};
const KIB: usize = 1024;
pub const MAX_QUESTION_BYTES: usize = 4 * KIB;
pub const MAX_LABEL_BYTES: usize = 128;
pub const MIN_LABELS: usize = 2;
pub const MAX_CLASSIFY_LABELS: usize = 32;
pub const MAX_RUBRIC_LEVELS: usize = 16;
pub const PROPOSITION_LABELS: [&str; 2] = ["false", "true"];
pub const MAX_CONTEXT_FIELDS: usize = 64;
pub const MAX_CONTEXT_VALUE_BYTES: usize = 64 * KIB;
pub const MAX_CANONICAL_BYTES: usize = 128 * KIB;
pub const MAX_NAME_BYTES: usize = 128;
pub const MAX_META_BYTES: usize = 256;
pub const MAX_DETAIL_BYTES: usize = 256;
pub const MAX_CAPABILITY_ENTRIES: usize = 64;
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum DynamicOperation {
Proposition,
Classify,
Rubric,
}
impl DynamicOperation {
pub fn name(self) -> &'static str {
match self {
DynamicOperation::Proposition => "proposition",
DynamicOperation::Classify => "classify",
DynamicOperation::Rubric => "rubric",
}
}
}
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum DynamicOutputKind {
Label,
Distribution,
OrdinalDistribution,
}
impl DynamicOutputKind {
pub fn name(self) -> &'static str {
match self {
DynamicOutputKind::Label => "label",
DynamicOutputKind::Distribution => "distribution",
DynamicOutputKind::OrdinalDistribution => "ordinal_distribution",
}
}
}
#[derive(
Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, schemars::JsonSchema,
)]
#[serde(deny_unknown_fields)]
pub struct TaskFamily {
pub name: String,
pub version: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum ContextValue {
Bool(bool),
Integer(i64),
Decimal(Decimal),
Text(String),
Timestamp(Timestamp),
}
impl ContextValue {
pub fn value_bytes(&self) -> usize {
match self {
ContextValue::Bool(b) => {
if *b {
4
} else {
5
}
}
ContextValue::Integer(i) => i.to_string().len(),
ContextValue::Decimal(d) => d.to_canonical_string().len(),
ContextValue::Text(t) => t.len(),
ContextValue::Timestamp(t) => t.as_ms().to_string().len(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ContextField {
pub name: String,
pub value: ContextValue,
pub provenance: ProvenanceClass,
pub source: String,
pub derivation: Derivation,
pub source_revision: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum AbstentionPolicy {
Thresholds {
min_top_mass: Decimal,
min_margin: Decimal,
},
LabelOnly,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicQuestion {
pub schema: String,
pub request_id: String,
pub scope: Scope,
pub operation: DynamicOperation,
pub task_family: TaskFamily,
pub question: String,
pub labels: Vec<String>,
pub context: Vec<ContextField>,
pub output: DynamicOutputKind,
pub deadline_ms: u64,
pub max_cost_units: u64,
pub artifact: ArtifactId,
pub privacy: Vec<String>,
pub abstention: AbstentionPolicy,
}
impl Document for DynamicQuestion {
const SCHEMA: &'static str = schema::DYNAMIC_QUESTION;
const LIMITS: ParseLimits = DYNAMIC_QUESTION_V1;
fn schema(&self) -> &str {
&self.schema
}
fn parse(bytes: &[u8]) -> Result<Self, DocumentError> {
parse_checked(bytes, DynamicQuestion::check_bounds)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BoundsError(pub String);
impl fmt::Display for BoundsError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for BoundsError {}
fn bounds<T>(msg: impl Into<String>) -> Result<T, BoundsError> {
Err(BoundsError(msg.into()))
}
fn short(s: &str) -> String {
s.chars().take(64).collect()
}
fn sized(what: &str, s: &str, min: usize, max: usize) -> Result<(), BoundsError> {
if s.len() < min || s.len() > max {
return bounds(format!(
"{what} is {} bytes; {min} to {max} are allowed",
s.len()
));
}
Ok(())
}
fn unique<'a>(what: &str, items: impl Iterator<Item = &'a str>) -> Result<(), BoundsError> {
let mut seen = BTreeSet::new();
for s in items {
if !seen.insert(s) {
return bounds(format!("duplicate {what} {:?}", short(s)));
}
}
Ok(())
}
fn parse_checked<T: Document>(
bytes: &[u8],
check: fn(&T) -> Result<(), BoundsError>,
) -> Result<T, DocumentError> {
let doc: T = parse_bounded(bytes, &T::LIMITS).map_err(DocumentError::Parse)?;
if doc.schema() != T::SCHEMA {
return Err(DocumentError::Schema {
expected: T::SCHEMA,
found: doc.schema().chars().take(80).collect(),
});
}
check(&doc).map_err(|e| DocumentError::Parse(ParseError::Schema { message: e.0 }))?;
Ok(doc)
}
impl DynamicQuestion {
pub fn check_bounds(&self) -> Result<(), BoundsError> {
sized("request_id", &self.request_id, 1, MAX_META_BYTES)?;
sized(
"task_family.name",
&self.task_family.name,
1,
MAX_NAME_BYTES,
)?;
sized(
"task_family.version",
&self.task_family.version,
1,
MAX_NAME_BYTES,
)?;
sized("question", &self.question, 1, MAX_QUESTION_BYTES)?;
let max_labels = MAX_CLASSIFY_LABELS.max(MAX_RUBRIC_LEVELS);
if self.labels.len() < MIN_LABELS || self.labels.len() > max_labels {
return bounds(format!(
"{} labels; {MIN_LABELS} to {max_labels} are allowed",
self.labels.len()
));
}
for l in &self.labels {
sized("a label", l, 1, MAX_LABEL_BYTES)?;
}
unique("label", self.labels.iter().map(String::as_str))?;
if self.context.len() > MAX_CONTEXT_FIELDS {
return bounds(format!(
"{} context fields; at most {MAX_CONTEXT_FIELDS} are allowed",
self.context.len()
));
}
let mut value_bytes = 0usize;
for c in &self.context {
sized("a context field name", &c.name, 1, MAX_NAME_BYTES)?;
sized("a context source", &c.source, 1, MAX_META_BYTES)?;
sized(
"a context source revision",
&c.source_revision,
1,
MAX_META_BYTES,
)?;
value_bytes = value_bytes.saturating_add(c.value.value_bytes());
}
unique(
"context field",
self.context.iter().map(|c| c.name.as_str()),
)?;
if value_bytes > MAX_CONTEXT_VALUE_BYTES {
return bounds(format!(
"context values are {value_bytes} bytes; at most {MAX_CONTEXT_VALUE_BYTES} are allowed"
));
}
if self.privacy.len() > MAX_CAPABILITY_ENTRIES {
return bounds("too many privacy options");
}
for p in &self.privacy {
sized("a privacy option", p, 1, MAX_NAME_BYTES)?;
}
unique("privacy option", self.privacy.iter().map(String::as_str))?;
let canonical = self
.canonical()
.map_err(|e| BoundsError(format!("no canonical form: {e}")))?;
if canonical.len() > MAX_CANONICAL_BYTES {
return bounds(format!(
"the request is {} canonical bytes; at most {MAX_CANONICAL_BYTES} are allowed",
canonical.len()
));
}
Ok(())
}
pub fn context_value_bytes(&self) -> usize {
self.context
.iter()
.map(|c| c.value.value_bytes())
.fold(0, usize::saturating_add)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct TemplateRef {
pub name: String,
pub version: String,
pub id: TemplateId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicOperationSupport {
pub operation: DynamicOperation,
pub outputs: Vec<DynamicOutputKind>,
pub max_labels: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum HiddenState {
None,
Present,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicCapability {
pub schema: String,
pub backend_id: String,
pub artifact: ArtifactId,
pub operations: Vec<DynamicOperationSupport>,
pub max_question_bytes: u64,
pub max_context_bytes: u64,
pub max_context_fields: u64,
pub task_families: Vec<TaskFamily>,
pub template: TemplateRef,
pub preprocessing: String,
pub truncation: InputExcess,
pub served: ServedTerms,
pub privacy: Vec<String>,
pub cancellation: CancelSupport,
pub cost: CostModel,
pub hidden_state: HiddenState,
}
impl Document for DynamicCapability {
const SCHEMA: &'static str = schema::DYNAMIC_CAPABILITY;
const LIMITS: ParseLimits = DESCRIPTOR_V1;
fn schema(&self) -> &str {
&self.schema
}
fn parse(bytes: &[u8]) -> Result<Self, DocumentError> {
parse_checked(bytes, DynamicCapability::check_bounds)
}
}
impl Identified for DynamicCapability {
type Id = DynamicCapabilityId;
fn wrap(digest: String) -> DynamicCapabilityId {
DynamicCapabilityId::from_digest(digest)
}
}
impl DynamicCapability {
pub fn check_bounds(&self) -> Result<(), BoundsError> {
sized("backend_id", &self.backend_id, 1, MAX_NAME_BYTES)?;
sized("template.name", &self.template.name, 1, MAX_NAME_BYTES)?;
sized(
"template.version",
&self.template.version,
1,
MAX_NAME_BYTES,
)?;
sized("preprocessing", &self.preprocessing, 1, MAX_NAME_BYTES)?;
if self.operations.is_empty() || self.operations.len() > 3 {
return bounds("one to three operations are declared");
}
let mut ops = BTreeSet::new();
for o in &self.operations {
if !ops.insert(o.operation) {
return bounds(format!("duplicate operation {}", o.operation.name()));
}
if o.outputs.is_empty() || o.outputs.len() > 3 {
return bounds("one to three output kinds per operation");
}
let set: BTreeSet<_> = o.outputs.iter().collect();
if set.len() != o.outputs.len() {
return bounds("duplicate output kind");
}
}
if self.task_families.is_empty() || self.task_families.len() > MAX_CAPABILITY_ENTRIES {
return bounds("one to 64 task families are declared");
}
for t in &self.task_families {
sized("a task family name", &t.name, 1, MAX_NAME_BYTES)?;
sized("a task family version", &t.version, 1, MAX_NAME_BYTES)?;
}
let set: BTreeSet<_> = self.task_families.iter().collect();
if set.len() != self.task_families.len() {
return bounds("duplicate task family");
}
if self.privacy.len() > MAX_CAPABILITY_ENTRIES {
return bounds("too many privacy options");
}
for p in &self.privacy {
sized("a privacy option", p, 1, MAX_NAME_BYTES)?;
}
unique("privacy option", self.privacy.iter().map(String::as_str))?;
let pinned = self.served.disclosure == Disclosure::Pinned;
match &self.served.identity {
Some(i) if pinned => sized("served.identity", i, 1, MAX_META_BYTES)?,
None if !pinned => {}
_ => return bounds("served.identity is present exactly when disclosure is pinned"),
}
Ok(())
}
pub fn support(&self, operation: DynamicOperation) -> Option<&DynamicOperationSupport> {
self.operations.iter().find(|o| o.operation == operation)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicIdentity {
pub scope: Scope,
pub operation: DynamicOperation,
pub task_family: TaskFamily,
pub question: String,
pub labels: Vec<String>,
pub context: Vec<ContextField>,
pub output: DynamicOutputKind,
pub artifact: ArtifactId,
pub privacy: Vec<String>,
pub template: TemplateRef,
pub preprocessing: String,
}
impl DynamicIdentity {
pub fn of(q: &DynamicQuestion, capability: &DynamicCapability) -> DynamicIdentity {
DynamicIdentity {
scope: q.scope.clone(),
operation: q.operation,
task_family: q.task_family.clone(),
question: q.question.clone(),
labels: q.labels.clone(),
context: q.context.clone(),
output: q.output,
artifact: q.artifact.clone(),
privacy: q.privacy.clone(),
template: capability.template.clone(),
preprocessing: capability.preprocessing.clone(),
}
}
pub fn id(&self) -> Result<DynamicRequestId, CanonicalError> {
Ok(DynamicRequestId::from_digest(tagged_digest(
schema::DYNAMIC_QUESTION,
&canonical_bytes(self)?,
)))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicProjection {
pub operation: DynamicOperation,
pub task_family: TaskFamily,
pub question: String,
pub labels: Vec<String>,
pub context: Vec<ProjectedField>,
pub output: DynamicOutputKind,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ProjectedField {
pub name: String,
pub value: ContextValue,
}
impl DynamicProjection {
pub fn of(q: &DynamicQuestion) -> DynamicProjection {
DynamicProjection {
operation: q.operation,
task_family: q.task_family.clone(),
question: q.question.clone(),
labels: q.labels.clone(),
context: q
.context
.iter()
.map(|c| ProjectedField {
name: c.name.clone(),
value: c.value.clone(),
})
.collect(),
output: q.output,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicRawOutput {
Label(String),
Distribution(BTreeMap<String, f64>),
Unrecognized {
shape: ResponseShape,
excerpt: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum ResponseShape {
ToolCall,
Instruction,
Schema,
UrlFetch,
Action,
FreeText,
NonFiniteMass,
Other,
}
impl ResponseShape {
pub fn name(self) -> &'static str {
match self {
ResponseShape::ToolCall => "tool_call",
ResponseShape::Instruction => "instruction",
ResponseShape::Schema => "schema",
ResponseShape::UrlFetch => "url_fetch",
ResponseShape::Action => "action",
ResponseShape::FreeText => "free_text",
ResponseShape::NonFiniteMass => "non_finite_mass",
ResponseShape::Other => "other",
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicOutcome {
pub schema: String,
pub request_id: String,
pub request: Option<DynamicRequestId>,
pub scope: Scope,
pub task_family: TaskFamily,
pub operation: DynamicOperation,
pub output: DynamicOutputKind,
pub backend_id: String,
pub requested_artifact: ArtifactId,
pub served: ServedIdentity,
pub capability: Option<DynamicCapabilityId>,
pub template: Option<TemplateRef>,
pub privacy: Vec<String>,
pub attempt: Option<DynamicAttempt>,
pub timing: DynamicTiming,
pub cost: CostSummary,
pub raw: DynamicRaw,
pub validation: DynamicValidation,
pub abstention: AbstentionDecision,
pub result: DynamicResult,
}
impl Document for DynamicOutcome {
const SCHEMA: &'static str = schema::DYNAMIC_OUTCOME;
const LIMITS: ParseLimits = RECORD_V1;
fn schema(&self) -> &str {
&self.schema
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicAttempt {
pub attempt_id: String,
pub dispatched_ms: u64,
pub ended_ms: u64,
pub end: AttemptEnd,
pub cancellation: Cancellation,
pub remote: RemoteState,
pub cost: AttemptCost,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicTiming {
pub elapsed_ms: u64,
pub deadline_ms: u64,
pub deadline_expired: bool,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicRaw {
NotDispatched,
Output(DynamicRawOutput),
Failed { class: FailureClass, detail: String },
NoAnswer,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicValidation {
NotRun,
Valid,
Invalid { detail: String },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum AbstentionRule {
MinTopMass,
MinMargin,
}
impl AbstentionRule {
pub fn name(self) -> &'static str {
match self {
AbstentionRule::MinTopMass => "min_top_mass",
AbstentionRule::MinMargin => "min_margin",
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum AbstentionDecision {
NotApplied,
LabelOnly,
Passed {
top_mass: f64,
margin: f64,
},
Abstained {
rule: AbstentionRule,
top_mass: f64,
margin: f64,
},
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicResult {
Proposal(DynamicProposal),
Unresolved(DynamicUnresolved),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicProposal {
pub answer: DynamicAnswer,
pub derivation: Derivation,
pub lineage: DynamicLineage,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct LabelMass {
pub label: String,
pub mass: f64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicAnswer {
Label {
label: String,
},
UncalibratedDistribution {
label: String,
masses: Vec<LabelMass>,
},
OrdinalDistribution {
level: String,
masses: Vec<LabelMass>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct DynamicLineage {
pub request: DynamicRequestId,
pub context: Vec<String>,
pub backend_id: String,
pub artifact: ArtifactId,
pub attempt_id: String,
pub exact_processing: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicUnresolved {
Refused {
refusal: DynamicRefusal,
},
BackendFailed {
class: FailureClass,
detail: String,
},
InvalidOutput {
detail: String,
},
ServedIdentityMismatch {
pinned: String,
served: Option<String>,
},
Abstained {
rule: AbstentionRule,
},
DeadlineExceeded,
Cancelled,
}
impl DynamicUnresolved {
pub fn name(&self) -> &'static str {
match self {
DynamicUnresolved::Refused { .. } => "refused",
DynamicUnresolved::BackendFailed { .. } => "backend_failed",
DynamicUnresolved::InvalidOutput { .. } => "invalid_output",
DynamicUnresolved::ServedIdentityMismatch { .. } => "served_identity_mismatch",
DynamicUnresolved::Abstained { .. } => "abstained",
DynamicUnresolved::DeadlineExceeded => "deadline_exceeded",
DynamicUnresolved::Cancelled => "cancelled",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum DynamicRefusal {
NoDynamicCapability,
CapabilityInvalid {
detail: String,
},
ArtifactMismatch {
declared: ArtifactId,
},
OperationUnsupported {
operation: DynamicOperation,
output: DynamicOutputKind,
},
TaskFamilyUndeclared,
ExceedsDeclaredLimit {
limit: String,
declared: u64,
requested: u64,
},
PrivacyUnsupported {
option: String,
},
CostUnenforceable {
model: CostModel,
},
CostUnknown,
CostAboveCap {
bound: u64,
cap: u64,
},
BudgetExhausted,
}