1mod questions;
8pub use questions::*;
9
10use std::{fmt, net::IpAddr, time::Duration};
11
12use bytes::Bytes;
13use http::{Method, Request, Uri};
14use http_body_util::{BodyExt, Full};
15use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
16use hyper_util::{
17 client::legacy::connect::HttpConnector,
18 rt::{TokioExecutor, TokioTimer},
19};
20
21use http::header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_TYPE, RETRY_AFTER, USER_AGENT};
22use serde::{de::DeserializeOwned, Deserialize, Serialize};
23use serde_json::{Map, Value};
24
25#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
27#[non_exhaustive]
28#[serde(rename_all = "lowercase")]
29pub enum QuestionType {
30 Choice,
32 Noul,
34 Score,
36 Multi,
38 Act,
40}
41
42#[derive(Clone, Debug, Serialize, Deserialize)]
45pub struct Question {
46 #[serde(rename = "type")]
48 pub kind: QuestionType,
49 #[serde(skip_serializing_if = "Option::is_none")]
51 pub instructions: Option<Value>,
52 #[serde(skip_serializing_if = "Option::is_none")]
54 pub criteria: Option<Value>,
55 #[serde(flatten)]
58 pub extra: Map<String, Value>,
59}
60
61impl Question {
62 pub fn new(kind: QuestionType, instructions: impl Into<String>) -> Self {
64 Self { kind, instructions: Some(Value::String(instructions.into())), criteria: None, extra: Map::new() }
65 }
66}
67
68#[derive(Clone, Debug, Serialize, Deserialize)]
70pub struct DecisionRequest {
71 pub state: Value,
73 pub questions: Map<String, Value>,
75 #[serde(skip_serializing_if = "Option::is_none")]
77 pub model: Option<String>,
78 #[serde(flatten)]
80 pub extra: Map<String, Value>,
81}
82
83impl DecisionRequest {
84 pub fn new(state: impl Into<Value>) -> Self {
86 Self { state: state.into(), questions: Map::new(), model: None, extra: Map::new() }
87 }
88
89 pub fn insert_question(&mut self, id: impl Into<String>, question: &Question) -> Result<(), Error> {
91 self.questions.insert(id.into(), serde_json::to_value(question).map_err(Error::Encode)?);
92 Ok(())
93 }
94}
95
96#[derive(Clone, Debug, Deserialize)]
98pub struct Answer {
99 #[serde(rename = "type")]
101 pub kind: QuestionType,
102 pub probabilities: indexmap::IndexMap<String, f64>,
104 pub confidence: Option<f64>,
106 pub choice: Option<String>,
108 pub noul: Option<f64>,
110 pub score: Option<f64>,
112 pub legend: Option<indexmap::IndexMap<String, String>>,
114 pub selected: Option<Vec<String>>,
116 pub threshold: Option<f64>,
118 pub set_confidence: Option<f64>,
120 pub action: Option<String>,
122 pub answer: Option<String>,
124 pub expected_costs: Option<indexmap::IndexMap<String, f64>>,
126 pub calibration: Option<String>,
128 pub refused: Option<String>,
130 pub max_error: Option<String>,
132 pub evidence: Option<Vec<EvidenceSpan>>,
134 #[serde(flatten)]
136 pub extra: Map<String, Value>,
137}
138
139#[derive(Clone, Debug, Deserialize)]
141pub struct EvidenceSpan {
142 pub text: String,
144 pub start: usize,
146 pub end: usize,
148 pub probability: f64,
150}
151
152#[derive(Clone, Debug, Deserialize)]
154pub struct Usage {
155 pub input_tokens: u64,
157 pub output_tokens: u64,
159 pub questions: Option<u64>,
161 pub branches: Option<u64>,
163 pub latency_ms: Option<f64>,
165 #[serde(flatten)]
167 pub extra: Map<String, Value>,
168}
169
170#[derive(Clone, Debug, Deserialize)]
172pub struct DecisionResponse {
173 pub model: String,
175 pub answers: indexmap::IndexMap<String, Answer>,
177 pub usage: Usage,
179 #[serde(flatten)]
181 pub extra: Map<String, Value>,
182}
183
184#[derive(Clone, Debug, Deserialize)]
186pub struct ModelMetadata {
187 pub name: String,
189 #[serde(flatten)]
191 pub extra: Map<String, Value>,
192}
193
194#[derive(Clone, Debug, Deserialize)]
196pub struct ModelsResponse {
197 pub models: Vec<ModelMetadata>,
199 #[serde(flatten)]
201 pub extra: Map<String, Value>,
202}
203
204#[non_exhaustive]
206pub enum Error {
207 InvalidConfig(&'static str),
209 InvalidRequest(&'static str),
211 InvalidResponse(&'static str),
213 Transport(Box<dyn std::error::Error + Send + Sync>),
215 Encode(serde_json::Error),
217 Protocol(serde_json::Error),
219 Timeout,
221 ResponseTooLarge,
223 Http {
225 status: u16,
227 body: Value,
229 retry_after: Option<String>,
231 },
232}
233
234impl fmt::Display for Error {
235 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236 match self {
237 Self::InvalidConfig(message) => write!(f, "invalid Basal client configuration: {message}"),
238 Self::InvalidRequest(message) => write!(f, "invalid Basal request: {message}"),
239 Self::InvalidResponse(message) => write!(f, "invalid Basal answer: {message}"),
240 Self::Transport(_) => f.write_str("Basal HTTP transport failed"),
241 Self::Encode(_) => f.write_str("could not encode Basal request"),
242 Self::Protocol(_) => f.write_str("invalid Basal JSON response"),
243 Self::Timeout => f.write_str("Basal HTTP request timed out"),
244 Self::ResponseTooLarge => f.write_str("Basal response exceeded the configured byte limit"),
245 Self::Http { status, .. } => write!(f, "Basal HTTP {status}"),
246 }
247 }
248}
249
250impl fmt::Debug for Error {
251 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
252 fmt::Display::fmt(self, f)
254 }
255}
256
257impl std::error::Error for Error {
258 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
259 match self {
260 Self::Transport(source) => Some(source.as_ref()),
261 Self::Encode(source) | Self::Protocol(source) => Some(source),
262 _ => None,
263 }
264 }
265}
266
267pub struct ClientOptions<'a> {
269 pub base_url: &'a str,
271 pub api_key: Option<&'a str>,
273 pub timeout: Duration,
275 pub max_response_bytes: usize,
277}
278
279impl Default for ClientOptions<'_> {
280 fn default() -> Self {
281 Self {
282 base_url: "http://127.0.0.1:8000",
283 api_key: None,
284 timeout: Duration::from_secs(120),
285 max_response_bytes: 8 * 1024 * 1024,
286 }
287 }
288}
289
290#[derive(Clone)]
293pub struct Client {
294 http: hyper_util::client::legacy::Client<HttpsConnector<HttpConnector>, Full<Bytes>>,
295 base_url: String,
296 headers: HeaderMap,
297 timeout: Duration,
298 max_response_bytes: usize,
299}
300
301impl Client {
302 pub fn new(base_url: &str) -> Result<Self, Error> {
304 Self::with_options(ClientOptions { base_url, ..ClientOptions::default() })
305 }
306
307 pub fn with_options(options: ClientOptions<'_>) -> Result<Self, Error> {
310 let url: Uri = options.base_url.parse().map_err(|_| Error::InvalidConfig("invalid URL"))?;
311 let scheme = url.scheme_str().unwrap_or_default();
312 let host = url.host().unwrap_or_default();
313 if !matches!(scheme, "http" | "https")
314 || host.is_empty()
315 || options.base_url.contains(['?', '#'])
316 || url.authority().is_some_and(|a| a.as_str().contains('@'))
317 {
318 return Err(Error::InvalidConfig("expected an HTTP(S) URL without credentials, query or fragment"));
319 }
320 let loopback = host.eq_ignore_ascii_case("localhost")
321 || host.trim_matches(['[', ']']).parse::<IpAddr>().is_ok_and(|ip| ip.is_loopback());
322 if scheme != "https" && !loopback {
323 return Err(Error::InvalidConfig("base_url requires HTTPS or loopback HTTP"));
324 }
325 if options.timeout.is_zero() || options.max_response_bytes == 0 {
326 return Err(Error::InvalidConfig("timeout and max_response_bytes must be positive"));
327 }
328 let mut headers = HeaderMap::new();
329 headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
330 headers.insert(USER_AGENT, HeaderValue::from_static(concat!("basal-client-rust/", env!("CARGO_PKG_VERSION"))));
331 if let Some(key) = options.api_key {
332 if key.is_empty() || key.contains(['\r', '\n']) {
333 return Err(Error::InvalidConfig("api_key must be a nonempty, single-line token"));
334 }
335 let mut value = HeaderValue::from_str(&format!("Bearer {key}"))
336 .map_err(|_| Error::InvalidConfig("invalid bearer token"))?;
337 value.set_sensitive(true);
338 headers.insert(AUTHORIZATION, value);
339 }
340 let connector = HttpsConnectorBuilder::new()
341 .with_provider_and_platform_verifier(rustls::crypto::ring::default_provider())
342 .map_err(|source| Error::Transport(Box::new(source)))?
343 .https_or_http()
344 .enable_http1()
345 .build();
346 let http = hyper_util::client::legacy::Client::builder(TokioExecutor::new())
347 .retry_canceled_requests(false)
348 .pool_max_idle_per_host(8)
349 .pool_timer(TokioTimer::new())
350 .pool_idle_timeout(Duration::from_secs(90))
351 .build(connector);
352 Ok(Self {
353 http,
354 base_url: format!("{}/", options.base_url.trim_end_matches('/')),
355 headers,
356 timeout: options.timeout,
357 max_response_bytes: options.max_response_bytes,
358 })
359 }
360
361 async fn send<T: DeserializeOwned>(&self, method: Method, path: &str, body: Option<Vec<u8>>) -> Result<T, Error> {
362 let work = async {
363 let mut request = Request::builder().method(method).uri(format!("{}{path}", self.base_url));
365 if let Some(headers) = request.headers_mut() {
366 *headers = self.headers.clone();
367 }
368 if body.is_some() {
369 request = request.header(CONTENT_TYPE, "application/json");
370 }
371 let request = request
372 .body(Full::new(Bytes::from(body.unwrap_or_default())))
373 .map_err(|source| Error::Transport(Box::new(source)))?;
374 let response = self.http.request(request).await.map_err(|source| Error::Transport(Box::new(source)))?;
375 let status = response.status();
376 let retry_after = response.headers().get(RETRY_AFTER).and_then(|v| v.to_str().ok()).map(String::from);
377 let mut body = response.into_body();
378 let mut bytes = Vec::new();
379 while let Some(frame) = body.frame().await {
380 let frame = frame.map_err(|source| Error::Transport(Box::new(source)))?;
381 if let Ok(data) = frame.into_data() {
382 if data.len() > self.max_response_bytes.saturating_sub(bytes.len()) {
383 return Err(Error::ResponseTooLarge);
384 }
385 bytes.extend_from_slice(&data);
386 }
387 }
388 if !status.is_success() {
389 let body = serde_json::from_slice(&bytes)
390 .unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&bytes).into_owned()));
391 return Err(Error::Http { status: status.as_u16(), body, retry_after });
392 }
393 serde_json::from_slice(&bytes).map_err(Error::Protocol)
394 };
395 tokio::time::timeout(self.timeout, work).await.map_err(|_| Error::Timeout)?
396 }
397
398 pub async fn decide(&self, state: impl Into<Value>, questions: &Questions) -> Result<TypedDecisionResponse, Error> {
400 self.decide_with_options(state, questions, DecisionOptions::default()).await
401 }
402
403 pub async fn decide_with_options(
405 &self,
406 state: impl Into<Value>,
407 questions: &Questions,
408 options: DecisionOptions<'_>,
409 ) -> Result<TypedDecisionResponse, Error> {
410 let request = questions.request(state.into(), options)?;
411 let response = self.decide_request(&request).await?;
412 questions.clone().parse(response)
413 }
414
415 pub async fn decide_request(&self, request: &DecisionRequest) -> Result<DecisionResponse, Error> {
417 self.send(Method::POST, "v1/basal", Some(serde_json::to_vec(request).map_err(Error::Encode)?)).await
418 }
419
420 pub async fn models(&self) -> Result<ModelsResponse, Error> {
422 self.send(Method::GET, "v1/models", None).await
423 }
424}