jev-sdk 0.1.0

//! Answers returned by the System One API.

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};

/// The response to a `system_one` evaluation.
///
/// Answers are keyed by the same ids used in the request's question map.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SystemOneResponse {
    /// The model that performed the evaluation.
    pub model: String,
    /// One [`Answer`] per question, keyed by question id.
    pub answers: IndexMap<String, Answer>,
    /// Token usage for the request.
    pub usage: Usage,
}

impl SystemOneResponse {
    /// Get the answer for a question id, regardless of its type.
    pub fn get(&self, id: &str) -> Option<&Answer> {
        self.answers.get(id)
    }

    /// Iterate over all [`Answer::Noul`] answers.
    pub fn nouls(&self) -> impl Iterator<Item = (&str, &NoulAnswer)> {
        self.answers.iter().filter_map(|(k, v)| match v {
            Answer::Noul(a) => Some((k.as_str(), a)),
            _ => None,
        })
    }

    /// Iterate over all [`Answer::Choice`] answers.
    pub fn choices(&self) -> impl Iterator<Item = (&str, &ChoiceAnswer)> {
        self.answers.iter().filter_map(|(k, v)| match v {
            Answer::Choice(a) => Some((k.as_str(), a)),
            _ => None,
        })
    }

    /// Iterate over all [`Answer::Score`] answers.
    pub fn scores(&self) -> impl Iterator<Item = (&str, &ScoreAnswer)> {
        self.answers.iter().filter_map(|(k, v)| match v {
            Answer::Score(a) => Some((k.as_str(), a)),
            _ => None,
        })
    }

    /// Get a [`NoulAnswer`] by question id, if it is a noul answer.
    pub fn noul(&self, id: &str) -> Option<&NoulAnswer> {
        self.answers.get(id).and_then(Answer::as_noul)
    }

    /// Get a [`ChoiceAnswer`] by question id, if it is a choice answer.
    pub fn choice(&self, id: &str) -> Option<&ChoiceAnswer> {
        self.answers.get(id).and_then(Answer::as_choice)
    }

    /// Get a [`ScoreAnswer`] by question id, if it is a score answer.
    pub fn score(&self, id: &str) -> Option<&ScoreAnswer> {
        self.answers.get(id).and_then(Answer::as_score)
    }
}

/// Token usage for a request.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
pub struct Usage {
    /// Tokens read from the input (state + questions).
    pub input_tokens: u64,
    /// Tokens produced as output.
    pub output_tokens: u64,
}

/// One typed answer. The variant matches the type of the question asked.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum Answer {
    /// Answer to a noul question.
    Noul(NoulAnswer),
    /// Answer to a choice question.
    Choice(ChoiceAnswer),
    /// Answer to a score question.
    Score(ScoreAnswer),
}

impl Answer {
    /// The noul answer, if this is one.
    pub fn as_noul(&self) -> Option<&NoulAnswer> {
        match self {
            Answer::Noul(a) => Some(a),
            _ => None,
        }
    }

    /// The choice answer, if this is one.
    pub fn as_choice(&self) -> Option<&ChoiceAnswer> {
        match self {
            Answer::Choice(a) => Some(a),
            _ => None,
        }
    }

    /// The score answer, if this is one.
    pub fn as_score(&self) -> Option<&ScoreAnswer> {
        match self {
            Answer::Score(a) => Some(a),
            _ => None,
        }
    }
}

/// Answer to a noul question: the probability that the answer is "yes".
#[derive(Debug, Clone, Copy, Deserialize, Serialize)]
pub struct NoulAnswer {
    /// The yes/no answer on a scale from 0 (no) to 1 (yes).
    pub noul: f64,
}

/// Answer to a choice question.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ChoiceAnswer {
    /// The highest-probability option.
    pub choice: String,
    /// Every option mapped to its probability. Sums to 1.
    #[serde(default)]
    pub probabilities: IndexMap<String, f64>,
    /// How certain the model is, derived from the probability distribution.
    pub confidence: f64,
}

/// Answer to a score question.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ScoreAnswer {
    /// The probability-weighted value across the levels; can land between
    /// levels.
    pub score: f64,
    /// Each level index (as a string) mapped back to its description.
    #[serde(default)]
    pub legend: IndexMap<String, String>,
    /// Each level index (as a string) mapped to its probability. Sums to 1.
    #[serde(default)]
    pub probabilities: IndexMap<String, f64>,
    /// How certain the model is, derived from the probability distribution.
    pub confidence: f64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn parses_full_response() {
        let raw = json!({
            "model": "jev-1.13.0",
            "answers": {
                "department": {
                    "type": "choice",
                    "choice": "technical",
                    "probabilities": {"billing": 0.159, "technical": 0.84, "sales": 0.001},
                    "confidence": 0.596,
                },
                "frustration": {
                    "type": "score",
                    "score": 1.035,
                    "legend": {"0": "Calm", "1": "Frustrated", "2": "Very angry"},
                    "probabilities": {"0": 0.1, "1": 0.6, "2": 0.3},
                    "confidence": 0.842,
                },
                "is_urgent": {"type": "noul", "noul": 0.999},
            },
            "usage": {"input_tokens": 312, "output_tokens": 48},
        });

        let resp: SystemOneResponse = serde_json::from_value(raw).unwrap();
        assert_eq!(resp.model, "jev-1.13.0");
        assert_eq!(resp.answers.len(), 3);
        assert_eq!(resp.usage.input_tokens, 312);

        let dept = resp.choice("department").unwrap();
        assert_eq!(dept.choice, "technical");
        assert_eq!(dept.probabilities["billing"], 0.159);
        assert_eq!(dept.confidence, 0.596);

        let fr = resp.score("frustration").unwrap();
        assert_eq!(fr.score, 1.035);
        assert_eq!(fr.legend["2"], "Very angry");
        assert_eq!(fr.probabilities["1"], 0.6);

        let urg = resp.noul("is_urgent").unwrap();
        assert_eq!(urg.noul, 0.999);

        assert_eq!(resp.nouls().count(), 1);
        assert_eq!(resp.choices().count(), 1);
        assert_eq!(resp.scores().count(), 1);
    }

    #[test]
    fn score_answer_tolerates_missing_probabilities() {
        let raw = json!({
            "model": "jev-latest",
            "answers": {
                "s": {"type": "score", "score": 0.0, "confidence": 1.0, "legend": {"0": "only"}}
            },
            "usage": {"input_tokens": 1, "output_tokens": 1},
        });
        let resp: SystemOneResponse = serde_json::from_value(raw).unwrap();
        assert!(resp.score("s").unwrap().probabilities.is_empty());
    }

    #[test]
    fn responses_round_trip_through_json() {
        // Tools can re-serialize a response (caching, passthrough output).
        let raw = json!({
            "model": "jev-1.13.0",
            "answers": {
                "ok": {"type": "noul", "noul": 0.87},
                "path": {
                    "type": "choice",
                    "choice": "a",
                    "probabilities": {"a": 0.9, "b": 0.1},
                    "confidence": 0.8,
                },
            },
            "usage": {"input_tokens": 10, "output_tokens": 4},
        });
        let resp: SystemOneResponse = serde_json::from_value(raw.clone()).unwrap();
        assert_eq!(serde_json::to_value(&resp).unwrap(), raw);
    }
}