polyc_judgment/

lib.rs

//! Provider-agnostic judgment trait and wire types for polychrome.
//!
//! A judgment is a typed decision over a state: a yes/no probability
//! ([`Question::Noul`]) or one option from a closed set
//! ([`Question::Choice`]). A judgment model returns probabilities, never
//! generated text, so code can threshold and combine the answers without
//! parsing prose.
//!
//! [`JudgmentProvider`] is the seam between a decision site and any concrete
//! backend, the same way [`polyc_llm::LlmProvider`] is for completions: one
//! implementation crate per backend, dispatched behind [`DynJudgment`]. The
//! core never names a vendor.
//!
//! [`polyc_llm::LlmProvider`]: https://docs.rs/polyc-llm

use std::{collections::BTreeMap, sync::Arc};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

pub mod fallback;
/// Canned [`JudgmentProvider`] stubs for wiring and tests.
///
/// The `test-fixtures` feature gates the module, so a normal build cannot
/// link a stub. The module doc says why that matters.
#[cfg(feature = "test-fixtures")]
pub mod stub;

/// What a yes and a no mean for a [`Question::Noul`], when the boundary needs
/// spelling out.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct NoulCriteria {
    /// What a value near 1 means.
    #[serde(rename = "true")]
    pub yes: String,
    /// What a value near 0 means.
    #[serde(rename = "false")]
    pub no: String,
}

/// One typed question over a state.
///
/// The variant sets the answer shape. Every variant carries `instructions`,
/// the judgment itself, phrased so the model can answer it from the state
/// alone. Question ids are chosen by the caller and never reach the model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Question {
    /// A yes/no question. The answer is the probability of yes.
    Noul {
        /// The yes/no question to evaluate.
        instructions: String,
        /// Optional descriptions of what yes and no mean.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        criteria: Option<NoulCriteria>,
    },
    /// One option from a closed set. The answer is the chosen option plus the
    /// full distribution.
    Choice {
        /// What the model decides.
        instructions: String,
        /// Option name to rubric description. `None` when an option needs no
        /// extra detail.
        criteria: BTreeMap<String, Option<String>>,
    },
}

/// One request: a state and the questions to answer over it.
///
/// Every question sees the same state and is answered independently, so a
/// caller asks every question it might need in one request.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct JudgmentRequest {
    /// The content to judge: a string, an object with named fields, or an
    /// array of records.
    pub state: serde_json::Value,
    /// Caller-chosen id to question. Answers come back under the same ids.
    pub questions: BTreeMap<String, Question>,
}

/// One typed answer, matching its question's variant.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Answer {
    /// The probability that the answer is yes, 0 to 1.
    Noul {
        /// Probability of yes.
        noul: f64,
    },
    /// The highest-probability option and the full distribution.
    Choice {
        /// The chosen option.
        choice: String,
        /// Every option to its probability. The values sum to 1.
        probabilities: BTreeMap<String, f64>,
        /// `confidence` is the maximum probability in `probabilities`: the
        /// probability mass on `choice`, 0 to 1.
        confidence: f64,
    },
}

/// Token accounting for one request.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct JudgmentUsage {
    /// Tokens the backend read.
    pub input_tokens: u64,
    /// Tokens the backend wrote.
    pub output_tokens: u64,
}

/// Which path answered a [`JudgmentRequest`].
///
/// Lets a caller or an evaluation tell a dedicated judgment backend's answer
/// apart from a [`fallback::FallbackJudgment`] answer produced by its
/// fallback instead. A wire response that omits this field (every judgment
/// backend's own wire shape does — the field is local to this seam) decodes
/// as [`Self::Primary`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JudgmentSource {
    /// A dedicated judgment backend answered directly.
    #[default]
    Primary,
    /// The primary backend failed in a way a fallback can fix, and a
    /// fallback backend answered instead.
    Fallback,
}

/// One response: an answer per question id, plus usage.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct JudgmentResponse {
    /// The backend model that answered.
    pub model: String,
    /// Question id to answer.
    pub answers: BTreeMap<String, Answer>,
    /// Token accounting.
    pub usage: JudgmentUsage,
    /// Which path answered. Defaults to [`JudgmentSource::Primary`] when a
    /// wire response omits it, which every backend's own response does.
    #[serde(default)]
    pub source: JudgmentSource,
}

impl JudgmentResponse {
    /// Returns the yes probability of the [`Answer::Noul`] under `id`.
    ///
    /// `None` when the id is absent or the answer is not a Noul. A caller
    /// that needs a hard decision thresholds the value itself, and treats
    /// `None` as its fail-closed branch.
    #[must_use]
    pub fn noul(&self, id: &str) -> Option<f64> {
        match self.answers.get(id) {
            Some(Answer::Noul { noul }) => Some(*noul),
            _ => None,
        }
    }
}

