mod media;
pub use media::*;
mod tasks;
pub use tasks::*;
mod questions;
pub use questions::*;
use std::{fmt, net::IpAddr, time::Duration};
use bytes::Bytes;
use http::{Method, Request, Uri};
use http_body_util::{BodyExt, Full};
use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
use hyper_util::{
client::legacy::connect::HttpConnector,
rt::{TokioExecutor, TokioTimer},
};
use http::header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_TYPE, RETRY_AFTER, USER_AGENT};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
#[non_exhaustive]
#[serde(rename_all = "lowercase")]
pub enum QuestionType {
Choice,
Noul,
Score,
Multi,
Act,
Tool,
Classify,
Verify,
Compress,
Extract,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Question {
#[serde(rename = "type")]
pub kind: QuestionType,
#[serde(skip_serializing_if = "Option::is_none")]
pub instructions: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub criteria: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
impl Question {
pub fn new(kind: QuestionType, instructions: impl Into<String>) -> Self {
Self { kind, instructions: Some(Value::String(instructions.into())), criteria: None, extra: Map::new() }
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct DecisionRequest {
pub state: Value,
pub questions: Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub lang: Option<Language>,
#[serde(skip_serializing_if = "Option::is_none")]
pub media: Option<Vec<MediaInput>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub detail: Option<MediaDetail>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
impl DecisionRequest {
pub fn new(state: impl Into<Value>) -> Self {
Self {
state: state.into(),
questions: Map::new(),
model: None,
lang: None,
media: None,
detail: None,
extra: Map::new(),
}
}
pub fn insert_question(&mut self, id: impl Into<String>, question: &Question) -> Result<(), Error> {
self.questions.insert(id.into(), serde_json::to_value(question).map_err(Error::Encode)?);
Ok(())
}
}
#[derive(Clone, Debug, Deserialize)]
pub struct Answer {
#[serde(rename = "type")]
pub kind: QuestionType,
#[serde(default)]
pub probabilities: indexmap::IndexMap<String, f64>,
pub confidence: Option<f64>,
pub choice: Option<String>,
pub noul: Option<f64>,
pub score: Option<f64>,
pub legend: Option<indexmap::IndexMap<String, String>>,
pub selected: Option<Vec<String>>,
pub threshold: Option<f64>,
pub set_confidence: Option<f64>,
pub action: Option<String>,
pub answer: Option<String>,
pub expected_costs: Option<indexmap::IndexMap<String, f64>>,
pub calibration: Option<String>,
pub refused: Option<String>,
pub max_error: Option<String>,
pub evidence: Option<Vec<EvidenceSpan>>,
pub objective: Option<MultiObjective>,
pub expected_score: Option<f64>,
pub constraints_applied: Option<Value>,
pub tool: Option<String>,
pub tools: Option<Vec<String>>,
pub call: Option<bool>,
pub multi: Option<bool>,
pub ranking: Option<Vec<ToolRanking>>,
pub deferred: Option<bool>,
pub experimental: Option<bool>,
pub label: Option<String>,
pub top_label: Option<String>,
pub labels: Option<Vec<String>>,
pub scores: Option<indexmap::IndexMap<String, f64>>,
pub automated: Option<Automation>,
pub claims: Option<indexmap::IndexMap<String, ClaimResult>>,
pub all_supported: Option<bool>,
pub not_supported: Option<Vec<String>>,
pub compressed: Option<String>,
pub kept: Option<Vec<usize>>,
pub units: Option<Vec<CompressionUnit>>,
pub unit: Option<CompressionUnitKind>,
pub original_chars: Option<usize>,
pub compressed_chars: Option<usize>,
pub ratio: Option<f64>,
pub data: Option<indexmap::IndexMap<String, Value>>,
pub fields: Option<indexmap::IndexMap<String, ExtractField>>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ToolRanking {
pub tool: String,
pub probability: f64,
}
#[derive(Clone, Debug, Deserialize)]
pub struct EvidenceSpan {
pub text: String,
pub start: usize,
pub end: usize,
pub probability: f64,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Usage {
pub input_tokens: u64,
pub output_tokens: u64,
pub questions: Option<u64>,
pub branches: Option<u64>,
pub latency_ms: Option<f64>,
pub longest_prompt_tokens: Option<u64>,
pub second_order_branches: Option<u64>,
pub media: Option<MediaUsage>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct DecisionResponse {
pub model: String,
pub answers: indexmap::IndexMap<String, Answer>,
pub usage: Usage,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct ModelMetadata {
pub name: String,
pub prompt_format: Option<u8>,
pub orders: Option<Value>,
pub media: Option<Vec<MediaKind>>,
pub temperature_per_prim: Option<indexmap::IndexMap<String, f64>>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct ModelsResponse {
pub models: Vec<ModelMetadata>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[non_exhaustive]
pub enum Error {
InvalidConfig(&'static str),
InvalidRequest(&'static str),
InvalidResponse(&'static str),
Transport(Box<dyn std::error::Error + Send + Sync>),
Encode(serde_json::Error),
Protocol(serde_json::Error),
Timeout,
ResponseTooLarge,
Http {
status: u16,
body: Value,
retry_after: Option<String>,
},
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidConfig(message) => write!(f, "invalid Basal client configuration: {message}"),
Self::InvalidRequest(message) => write!(f, "invalid Basal request: {message}"),
Self::InvalidResponse(message) => write!(f, "invalid Basal answer: {message}"),
Self::Transport(_) => f.write_str("Basal HTTP transport failed"),
Self::Encode(_) => f.write_str("could not encode Basal request"),
Self::Protocol(_) => f.write_str("invalid Basal JSON response"),
Self::Timeout => f.write_str("Basal HTTP request timed out"),
Self::ResponseTooLarge => f.write_str("Basal response exceeded the configured byte limit"),
Self::Http { status, .. } => write!(f, "Basal HTTP {status}"),
}
}
}
impl fmt::Debug for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(self, f)
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Transport(source) => Some(source.as_ref()),
Self::Encode(source) | Self::Protocol(source) => Some(source),
_ => None,
}
}
}
pub struct ClientOptions<'a> {
pub base_url: &'a str,
pub api_key: Option<&'a str>,
pub timeout: Duration,
pub max_response_bytes: usize,
}
impl Default for ClientOptions<'_> {
fn default() -> Self {
Self {
base_url: "http://127.0.0.1:8000",
api_key: None,
timeout: Duration::from_secs(120),
max_response_bytes: 8 * 1024 * 1024,
}
}
}
#[derive(Clone)]
pub struct Client {
http: hyper_util::client::legacy::Client<HttpsConnector<HttpConnector>, Full<Bytes>>,
base_url: String,
headers: HeaderMap,
timeout: Duration,
max_response_bytes: usize,
}
impl Client {
pub fn new(base_url: &str) -> Result<Self, Error> {
Self::with_options(ClientOptions { base_url, ..ClientOptions::default() })
}
pub fn with_options(options: ClientOptions<'_>) -> Result<Self, Error> {
let url: Uri = options.base_url.parse().map_err(|_| Error::InvalidConfig("invalid URL"))?;
let scheme = url.scheme_str().unwrap_or_default();
let host = url.host().unwrap_or_default();
if !matches!(scheme, "http" | "https")
|| host.is_empty()
|| options.base_url.contains(['?', '#'])
|| url.authority().is_some_and(|a| a.as_str().contains('@'))
{
return Err(Error::InvalidConfig("expected an HTTP(S) URL without credentials, query or fragment"));
}
let loopback = host.eq_ignore_ascii_case("localhost")
|| host.trim_matches(['[', ']']).parse::<IpAddr>().is_ok_and(|ip| ip.is_loopback());
if scheme != "https" && !loopback {
return Err(Error::InvalidConfig("base_url requires HTTPS or loopback HTTP"));
}
if options.timeout.is_zero() || options.max_response_bytes == 0 {
return Err(Error::InvalidConfig("timeout and max_response_bytes must be positive"));
}
let mut headers = HeaderMap::new();
headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
headers.insert(USER_AGENT, HeaderValue::from_static(concat!("basal-client-rust/", env!("CARGO_PKG_VERSION"))));
if let Some(key) = options.api_key {
if key.is_empty() || key.contains(['\r', '\n']) {
return Err(Error::InvalidConfig("api_key must be a nonempty, single-line token"));
}
let mut value = HeaderValue::from_str(&format!("Bearer {key}"))
.map_err(|_| Error::InvalidConfig("invalid bearer token"))?;
value.set_sensitive(true);
headers.insert(AUTHORIZATION, value);
}
let connector = HttpsConnectorBuilder::new()
.with_provider_and_platform_verifier(rustls::crypto::ring::default_provider())
.map_err(|source| Error::Transport(Box::new(source)))?
.https_or_http()
.enable_http1()
.build();
let http = hyper_util::client::legacy::Client::builder(TokioExecutor::new())
.retry_canceled_requests(false)
.pool_max_idle_per_host(8)
.pool_timer(TokioTimer::new())
.pool_idle_timeout(Duration::from_secs(90))
.build(connector);
Ok(Self {
http,
base_url: format!("{}/", options.base_url.trim_end_matches('/')),
headers,
timeout: options.timeout,
max_response_bytes: options.max_response_bytes,
})
}
async fn send<T: DeserializeOwned>(&self, method: Method, path: &str, body: Option<Vec<u8>>) -> Result<T, Error> {
let work = async {
let mut request = Request::builder().method(method).uri(format!("{}{path}", self.base_url));
if let Some(headers) = request.headers_mut() {
*headers = self.headers.clone();
}
if body.is_some() {
request = request.header(CONTENT_TYPE, "application/json");
}
let request = request
.body(Full::new(Bytes::from(body.unwrap_or_default())))
.map_err(|source| Error::Transport(Box::new(source)))?;
let response = self.http.request(request).await.map_err(|source| Error::Transport(Box::new(source)))?;
let status = response.status();
let retry_after = response.headers().get(RETRY_AFTER).and_then(|v| v.to_str().ok()).map(String::from);
let mut body = response.into_body();
let mut bytes = Vec::new();
while let Some(frame) = body.frame().await {
let frame = frame.map_err(|source| Error::Transport(Box::new(source)))?;
if let Ok(data) = frame.into_data() {
if data.len() > self.max_response_bytes.saturating_sub(bytes.len()) {
return Err(Error::ResponseTooLarge);
}
bytes.extend_from_slice(&data);
}
}
if !status.is_success() {
let body = serde_json::from_slice(&bytes)
.unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&bytes).into_owned()));
return Err(Error::Http { status: status.as_u16(), body, retry_after });
}
serde_json::from_slice(&bytes).map_err(Error::Protocol)
};
tokio::time::timeout(self.timeout, work).await.map_err(|_| Error::Timeout)?
}
pub async fn decide<Q: QuestionSet>(
&self,
state: impl Into<Value>,
questions: &Q,
) -> Result<TypedDecisionResponse<Q::Schema>, Error> {
self.decide_with_options(state, questions, DecisionOptions::default()).await
}
pub async fn decide_with_options<Q: QuestionSet>(
&self,
state: impl Into<Value>,
questions: &Q,
options: DecisionOptions<'_>,
) -> Result<TypedDecisionResponse<Q::Schema>, Error> {
let request = questions.as_questions().request(state.into(), options)?;
let response = self.decide_request(&request).await?;
questions.as_questions().clone().parse(response, questions.schema())
}
pub async fn decide_request(&self, request: &DecisionRequest) -> Result<DecisionResponse, Error> {
self.send(Method::POST, "v1/basal", Some(serde_json::to_vec(request).map_err(Error::Encode)?)).await
}
pub async fn models(&self) -> Result<ModelsResponse, Error> {
self.send(Method::GET, "v1/models", None).await
}
}