gemini_adk_rs/decision/
mod.rs

1//! Decision models: typed questions about shared state, answered with
2//! calibrated probabilities instead of text.
3//!
4//! A decision model (a "System One" model, such as TypeSafe AI's Jev on
5//! Vercel AI Gateway) takes a `state` (a string, an object or an array) and a
6//! map of named questions. Each answer comes back typed:
7//!
8//! | Question | Answer |
9//! |---|---|
10//! | [`Question::Boolean`] | `probability` that the statement is true |
11//! | [`Question::Choice`] | the chosen option and each option's probability |
12//! | [`Question::Score`] | a probability-weighted position on an ordered scale |
13//!
14//! A decision model writes no text, so it cannot fill an open slot such as a
15//! name or a date. It can pick among options the session already holds (the
16//! time slots a tool offered) and decide categorical things (a confirmation,
17//! an intent, which stage comes next), usually far faster than a language
18//! model. [`Decisions`] holds the questions a session asks and asks them
19//! against the rolling conversation; guards name them with the `decided` atom
20//! (see [`crate::flow::decided`]). [`GatewayDecisionModel`] (feature
21//! `ai-gateway`) calls Vercel AI Gateway.
22//!
23//! ```
24//! use gemini_adk_rs::decision::{Answer, DecisionRequest, MockDecisionModel, Question, DecisionModel};
25//! use serde_json::json;
26//!
27//! # tokio_test::block_on(async {
28//! let model = MockDecisionModel::new(|_| {
29//!     Ok([("confirmed".to_string(), Answer::boolean(0.97))].into())
30//! });
31//! let response = model
32//!     .decide(DecisionRequest::new(
33//!         json!("Caller: yes, book it."),
34//!         [("confirmed", Question::boolean("Did the caller agree to book?"))],
35//!     ))
36//!     .await
37//!     .unwrap();
38//! assert_eq!(response.answers["confirmed"].probability(), Some(0.97));
39//! # });
40//! ```
41
42#[cfg(feature = "ai-gateway")]
43mod gateway;
44mod service;
45
46use std::collections::BTreeMap;
47use std::sync::Arc;
48use std::time::Duration;
49
50use async_trait::async_trait;
51use serde::{Deserialize, Serialize};
52use serde_json::{Value, json};
53
54use crate::llm::LlmError;
55
56#[cfg(feature = "ai-gateway")]
57pub use gateway::GatewayDecisionModel;
58pub use service::{
59    Conversation, DEFAULT_BOOLEAN_THRESHOLD, DEFAULT_CERTAINTY_THRESHOLD, DEFAULT_DECISION_TIMEOUT,
60    DEFAULT_HISTORY_TURNS, Decision, Decisions, NONE_OPTION, Round,
61};
62
63/// The most options a choice question may have.
64pub const MAX_CHOICE_OPTIONS: usize = 255;
65/// The fewest and most levels a score question may have.
66pub const SCORE_LEVELS: std::ops::RangeInclusive<usize> = 2..=10;
67
68/// What the true and false cases of a boolean question mean.
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
70pub struct BooleanCriteria {
71    /// What makes the statement true.
72    #[serde(rename = "true")]
73    pub when_true: String,
74    /// What makes it false.
75    #[serde(rename = "false")]
76    pub when_false: String,
77}
78
79/// One typed question, as the decision API takes it.
80#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
81#[serde(tag = "type", rename_all = "lowercase")]
82pub enum Question {
83    /// Answered with the probability that the statement is true.
84    Boolean {
85        /// The question.
86        instructions: String,
87        /// Optional definitions of the true and false cases.
88        #[serde(default, skip_serializing_if = "Option::is_none")]
89        criteria: Option<BooleanCriteria>,
90    },
91    /// Answered with one option and each option's probability.
92    Choice {
93        /// The question.
94        instructions: String,
95        /// Option name to description, 1 to 255 options.
96        criteria: BTreeMap<String, String>,
97    },
98    /// Answered with a position on an ordered scale.
99    Score {
100        /// The question.
101        instructions: String,
102        /// Level descriptions, lowest first, 2 to 10 levels.
103        criteria: Vec<String>,
104    },
105}
106
107impl Question {
108    /// A boolean question.
109    pub fn boolean(instructions: impl Into<String>) -> Self {
110        Self::Boolean {
111            instructions: instructions.into(),
112            criteria: None,
113        }
114    }
115
116    /// A choice among named options, each with a description.
117    pub fn choice<K: Into<String>, D: Into<String>>(
118        instructions: impl Into<String>,
119        options: impl IntoIterator<Item = (K, D)>,
120    ) -> Self {
121        Self::Choice {
122            instructions: instructions.into(),
123            criteria: options
124                .into_iter()
125                .map(|(k, d)| (k.into(), d.into()))
126                .collect(),
127        }
128    }
129
130    /// A score on an ordered scale, lowest level first.
131    pub fn score<L: Into<String>>(
132        instructions: impl Into<String>,
133        levels: impl IntoIterator<Item = L>,
134    ) -> Self {
135        Self::Score {
136            instructions: instructions.into(),
137            criteria: levels.into_iter().map(Into::into).collect(),
138        }
139    }
140
141    /// Define what the true and false cases of a boolean question mean. A
142    /// no-op on other question types.
143    #[must_use]
144    pub fn when(mut self, when_true: impl Into<String>, when_false: impl Into<String>) -> Self {
145        if let Self::Boolean { criteria, .. } = &mut self {
146            *criteria = Some(BooleanCriteria {
147                when_true: when_true.into(),
148                when_false: when_false.into(),
149            });
150        }
151        self
152    }
153
154    /// The question's instructions.
155    pub fn instructions(&self) -> &str {
156        match self {
157            Self::Boolean { instructions, .. }
158            | Self::Choice { instructions, .. }
159            | Self::Score { instructions, .. } => instructions,
160        }
161    }
162
163    /// Check the limits the decision API enforces, so a bad question fails
164    /// when it is built rather than on every call.
165    pub fn validate(&self) -> Result<(), String> {
166        if self.instructions().trim().is_empty() {
167            return Err("a question needs instructions".into());
168        }
169        match self {
170            Self::Boolean { .. } => Ok(()),
171            Self::Choice { criteria, .. } if criteria.is_empty() => {
172                Err("a choice question needs at least one option".into())
173            }
174            Self::Choice { criteria, .. } if criteria.len() > MAX_CHOICE_OPTIONS => Err(format!(
175                "a choice question takes at most {MAX_CHOICE_OPTIONS} options, not {}",
176                criteria.len()
177            )),
178            Self::Choice { .. } => Ok(()),
179            Self::Score { criteria, .. } if !SCORE_LEVELS.contains(&criteria.len()) => {
180                Err(format!(
181                    "a score question takes {} to {} levels, not {}",
182                    SCORE_LEVELS.start(),
183                    SCORE_LEVELS.end(),
184                    criteria.len()
185                ))
186            }
187            Self::Score { .. } => Ok(()),
188        }
189    }
190}
191
192/// One answer, typed like its question.
193///
194/// Probabilities and scores arrive rounded to two decimals, so a
195/// distribution may not sum to exactly 1; it is not renormalized.
196#[derive(Debug, Clone, PartialEq)]
197pub enum Answer {
198    /// The probability that the statement is true (not a confidence in a
199    /// yes/no answer).
200    Boolean {
201        /// `P(true)`, 0 to 1.
202        probability: f64,
203    },
204    /// The chosen option.
205    Choice {
206        /// The selected option name.
207        choice: String,
208        /// Each option's probability.
209        probabilities: BTreeMap<String, f64>,
210        /// How concentrated the distribution is, 0 (even) to 1 (all on one
211        /// option), when the model reports it.
212        confidence: Option<f64>,
213    },
214    /// A position on the question's scale.
215    Score {
216        /// The probability-weighted level, indexed from 0.
217        score: f64,
218        /// Each level's probability, keyed `"0"`, `"1"`, …
219        probabilities: BTreeMap<String, f64>,
220        /// As for [`Answer::Choice`].
221        confidence: Option<f64>,
222    },
223    /// The model declined to answer.
224    Refusal(Value),
225    /// An answer type this client does not know, kept as received.
226    Other(Value),
227}
228
229impl Answer {
230    /// A boolean answer.
231    pub fn boolean(probability: f64) -> Self {
232        Self::Boolean { probability }
233    }
234
235    /// A choice answer with no distribution.
236    pub fn choice(choice: impl Into<String>, confidence: f64) -> Self {
237        Self::Choice {
238            choice: choice.into(),
239            probabilities: BTreeMap::new(),
240            confidence: Some(confidence),
241        }
242    }
243
244    /// A score answer with no distribution.
245    pub fn score(score: f64, confidence: f64) -> Self {
246        Self::Score {
247            score,
248            probabilities: BTreeMap::new(),
249            confidence: Some(confidence),
250        }
251    }
252
253    /// `P(true)` of a boolean answer.
254    pub fn probability(&self) -> Option<f64> {
255        match self {
256            Self::Boolean { probability } => Some(*probability),
257            _ => None,
258        }
259    }
260
261    /// How sure a choice or score answer is: the model's reported confidence,
262    /// or else the probability of the chosen option (the most likely level
263    /// for a score).
264    pub fn certainty(&self) -> Option<f64> {
265        match self {
266            Self::Choice {
267                choice,
268                probabilities,
269                confidence,
270            } => confidence.or_else(|| probabilities.get(choice).copied()),
271            // A score between two levels (2.5 on a 0 to 3 scale) is not
272            // unsure about the direction: its certainty is the probability of
273            // the two levels it lies between.
274            Self::Score {
275                score,
276                probabilities,
277                confidence,
278            } if !probabilities.is_empty() => {
279                let level = |i: f64| {
280                    probabilities
281                        .get(&format!("{}", i as u64))
282                        .copied()
283                        .unwrap_or(0.0)
284                };
285                let (lo, hi) = (score.floor(), score.ceil());
286                let mass = if lo == hi {
287                    level(lo)
288                } else {
289                    level(lo) + level(hi)
290                };
291                Some(confidence.map_or(mass, |c| c.max(mass)))
292            }
293            Self::Score { confidence, .. } => *confidence,
294            _ => None,
295        }
296    }
297
298    /// Read an answer from the decision API's JSON.
299    pub fn from_value(value: Value) -> Self {
300        let number = |v: &Value, key: &str| v.get(key).and_then(Value::as_f64);
301        let distribution = |v: &Value| -> BTreeMap<String, f64> {
302            v.get("probabilities")
303                .and_then(Value::as_object)
304                .map(|m| {
305                    m.iter()
306                        .filter_map(|(k, p)| p.as_f64().map(|p| (k.clone(), p)))
307                        .collect()
308                })
309                .unwrap_or_default()
310        };
311        match value.get("type").and_then(Value::as_str) {
312            Some("boolean") => match number(&value, "probability") {
313                Some(probability) => Self::Boolean { probability },
314                None => Self::Other(value),
315            },
316            Some("choice") => match value.get("choice").and_then(Value::as_str) {
317                Some(choice) => Self::Choice {
318                    choice: choice.to_string(),
319                    probabilities: distribution(&value),
320                    confidence: number(&value, "confidence"),
321                },
322                None => Self::Other(value),
323            },
324            Some("score") => match number(&value, "score") {
325                Some(score) => Self::Score {
326                    score,
327                    probabilities: distribution(&value),
328                    confidence: number(&value, "confidence"),
329                },
330                None => Self::Other(value),
331            },
332            Some("refusal") => Self::Refusal(value),
333            _ => Self::Other(value),
334        }
335    }
336
337    /// The answer as the decision API writes it.
338    pub fn to_value(&self) -> Value {
339        match self {
340            Self::Boolean { probability } => {
341                json!({ "type": "boolean", "probability": probability })
342            }
343            Self::Choice {
344                choice,
345                probabilities,
346                confidence,
347            } => {
348                let mut v =
349                    json!({ "type": "choice", "choice": choice, "probabilities": probabilities });
350                if let Some(c) = confidence {
351                    v["confidence"] = json!(c);
352                }
353                v
354            }
355            Self::Score {
356                score,
357                probabilities,
358                confidence,
359            } => {
360                let mut v =
361                    json!({ "type": "score", "score": score, "probabilities": probabilities });
362                if let Some(c) = confidence {
363                    v["confidence"] = json!(c);
364                }
365                v
366            }
367            Self::Refusal(v) | Self::Other(v) => v.clone(),
368        }
369    }
370}
371
372/// When AI Gateway should rerun a decision with another model.
373///
374/// Serializes to the `when` object of a Gateway decision fallback.
375#[derive(Debug, Clone, PartialEq)]
376pub enum FallbackWhen {
377    /// A choice or score answer's confidence is below `below` (all such
378    /// questions when `question` is `None`).
379    ConfidenceBelow {
380        /// The question to check, or every choice and score question.
381        question: Option<String>,
382        /// The threshold, 0 to 1.
383        below: f64,
384    },
385    /// A boolean answer's `P(true)` is inside `[low, high]`.
386    ProbabilityBetween {
387        /// The question to check, or every boolean question.
388        question: Option<String>,
389        /// Inclusive lower bound.
390        low: f64,
391        /// Inclusive upper bound.
392        high: f64,
393    },
394    /// At least one condition matches.
395    Any(Vec<FallbackWhen>),
396    /// Every condition matches.
397    All(Vec<FallbackWhen>),
398    /// At least `count` conditions match.
399    AtLeast {
400        /// How many must match.
401        count: usize,
402        /// The conditions.
403        conditions: Vec<FallbackWhen>,
404    },
405}
406
407impl FallbackWhen {
408    /// The Gateway's JSON form.
409    pub fn to_value(&self) -> Value {
410        let with_question = |mut v: Value, question: &Option<String>| {
411            if let Some(q) = question {
412                v["question"] = json!(q);
413            }
414            v
415        };
416        match self {
417            Self::ConfidenceBelow { question, below } => {
418                with_question(json!({ "confidenceBelow": below }), question)
419            }
420            Self::ProbabilityBetween {
421                question,
422                low,
423                high,
424            } => with_question(json!({ "probabilityBetween": [low, high] }), question),
425            Self::Any(c) => json!({ "any": c.iter().map(Self::to_value).collect::<Vec<_>>() }),
426            Self::All(c) => json!({ "all": c.iter().map(Self::to_value).collect::<Vec<_>>() }),
427            Self::AtLeast { count, conditions } => json!({
428                "atLeast": {
429                    "count": count,
430                    "conditions": conditions.iter().map(Self::to_value).collect::<Vec<_>>(),
431                }
432            }),
433        }
434    }
435}
436
437/// A decision request: one state and its questions.
438#[derive(Debug, Clone, PartialEq)]
439pub struct DecisionRequest {
440    /// The shared state every question is asked about.
441    pub state: Value,
442    /// Named questions; each name keys its answer.
443    pub questions: BTreeMap<String, Question>,
444    /// Provider options passed through unchanged (for AI Gateway,
445    /// `{"gateway": {...}}`). Merged over the model's defaults.
446    pub provider_options: Option<Value>,
447}
448
449impl DecisionRequest {
450    /// A request with no provider options.
451    pub fn new<K: Into<String>>(
452        state: Value,
453        questions: impl IntoIterator<Item = (K, Question)>,
454    ) -> Self {
455        Self {
456            state,
457            questions: questions.into_iter().map(|(k, q)| (k.into(), q)).collect(),
458            provider_options: None,
459        }
460    }
461}
462
463/// Token usage of one decision.
464#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
465pub struct DecisionUsage {
466    /// Input tokens billed.
467    pub input_tokens: u64,
468    /// Output tokens reported.
469    pub output_tokens: u64,
470}
471
472/// A decision model's answers and what it cost.
473#[derive(Debug, Clone, PartialEq)]
474pub struct DecisionResponse {
475    /// The model that produced the answers (a fallback model when one ran).
476    pub model: String,
477    /// One answer per question.
478    pub answers: BTreeMap<String, Answer>,
479    /// Token usage, when reported.
480    pub usage: Option<DecisionUsage>,
481    /// The cost in US dollars, as the provider reports it.
482    pub cost: Option<String>,
483    /// Questions that triggered a Gateway fallback, with the reason
484    /// (`confidence_below`, `probability_between`, `refused`, …). Empty when
485    /// the primary model's answers were returned.
486    pub fallback_triggered_by: Vec<(String, String)>,
487    /// Wall-clock time of the call, as the client measured it.
488    pub latency: Duration,
489    /// Provider metadata as received.
490    pub provider_metadata: Value,
491}
492
493impl DecisionResponse {
494    /// A response with only answers, for mocks and tests.
495    pub fn with_answers(model: impl Into<String>, answers: BTreeMap<String, Answer>) -> Self {
496        Self {
497            model: model.into(),
498            answers,
499            usage: None,
500            cost: None,
501            fallback_triggered_by: Vec::new(),
502            latency: Duration::ZERO,
503            provider_metadata: Value::Null,
504        }
505    }
506}
507
508/// A model that answers typed questions about a state.
509#[async_trait]
510pub trait DecisionModel: Send + Sync {
511    /// The model's identifier, such as `typesafe-ai/jev`.
512    fn model_id(&self) -> &str;
513
514    /// Answer `request`'s questions.
515    async fn decide(&self, request: DecisionRequest) -> Result<DecisionResponse, LlmError>;
516}
517
518type MockAnswerFn =
519    dyn Fn(&DecisionRequest) -> Result<BTreeMap<String, Answer>, LlmError> + Send + Sync;
520
521/// A [`DecisionModel`] whose answers come from a closure, for offline tests.
522/// It records every request it receives.
523pub struct MockDecisionModel {
524    answer: Arc<MockAnswerFn>,
525    requests: parking_lot::Mutex<Vec<DecisionRequest>>,
526    latency: Duration,
527}
528
529impl MockDecisionModel {
530    /// A mock answering with `answer`.
531    pub fn new(
532        answer: impl Fn(&DecisionRequest) -> Result<BTreeMap<String, Answer>, LlmError>
533        + Send
534        + Sync
535        + 'static,
536    ) -> Self {
537        Self {
538            answer: Arc::new(answer),
539            requests: parking_lot::Mutex::new(Vec::new()),
540            latency: Duration::ZERO,
541        }
542    }
543
544    /// Wait `latency` before answering.
545    #[must_use]
546    pub fn with_latency(mut self, latency: Duration) -> Self {
547        self.latency = latency;
548        self
549    }
550
551    /// Every request received so far.
552    pub fn requests(&self) -> Vec<DecisionRequest> {
553        self.requests.lock().clone()
554    }
555}
556
557#[async_trait]
558impl DecisionModel for MockDecisionModel {
559    fn model_id(&self) -> &str {
560        "mock-decision"
561    }
562
563    async fn decide(&self, request: DecisionRequest) -> Result<DecisionResponse, LlmError> {
564        self.requests.lock().push(request.clone());
565        if !self.latency.is_zero() {
566            tokio::time::sleep(self.latency).await;
567        }
568        let answers = (self.answer)(&request)?;
569        let mut response = DecisionResponse::with_answers("mock-decision", answers);
570        response.latency = self.latency;
571        Ok(response)
572    }
573}
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578
579    #[test]
580    fn questions_serialize_as_the_decision_api_takes_them() {
581        let q = Question::boolean("Did the caller agree?").when("said yes", "anything else");
582        assert_eq!(
583            serde_json::to_value(&q).unwrap(),
584            json!({
585                "type": "boolean",
586                "instructions": "Did the caller agree?",
587                "criteria": { "true": "said yes", "false": "anything else" }
588            })
589        );
590        let q = Question::choice(
591            "Route it.",
592            [("billing", "charges"), ("shipping", "delivery")],
593        );
594        assert_eq!(
595            serde_json::to_value(&q).unwrap(),
596            json!({
597                "type": "choice",
598                "instructions": "Route it.",
599                "criteria": { "billing": "charges", "shipping": "delivery" }
600            })
601        );
602        let q = Question::score("How urgent?", ["low", "medium", "high"]);
603        assert_eq!(
604            serde_json::to_value(&q).unwrap(),
605            json!({ "type": "score", "instructions": "How urgent?", "criteria": ["low", "medium", "high"] })
606        );
607    }
608
609    #[test]
610    fn question_limits_are_checked() {
611        assert!(Question::boolean("").validate().is_err());
612        let none: [(&str, &str); 0] = [];
613        assert!(Question::choice("Pick.", none).validate().is_err());
614        let many = (0..256).map(|i| (format!("o{i}"), "x".to_string()));
615        assert!(Question::choice("Pick.", many).validate().is_err());
616        assert!(Question::score("Rate.", ["only one"]).validate().is_err());
617        assert!(Question::score("Rate.", ["a", "b"]).validate().is_ok());
618    }
619
620    #[test]
621    fn answers_are_read_from_the_decision_api() {
622        let a = Answer::from_value(json!({ "type": "boolean", "probability": 0.98 }));
623        assert_eq!(a.probability(), Some(0.98));
624
625        let a = Answer::from_value(json!({
626            "type": "choice", "choice": "billing",
627            "probabilities": { "billing": 0.9, "shipping": 0.1 }
628        }));
629        assert_eq!(
630            a.certainty(),
631            Some(0.9),
632            "no confidence: the chosen option's probability"
633        );
634        let a = Answer::from_value(json!({
635            "type": "choice", "choice": "billing", "confidence": 0.4,
636            "probabilities": { "billing": 0.9, "shipping": 0.1 }
637        }));
638        assert_eq!(a.certainty(), Some(0.4), "the reported confidence wins");
639
640        let a = Answer::from_value(json!({
641            "type": "score", "score": 2.97,
642            "probabilities": { "0": 0, "1": 0, "2": 0.02, "3": 0.98 }
643        }));
644        assert!(matches!(a, Answer::Score { score, .. } if (score - 2.97).abs() < 1e-9));
645        assert_eq!(
646            a.certainty(),
647            Some(1.0),
648            "the two levels the score lies between hold all the probability"
649        );
650
651        assert!(matches!(
652            Answer::from_value(json!({ "type": "refusal" })),
653            Answer::Refusal(_)
654        ));
655        assert!(matches!(
656            Answer::from_value(json!({ "type": "boolean" })),
657            Answer::Other(_)
658        ));
659        let round = json!({ "type": "choice", "choice": "a", "probabilities": { "a": 1.0 }, "confidence": 1.0 });
660        assert_eq!(Answer::from_value(round.clone()).to_value(), round);
661    }
662
663    #[test]
664    fn fallback_conditions_take_the_gateway_shape() {
665        let when = FallbackWhen::Any(vec![
666            FallbackWhen::ConfidenceBelow {
667                question: Some("route".into()),
668                below: 0.6,
669            },
670            FallbackWhen::ProbabilityBetween {
671                question: None,
672                low: 0.3,
673                high: 0.8,
674            },
675            FallbackWhen::AtLeast {
676                count: 1,
677                conditions: vec![FallbackWhen::All(vec![FallbackWhen::ConfidenceBelow {
678                    question: None,
679                    below: 0.5,
680                }])],
681            },
682        ]);
683        assert_eq!(
684            when.to_value(),
685            json!({ "any": [
686                { "question": "route", "confidenceBelow": 0.6 },
687                { "probabilityBetween": [0.3, 0.8] },
688                { "atLeast": { "count": 1, "conditions": [ { "all": [ { "confidenceBelow": 0.5 } ] } ] } }
689            ] })
690        );
691    }
692
693    #[tokio::test]
694    async fn the_mock_records_requests() {
695        let model = MockDecisionModel::new(|req| {
696            Ok(req
697                .questions
698                .keys()
699                .map(|k| (k.clone(), Answer::boolean(0.5)))
700                .collect())
701        });
702        let response = model
703            .decide(DecisionRequest::new(
704                json!("x"),
705                [("a", Question::boolean("A?"))],
706            ))
707            .await
708            .unwrap();
709        assert_eq!(response.answers.len(), 1);
710        assert_eq!(model.requests().len(), 1);
711    }
712}
gemini_adk_rs/decision/mod.rs

gemini_adk_rs/decision/
mod.rs
