gemini_adk_rs/decision/
gateway.rs

1//! Vercel AI Gateway's decision API (`POST /v1/evaluate`).
2
3use std::time::{Duration, Instant};
4
5use async_trait::async_trait;
6use serde_json::{Map, Value, json};
7
8use super::{
9    Answer, DecisionModel, DecisionRequest, DecisionResponse, DecisionUsage, FallbackWhen,
10};
11use crate::llm::LlmError;
12
13/// A decision model served by Vercel AI Gateway, such as TypeSafe AI's Jev.
14///
15/// Authenticates with `AI_GATEWAY_API_KEY`, or with a `VERCEL_OIDC_TOKEN`
16/// (what `vercel env pull` writes and what deployments on Vercel receive).
17///
18/// ```no_run
19/// use gemini_adk_rs::decision::{DecisionModel, DecisionRequest, GatewayDecisionModel, Question};
20/// use serde_json::json;
21///
22/// # async fn run() -> Result<(), gemini_adk_rs::llm::LlmError> {
23/// let jev = GatewayDecisionModel::from_env()?;
24/// let response = jev
25///     .decide(DecisionRequest::new(
26///         json!({ "caller": "Yes, please book it." }),
27///         [("confirmed", Question::boolean("Did the caller agree to book?"))],
28///     ))
29///     .await?;
30/// println!("{:?} in {:?}", response.answers["confirmed"], response.latency);
31/// # Ok(()) }
32/// ```
33#[derive(Clone)]
34pub struct GatewayDecisionModel {
35    client: reqwest::Client,
36    base_url: String,
37    model: String,
38    token: String,
39    timeout: Duration,
40    options: Value,
41}
42
43impl std::fmt::Debug for GatewayDecisionModel {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        // Never print the token.
46        f.debug_struct("GatewayDecisionModel")
47            .field("base_url", &self.base_url)
48            .field("model", &self.model)
49            .field("timeout", &self.timeout)
50            .field("options", &self.options)
51            .finish_non_exhaustive()
52    }
53}
54
55impl GatewayDecisionModel {
56    /// AI Gateway's address.
57    pub const DEFAULT_BASE_URL: &'static str = "https://ai-gateway.vercel.sh";
58    /// TypeSafe AI's Jev, the default model.
59    pub const JEV: &'static str = "typesafe-ai/jev";
60    /// How long a decision may take before the call fails.
61    pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
62
63    /// A client for `model`, authenticating with `token`.
64    pub fn new(model: impl Into<String>, token: impl Into<String>) -> Self {
65        Self {
66            client: reqwest::Client::new(),
67            base_url: Self::DEFAULT_BASE_URL.to_string(),
68            model: model.into(),
69            token: token.into(),
70            timeout: Self::DEFAULT_TIMEOUT,
71            options: Value::Null,
72        }
73    }
74
75    /// A client configured from the environment:
76    ///
77    /// | Variable | Use |
78    /// |---|---|
79    /// | `AI_GATEWAY_API_KEY`, else `VERCEL_OIDC_TOKEN` | Credential (required) |
80    /// | `AI_GATEWAY_DECISION_MODEL` | Model, default `typesafe-ai/jev` |
81    /// | `AI_GATEWAY_BASE_URL` | Gateway address, default `https://ai-gateway.vercel.sh` |
82    pub fn from_env() -> Result<Self, LlmError> {
83        let var = |name: &str| std::env::var(name).ok().filter(|v| !v.trim().is_empty());
84        let token = var("AI_GATEWAY_API_KEY")
85            .or_else(|| var("VERCEL_OIDC_TOKEN"))
86            .ok_or_else(|| {
87                LlmError::Auth(
88                    "no AI Gateway credential: set AI_GATEWAY_API_KEY (for example in \
89                     .env.local), or run `vercel env pull` for a VERCEL_OIDC_TOKEN"
90                        .into(),
91                )
92            })?;
93        let model = var("AI_GATEWAY_DECISION_MODEL").unwrap_or_else(|| Self::JEV.to_string());
94        let mut this = Self::new(model, token);
95        if let Some(base) = var("AI_GATEWAY_BASE_URL") {
96            this.base_url = base.trim_end_matches('/').to_string();
97        }
98        Ok(this)
99    }
100
101    /// Use another Gateway address (a proxy, or a local server in tests).
102    #[must_use]
103    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
104        self.base_url = base_url.into().trim_end_matches('/').to_string();
105        self
106    }
107
108    /// Fail a decision that takes longer than `timeout`.
109    #[must_use]
110    pub fn with_timeout(mut self, timeout: Duration) -> Self {
111        self.timeout = timeout;
112        self
113    }
114
115    /// Ask the Gateway to route only to providers with zero data retention.
116    #[must_use]
117    pub fn zero_data_retention(mut self, on: bool) -> Self {
118        merge_json(
119            &mut self.options,
120            &object("gateway", json!({ "zeroDataRetention": on })),
121        );
122        self
123    }
124
125    /// Restrict which providers may serve the model.
126    #[must_use]
127    pub fn only_providers<S: Into<String>>(
128        mut self,
129        providers: impl IntoIterator<Item = S>,
130    ) -> Self {
131        let only: Vec<String> = providers.into_iter().map(Into::into).collect();
132        merge_json(
133            &mut self.options,
134            &object("gateway", json!({ "only": only })),
135        );
136        self
137    }
138
139    /// Rerun the decision with `model` when `when` matches the primary
140    /// answers (a Gateway decision fallback). `model` may be a language
141    /// model, which answers through structured output and reports no
142    /// confidence. Both stages are billed and their latencies add.
143    #[must_use]
144    pub fn fallback(mut self, model: impl Into<String>, when: FallbackWhen) -> Self {
145        let entry = json!({ "model": model.into(), "when": when.to_value() });
146        merge_json(
147            &mut self.options,
148            &object("gateway", json!({ "models": [entry] })),
149        );
150        self
151    }
152
153    /// The request body for `request`.
154    pub fn body(&self, request: &DecisionRequest) -> Value {
155        let mut body = json!({
156            "model": self.model,
157            "state": request.state,
158            "questions": request.questions,
159        });
160        let mut options = self.options.clone();
161        if let Some(extra) = &request.provider_options {
162            merge_json(&mut options, extra);
163        }
164        if !options.is_null() {
165            body["providerOptions"] = options;
166        }
167        body
168    }
169
170    /// Read a `/v1/evaluate` response.
171    pub fn parse(status: u16, body: &str, latency: Duration) -> Result<DecisionResponse, LlmError> {
172        let value: Value = serde_json::from_str(body).map_err(|e| LlmError::Api {
173            status,
174            message: format!("unreadable response ({e}): {}", truncate(body, 300)),
175        })?;
176        if !(200..300).contains(&status) {
177            let message = value["error"]["message"]
178                .as_str()
179                .map_or_else(|| truncate(body, 300), str::to_string);
180            return Err(match status {
181                401 => LlmError::Auth(format!(
182                    "AI Gateway rejected the credential: {message}. Check AI_GATEWAY_API_KEY"
183                )),
184                _ => LlmError::Api { status, message },
185            });
186        }
187        let answers = value["answers"]
188            .as_object()
189            .ok_or_else(|| LlmError::Api {
190                status,
191                message: format!("response has no answers: {}", truncate(body, 300)),
192            })?
193            .iter()
194            .map(|(k, v)| (k.clone(), Answer::from_value(v.clone())))
195            .collect();
196        let usage = value.get("usage").map(|u| DecisionUsage {
197            input_tokens: u["inputTokens"].as_u64().unwrap_or(0),
198            output_tokens: u["outputTokens"].as_u64().unwrap_or(0),
199        });
200        let gateway = &value["providerMetadata"]["gateway"];
201        let fallback_triggered_by = gateway["routing"]["modelAttempts"]
202            .as_array()
203            .into_iter()
204            .flatten()
205            .filter_map(|attempt| attempt["triggeredBy"].as_array())
206            .flatten()
207            .map(|t| {
208                (
209                    t["question"].as_str().unwrap_or_default().to_string(),
210                    t["reason"].as_str().unwrap_or_default().to_string(),
211                )
212            })
213            .collect();
214        Ok(DecisionResponse {
215            model: value["model"].as_str().unwrap_or_default().to_string(),
216            answers,
217            usage,
218            cost: gateway["cost"].as_str().map(str::to_string),
219            fallback_triggered_by,
220            latency,
221            provider_metadata: value["providerMetadata"].clone(),
222        })
223    }
224}
225
226/// Merge `overlay` into `base`, objects recursively; other values replace.
227fn merge_json(base: &mut Value, overlay: &Value) {
228    match (base, overlay) {
229        (Value::Object(b), Value::Object(o)) => {
230            for (k, v) in o {
231                merge_json(b.entry(k.clone()).or_insert(Value::Null), v);
232            }
233        }
234        (b, o) => *b = o.clone(),
235    }
236}
237
238/// A JSON object with one key, for building nested options.
239fn object(key: &str, value: Value) -> Value {
240    let mut m = Map::new();
241    m.insert(key.to_string(), value);
242    Value::Object(m)
243}
244
245fn truncate(s: &str, max: usize) -> String {
246    match s.char_indices().nth(max) {
247        Some((i, _)) => format!("{}…", &s[..i]),
248        None => s.to_string(),
249    }
250}
251
252#[async_trait]
253impl DecisionModel for GatewayDecisionModel {
254    fn model_id(&self) -> &str {
255        &self.model
256    }
257
258    async fn decide(&self, request: DecisionRequest) -> Result<DecisionResponse, LlmError> {
259        for (id, question) in &request.questions {
260            question
261                .validate()
262                .map_err(|e| LlmError::Config(format!("question '{id}': {e}")))?;
263        }
264        let started = Instant::now();
265        let response = self
266            .client
267            .post(format!("{}/v1/evaluate", self.base_url))
268            .bearer_auth(&self.token)
269            .timeout(self.timeout)
270            .json(&self.body(&request))
271            .send()
272            .await
273            .map_err(|e| LlmError::Transport(e.to_string()))?;
274        let status = response.status().as_u16();
275        let text = response
276            .text()
277            .await
278            .map_err(|e| LlmError::Transport(e.to_string()))?;
279        let latency = started.elapsed();
280        let parsed = Self::parse(status, &text, latency);
281        match &parsed {
282            Ok(r) => tracing::info!(
283                model = %r.model,
284                ms = latency.as_millis() as u64,
285                questions = r.answers.len(),
286                fallback = !r.fallback_triggered_by.is_empty(),
287                "decision"
288            ),
289            Err(e) => {
290                tracing::warn!(model = %self.model, ms = latency.as_millis() as u64, "decision failed: {e}");
291            }
292        }
293        parsed
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::decision::Question;
301
302    fn jev() -> GatewayDecisionModel {
303        GatewayDecisionModel::new(GatewayDecisionModel::JEV, "test-token")
304    }
305
306    #[test]
307    fn the_body_carries_state_questions_and_options() {
308        let model = jev().zero_data_retention(true).fallback(
309            "google/gemini-3.8-flash",
310            FallbackWhen::ProbabilityBetween {
311                question: Some("confirmed".into()),
312                low: 0.3,
313                high: 0.8,
314            },
315        );
316        let mut request = DecisionRequest::new(
317            json!({ "caller": "yes" }),
318            [("confirmed", Question::boolean("Agreed?"))],
319        );
320        request.provider_options = Some(json!({ "gateway": { "only": ["typesafe-ai"] } }));
321        assert_eq!(
322            model.body(&request),
323            json!({
324                "model": "typesafe-ai/jev",
325                "state": { "caller": "yes" },
326                "questions": { "confirmed": { "type": "boolean", "instructions": "Agreed?" } },
327                "providerOptions": { "gateway": {
328                    "zeroDataRetention": true,
329                    "only": ["typesafe-ai"],
330                    "models": [ { "model": "google/gemini-3.8-flash",
331                                  "when": { "question": "confirmed", "probabilityBetween": [0.3, 0.8] } } ]
332                } }
333            })
334        );
335        assert!(
336            jev().body(&request).get("providerOptions").is_some(),
337            "per-request options alone"
338        );
339        let plain = DecisionRequest::new(json!("x"), [("a", Question::boolean("A?"))]);
340        assert!(jev().body(&plain).get("providerOptions").is_none());
341    }
342
343    #[test]
344    fn a_documented_response_is_read() {
345        let body = r#"{
346          "model": "typesafe-ai/jev",
347          "answers": { "refund": { "type": "boolean", "probability": 0.98 } },
348          "usage": { "inputTokens": 275, "outputTokens": 20 },
349          "providerMetadata": { "gateway": {
350            "routing": { "originalModelId": "typesafe-ai/jev", "finalProvider": "typesafe-ai" },
351            "cost": "0.00001155", "generationId": "gen_1" } }
352        }"#;
353        let r = GatewayDecisionModel::parse(200, body, Duration::from_millis(120)).unwrap();
354        assert_eq!(r.model, "typesafe-ai/jev");
355        assert_eq!(r.answers["refund"].probability(), Some(0.98));
356        assert_eq!(
357            r.usage,
358            Some(DecisionUsage {
359                input_tokens: 275,
360                output_tokens: 20
361            })
362        );
363        assert_eq!(r.cost.as_deref(), Some("0.00001155"));
364        assert!(r.fallback_triggered_by.is_empty());
365    }
366
367    #[test]
368    fn a_triggered_fallback_is_reported() {
369        let body = r#"{
370          "model": "openai/gpt-6-astra",
371          "answers": { "intent": { "type": "choice", "choice": "billing" } },
372          "providerMetadata": { "gateway": { "routing": { "modelAttempts": [
373            { "canonicalSlug": "typesafe-ai/jev", "success": true },
374            { "canonicalSlug": "openai/gpt-6-astra", "success": true,
375              "triggeredBy": [ { "question": "intent", "reason": "confidence_below" } ] }
376          ] } } }
377        }"#;
378        let r = GatewayDecisionModel::parse(200, body, Duration::ZERO).unwrap();
379        assert_eq!(r.model, "openai/gpt-6-astra");
380        assert_eq!(
381            r.fallback_triggered_by,
382            vec![("intent".to_string(), "confidence_below".to_string())]
383        );
384        assert_eq!(
385            r.answers["intent"].certainty(),
386            None,
387            "a language model reports none"
388        );
389    }
390
391    #[test]
392    fn errors_keep_the_gateway_message() {
393        let restricted = r#"{"error":{"message":"Free tier users do not have access to this model.","type":"no_providers_available"}}"#;
394        match GatewayDecisionModel::parse(403, restricted, Duration::ZERO) {
395            Err(LlmError::Api {
396                status: 403,
397                message,
398            }) => {
399                assert!(message.contains("Free tier"), "{message}");
400            }
401            other => panic!("{other:?}"),
402        }
403        let auth = r#"{"error":{"message":"Authentication failed","type":"authentication_error"}}"#;
404        assert!(matches!(
405            GatewayDecisionModel::parse(401, auth, Duration::ZERO),
406            Err(LlmError::Auth(_))
407        ));
408        let busy = GatewayDecisionModel::parse(503, "upstream busy", Duration::ZERO).unwrap_err();
409        assert!(busy.is_retryable(), "{busy:?}");
410    }
411
412    #[test]
413    fn provider_options_merge_recursively() {
414        let mut base = json!({ "gateway": { "zeroDataRetention": true, "only": ["typesafe-ai"] } });
415        merge_json(
416            &mut base,
417            &json!({ "gateway": { "only": ["digitalocean"] }, "openai": { "x": 1 } }),
418        );
419        assert_eq!(
420            base,
421            json!({ "gateway": { "zeroDataRetention": true, "only": ["digitalocean"] }, "openai": { "x": 1 } })
422        );
423    }
424
425    #[test]
426    fn the_token_is_never_printed() {
427        let printed = format!("{:?}", jev());
428        assert!(!printed.contains("test-token"), "{printed}");
429    }
430
431    #[tokio::test]
432    async fn an_invalid_question_fails_before_the_call() {
433        let none: [(&str, &str); 0] = [];
434        let err = jev()
435            .with_base_url("http://127.0.0.1:9")
436            .decide(DecisionRequest::new(
437                json!("x"),
438                [("pick", Question::choice("Pick.", none))],
439            ))
440            .await
441            .unwrap_err();
442        assert!(matches!(err, LlmError::Config(_)), "{err:?}");
443    }
444
445    #[tokio::test]
446    async fn a_call_reaches_the_gateway_and_reads_the_answers() {
447        use tokio::io::{AsyncReadExt, AsyncWriteExt};
448        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
449        let addr = listener.local_addr().unwrap();
450        let server = tokio::spawn(async move {
451            let (mut socket, _) = listener.accept().await.unwrap();
452            let mut seen = Vec::new();
453            let mut buf = [0u8; 4096];
454            // Read until the JSON body has arrived.
455            while !String::from_utf8_lossy(&seen).contains("\"questions\"") {
456                let n = socket.read(&mut buf).await.unwrap();
457                if n == 0 {
458                    break;
459                }
460                seen.extend_from_slice(&buf[..n]);
461            }
462            let body = r#"{"model":"typesafe-ai/jev","answers":{"confirmed":{"type":"boolean","probability":0.93}}}"#;
463            let reply = format!(
464                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
465                body.len()
466            );
467            socket.write_all(reply.as_bytes()).await.unwrap();
468            String::from_utf8_lossy(&seen).to_string()
469        });
470        let model = jev().with_base_url(format!("http://{addr}"));
471        let response = model
472            .decide(DecisionRequest::new(
473                json!("Caller: yes"),
474                [("confirmed", Question::boolean("Agreed?"))],
475            ))
476            .await
477            .unwrap();
478        assert_eq!(response.answers["confirmed"].probability(), Some(0.93));
479        let request = server.await.unwrap();
480        assert!(request.starts_with("POST /v1/evaluate"), "{request}");
481        assert!(
482            request
483                .to_lowercase()
484                .contains("authorization: bearer test-token"),
485            "{request}"
486        );
487    }
488}
gemini_adk_rs/decision/gateway.rs

gemini_adk_rs/decision/
gateway.rs
