audacity_sdk/
client.rs

1//! The main `Client` — fluent builder-style API mirroring aws-sdk-bedrockruntime.
2
3use std::time::{Duration, Instant};
4
5use rand::Rng;
6use serde_json::Value;
7
8use crate::config::Config;
9use crate::error::{parse_error, parse_retry_after};
10use crate::mapping::{build_request_body, parse_response};
11use crate::stream::{spawn_stream_driver, ConverseStreamOutputHandle};
12use crate::types::{
13    ConverseOutput, InferenceConfiguration, Message, SystemContentBlock, ToolConfiguration,
14};
15use crate::Error;
16
17const USER_AGENT: &str = concat!("audacity-sdk-rust/", env!("CARGO_PKG_VERSION"));
18
19/// The Audacity SDK client.
20#[derive(Clone)]
21pub struct Client {
22    config: Config,
23    http: reqwest::Client,
24}
25
26impl Client {
27    /// Create a new client from an explicit [`Config`].
28    pub fn new(config: &Config) -> Self {
29        let http = reqwest::Client::builder()
30            .timeout(config.timeout)
31            .build()
32            .expect("failed to build reqwest client");
33        Self {
34            config: config.clone(),
35            http,
36        }
37    }
38
39    /// Create a client from the environment (`AUDACITY_API_KEY`, `AUDACITY_BASE_URL`).
40    /// Returns `Err(Error::MissingApiKey)` if no key is found.
41    pub fn from_env() -> Result<Self, Error> {
42        let config = Config::builder().build()?;
43        Ok(Self::new(&config))
44    }
45
46    /// Start building a `converse` request.
47    pub fn converse(&self) -> ConverseBuilder {
48        ConverseBuilder::new(self.clone())
49    }
50
51    /// Start building a `converse_stream` request.
52    pub fn converse_stream(&self) -> ConverseStreamBuilder {
53        ConverseStreamBuilder::new(self.clone())
54    }
55
56    // ── internal HTTP helpers ─────────────────────────────────────────────────
57
58    pub(crate) async fn send_converse(
59        &self,
60        input: &ConverseInput,
61    ) -> Result<ConverseOutput, Error> {
62        let body = build_request_body(input, false);
63        let (resp, start) = self.request_with_retries(&body, "application/json").await?;
64
65        let latency = start.elapsed().as_millis() as u64;
66        let raw: Value = resp
67            .json()
68            .await
69            .map_err(|e| Error::Sdk(format!("JSON decode error: {e}")))?;
70        parse_response(&raw, latency)
71    }
72
73    pub(crate) async fn send_converse_stream(
74        &self,
75        input: &ConverseInput,
76    ) -> Result<ConverseStreamOutputHandle, Error> {
77        let body = build_request_body(input, true);
78        let (resp, start) = self
79            .request_with_retries(&body, "text/event-stream")
80            .await?;
81
82        // Once we have the 200, never retry — spec §4
83        Ok(spawn_stream_driver(resp, start))
84    }
85
86    /// POST the request body with the spec §4 retry policy (attempts, jittered
87    /// backoff, `Retry-After`), returning the 200 response + request start time.
88    /// Retryability is decided solely by [`Error::is_retryable`].
89    async fn request_with_retries(
90        &self,
91        body: &Value,
92        accept: &str,
93    ) -> Result<(reqwest::Response, Instant), Error> {
94        let url = format!("{}/v1/chat/completions", self.config.base_url);
95        let max_attempts = self.config.max_retries + 1;
96        let mut attempt = 0u32;
97
98        loop {
99            let start = Instant::now();
100            let result = self
101                .http
102                .post(&url)
103                .header("Authorization", format!("Bearer {}", self.config.api_key))
104                .header("Content-Type", "application/json")
105                .header("Accept", accept)
106                .header("User-Agent", USER_AGENT)
107                .json(body)
108                .send()
109                .await;
110
111            let err = match result {
112                Ok(resp) if resp.status().as_u16() == 200 => return Ok((resp, start)),
113                Ok(resp) => {
114                    let status = resp.status().as_u16();
115                    let retry_after = resp
116                        .headers()
117                        .get("retry-after")
118                        .and_then(|v| v.to_str().ok())
119                        .and_then(parse_retry_after);
120                    let body_text = resp.text().await.unwrap_or_default();
121                    parse_error(status, &body_text, retry_after)
122                }
123                Err(e) => Error::Sdk(e.to_string()),
124            };
125
126            attempt += 1;
127            if attempt >= max_attempts || !err.is_retryable() {
128                return Err(err);
129            }
130            backoff(attempt, err.retry_after_seconds()).await;
131        }
132    }
133}
134
135/// Jittered exponential backoff per spec: `random(0, min(20s, 0.5s * 2^attempt))`.
136/// If `retry_after` is set, sleep at least that long.
137async fn backoff(attempt: u32, retry_after: Option<f64>) {
138    let cap = 20.0_f64;
139    let base = 0.5_f64 * 2f64.powi(attempt as i32);
140    let ceiling = base.min(cap);
141    let jittered: f64 = rand::thread_rng().gen_range(0.0..=ceiling);
142    let sleep_secs = if let Some(ra) = retry_after {
143        jittered.max(ra)
144    } else {
145        jittered
146    };
147    tokio::time::sleep(Duration::from_secs_f64(sleep_secs)).await;
148}
149
150// ── Input type ────────────────────────────────────────────────────────────────
151
152/// The canonical input for both `converse` and `converse_stream`.
153#[derive(Default)]
154pub struct ConverseInput {
155    pub model_id: String,
156    pub messages: Vec<Message>,
157    pub system: Vec<SystemContentBlock>,
158    pub inference_config: Option<InferenceConfiguration>,
159    pub tool_config: Option<ToolConfiguration>,
160    pub additional_model_request_fields: Option<Value>,
161}
162
163impl ConverseInput {
164    /// Validate the required fields before sending (shared by both builders).
165    fn validate(&self) -> Result<(), Error> {
166        if self.model_id.is_empty() {
167            return Err(Error::Sdk("model_id is required".into()));
168        }
169        if self.messages.is_empty() {
170            return Err(Error::Sdk("messages must not be empty".into()));
171        }
172        Ok(())
173    }
174}
175
176// ── Fluent builders ───────────────────────────────────────────────────────────
177
178/// Generates the shared constructor + input setters for a fluent builder
179/// wrapping [`ConverseInput`]. Both builder types stay public (aws-sdk parity);
180/// only the method bodies are defined once here.
181macro_rules! impl_converse_input_builder {
182    ($builder:ident) => {
183        impl $builder {
184            fn new(client: Client) -> Self {
185                Self {
186                    client,
187                    input: ConverseInput::default(),
188                }
189            }
190
191            pub fn model_id(mut self, id: impl Into<String>) -> Self {
192                self.input.model_id = id.into();
193                self
194            }
195
196            pub fn messages(mut self, msg: Message) -> Self {
197                self.input.messages.push(msg);
198                self
199            }
200
201            pub fn set_messages(mut self, msgs: Vec<Message>) -> Self {
202                self.input.messages = msgs;
203                self
204            }
205
206            pub fn system(mut self, block: SystemContentBlock) -> Self {
207                self.input.system.push(block);
208                self
209            }
210
211            pub fn set_system(mut self, blocks: Vec<SystemContentBlock>) -> Self {
212                self.input.system = blocks;
213                self
214            }
215
216            pub fn inference_config(mut self, cfg: InferenceConfiguration) -> Self {
217                self.input.inference_config = Some(cfg);
218                self
219            }
220
221            pub fn tool_config(mut self, tc: ToolConfiguration) -> Self {
222                self.input.tool_config = Some(tc);
223                self
224            }
225
226            pub fn additional_model_request_fields(mut self, v: Value) -> Self {
227                self.input.additional_model_request_fields = Some(v);
228                self
229            }
230        }
231    };
232}
233
234/// Fluent builder for `converse`.
235pub struct ConverseBuilder {
236    client: Client,
237    input: ConverseInput,
238}
239
240impl_converse_input_builder!(ConverseBuilder);
241
242impl ConverseBuilder {
243    pub async fn send(self) -> Result<ConverseOutput, Error> {
244        self.input.validate()?;
245        self.client.send_converse(&self.input).await
246    }
247}
248
249/// Fluent builder for `converse_stream`.
250pub struct ConverseStreamBuilder {
251    client: Client,
252    input: ConverseInput,
253}
254
255impl_converse_input_builder!(ConverseStreamBuilder);
256
257impl ConverseStreamBuilder {
258    pub async fn send(self) -> Result<ConverseStreamOutputHandle, Error> {
259        self.input.validate()?;
260        self.client.send_converse_stream(&self.input).await
261    }
262}
audacity_sdk/client.rs

audacity_sdk/
client.rs
