Skip to main content

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