jevapi 0.1.0

//! Typed representations of the evaluation API's JSON schema.

use crate::{DEFAULT_MODEL, ValidationError};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use std::collections::BTreeMap;

/// Text or structured content. Top-level null, booleans, and numbers are excluded.
/// Nested object/array values may contain any JSON value.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Content {
    /// Plain text.
    Text(String),
    /// A JSON object.
    Object(Map<String, Value>),
    /// A JSON array.
    Array(Vec<Value>),
}

impl From<&str> for Content {
    fn from(value: &str) -> Self {
        Self::Text(value.into())
    }
}
impl From<String> for Content {
    fn from(value: String) -> Self {
        Self::Text(value)
    }
}
impl From<Map<String, Value>> for Content {
    fn from(value: Map<String, Value>) -> Self {
        Self::Object(value)
    }
}
impl From<Vec<Value>> for Content {
    fn from(value: Vec<Value>) -> Self {
        Self::Array(value)
    }
}
impl TryFrom<Value> for Content {
    type Error = ValidationError;

    fn try_from(value: Value) -> Result<Self, Self::Error> {
        match value {
            Value::String(value) => Ok(Self::Text(value)),
            Value::Object(value) => Ok(Self::Object(value)),
            Value::Array(value) => Ok(Self::Array(value)),
            _ => Err(ValidationError::new(
                "content",
                "expected a string, object, or array",
            )),
        }
    }
}

/// Optional descriptions of what yes and no mean for a Noul question.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct NoulCriteria {
    /// Meaning of a value near one; serialized as `true`.
    #[serde(rename = "true", skip_serializing_if = "Option::is_none")]
    pub yes: Option<Content>,
    /// Meaning of a value near zero; serialized as `false`.
    #[serde(rename = "false", skip_serializing_if = "Option::is_none")]
    pub no: Option<Content>,
}

impl NoulCriteria {
    /// Set the description for yes.
    #[must_use]
    pub fn yes(mut self, description: impl Into<Content>) -> Self {
        self.yes = Some(description.into());
        self
    }

    /// Set the description for no.
    #[must_use]
    pub fn no(mut self, description: impl Into<Content>) -> Self {
        self.no = Some(description.into());
        self
    }
}

/// A typed question, serialized with the API's `type` discriminator.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum Question {
    /// A yes/no question returning the probability of yes.
    Noul {
        /// The question and optional structured context.
        instructions: Content,
        /// Optional descriptions of yes and no.
        #[serde(skip_serializing_if = "Option::is_none")]
        criteria: Option<NoulCriteria>,
    },
    /// Choose one of at most 255 named options.
    Choice {
        /// What the model should decide.
        instructions: Content,
        /// Options and optional descriptions. `None` is serialized as JSON null.
        criteria: BTreeMap<String, Option<Content>>,
    },
    /// Score against two to ten ordered levels.
    Score {
        /// What the model should rate.
        instructions: Content,
        /// Level descriptions, in increasing score order.
        criteria: Vec<Content>,
    },
}

impl Question {
    /// Construct a yes/no question without a rubric.
    pub fn noul(instructions: impl Into<Content>) -> Self {
        Self::Noul {
            instructions: instructions.into(),
            criteria: None,
        }
    }

    /// Construct a yes/no question with descriptions of yes and/or no.
    pub fn noul_with_criteria(instructions: impl Into<Content>, criteria: NoulCriteria) -> Self {
        Self::Noul {
            instructions: instructions.into(),
            criteria: Some(criteria),
        }
    }

    /// Construct a choice question. Use `None` for an option without a description.
    /// Repeated option names replace earlier entries, as with a map.
    pub fn choice<K: Into<String>>(
        instructions: impl Into<Content>,
        criteria: impl IntoIterator<Item = (K, Option<Content>)>,
    ) -> Self {
        Self::Choice {
            instructions: instructions.into(),
            criteria: criteria
                .into_iter()
                .map(|(key, value)| (key.into(), value))
                .collect(),
        }
    }

    /// Construct a score question from ordered level descriptions.
    pub fn score<C: Into<Content>>(
        instructions: impl Into<Content>,
        criteria: impl IntoIterator<Item = C>,
    ) -> Self {
        Self::Score {
            instructions: instructions.into(),
            criteria: criteria.into_iter().map(Into::into).collect(),
        }
    }

    /// Check rubric sizes. This also runs automatically before sending a request.
    pub fn validate(&self) -> Result<(), ValidationError> {
        match self {
            Self::Choice { criteria, .. } if !(1..=255).contains(&criteria.len()) => Err(
                ValidationError::new("criteria", "choice requires between 1 and 255 options"),
            ),
            Self::Score { criteria, .. } if !(2..=10).contains(&criteria.len()) => Err(
                ValidationError::new("criteria", "score requires between 2 and 10 levels"),
            ),
            _ => Ok(()),
        }
    }
}

/// Input to `POST /v1/systemone`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EvaluationRequest {
    /// Content to evaluate.
    pub state: Content,
    /// Model name or alias; constructors default to `jev-latest`.
    pub model: String,
    /// Questions keyed by caller-defined identifiers.
    pub questions: BTreeMap<String, Question>,
}

impl EvaluationRequest {
    /// Start a request with its required state and the default model.
    pub fn new(state: impl Into<Content>) -> Self {
        Self {
            state: state.into(),
            model: DEFAULT_MODEL.into(),
            questions: BTreeMap::new(),
        }
    }

    /// Override the model. Arbitrary version strings are supported.
    #[must_use]
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = model.into();
        self
    }

    /// Add a question, replacing any previous question with the same identifier.
    #[must_use]
    pub fn question(mut self, id: impl Into<String>, question: Question) -> Self {
        self.questions.insert(id.into(), question);
        self
    }

    /// Check that the model and question set are nonempty and rubrics have valid sizes.
    /// Deserialized or directly modified requests are checked by `Client::evaluate` too.
    pub fn validate(&self) -> Result<(), ValidationError> {
        if self.model.trim().is_empty() {
            return Err(ValidationError::new("model", "must not be empty"));
        }
        if self.questions.is_empty() {
            return Err(ValidationError::new(
                "questions",
                "at least one question is required",
            ));
        }
        for (id, question) in &self.questions {
            question.validate().map_err(|mut error| {
                error.path = format!("questions[{id:?}].{}", error.path);
                error
            })?;
        }
        Ok(())
    }
}

/// A result whose variant matches the question's `type`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum Answer {
    /// Probability of yes.
    Noul {
        /// Value from zero (no) to one (yes).
        noul: f64,
    },
    /// Selected option and its full distribution.
    Choice {
        /// The highest-probability option.
        choice: String,
        /// Probability for every option.
        probabilities: BTreeMap<String, f64>,
        /// Confidence from zero to one.
        confidence: f64,
    },
    /// A probability-weighted score over the rubric levels.
    Score {
        /// Weighted score, potentially between levels.
        score: f64,
        /// String level indices mapped to their descriptions.
        legend: BTreeMap<String, String>,
        /// Probability for every level, keyed by string level indices.
        probabilities: BTreeMap<String, f64>,
        /// Confidence from zero to one.
        confidence: f64,
    },
}

/// Token counts reported by the API.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    /// Input tokens consumed.
    pub input_tokens: u64,
    /// Output tokens produced.
    pub output_tokens: u64,
}

/// The complete evaluation response.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EvaluationResponse {
    /// The model version that performed the evaluation.
    pub model: String,
    /// Answers keyed by the original question identifiers.
    pub answers: BTreeMap<String, Answer>,
    /// Token usage for this evaluation.
    pub usage: Usage,
}