typesafe/
client.rs

1//! The asynchronous client.
2
3use std::future::{Future, IntoFuture};
4use std::path::PathBuf;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8
9use http::header::{
10    ACCEPT, AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue, USER_AGENT,
11};
12use http::{Method, StatusCode};
13use serde::Serialize;
14use serde_json::{Map, Value};
15
16use crate::cassette;
17use crate::constants::*;
18use crate::error::{ApiError, Error, ResponseValidationError, Result, lenient_body};
19use crate::question::Questions;
20use crate::response::{
21    DecodeFailure, DecodedSystemOne, ListModelsResponse, ResponseMeta, SystemOneResponse,
22    decode_models, decode_system_one,
23};
24use crate::retry::RetryPolicy;
25
26pub(crate) type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
27
28/// Builder for [`Client`]. Explicit settings win over environment variables; empty or
29/// whitespace-only environment values are ignored.
30#[derive(Debug, Default)]
31pub struct ClientBuilder {
32    api_key: Option<String>,
33    base_url: Option<String>,
34    model: Option<String>,
35    timeout: Option<Duration>,
36    retry: Option<RetryPolicy>,
37    headers: HeaderMap,
38    http: Option<reqwest::Client>,
39    record: Option<PathBuf>,
40    replay: Option<PathBuf>,
41}
42
43impl ClientBuilder {
44    /// API key (else `TYPESAFE_API_KEY`). Surrounding whitespace is trimmed; an empty key, or one
45    /// with whitespace, control or non-ASCII characters inside it, is rejected by `build`.
46    pub fn api_key(mut self, key: impl Into<String>) -> Self {
47        self.api_key = Some(key.into());
48        self
49    }
50
51    /// API root (else `TYPESAFE_BASE_URL`, else `https://api.typesafe.ai`).
52    pub fn base_url(mut self, url: impl Into<String>) -> Self {
53        self.base_url = Some(url.into());
54        self
55    }
56
57    /// Default model (else `TYPESAFE_DEFAULT_MODEL`, else `jev-latest`).
58    pub fn model(mut self, model: impl Into<String>) -> Self {
59        self.model = Some(model.into());
60        self
61    }
62
63    /// Per-attempt timeout (default 10s).
64    pub fn timeout(mut self, timeout: Duration) -> Self {
65        self.timeout = Some(timeout);
66        self
67    }
68
69    /// Default retry policy.
70    pub fn retry(mut self, policy: RetryPolicy) -> Self {
71        self.retry = Some(policy);
72        self
73    }
74
75    /// Extra header sent with every request. Authentication and SDK identification headers
76    /// cannot be overridden.
77    pub fn header(mut self, name: HeaderName, value: HeaderValue) -> Self {
78        self.headers.insert(name, value);
79        self
80    }
81
82    /// Write each successful System One response to `<dir>/<key>.json` (else `TYPESAFE_RECORD`).
83    /// The directory is created if needed. See [`cassette`](crate::cassette).
84    pub fn record(mut self, dir: impl Into<PathBuf>) -> Self {
85        self.record = Some(dir.into());
86        self
87    }
88
89    /// Answer System One calls from `<dir>/<key>.json` instead of the network (else
90    /// `TYPESAFE_REPLAY`); no API key is needed. A request that was never recorded fails with
91    /// [`Error::ReplayMiss`]. See [`cassette`](crate::cassette).
92    pub fn replay(mut self, dir: impl Into<PathBuf>) -> Self {
93        self.replay = Some(dir.into());
94        self
95    }
96
97    /// Use your own `reqwest::Client` (proxies, TLS, connection pools…). Requires the
98    /// `reqwest-client` feature.
99    #[cfg(feature = "reqwest-client")]
100    pub fn http_client(mut self, client: reqwest::Client) -> Self {
101        self.http = Some(client);
102        self
103    }
104
105    /// Build the client.
106    pub fn build(self) -> Result<Client> {
107        let cassette = match (
108            resolve_path(self.record, RECORD_ENV),
109            resolve_path(self.replay, REPLAY_ENV),
110        ) {
111            (Some(_), Some(_)) => {
112                return Err(Error::Config(format!(
113                    "Both record and replay are set; a client does one or the other. Unset \
114                     {RECORD_ENV} or {REPLAY_ENV}, or drop one of the builder calls."
115                )));
116            }
117            (Some(dir), None) => {
118                std::fs::create_dir_all(&dir).map_err(|e| {
119                    Error::Config(format!("Cannot record into {}: {e}.", dir.display()))
120                })?;
121                Some(Cassette::Record(dir))
122            }
123            (None, Some(dir)) => Some(Cassette::Replay(dir)),
124            (None, None) => None,
125        };
126        let replaying = matches!(cassette, Some(Cassette::Replay(_)));
127        // A replaying client never sends anything, so it has no use for a key.
128        let api_key = resolve(self.api_key, API_KEY_ENV, None)
129            .map(|k| k.trim().to_owned())
130            .filter(|k| !k.is_empty());
131        if api_key.is_none() && !replaying {
132            return Err(Error::Config(format!(
133                "No API key was provided. Pass api_key or set the {API_KEY_ENV} environment variable."
134            )));
135        }
136        if api_key
137            .as_deref()
138            .is_some_and(|k| !k.bytes().all(|b| b.is_ascii_graphic()))
139        {
140            return Err(Error::Config(
141                "API key must contain only printable ASCII characters without whitespace.".into(),
142            ));
143        }
144        let base_url = resolve(self.base_url, BASE_URL_ENV, Some(DEFAULT_BASE_URL))
145            .unwrap_or_default()
146            .trim_end_matches('/')
147            .to_owned();
148        check_base_url(&base_url)?;
149        let model = resolve(self.model, DEFAULT_MODEL_ENV, Some(DEFAULT_MODEL)).unwrap_or_default();
150        let timeout = check_timeout(self.timeout.unwrap_or(DEFAULT_TIMEOUT))?;
151        let retry = self.retry.unwrap_or_default();
152        retry.validate()?;
153
154        let mut protected = HeaderMap::new();
155        if let Some(api_key) = api_key {
156            let mut auth = HeaderValue::from_str(&format!("Bearer {api_key}"))
157                .expect("a printable ASCII key is a valid header value");
158            auth.set_sensitive(true);
159            protected.insert(AUTHORIZATION, auth);
160        }
161        protected.insert(ACCEPT, HeaderValue::from_static("application/json"));
162        let ident = HeaderValue::from_str(&format!("{SDK_NAME}/{VERSION}")).expect("ascii");
163        protected.insert(USER_AGENT, ident.clone());
164        protected.insert(HeaderName::from_static(SDK_HEADER), ident);
165        protected.insert(
166            HeaderName::from_static(RUNTIME_HEADER),
167            HeaderValue::from_str(&format!(
168                "rust ({}; {})",
169                std::env::consts::OS,
170                std::env::consts::ARCH
171            ))
172            .expect("ascii"),
173        );
174
175        Ok(Client {
176            inner: Arc::new(Inner {
177                http: self.http.unwrap_or_default(),
178                base_url,
179                model,
180                timeout,
181                retry,
182                default_headers: self.headers,
183                protected,
184                cassette,
185            }),
186        })
187    }
188}
189
190fn resolve(explicit: Option<String>, env: &str, default: Option<&str>) -> Option<String> {
191    explicit
192        .or_else(|| {
193            std::env::var(env)
194                .ok()
195                .map(|v| v.trim().to_owned())
196                .filter(|v| !v.is_empty())
197        })
198        .or_else(|| default.map(str::to_owned))
199}
200
201fn resolve_path(explicit: Option<PathBuf>, env: &str) -> Option<PathBuf> {
202    explicit.or_else(|| resolve(None, env, None).map(PathBuf::from))
203}
204
205fn check_base_url(url: &str) -> Result<()> {
206    match reqwest::Url::parse(url) {
207        Ok(u) if matches!(u.scheme(), "http" | "https") && u.has_host() => Ok(()),
208        Ok(_) => Err(Error::Config(format!(
209            "base_url must be an http(s) URL with a host, got {url:?}."
210        ))),
211        Err(e) => Err(Error::Config(format!(
212            "base_url {url:?} is not a valid URL: {e}."
213        ))),
214    }
215}
216
217fn check_timeout(t: Duration) -> Result<Duration> {
218    if t.is_zero() {
219        Err(Error::Config("timeout must be a positive duration.".into()))
220    } else {
221        Ok(t)
222    }
223}
224
225#[derive(Debug)]
226struct Inner {
227    http: reqwest::Client,
228    base_url: String,
229    model: String,
230    timeout: Duration,
231    retry: RetryPolicy,
232    default_headers: HeaderMap,
233    protected: HeaderMap,
234    cassette: Option<Cassette>,
235}
236
237/// Where System One responses are recorded to or replayed from.
238#[derive(Debug)]
239enum Cassette {
240    Record(PathBuf),
241    Replay(PathBuf),
242}
243
244/// TypeSafe API client. Cheap to clone; clones share the connection pool.
245///
246/// ```no_run
247/// use typesafe::{Client, Choice, Noul, Questions, Score};
248///
249/// # async fn run() -> typesafe::Result<()> {
250/// let client = Client::from_env()?;
251/// let res = client
252///     .system_one(
253///         "I've been trying to connect Stripe for 3 days. Please help ASAP.",
254///         Questions::new()
255///             .with("department", Choice::new("Which team should handle this")
256///                 .option("billing", "Payment or subscription issues")
257///                 .option("technical", "Bugs or integration problems"))
258///             .with("frustration", Score::new("How frustrated", ["Calm", "Frustrated", "Very angry"]))
259///             .with("is_urgent", Noul::new("The message conveys urgency")),
260///     )
261///     .await?;
262/// println!("{}", res.choice("department").unwrap().choice);
263/// # Ok(()) }
264/// ```
265#[derive(Debug, Clone)]
266pub struct Client {
267    inner: Arc<Inner>,
268}
269
270impl Client {
271    /// Start configuring a client.
272    pub fn builder() -> ClientBuilder {
273        ClientBuilder::default()
274    }
275
276    /// A client configured entirely from the environment.
277    pub fn from_env() -> Result<Self> {
278        Self::builder().build()
279    }
280
281    /// The default model.
282    pub fn default_model(&self) -> &str {
283        &self.inner.model
284    }
285
286    /// Ask typed questions about `state` (a string, or anything `Serialize` that becomes a JSON
287    /// object/array). Returns a request builder: `.await` it directly or set per-call options first.
288    pub fn system_one<S: Serialize>(
289        &self,
290        state: S,
291        questions: impl Into<Questions>,
292    ) -> SystemOneRequest {
293        SystemOneRequest {
294            client: self.clone(),
295            state: serde_json::to_value(state)
296                .map_err(|e| format!("The state could not be encoded as JSON: {e}")),
297            questions: questions.into(),
298            model: None,
299            extra_body: Map::new(),
300            opts: CallOptions::default(),
301        }
302    }
303
304    /// The Models resource.
305    pub fn models(&self) -> Models {
306        Models {
307            client: self.clone(),
308        }
309    }
310
311    async fn execute(
312        &self,
313        method: Method,
314        path: &str,
315        body: Option<Vec<u8>>,
316        opts: CallOptions,
317    ) -> Result<(Vec<u8>, ResponseMeta, String)> {
318        let inner = &self.inner;
319        let retry = opts.retry.as_ref().unwrap_or(&inner.retry);
320        retry.validate()?;
321        let timeout = check_timeout(opts.timeout.unwrap_or(inner.timeout))?;
322        let url = format!("{}{}", inner.base_url, path);
323        let endpoint = format!("{method} {}", redact_url(&url));
324
325        let mut headers = inner.default_headers.clone();
326        headers.extend(opts.headers);
327        headers.remove(RETRY_COUNT_HEADER);
328        headers.extend(inner.protected.clone());
329        if body.is_some() {
330            headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
331        }
332
333        let started = Instant::now();
334        let mut attempts: u32 = 0;
335        loop {
336            let mut h = headers.clone();
337            if attempts > 0 {
338                h.insert(
339                    HeaderName::from_static(RETRY_COUNT_HEADER),
340                    HeaderValue::from(attempts),
341                );
342                tracing::info!(%endpoint, retry = attempts, "retrying");
343            }
344            attempts += 1;
345            let result = self
346                .attempt(&method, &url, &endpoint, h, body.clone(), timeout)
347                .await;
348            match result {
349                Ok((bytes, status, resp_headers)) => {
350                    let meta = ResponseMeta {
351                        status,
352                        headers: resp_headers,
353                        attempts,
354                    };
355                    return Ok((bytes, meta, endpoint));
356                }
357                Err(err) => {
358                    if !retry.is_retryable(&err) {
359                        return Err(err);
360                    }
361                    let delay = retry.delay(attempts, &err);
362                    if retry.should_stop(attempts, started.elapsed(), delay) {
363                        return Err(err);
364                    }
365                    if !delay.is_zero() {
366                        tokio::time::sleep(delay).await;
367                    }
368                }
369            }
370        }
371    }
372
373    async fn attempt(
374        &self,
375        method: &Method,
376        url: &str,
377        endpoint: &str,
378        headers: HeaderMap,
379        body: Option<Vec<u8>>,
380        timeout: Duration,
381    ) -> Result<(Vec<u8>, u16, HeaderMap)> {
382        let t0 = Instant::now();
383        tracing::debug!(%endpoint, "->");
384        if tracing::enabled!(tracing::Level::TRACE) {
385            tracing::trace!(%endpoint, headers = ?redacted(&headers),
386                body = %body.as_deref().map(String::from_utf8_lossy).unwrap_or_default(), "->");
387        }
388        let mut req = self
389            .inner
390            .http
391            .request(method.clone(), url)
392            .headers(headers)
393            .timeout(timeout);
394        if let Some(b) = body {
395            req = req.body(b);
396        }
397        let map_err = |e: reqwest::Error| {
398            tracing::debug!(%endpoint, error = %e, "<- transport error");
399            if e.is_timeout() {
400                Error::Timeout(timeout)
401            } else {
402                Error::Connection(Box::new(e))
403            }
404        };
405        let resp = req.send().await.map_err(map_err)?;
406        let status = resp.status();
407        let resp_headers = resp.headers().clone();
408        let bytes = resp.bytes().await.map_err(map_err)?;
409
410        tracing::debug!(
411            %endpoint,
412            status = status.as_u16(),
413            elapsed_ms = t0.elapsed().as_millis() as u64,
414            request_id = resp_headers.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()).unwrap_or("-"),
415            "<-"
416        );
417        if tracing::enabled!(tracing::Level::TRACE) {
418            tracing::trace!(%endpoint, headers = ?redacted(&resp_headers),
419                body = %String::from_utf8_lossy(&bytes), "<-");
420        }
421
422        if !status.is_success() {
423            return Err(Error::Api(Box::new(ApiError::new(
424                status.as_u16(),
425                lenient_body(&bytes),
426                resp_headers,
427                Some(endpoint.to_owned()),
428            ))));
429        }
430        Ok((bytes.to_vec(), status.as_u16(), resp_headers))
431    }
432}
433
434fn validation_error(
435    status: StatusCode,
436    body: Option<Value>,
437    headers: HeaderMap,
438    endpoint: &str,
439    f: DecodeFailure,
440) -> Error {
441    Error::ResponseValidation(Box::new(ResponseValidationError {
442        status: status.as_u16(),
443        field_path: f.path,
444        detail: f.detail,
445        body,
446        headers,
447        endpoint: Some(endpoint.to_owned()),
448    }))
449}
450
451fn redact_url(url: &str) -> String {
452    match reqwest::Url::parse(url) {
453        Ok(mut u) => {
454            let _ = u.set_username("");
455            let _ = u.set_password(None);
456            u.set_query(None);
457            u.set_fragment(None);
458            u.to_string()
459        }
460        Err(_) => url.to_owned(),
461    }
462}
463
464/// A credential-bearing header: the known names, plus anything that says it carries one, as an AI
465/// gateway's own key does (`cf-aig-authorization`, `x-portkey-api-key`, `x-gateway-token`).
466fn is_secret(name: &str) -> bool {
467    let name = name.to_ascii_lowercase();
468    SECRET_HEADERS.contains(&name.as_str())
469        || ["authorization", "api-key", "token", "secret"]
470            .iter()
471            .any(|part| name.contains(part))
472}
473
474fn redacted(headers: &HeaderMap) -> Vec<(String, String)> {
475    headers
476        .iter()
477        .map(|(k, v)| {
478            let value = if is_secret(k.as_str()) {
479                "[REDACTED]".to_owned()
480            } else {
481                v.to_str().unwrap_or("<binary>").to_owned()
482            };
483            (k.as_str().to_owned(), value)
484        })
485        .collect()
486}
487
488#[derive(Debug, Default, Clone)]
489struct CallOptions {
490    retry: Option<RetryPolicy>,
491    timeout: Option<Duration>,
492    headers: HeaderMap,
493}
494
495macro_rules! call_option_methods {
496    () => {
497        /// Override the retry policy for this call.
498        pub fn retry(mut self, policy: RetryPolicy) -> Self {
499            self.opts.retry = Some(policy);
500            self
501        }
502
503        /// Override the per-attempt timeout for this call.
504        pub fn timeout(mut self, timeout: Duration) -> Self {
505            self.opts.timeout = Some(timeout);
506            self
507        }
508
509        /// Add a header for this call (protected headers still win).
510        pub fn header(mut self, name: HeaderName, value: HeaderValue) -> Self {
511            self.opts.headers.insert(name, value);
512            self
513        }
514    };
515}
516
517/// A pending `POST /v1/systemone`. Configure it, then `.await` it (or call [`send`](Self::send)).
518#[must_use = "requests do nothing until awaited"]
519#[derive(Debug)]
520pub struct SystemOneRequest {
521    client: Client,
522    state: std::result::Result<Value, String>,
523    questions: Questions,
524    model: Option<String>,
525    extra_body: Map<String, Value>,
526    opts: CallOptions,
527}
528
529impl SystemOneRequest {
530    call_option_methods!();
531
532    /// Override the model for this call.
533    pub fn model(mut self, model: impl Into<String>) -> Self {
534        self.model = Some(model.into());
535        self
536    }
537
538    /// Add a top-level body field. A key named `state`, `model` or `questions` replaces the
539    /// standard field. Useful for API fields this SDK version does not model yet.
540    pub fn extra_body(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
541        self.extra_body.insert(key.into(), value.into());
542        self
543    }
544
545    /// Send the request.
546    pub async fn send(self) -> Result<SystemOneResponse> {
547        let state = self.state.map_err(Error::InvalidRequest)?;
548        self.questions.validate()?;
549        let model = self
550            .model
551            .unwrap_or_else(|| self.client.inner.model.clone());
552        // Serialized as a struct (not via `serde_json::Value`) so question order reaches the wire
553        // unchanged. `extra_body` keys replace the standard field of the same name.
554        let extra = &self.extra_body;
555        let body = SystemOneBody {
556            state: (!extra.contains_key("state")).then_some(&state),
557            model: (!extra.contains_key("model")).then_some(&model),
558            questions: (!extra.contains_key("questions")).then_some(&self.questions),
559            extra,
560        };
561        let bytes = serde_json::to_vec(&body).map_err(|e| {
562            Error::InvalidRequest(format!(
563                "The request body could not be encoded as JSON: {e}"
564            ))
565        })?;
566
567        let key = cassette::body_key(&bytes);
568        let (bytes, meta, endpoint) = match &self.client.inner.cassette {
569            Some(Cassette::Replay(dir)) => replay(dir, &key)?,
570            _ => {
571                self.client
572                    .execute(Method::POST, SYSTEM_ONE_PATH, Some(bytes), self.opts)
573                    .await?
574            }
575        };
576        let decoded = decode_system_one(&bytes);
577        if let (Ok(_), Some(Cassette::Record(dir))) = (&decoded, &self.client.inner.cassette)
578            && let Err(e) = cassette::write(dir, &key, &bytes)
579        {
580            tracing::warn!(dir = %dir.display(), %key, error = %e, "could not record the response");
581        }
582        match decoded {
583            Ok(DecodedSystemOne {
584                model,
585                usage,
586                answers,
587                raw,
588            }) => Ok(SystemOneResponse {
589                model,
590                usage,
591                answers,
592                raw,
593                meta,
594            }),
595            Err(f) => Err(validation_error(
596                StatusCode::from_u16(meta.status).unwrap_or(StatusCode::OK),
597                lenient_body(&bytes),
598                meta.headers,
599                &endpoint,
600                f,
601            )),
602        }
603    }
604}
605
606/// The recorded response for `key`, in the shape a live call returns: no headers, no attempts.
607fn replay(dir: &std::path::Path, key: &str) -> Result<(Vec<u8>, ResponseMeta, String)> {
608    let path = cassette::path(dir, key);
609    let bytes = std::fs::read(&path).map_err(|_| Error::ReplayMiss {
610        key: key.to_owned(),
611        path: path.clone(),
612    })?;
613    tracing::debug!(path = %path.display(), "<- replayed");
614    let meta = ResponseMeta {
615        status: 200,
616        headers: HeaderMap::new(),
617        attempts: 0,
618    };
619    Ok((bytes, meta, format!("replay {}", path.display())))
620}
621
622#[derive(Serialize)]
623struct SystemOneBody<'a> {
624    #[serde(skip_serializing_if = "Option::is_none")]
625    state: Option<&'a Value>,
626    #[serde(skip_serializing_if = "Option::is_none")]
627    model: Option<&'a str>,
628    #[serde(skip_serializing_if = "Option::is_none")]
629    questions: Option<&'a Questions>,
630    #[serde(flatten)]
631    extra: &'a Map<String, Value>,
632}
633
634impl IntoFuture for SystemOneRequest {
635    type Output = Result<SystemOneResponse>;
636    type IntoFuture = BoxFuture<'static, Self::Output>;
637
638    fn into_future(self) -> Self::IntoFuture {
639        Box::pin(self.send())
640    }
641}
642
643/// The Models resource.
644#[derive(Debug, Clone)]
645pub struct Models {
646    client: Client,
647}
648
649impl Models {
650    /// `GET /v1/models`.
651    pub fn list(&self) -> ListModelsRequest {
652        ListModelsRequest {
653            client: self.client.clone(),
654            opts: CallOptions::default(),
655        }
656    }
657}
658
659/// A pending `GET /v1/models`.
660#[must_use = "requests do nothing until awaited"]
661#[derive(Debug)]
662pub struct ListModelsRequest {
663    client: Client,
664    opts: CallOptions,
665}
666
667impl ListModelsRequest {
668    call_option_methods!();
669
670    /// Send the request.
671    pub async fn send(self) -> Result<ListModelsResponse> {
672        if let Some(Cassette::Replay(dir)) = &self.client.inner.cassette {
673            return Err(Error::Config(format!(
674                "Listing models is not recorded, so a client replaying from {} cannot answer it.",
675                dir.display()
676            )));
677        }
678        let (bytes, meta, endpoint) = self
679            .client
680            .execute(Method::GET, MODELS_PATH, None, self.opts)
681            .await?;
682        match decode_models(&bytes) {
683            Ok((models, raw)) => Ok(ListModelsResponse { models, raw, meta }),
684            Err(f) => Err(validation_error(
685                StatusCode::from_u16(meta.status).unwrap_or(StatusCode::OK),
686                lenient_body(&bytes),
687                meta.headers,
688                &endpoint,
689                f,
690            )),
691        }
692    }
693}
694
695impl IntoFuture for ListModelsRequest {
696    type Output = Result<ListModelsResponse>;
697    type IntoFuture = BoxFuture<'static, Self::Output>;
698
699    fn into_future(self) -> Self::IntoFuture {
700        Box::pin(self.send())
701    }
702}
703
704#[cfg(test)]
705mod tests {
706    use super::is_secret;
707
708    #[test]
709    fn masks_a_gateways_key_as_well_as_the_apis() {
710        for name in [
711            "Authorization",
712            "cookie",
713            "cf-aig-authorization",
714            "x-portkey-api-key",
715            "x-gateway-token",
716            "x-client-secret",
717        ] {
718            assert!(is_secret(name), "{name}");
719        }
720        for name in ["content-type", "x-typesafe-request-id", "retry-after"] {
721            assert!(!is_secret(name), "{name}");
722        }
723    }
724}
typesafe/client.rs

typesafe/
client.rs
