1use 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, status_error};
10use crate::mapping::{build_request_body, parse_response};
11use crate::passthrough::{ChatCompletions, Messages};
12use crate::stream::{spawn_stream_driver, ConverseStreamOutputHandle};
13use crate::types::{
14 ConverseOutput, GenerateImageOutput, InferenceConfiguration, MediaResolution, Message,
15 SystemContentBlock, ToolConfiguration, UploadFileOutput,
16};
17use crate::Error;
18
19const USER_AGENT: &str = concat!("audacity-sdk-rust/", env!("CARGO_PKG_VERSION"));
20
21const DEFAULT_UPLOAD_CHUNK_SIZE: usize = 8 * 1024 * 1024;
24
25const MAX_UPLOAD_RECOVERY_ATTEMPTS: u32 = 5;
27
28#[derive(Clone)]
30pub struct Client {
31 config: Config,
32 http: reqwest::Client,
33}
34
35impl Client {
36 pub fn new(config: &Config) -> Result<Self, Error> {
44 let http = reqwest::Client::builder()
45 .build()
46 .map_err(|e| Error::sdk(format!("failed to build HTTP client: {e}")))?;
47 Ok(Self {
48 config: config.clone(),
49 http,
50 })
51 }
52
53 pub fn from_env() -> Result<Self, Error> {
56 let config = Config::builder().build()?;
57 Self::new(&config)
58 }
59
60 pub fn converse(&self) -> ConverseBuilder {
62 ConverseBuilder::new(self.clone())
63 }
64
65 pub fn converse_stream(&self) -> ConverseStreamBuilder {
67 ConverseStreamBuilder::new(self.clone())
68 }
69
70 pub fn upload_file(&self) -> UploadFileBuilder {
74 UploadFileBuilder::new(self.clone())
75 }
76
77 pub fn generate_image(&self) -> GenerateImageBuilder {
80 GenerateImageBuilder::new(self.clone())
81 }
82
83 pub fn chat_completions(&self) -> ChatCompletions {
86 ChatCompletions::new(self.clone())
87 }
88
89 pub fn messages(&self) -> Messages {
93 Messages::new(self.clone())
94 }
95
96 pub(crate) async fn send_converse(
99 &self,
100 input: &ConverseInput,
101 ) -> Result<ConverseOutput, Error> {
102 let body = build_request_body(input, false);
103 let start = Instant::now();
106 let text = self.post_converse(&body).await?;
107 let raw: Value = serde_json::from_str(&text)
108 .map_err(|e| Error::sdk(format!("JSON decode error: {e}")))?;
109 let latency = start.elapsed().as_millis() as u64;
110 parse_response(&raw, latency)
111 }
112
113 pub(crate) async fn send_converse_stream(
114 &self,
115 input: &ConverseInput,
116 ) -> Result<ConverseStreamOutputHandle, Error> {
117 let body = build_request_body(input, true);
118 let start = Instant::now();
119 let resp = self.post_converse_stream(&body).await?;
120
121 Ok(spawn_stream_driver(resp, start))
123 }
124
125 pub(crate) async fn send_upload_file(
126 &self,
127 data: Vec<u8>,
128 content_type: String,
129 chunk_size: usize,
130 ) -> Result<UploadFileOutput, Error> {
131 let url = format!("{}/v1/files", self.config.base_url);
134 let body = serde_json::json!({
135 "content_type": content_type,
136 "size_bytes": data.len(),
137 });
138 let text = match self
139 .request_with_retries(&url, &body, RequestMode::Converse)
140 .await?
141 {
142 Attempted::Body(text) => text,
143 Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
144 };
145 let ticket: UploadFileOutput = serde_json::from_str(&text)
146 .map_err(|e| Error::sdk(format!("JSON decode error on /v1/files response: {e}")))?;
147
148 let session_uri = self
153 .initiate_resumable_session(&ticket.upload_url, &content_type)
154 .await?;
155 self.upload_chunks(&session_uri, &data, chunk_size).await?;
156
157 Ok(ticket)
158 }
159
160 async fn initiate_resumable_session(
167 &self,
168 upload_url: &str,
169 content_type: &str,
170 ) -> Result<String, Error> {
171 let mut attempt = 0u32;
172 loop {
173 let outcome = self
174 .http
175 .post(upload_url)
176 .header("x-goog-resumable", "start")
177 .header("Content-Type", content_type)
178 .header("User-Agent", USER_AGENT)
179 .timeout(self.config.timeout)
180 .send()
181 .await;
182
183 let (err, retry_after) = match outcome {
184 Ok(resp) if resp.status().is_success() => {
185 return resp
186 .headers()
187 .get("location")
188 .and_then(|v| v.to_str().ok())
189 .map(str::to_owned)
190 .ok_or_else(|| {
191 Error::sdk(
192 "resumable upload initiation response missing Location header",
193 )
194 });
195 }
196 Ok(resp) => {
197 let status = resp.status().as_u16();
198 let retry_after = header_retry_after(&resp);
199 let transient = status == 429 || status >= 500;
200 let body_text = resp.text().await.unwrap_or_default();
201 let err = status_error(status, &body_text);
202 if !transient {
203 return Err(err);
204 }
205 (err, retry_after)
206 }
207 Err(e) => (
208 Error::sdk_network(format!("network error initiating resumable upload: {e}")),
209 None,
210 ),
211 };
212
213 attempt += 1;
214 if attempt > MAX_UPLOAD_RECOVERY_ATTEMPTS {
215 return Err(err);
216 }
217 backoff(attempt, retry_after).await;
218 }
219 }
220
221 async fn upload_chunks(
230 &self,
231 session_uri: &str,
232 data: &[u8],
233 chunk_size: usize,
234 ) -> Result<(), Error> {
235 let total = data.len();
236 let mut offset = 0usize;
237 let mut failures = 0u32;
238 let mut querying = total == 0;
241
242 loop {
243 let req = if querying {
244 self.http
245 .put(session_uri)
246 .header("Content-Range", format!("bytes */{total}"))
247 } else {
248 let end = (offset + chunk_size).min(total);
249 self.http
250 .put(session_uri)
251 .header(
252 "Content-Range",
253 format!("bytes {offset}-{}/{total}", end - 1),
254 )
255 .body(data[offset..end].to_vec())
256 };
257 let outcome = req
258 .header("User-Agent", USER_AGENT)
259 .timeout(self.config.timeout)
260 .send()
261 .await;
262
263 let (err, retry_after) = match outcome {
264 Ok(resp) if resp.status().is_success() => return Ok(()),
265 Ok(resp) if resp.status().as_u16() == 308 => {
266 let confirmed = confirmed_offset(resp.headers());
271 if confirmed > offset {
272 failures = 0; }
274 offset = confirmed;
275 querying = false;
276 continue;
277 }
278 Ok(resp) => {
279 let status = resp.status().as_u16();
280 let retry_after = header_retry_after(&resp);
281 let transient = status == 429 || status >= 500;
282 let body_text = resp.text().await.unwrap_or_default();
283 let err = status_error(status, &body_text);
284 if !transient {
285 return Err(err);
286 }
287 (err, retry_after)
288 }
289 Err(e) => (
290 Error::sdk_network(format!("network error uploading file chunk: {e}")),
291 None,
292 ),
293 };
294
295 failures += 1;
296 if failures > MAX_UPLOAD_RECOVERY_ATTEMPTS {
297 return Err(err);
298 }
299 backoff(failures, retry_after).await;
300 querying = true;
301 }
302 }
303
304 pub(crate) async fn send_generate_image(
305 &self,
306 body: &Value,
307 ) -> Result<GenerateImageOutput, Error> {
308 let url = format!("{}/v1/images/generations", self.config.base_url);
309 let text = match self
310 .request_with_retries(&url, body, RequestMode::Converse)
311 .await?
312 {
313 Attempted::Body(text) => text,
314 Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
315 };
316 serde_json::from_str(&text).map_err(|e| {
320 Error::sdk(format!(
321 "JSON decode error on /v1/images/generations response: {e}"
322 ))
323 })
324 }
325
326 async fn post_converse(&self, body: &Value) -> Result<String, Error> {
328 let url = format!("{}/v1/chat/completions", self.config.base_url);
329 match self
330 .request_with_retries(&url, body, RequestMode::Converse)
331 .await?
332 {
333 Attempted::Body(text) => Ok(text),
334 Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
335 }
336 }
337
338 async fn post_converse_stream(&self, body: &Value) -> Result<reqwest::Response, Error> {
340 let url = format!("{}/v1/chat/completions", self.config.base_url);
341 match self
342 .request_with_retries(&url, body, RequestMode::Stream)
343 .await?
344 {
345 Attempted::Stream(resp) => Ok(resp),
346 Attempted::Body(_) => unreachable!("Stream mode always yields Stream"),
347 }
348 }
349
350 pub(crate) async fn post_passthrough_json(
354 &self,
355 path: &str,
356 body: &Value,
357 extra_headers: &[(&str, &str)],
358 ) -> Result<Value, Error> {
359 let url = format!("{}{path}", self.config.base_url);
360 let text = match self
361 .request_with_retries_with_headers(&url, body, RequestMode::Converse, extra_headers)
362 .await?
363 {
364 Attempted::Body(text) => text,
365 Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
366 };
367 let parsed: Value = serde_json::from_str(&text)
368 .map_err(|e| Error::sdk(format!("failed to decode {path} response as JSON: {e}")))?;
369 if !parsed.is_object() {
370 return Err(Error::sdk(format!(
371 "malformed {path} response: body is not a JSON object"
372 )));
373 }
374 Ok(parsed)
375 }
376
377 pub(crate) async fn post_passthrough_stream(
380 &self,
381 path: &str,
382 body: &Value,
383 extra_headers: &[(&str, &str)],
384 ) -> Result<reqwest::Response, Error> {
385 let url = format!("{}{path}", self.config.base_url);
386 match self
387 .request_with_retries_with_headers(&url, body, RequestMode::Stream, extra_headers)
388 .await?
389 {
390 Attempted::Stream(resp) => Ok(resp),
391 Attempted::Body(_) => unreachable!("Stream mode always yields Stream"),
392 }
393 }
394
395 async fn request_with_retries(
404 &self,
405 url: &str,
406 body: &Value,
407 mode: RequestMode,
408 ) -> Result<Attempted, Error> {
409 self.request_with_retries_with_headers(url, body, mode, &[])
410 .await
411 }
412
413 async fn request_with_retries_with_headers(
416 &self,
417 url: &str,
418 body: &Value,
419 mode: RequestMode,
420 extra_headers: &[(&str, &str)],
421 ) -> Result<Attempted, Error> {
422 let max_attempts = self.config.max_retries + 1;
423 let mut attempt = 0u32;
424
425 loop {
426 let accept = match mode {
427 RequestMode::Converse => "application/json",
428 RequestMode::Stream => "text/event-stream",
429 };
430 let mut req = self
431 .http
432 .post(url)
433 .header("Authorization", format!("Bearer {}", self.config.api_key))
434 .header("Content-Type", "application/json")
435 .header("Accept", accept)
436 .header("User-Agent", USER_AGENT)
437 .json(body);
438 for (name, value) in extra_headers {
439 req = req.header(*name, *value);
440 }
441
442 let outcome: Result<reqwest::Response, Error> = match mode {
443 RequestMode::Converse => req
444 .timeout(self.config.timeout)
445 .send()
446 .await
447 .map_err(|e| Error::sdk_network(format!("network error: {e}"))),
448 RequestMode::Stream => {
449 match tokio::time::timeout(self.config.timeout, req.send()).await {
450 Ok(Ok(resp)) => Ok(resp),
451 Ok(Err(e)) => Err(Error::sdk_network(format!("network error: {e}"))),
452 Err(_) => Err(Error::sdk_network("timed out waiting for response headers")),
453 }
454 }
455 };
456
457 let err = match outcome {
458 Ok(resp) if resp.status().as_u16() == 200 => match mode {
459 RequestMode::Stream => return Ok(Attempted::Stream(resp)),
460 RequestMode::Converse => match resp.text().await {
461 Ok(text) => return Ok(Attempted::Body(text)),
462 Err(e) => {
463 Error::sdk_network(format!("network error reading response body: {e}"))
464 }
465 },
466 },
467 Ok(resp) => {
468 let status = resp.status().as_u16();
469 let retry_after = resp
470 .headers()
471 .get("retry-after")
472 .and_then(|v| v.to_str().ok())
473 .and_then(parse_retry_after);
474 let body_text = resp.text().await.unwrap_or_default();
475 parse_error(status, &body_text, retry_after)
476 }
477 Err(e) => e,
478 };
479
480 attempt += 1;
481 if attempt >= max_attempts || !err.is_retryable() {
482 return Err(err);
483 }
484 backoff(attempt, err.retry_after_seconds()).await;
485 }
486 }
487}
488
489#[derive(Clone, Copy)]
491enum RequestMode {
492 Converse,
493 Stream,
494}
495
496enum Attempted {
498 Body(String),
500 Stream(reqwest::Response),
502}
503
504fn confirmed_offset(headers: &reqwest::header::HeaderMap) -> usize {
507 headers
508 .get("range")
509 .and_then(|v| v.to_str().ok())
510 .and_then(|s| s.trim().strip_prefix("bytes=0-"))
511 .and_then(|n| n.parse::<usize>().ok())
512 .map(|n| n + 1)
513 .unwrap_or(0)
514}
515
516fn header_retry_after(resp: &reqwest::Response) -> Option<f64> {
518 resp.headers()
519 .get("retry-after")
520 .and_then(|v| v.to_str().ok())
521 .and_then(parse_retry_after)
522}
523
524async fn backoff(attempt: u32, retry_after: Option<f64>) {
527 let cap = 20.0_f64;
528 let base = 0.5_f64 * 2f64.powi(attempt as i32);
529 let ceiling = base.min(cap);
530 let jittered: f64 = rand::thread_rng().gen_range(0.0..=ceiling);
531 let sleep_secs = if let Some(ra) = retry_after {
532 jittered.max(ra)
533 } else {
534 jittered
535 };
536 tokio::time::sleep(Duration::from_secs_f64(sleep_secs)).await;
537}
538
539#[derive(Default)]
543pub struct ConverseInput {
544 pub model_id: String,
545 pub messages: Vec<Message>,
546 pub system: Vec<SystemContentBlock>,
547 pub inference_config: Option<InferenceConfiguration>,
548 pub media_resolution: Option<MediaResolution>,
549 pub tool_config: Option<ToolConfiguration>,
550 pub additional_model_request_fields: Option<Value>,
551}
552
553impl ConverseInput {
554 fn validate(&self) -> Result<(), Error> {
556 if self.model_id.is_empty() {
557 return Err(Error::client_validation("model_id is required"));
558 }
559 if self.messages.is_empty() {
560 return Err(Error::client_validation("messages must not be empty"));
561 }
562 if let Some(v) = &self.additional_model_request_fields {
563 if !v.is_object() {
564 return Err(Error::client_validation(
565 "additional_model_request_fields must be a JSON object",
566 ));
567 }
568 }
569 Ok(())
570 }
571}
572
573macro_rules! impl_converse_input_builder {
579 ($builder:ident) => {
580 impl $builder {
581 fn new(client: Client) -> Self {
582 Self {
583 client,
584 input: ConverseInput::default(),
585 }
586 }
587
588 pub fn model_id(mut self, id: impl Into<String>) -> Self {
589 self.input.model_id = id.into();
590 self
591 }
592
593 pub fn messages(mut self, msg: Message) -> Self {
594 self.input.messages.push(msg);
595 self
596 }
597
598 pub fn set_messages(mut self, msgs: Vec<Message>) -> Self {
599 self.input.messages = msgs;
600 self
601 }
602
603 pub fn system(mut self, block: SystemContentBlock) -> Self {
604 self.input.system.push(block);
605 self
606 }
607
608 pub fn set_system(mut self, blocks: Vec<SystemContentBlock>) -> Self {
609 self.input.system = blocks;
610 self
611 }
612
613 pub fn inference_config(mut self, cfg: InferenceConfiguration) -> Self {
614 self.input.inference_config = Some(cfg);
615 self
616 }
617
618 pub fn media_resolution(mut self, v: MediaResolution) -> Self {
619 self.input.media_resolution = Some(v);
620 self
621 }
622
623 pub fn tool_config(mut self, tc: ToolConfiguration) -> Self {
624 self.input.tool_config = Some(tc);
625 self
626 }
627
628 pub fn additional_model_request_fields(mut self, v: Value) -> Self {
629 self.input.additional_model_request_fields = Some(v);
630 self
631 }
632 }
633 };
634}
635
636pub struct ConverseBuilder {
638 client: Client,
639 input: ConverseInput,
640}
641
642impl_converse_input_builder!(ConverseBuilder);
643
644impl ConverseBuilder {
645 pub async fn send(self) -> Result<ConverseOutput, Error> {
646 self.input.validate()?;
647 self.client.send_converse(&self.input).await
648 }
649}
650
651pub struct ConverseStreamBuilder {
653 client: Client,
654 input: ConverseInput,
655}
656
657impl_converse_input_builder!(ConverseStreamBuilder);
658
659impl ConverseStreamBuilder {
660 pub async fn send(self) -> Result<ConverseStreamOutputHandle, Error> {
661 self.input.validate()?;
662 self.client.send_converse_stream(&self.input).await
663 }
664}
665
666pub struct UploadFileBuilder {
668 client: Client,
669 data: Option<Vec<u8>>,
670 content_type: Option<String>,
671 chunk_size: usize,
672}
673
674impl UploadFileBuilder {
675 fn new(client: Client) -> Self {
676 Self {
677 client,
678 data: None,
679 content_type: None,
680 chunk_size: DEFAULT_UPLOAD_CHUNK_SIZE,
681 }
682 }
683
684 pub fn data(mut self, data: impl Into<Vec<u8>>) -> Self {
686 self.data = Some(data.into());
687 self
688 }
689
690 pub fn content_type(mut self, content_type: impl Into<String>) -> Self {
693 self.content_type = Some(content_type.into());
694 self
695 }
696
697 #[doc(hidden)]
700 pub fn chunk_size(mut self, n: usize) -> Self {
701 self.chunk_size = n;
702 self
703 }
704
705 pub async fn send(self) -> Result<UploadFileOutput, Error> {
706 let data = self
707 .data
708 .ok_or_else(|| Error::client_validation("upload_file data is required"))?;
709 let content_type = self
710 .content_type
711 .ok_or_else(|| Error::client_validation("upload_file content_type is required"))?;
712 self.client
713 .send_upload_file(data, content_type, self.chunk_size)
714 .await
715 }
716}
717
718pub struct GenerateImageBuilder {
720 client: Client,
721 model: Option<String>,
722 prompt: Option<String>,
723 n: Option<u32>,
724 size: Option<String>,
725 quality: Option<String>,
726 response_format: Option<String>,
727 user: Option<String>,
728}
729
730impl GenerateImageBuilder {
731 fn new(client: Client) -> Self {
732 Self {
733 client,
734 model: None,
735 prompt: None,
736 n: None,
737 size: None,
738 quality: None,
739 response_format: None,
740 user: None,
741 }
742 }
743
744 pub fn model(mut self, model: impl Into<String>) -> Self {
747 self.model = Some(model.into());
748 self
749 }
750
751 pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
753 self.prompt = Some(prompt.into());
754 self
755 }
756
757 pub fn n(mut self, n: u32) -> Self {
759 self.n = Some(n);
760 self
761 }
762
763 pub fn size(mut self, size: impl Into<String>) -> Self {
766 self.size = Some(size.into());
767 self
768 }
769
770 pub fn quality(mut self, quality: impl Into<String>) -> Self {
772 self.quality = Some(quality.into());
773 self
774 }
775
776 pub fn response_format(mut self, response_format: impl Into<String>) -> Self {
779 self.response_format = Some(response_format.into());
780 self
781 }
782
783 pub fn user(mut self, user: impl Into<String>) -> Self {
785 self.user = Some(user.into());
786 self
787 }
788
789 pub async fn send(self) -> Result<GenerateImageOutput, Error> {
790 let model = self
791 .model
792 .filter(|m| !m.is_empty())
793 .ok_or_else(|| Error::client_validation("generate_image model is required"))?;
794 let prompt = self
795 .prompt
796 .filter(|p| !p.is_empty())
797 .ok_or_else(|| Error::client_validation("generate_image prompt is required"))?;
798
799 let mut body = serde_json::Map::new();
801 body.insert("model".into(), Value::String(model));
802 body.insert("prompt".into(), Value::String(prompt));
803 if let Some(n) = self.n {
804 body.insert("n".into(), Value::from(n));
805 }
806 if let Some(size) = self.size {
807 body.insert("size".into(), Value::String(size));
808 }
809 if let Some(quality) = self.quality {
810 body.insert("quality".into(), Value::String(quality));
811 }
812 if let Some(response_format) = self.response_format {
813 body.insert("response_format".into(), Value::String(response_format));
814 }
815 if let Some(user) = self.user {
816 body.insert("user".into(), Value::String(user));
817 }
818
819 self.client.send_generate_image(&Value::Object(body)).await
820 }
821}