Skip to main content

rig_typesafeai/
types.rs

1//! TypeSafe's JSON protocol. Prefer the typed question constructors for application code.
2use rig_core::completion::Cost;
3use rig_core::json_utils::Lenient;
4use serde::{Deserialize, Deserializer, Serialize, Serializer};
5use serde_json::Value;
6use std::collections::BTreeMap;
7
8/// Model-independent input to an evaluation. [`crate::Jev`] supplies the model
9/// from its configuration when encoding the full HTTP request.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct Request {
12    /// The shared state evaluated independently by every question.
13    pub state: Value,
14    /// Serialized question object. Keeping the JSON intact preserves duplicate
15    /// keys for validation and allows named structs to supply the shape.
16    pub questions: Box<serde_json::value::RawValue>,
17}
18
19/// A question as represented by the HTTP API.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21#[serde(tag = "type", rename_all = "lowercase")]
22pub enum Question {
23    Choice {
24        instructions: Value,
25        criteria: BTreeMap<String, Option<Value>>,
26    },
27    Score {
28        instructions: Value,
29        criteria: Vec<Value>,
30    },
31    Noul {
32        instructions: Value,
33        #[serde(skip_serializing_if = "Option::is_none")]
34        criteria: Option<BTreeMap<String, Value>>,
35    },
36}
37
38/// An answer as represented by the HTTP API.
39#[derive(Debug, Clone, Serialize, Deserialize)]
40#[serde(tag = "type", rename_all = "lowercase")]
41pub enum Answer {
42    Choice {
43        choice: String,
44        probabilities: BTreeMap<String, f64>,
45        confidence: f64,
46    },
47    Score {
48        score: f64,
49        probabilities: BTreeMap<String, f64>,
50        legend: BTreeMap<String, Value>,
51        confidence: f64,
52    },
53    Noul {
54        noul: f64,
55    },
56}
57
58/// Rig token accounting; counters absent from Jev remain unknown.
59pub use rig_core::completion::Usage;
60
61/// Raw response, including the transport's request identifier. It
62/// serializes in Jev's shape, which reads back as the same response, without
63/// the transport's request identifier or the whole reply document.
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct Response {
66    /// The model identifier returned by the provider.
67    pub model: String,
68    /// Answer object kept intact for validation and typed Serde decoding.
69    pub answers: Box<serde_json::value::RawValue>,
70    /// Token accounting when reported by the provider.
71    #[serde(
72        default,
73        deserialize_with = "reported_usage",
74        serialize_with = "jev_usage"
75    )]
76    pub usage: Option<Usage>,
77    #[serde(skip)]
78    /// Transport request identifier for diagnostics.
79    pub provider_request_id: Option<String>,
80    #[serde(skip)]
81    /// The whole reply document, including fields this type does not model.
82    pub raw: Value,
83}
84
85/// Jev's `usage`: its token counters and, through OpenRouter, a `cost` in
86/// USD. A counter or cost of another shape is unreported rather than failing
87/// the reply, whose answers do not depend on it.
88fn reported_usage<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<Usage>, D::Error> {
89    let usage = Option::<Value>::deserialize(deserializer)?;
90    Ok(usage.filter(Value::is_object).map(|usage| {
91        Usage::new()
92            .input_tokens(usage.u64("input_tokens"))
93            .output_tokens(usage.u64("output_tokens"))
94            .cost(
95                usage
96                    .f64("cost")
97                    .filter(|cost| cost.is_finite())
98                    .map(Cost::from_total),
99            )
100    }))
101}
102
103/// `usage` in Jev's shape: its token counters and its cost's total.
104fn jev_usage<S: Serializer>(usage: &Option<Usage>, serializer: S) -> Result<S::Ok, S::Error> {
105    #[derive(Serialize)]
106    struct JevUsage {
107        #[serde(skip_serializing_if = "Option::is_none")]
108        input_tokens: Option<u64>,
109        #[serde(skip_serializing_if = "Option::is_none")]
110        output_tokens: Option<u64>,
111        #[serde(skip_serializing_if = "Option::is_none")]
112        cost: Option<f64>,
113    }
114    usage
115        .map(|usage| JevUsage {
116            input_tokens: usage.input_tokens,
117            output_tokens: usage.output_tokens,
118            cost: usage.cost.map(|cost| cost.total),
119        })
120        .serialize(serializer)
121}