clawmates/crates/cm-decide/src/jev.rs

//! TypeSafe's Jev over `POST /v1/systemone`.
//!
//! Text-only, 64 K context (32 K for the state), no tools, no reasoning.
//! Priced on input tokens only. Not trained on customer requests. The key is
//! `TYPESAFE_API_KEY` and lives in the server's environment — it is never
//! handed to a mission container, and this client is only ever called from
//! the server.

use std::collections::BTreeMap;
use std::time::{Duration, Instant};

use serde::Deserialize;

use crate::{Answer, DecideError, Decider, Decision, Question, Usage};

pub const ENDPOINT: &str = "https://api.typesafe.ai/v1/systemone";
pub const DEFAULT_MODEL: &str = "jev-latest";

pub struct Jev {
    client: reqwest::Client,
    key: String,
    model: String,
}

impl Jev {
    /// From `TYPESAFE_API_KEY` (and `TYPESAFE_MODEL`, default `jev-latest`).
    /// `None` when the key is unset: the caller decides whether that means
    /// "skip the decision" or "use another backend" — it never means guess.
    pub fn from_env() -> Option<Self> {
        let key = std::env::var("TYPESAFE_API_KEY").ok()?;
        let key = key.trim().to_string();
        if key.is_empty() {
            return None;
        }
        let model = std::env::var("TYPESAFE_MODEL")
            .ok()
            .filter(|m| !m.trim().is_empty())
            .unwrap_or_else(|| DEFAULT_MODEL.to_string());
        Some(Self::new(key, model))
    }

    pub fn new(key: String, model: String) -> Self {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(30))
            .build()
            .expect("reqwest client");
        Self { client, key, model }
    }
}

#[derive(Deserialize)]
struct Reply {
    model: String,
    answers: BTreeMap<String, RawAnswer>,
    #[serde(default)]
    usage: Option<RawUsage>,
}

#[derive(Deserialize)]
struct RawUsage {
    #[serde(default)]
    input_tokens: u64,
    #[serde(default)]
    output_tokens: u64,
}

/// Their answer shapes. Score keys `probabilities` by level number as a
/// STRING ("0", "1", …); a `legend` mirrors the criteria and is dropped.
#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
enum RawAnswer {
    Choice {
        choice: String,
        probabilities: BTreeMap<String, f64>,
        confidence: f64,
    },
    Score {
        score: f64,
        probabilities: BTreeMap<String, f64>,
        confidence: f64,
    },
    Noul {
        noul: f64,
    },
}

impl RawAnswer {
    fn into_answer(self) -> Result<Answer, DecideError> {
        Ok(match self {
            RawAnswer::Choice { choice, probabilities, confidence } => Answer::Choice {
                choice,
                probabilities,
                confidence,
            },
            RawAnswer::Score { score, probabilities, confidence } => {
                let mut levels: Vec<(usize, f64)> = probabilities
                    .into_iter()
                    .map(|(k, v)| {
                        k.parse::<usize>()
                            .map(|i| (i, v))
                            .map_err(|_| DecideError::Shape(format!("score level key {k:?}")))
                    })
                    .collect::<Result<_, _>>()?;
                levels.sort_by_key(|(i, _)| *i);
                Answer::Score {
                    score,
                    probabilities: levels.into_iter().map(|(_, p)| p).collect(),
                    confidence,
                }
            }
            RawAnswer::Noul { noul } => Answer::Noul { noul },
        })
    }
}

#[async_trait::async_trait]
impl Decider for Jev {
    fn name(&self) -> &str {
        "jev"
    }

    async fn decide(
        &self,
        state: &str,
        questions: &BTreeMap<String, Question>,
    ) -> Result<Decision, DecideError> {
        let body = serde_json::json!({
            "state": state,
            "model": self.model,
            "questions": questions,
        });
        let started = Instant::now();
        let resp = self
            .client
            .post(ENDPOINT)
            .bearer_auth(&self.key)
            .json(&body)
            .send()
            .await
            .map_err(|e| DecideError::Transport(e.to_string()))?;
        let status = resp.status();
        let text = resp
            .text()
            .await
            .map_err(|e| DecideError::Transport(e.to_string()))?;
        if !status.is_success() {
            return Err(DecideError::Model(format!(
                "HTTP {status}: {}",
                text.chars().take(300).collect::<String>()
            )));
        }
        let reply: Reply =
            serde_json::from_str(&text).map_err(|e| DecideError::Shape(e.to_string()))?;
        let mut answers = BTreeMap::new();
        for (id, raw) in reply.answers {
            answers.insert(id, raw.into_answer()?);
        }
        Ok(Decision {
            model: reply.model,
            answers,
            usage: reply.usage.map(|u| Usage {
                input_tokens: u.input_tokens,
                output_tokens: u.output_tokens,
            }),
            latency: started.elapsed(),
        })
    }
}

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

    /// The documented response, read back into our shapes — including the
    /// string-keyed Score levels, which arrive unordered in a JSON object.
    #[test]
    fn the_documented_reply_parses() {
        let text = r#"{"model":"jev-1.13.0","answers":{
            "department":{"type":"choice","choice":"technical","confidence":0.78,
                "probabilities":{"technical":0.85,"sales":0.0,"billing":0.15}},
            "frustration":{"type":"score","score":1.0,"confidence":1.0,
                "legend":{"0":"calm","1":"frustrated","2":"angry"},
                "probabilities":{"2":0.0,"0":0.0,"1":1.0}},
            "is_urgent":{"type":"noul","noul":1.0}},
            "usage":{"input_tokens":392,"output_tokens":65}}"#;
        let reply: Reply = serde_json::from_str(text).unwrap();
        let a = reply.answers.get("frustration").unwrap();
        let RawAnswer::Score { .. } = a else { panic!() };
        let converted: BTreeMap<String, Answer> = reply
            .answers
            .into_iter()
            .map(|(k, v)| (k, v.into_answer().unwrap()))
            .collect();
        match &converted["frustration"] {
            Answer::Score { probabilities, score, .. } => {
                assert_eq!(probabilities, &vec![0.0, 1.0, 0.0]);
                assert_eq!(*score, 1.0);
            }
            other => panic!("{other:?}"),
        }
        match &converted["department"] {
            Answer::Choice { choice, .. } => assert_eq!(choice, "technical"),
            other => panic!("{other:?}"),
        }
    }
}