/// The failure classes a judgment backend reports.
///
/// The classes are stable across backends so a decision site can fail closed
/// the same way whichever backend answers. No variant carries a request body,
/// a response body, or a credential.
#[derive(Debug, thiserror::Error)]
pub enum JudgmentError {
    /// The backend rejected the request shape. Retrying the same request
    /// cannot succeed.
    #[error("judgment request rejected: {0}")]
    Invalid(String),
    /// The backend refused the credential.
    #[error("judgment backend refused the credential")]
    Unauthorized,
    /// The backend refused for lack of credit or quota.
    #[error("judgment backend refused for lack of credit or quota")]
    Exhausted,
    /// The backend asked the caller to back off.
    #[error("judgment backend rate limited (retry after {retry_after:?})")]
    RateLimited {
        /// The backoff the backend suggested, when it sent one.
        retry_after: Option<std::time::Duration>,
    },
    /// The backend is temporarily unavailable.
    #[error("judgment backend unavailable (status {status})")]
    Unavailable {
        /// The HTTP status the backend answered with.
        status: u16,
    },
    /// The request did not complete on the transport.
    #[error("judgment transport failed: {0}")]
    Transport(Box<dyn std::error::Error + Send + Sync + 'static>),
    /// The response did not decode into the answer types.
    #[error("judgment response malformed: {0}")]
    Malformed(String),
}

/// The seam between a decision site and any concrete judgment backend.
///
/// One implementation per backend. A decision site holds a
/// [`DynJudgment`] and never names the backend.
#[async_trait]
pub trait JudgmentProvider: Send + Sync + 'static {
    /// The backend's concrete error type.
    type Error: std::error::Error + Send + Sync + 'static;

    /// Answers every question in `request` over its state.
    ///
    /// # Errors
    ///
    /// Returns the backend's error when the request is rejected, the
    /// credential is refused, the transport fails, or the response does not
    /// decode.
    async fn judge(&self, request: JudgmentRequest) -> Result<JudgmentResponse, Self::Error>;
}

/// A backend error erased to one concrete type.
///
/// Transparent wrapper: [`Display`](std::fmt::Display) and
/// [`source`](std::error::Error::source) delegate to the inner error.
#[derive(Debug)]
pub struct BoxError(Box<dyn std::error::Error + Send + Sync + 'static>);

impl BoxError {
    /// Erases any backend error.
    #[must_use]
    pub fn new<E: std::error::Error + Send + Sync + 'static>(err: E) -> Self {
        Self(Box::new(err))
    }
}

impl std::fmt::Display for BoxError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.0, f)
    }
}

impl std::error::Error for BoxError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(&*self.0)
    }
}

/// Wraps a concrete [`JudgmentProvider`] and erases its error to
/// [`BoxError`], so the wrapped value coerces to [`DynJudgment`].
pub struct ErasedJudgment<P>(P);

#[async_trait]
impl<P: JudgmentProvider> JudgmentProvider for ErasedJudgment<P> {
    type Error = BoxError;

    async fn judge(&self, request: JudgmentRequest) -> Result<JudgmentResponse, Self::Error> {
        self.0.judge(request).await.map_err(BoxError::new)
    }
}

/// The single trait-object type a decision site stores.
pub type DynJudgment = dyn JudgmentProvider<Error = BoxError>;

/// Erases a concrete backend and wraps it in an `Arc` as a [`DynJudgment`].
#[must_use]
pub fn into_dyn<P: JudgmentProvider>(provider: P) -> Arc<DynJudgment> {
    Arc::new(ErasedJudgment(provider))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use super::*;

    #[test]
    fn noul_question_serializes_to_the_wire_shape() {
        let q = Question::Noul {
            instructions: "Is it urgent?".to_owned(),
            criteria: Some(NoulCriteria {
                yes: "Time-sensitive".to_owned(),
                no: "No urgency".to_owned(),
            }),
        };
        let json = serde_json::to_value(&q).expect("serializable");
        assert_eq!(
            json,
            serde_json::json!({
                "type": "noul",
                "instructions": "Is it urgent?",
                "criteria": {"true": "Time-sensitive", "false": "No urgency"}
            })
        );
    }

    #[test]
    fn noul_question_without_criteria_omits_the_field() {
        let q = Question::Noul {
            instructions: "Is it urgent?".to_owned(),
            criteria: None,
        };
        let json = serde_json::to_value(&q).expect("serializable");
        assert!(json.get("criteria").is_none(), "{json}");
    }

    #[test]
    fn choice_answer_round_trips() {
        let raw = serde_json::json!({
            "type": "choice",
            "choice": "billing",
            "probabilities": {"billing": 0.9, "sales": 0.1},
            "confidence": 0.85
        });
        let answer: Answer = serde_json::from_value(raw.clone()).expect("decodes");
        assert!(matches!(&answer, Answer::Choice { choice, .. } if choice == "billing"));
        assert_eq!(serde_json::to_value(&answer).expect("encodes"), raw);
    }

    #[test]
    fn response_noul_accessor_is_none_for_other_shapes() {
        let mut answers = BTreeMap::new();
        answers.insert("yes".to_owned(), Answer::Noul { noul: 0.7 });
        answers.insert(
            "pick".to_owned(),
            Answer::Choice {
                choice: "a".to_owned(),
                probabilities: BTreeMap::new(),
                confidence: 1.0,
            },
        );
        let response = JudgmentResponse {
            model: "m".to_owned(),
            answers,
            usage: JudgmentUsage::default(),
            source: JudgmentSource::Primary,
        };
        assert_eq!(response.noul("yes"), Some(0.7));
        assert_eq!(response.noul("pick"), None);
        assert_eq!(response.noul("absent"), None);
    }
}