1use 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#[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 pub fn api_key(mut self, key: impl Into<String>) -> Self {
46 self.api_key = Some(key.into());
47 self
48 }
49
50 pub fn base_url(mut self, url: impl Into<String>) -> Self {
52 self.base_url = Some(url.into());
53 self
54 }
55
56 pub fn model(mut self, model: impl Into<String>) -> Self {
58 self.model = Some(model.into());
59 self
60 }
61
62 pub fn timeout(mut self, timeout: Duration) -> Self {
64 self.timeout = Some(timeout);
65 self
66 }
67
68 pub fn retry(mut self, policy: RetryPolicy) -> Self {
70 self.retry = Some(policy);
71 self
72 }
73
74 pub fn header(mut self, name: HeaderName, value: HeaderValue) -> Self {
77 self.headers.insert(name, value);
78 self
79 }
80
81 pub fn record(mut self, dir: impl Into<PathBuf>) -> Self {
84 self.record = Some(dir.into());
85 self
86 }
87
88 pub fn replay(mut self, dir: impl Into<PathBuf>) -> Self {
92 self.replay = Some(dir.into());
93 self
94 }
95
96 #[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 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 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#[derive(Debug)]
229enum Cassette {
230 Record(PathBuf),
231 Replay(PathBuf),
232}
233
234#[derive(Debug, Clone)]
256pub struct Client {
257 inner: Arc<Inner>,
258}
259
260impl Client {
261 pub fn builder() -> ClientBuilder {
263 ClientBuilder::default()
264 }
265
266 pub fn from_env() -> Result<Self> {
268 Self::builder().build()
269 }
270
271 pub fn default_model(&self) -> &str {
273 &self.inner.model
274 }
275
276 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 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 pub fn retry(mut self, policy: RetryPolicy) -> Self {
479 self.opts.retry = Some(policy);
480 self
481 }
482
483 pub fn timeout(mut self, timeout: Duration) -> Self {
485 self.opts.timeout = Some(timeout);
486 self
487 }
488
489 pub fn header(mut self, name: HeaderName, value: HeaderValue) -> Self {
491 self.opts.headers.insert(name, value);
492 self
493 }
494 };
495}
496
497#[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 pub fn model(mut self, model: impl Into<String>) -> Self {
514 self.model = Some(model.into());
515 self
516 }
517
518 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 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 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
586fn 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#[derive(Debug, Clone)]
625pub struct Models {
626 client: Client,
627}
628
629impl Models {
630 pub fn list(&self) -> ListModelsRequest {
632 ListModelsRequest {
633 client: self.client.clone(),
634 opts: CallOptions::default(),
635 }
636 }
637}
638
639#[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 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}