audacity_sdk/

client.rs

//! The main `Client` — fluent builder-style API mirroring aws-sdk-bedrockruntime.

use std::time::{Duration, Instant};

use rand::Rng;
use serde_json::Value;

use crate::config::Config;
use crate::error::{parse_error, parse_retry_after, status_error};
use crate::mapping::{build_request_body, parse_response};
use crate::stream::{spawn_stream_driver, ConverseStreamOutputHandle};
use crate::types::{
    ConverseOutput, GenerateImageOutput, InferenceConfiguration, MediaResolution, Message,
    SystemContentBlock, ToolConfiguration, UploadFileOutput,
};
use crate::Error;

const USER_AGENT: &str = concat!("audacity-sdk-rust/", env!("CARGO_PKG_VERSION"));

/// Default resumable-upload chunk size (spec §6): 8 MiB, a multiple of the
/// protocol's required 256 KiB granularity.
const DEFAULT_UPLOAD_CHUNK_SIZE: usize = 8 * 1024 * 1024;

/// Max §6 recovery attempts (offset query + resume) without confirmed progress.
const MAX_UPLOAD_RECOVERY_ATTEMPTS: u32 = 5;

/// The Audacity SDK client.
#[derive(Clone)]
pub struct Client {
    config: Config,
    http: reqwest::Client,
}

impl Client {
    /// Create a new client from an explicit [`Config`].
    ///
    /// No client-wide timeout is set: per spec §1 the configured timeout bounds
    /// each `converse` attempt (including the body read) via a per-request
    /// timeout, and a `converse_stream` attempt only up to the response
    /// headers — the SSE body read is unbounded so long generations are never
    /// killed mid-stream.
    pub fn new(config: &Config) -> Result<Self, Error> {
        let http = reqwest::Client::builder()
            .build()
            .map_err(|e| Error::sdk(format!("failed to build HTTP client: {e}")))?;
        Ok(Self {
            config: config.clone(),
            http,
        })
    }

    /// Create a client from the environment (`AUDACITY_API_KEY`, `AUDACITY_BASE_URL`).
    /// Returns `Err(Error::MissingApiKey)` if no key is found.
    pub fn from_env() -> Result<Self, Error> {
        let config = Config::builder().build()?;
        Self::new(&config)
    }

    /// Start building a `converse` request.
    pub fn converse(&self) -> ConverseBuilder {
        ConverseBuilder::new(self.clone())
    }

    /// Start building a `converse_stream` request.
    pub fn converse_stream(&self) -> ConverseStreamBuilder {
        ConverseStreamBuilder::new(self.clone())
    }

    /// Start building a file upload (spec §6) — for videos too large to send
    /// inline. Returns an [`UploadFileOutput`] whose `uri` can be referenced
    /// in [`VideoSource::Uri`](crate::types::VideoSource::Uri).
    pub fn upload_file(&self) -> UploadFileBuilder {
        UploadFileBuilder::new(self.clone())
    }

    /// Start building an image generation (spec §8) — a text prompt in,
    /// generated image(s) out (signed URL or inline base64).
    pub fn generate_image(&self) -> GenerateImageBuilder {
        GenerateImageBuilder::new(self.clone())
    }

    // ── internal HTTP helpers ─────────────────────────────────────────────────

    pub(crate) async fn send_converse(
        &self,
        input: &ConverseInput,
    ) -> Result<ConverseOutput, Error> {
        let body = build_request_body(input, false);
        // Spec §3: latencyMs is measured from request start (before the first
        // attempt) to the response being fully parsed.
        let start = Instant::now();
        let text = self.post_converse(&body).await?;
        let raw: Value =
            serde_json::from_str(&text).map_err(|e| Error::sdk(format!("JSON decode error: {e}")))?;
        let latency = start.elapsed().as_millis() as u64;
        parse_response(&raw, latency)
    }

    pub(crate) async fn send_converse_stream(
        &self,
        input: &ConverseInput,
    ) -> Result<ConverseStreamOutputHandle, Error> {
        let body = build_request_body(input, true);
        let start = Instant::now();
        let resp = self.post_converse_stream(&body).await?;

        // Once we have the 200, never retry — spec §4
        Ok(spawn_stream_driver(resp, start))
    }

    pub(crate) async fn send_upload_file(
        &self,
        data: Vec<u8>,
        content_type: String,
        chunk_size: usize,
    ) -> Result<UploadFileOutput, Error> {
        // Step 1: request an upload ticket — same auth/error/retry machinery
        // as converse. The response is flat JSON (not enveloped).
        let url = format!("{}/v1/files", self.config.base_url);
        let body = serde_json::json!({
            "content_type": content_type,
            "size_bytes": data.len(),
        });
        let text = match self
            .request_with_retries(&url, &body, RequestMode::Converse)
            .await?
        {
            Attempted::Body(text) => text,
            Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
        };
        let ticket: UploadFileOutput = serde_json::from_str(&text)
            .map_err(|e| Error::sdk(format!("JSON decode error on /v1/files response: {e}")))?;

        // Step 2: GCS resumable-session protocol (spec §6) — initiate a
        // session at the presigned URL, then PUT the bytes in Content-Range
        // chunks with bounded offset-query recovery. Errors map by HTTP
        // status only (storage-provider bodies are XML, not gateway payloads).
        let session_uri = self
            .initiate_resumable_session(&ticket.upload_url, &content_type)
            .await?;
        self.upload_chunks(&session_uri, &data, chunk_size).await?;

        Ok(ticket)
    }

    /// Spec §6 step 2a: `POST <upload_url>` with `x-goog-resumable: start`
    /// and the Content-Type bound into the signature, an empty body, and no
    /// Authorization header (the URL is self-authorizing). A 2xx response
    /// carries the session URI in its `Location` header. Transient failures
    /// (network, 5xx, 429) retry with the §4 backoff under the same bound as
    /// chunk recovery; other non-2xx fail immediately, mapped by HTTP status.
    async fn initiate_resumable_session(
        &self,
        upload_url: &str,
        content_type: &str,
    ) -> Result<String, Error> {
        let mut attempt = 0u32;
        loop {
            let outcome = self
                .http
                .post(upload_url)
                .header("x-goog-resumable", "start")
                .header("Content-Type", content_type)
                .header("User-Agent", USER_AGENT)
                .timeout(self.config.timeout)
                .send()
                .await;

            let (err, retry_after) = match outcome {
                Ok(resp) if resp.status().is_success() => {
                    return resp
                        .headers()
                        .get("location")
                        .and_then(|v| v.to_str().ok())
                        .map(str::to_owned)
                        .ok_or_else(|| {
                            Error::sdk(
                                "resumable upload initiation response missing Location header",
                            )
                        });
                }
                Ok(resp) => {
                    let status = resp.status().as_u16();
                    let retry_after = header_retry_after(&resp);
                    let transient = status == 429 || status >= 500;
                    let body_text = resp.text().await.unwrap_or_default();
                    let err = status_error(status, &body_text);
                    if !transient {
                        return Err(err);
                    }
                    (err, retry_after)
                }
                Err(e) => (
                    Error::sdk_network(format!("network error initiating resumable upload: {e}")),
                    None,
                ),
            };

            attempt += 1;
            if attempt > MAX_UPLOAD_RECOVERY_ATTEMPTS {
                return Err(err);
            }
            backoff(attempt, retry_after).await;
        }
    }

    /// Spec §6 steps 2b–2d: PUT the file to the session URI in order with
    /// `Content-Range: bytes <first>-<last>/<total>` chunks. A 308 confirms
    /// persisted bytes via its `Range: bytes=0-N` header (the next chunk
    /// starts at `N+1`); 200/201 completes the upload. On a transient failure
    /// (network, 5xx, 429) the confirmed offset is re-queried with
    /// `Content-Range: bytes */<total>` and the upload resumes from there —
    /// at most 5 recovery attempts with the §4 backoff, the counter resetting
    /// on confirmed progress. Other 4xx fail immediately by HTTP status.
    async fn upload_chunks(
        &self,
        session_uri: &str,
        data: &[u8],
        chunk_size: usize,
    ) -> Result<(), Error> {
        let total = data.len();
        let mut offset = 0usize;
        let mut failures = 0u32;
        // When true the next request is an offset query rather than a chunk.
        // An empty file finalizes with the query form (`bytes */0`).
        let mut querying = total == 0;

        loop {
            let req = if querying {
                self.http
                    .put(session_uri)
                    .header("Content-Range", format!("bytes */{total}"))
            } else {
                let end = (offset + chunk_size).min(total);
                self.http
                    .put(session_uri)
                    .header(
                        "Content-Range",
                        format!("bytes {offset}-{}/{total}", end - 1),
                    )
                    .body(data[offset..end].to_vec())
            };
            let outcome = req
                .header("User-Agent", USER_AGENT)
                .timeout(self.config.timeout)
                .send()
                .await;

            let (err, retry_after) = match outcome {
                Ok(resp) if resp.status().is_success() => return Ok(()),
                Ok(resp) if resp.status().as_u16() == 308 => {
                    // Resume Incomplete: the Range header confirms how many
                    // bytes are persisted; no Range means none are. A 308
                    // showing less progress than sent just means the upload
                    // continues from the confirmed offset.
                    let confirmed = confirmed_offset(resp.headers());
                    if confirmed > offset {
                        failures = 0; // progress resets the recovery budget
                    }
                    offset = confirmed;
                    querying = false;
                    continue;
                }
                Ok(resp) => {
                    let status = resp.status().as_u16();
                    let retry_after = header_retry_after(&resp);
                    let transient = status == 429 || status >= 500;
                    let body_text = resp.text().await.unwrap_or_default();
                    let err = status_error(status, &body_text);
                    if !transient {
                        return Err(err);
                    }
                    (err, retry_after)
                }
                Err(e) => (
                    Error::sdk_network(format!("network error uploading file chunk: {e}")),
                    None,
                ),
            };

            failures += 1;
            if failures > MAX_UPLOAD_RECOVERY_ATTEMPTS {
                return Err(err);
            }
            backoff(failures, retry_after).await;
            querying = true;
        }
    }

    pub(crate) async fn send_generate_image(
        &self,
        body: &Value,
    ) -> Result<GenerateImageOutput, Error> {
        let url = format!("{}/v1/images/generations", self.config.base_url);
        let text = match self
            .request_with_retries(&url, body, RequestMode::Converse)
            .await?
        {
            Attempted::Body(text) => text,
            Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
        };
        // NOTE: the §1 defensive data-envelope unwrap is deliberately NOT
        // applied here — `data` is a legitimate top-level key of the images
        // response (the list of generated images).
        serde_json::from_str(&text).map_err(|e| {
            Error::sdk(format!(
                "JSON decode error on /v1/images/generations response: {e}"
            ))
        })
    }

    /// Retried POST for `Converse` — resolves to the fully-read 200 body.
    async fn post_converse(&self, body: &Value) -> Result<String, Error> {
        let url = format!("{}/v1/chat/completions", self.config.base_url);
        match self.request_with_retries(&url, body, RequestMode::Converse).await? {
            Attempted::Body(text) => Ok(text),
            Attempted::Stream(_) => unreachable!("Converse mode always yields Body"),
        }
    }

    /// Retried POST for `ConverseStream` — resolves to the unread 200 response.
    async fn post_converse_stream(&self, body: &Value) -> Result<reqwest::Response, Error> {
        let url = format!("{}/v1/chat/completions", self.config.base_url);
        match self.request_with_retries(&url, body, RequestMode::Stream).await? {
            Attempted::Stream(resp) => Ok(resp),
            Attempted::Body(_) => unreachable!("Stream mode always yields Stream"),
        }
    }

    /// POST the request body with the spec §4 retry policy (attempts, jittered
    /// backoff, `Retry-After`). Retryability is decided solely by
    /// [`Error::is_retryable`].
    ///
    /// Timeout semantics (spec §1): a `Converse` attempt is bounded end to end
    /// — including reading the response body, which happens inside this loop so
    /// a network failure mid-body is still retryable. A `Stream` attempt is
    /// bounded only until status + headers arrive; the SSE body is unbounded.
    async fn request_with_retries(
        &self,
        url: &str,
        body: &Value,
        mode: RequestMode,
    ) -> Result<Attempted, Error> {
        let max_attempts = self.config.max_retries + 1;
        let mut attempt = 0u32;

        loop {
            let accept = match mode {
                RequestMode::Converse => "application/json",
                RequestMode::Stream => "text/event-stream",
            };
            let req = self
                .http
                .post(url)
                .header("Authorization", format!("Bearer {}", self.config.api_key))
                .header("Content-Type", "application/json")
                .header("Accept", accept)
                .header("User-Agent", USER_AGENT)
                .json(body);

            let outcome: Result<reqwest::Response, Error> = match mode {
                RequestMode::Converse => req
                    .timeout(self.config.timeout)
                    .send()
                    .await
                    .map_err(|e| Error::sdk_network(format!("network error: {e}"))),
                RequestMode::Stream => {
                    match tokio::time::timeout(self.config.timeout, req.send()).await {
                        Ok(Ok(resp)) => Ok(resp),
                        Ok(Err(e)) => Err(Error::sdk_network(format!("network error: {e}"))),
                        Err(_) => Err(Error::sdk_network(
                            "timed out waiting for response headers",
                        )),
                    }
                }
            };

            let err = match outcome {
                Ok(resp) if resp.status().as_u16() == 200 => match mode {
                    RequestMode::Stream => return Ok(Attempted::Stream(resp)),
                    RequestMode::Converse => match resp.text().await {
                        Ok(text) => return Ok(Attempted::Body(text)),
                        Err(e) => {
                            Error::sdk_network(format!("network error reading response body: {e}"))
                        }
                    },
                },
                Ok(resp) => {
                    let status = resp.status().as_u16();
                    let retry_after = resp
                        .headers()
                        .get("retry-after")
                        .and_then(|v| v.to_str().ok())
                        .and_then(parse_retry_after);
                    let body_text = resp.text().await.unwrap_or_default();
                    parse_error(status, &body_text, retry_after)
                }
                Err(e) => e,
            };

            attempt += 1;
            if attempt >= max_attempts || !err.is_retryable() {
                return Err(err);
            }
            backoff(attempt, err.retry_after_seconds()).await;
        }
    }
}

/// Which operation a retried request is for (drives timeout + body handling).
#[derive(Clone, Copy)]
enum RequestMode {
    Converse,
    Stream,
}

/// Successful outcome of [`Client::request_with_retries`].
enum Attempted {
    /// Fully read 200 response body (`Converse`).
    Body(String),
    /// Unread 200 response, headers received (`Stream`).
    Stream(reqwest::Response),
}

/// Parse the confirmed byte count from a 308 `Range: bytes=0-N` header —
/// `N+1` bytes are persisted. A missing/unparsable Range means none are.
fn confirmed_offset(headers: &reqwest::header::HeaderMap) -> usize {
    headers
        .get("range")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.trim().strip_prefix("bytes=0-"))
        .and_then(|n| n.parse::<usize>().ok())
        .map(|n| n + 1)
        .unwrap_or(0)
}

/// Extract a `Retry-After` header value in seconds, if present.
fn header_retry_after(resp: &reqwest::Response) -> Option<f64> {
    resp.headers()
        .get("retry-after")
        .and_then(|v| v.to_str().ok())
        .and_then(parse_retry_after)
}

/// Jittered exponential backoff per spec: `random(0, min(20s, 0.5s * 2^attempt))`.
/// If `retry_after` is set, sleep at least that long.
async fn backoff(attempt: u32, retry_after: Option<f64>) {
    let cap = 20.0_f64;
    let base = 0.5_f64 * 2f64.powi(attempt as i32);
    let ceiling = base.min(cap);
    let jittered: f64 = rand::thread_rng().gen_range(0.0..=ceiling);
    let sleep_secs = if let Some(ra) = retry_after {
        jittered.max(ra)
    } else {
        jittered
    };
    tokio::time::sleep(Duration::from_secs_f64(sleep_secs)).await;
}

// ── Input type ────────────────────────────────────────────────────────────────

/// The canonical input for both `converse` and `converse_stream`.
#[derive(Default)]
pub struct ConverseInput {
    pub model_id: String,
    pub messages: Vec<Message>,
    pub system: Vec<SystemContentBlock>,
    pub inference_config: Option<InferenceConfiguration>,
    pub media_resolution: Option<MediaResolution>,
    pub tool_config: Option<ToolConfiguration>,
    pub additional_model_request_fields: Option<Value>,
}

impl ConverseInput {
    /// Validate the required fields before sending (shared by both builders).
    fn validate(&self) -> Result<(), Error> {
        if self.model_id.is_empty() {
            return Err(Error::client_validation("model_id is required"));
        }
        if self.messages.is_empty() {
            return Err(Error::client_validation("messages must not be empty"));
        }
        if let Some(v) = &self.additional_model_request_fields {
            if !v.is_object() {
                return Err(Error::client_validation(
                    "additional_model_request_fields must be a JSON object",
                ));
            }
        }
        Ok(())
    }
}

// ── Fluent builders ───────────────────────────────────────────────────────────

/// Generates the shared constructor + input setters for a fluent builder
/// wrapping [`ConverseInput`]. Both builder types stay public (aws-sdk parity);
/// only the method bodies are defined once here.
macro_rules! impl_converse_input_builder {
    ($builder:ident) => {
        impl $builder {
            fn new(client: Client) -> Self {
                Self {
                    client,
                    input: ConverseInput::default(),
                }
            }

            pub fn model_id(mut self, id: impl Into<String>) -> Self {
                self.input.model_id = id.into();
                self
            }

            pub fn messages(mut self, msg: Message) -> Self {
                self.input.messages.push(msg);
                self
            }

            pub fn set_messages(mut self, msgs: Vec<Message>) -> Self {
                self.input.messages = msgs;
                self
            }

            pub fn system(mut self, block: SystemContentBlock) -> Self {
                self.input.system.push(block);
                self
            }

            pub fn set_system(mut self, blocks: Vec<SystemContentBlock>) -> Self {
                self.input.system = blocks;
                self
            }

            pub fn inference_config(mut self, cfg: InferenceConfiguration) -> Self {
                self.input.inference_config = Some(cfg);
                self
            }

            pub fn media_resolution(mut self, v: MediaResolution) -> Self {
                self.input.media_resolution = Some(v);
                self
            }

            pub fn tool_config(mut self, tc: ToolConfiguration) -> Self {
                self.input.tool_config = Some(tc);
                self
            }

            pub fn additional_model_request_fields(mut self, v: Value) -> Self {
                self.input.additional_model_request_fields = Some(v);
                self
            }
        }
    };
}

/// Fluent builder for `converse`.
pub struct ConverseBuilder {
    client: Client,
    input: ConverseInput,
}

impl_converse_input_builder!(ConverseBuilder);

impl ConverseBuilder {
    pub async fn send(self) -> Result<ConverseOutput, Error> {
        self.input.validate()?;
        self.client.send_converse(&self.input).await
    }
}

/// Fluent builder for `converse_stream`.
pub struct ConverseStreamBuilder {
    client: Client,
    input: ConverseInput,
}

impl_converse_input_builder!(ConverseStreamBuilder);

impl ConverseStreamBuilder {
    pub async fn send(self) -> Result<ConverseStreamOutputHandle, Error> {
        self.input.validate()?;
        self.client.send_converse_stream(&self.input).await
    }
}

/// Fluent builder for `upload_file` (spec §6).
pub struct UploadFileBuilder {
    client: Client,
    data: Option<Vec<u8>>,
    content_type: Option<String>,
    chunk_size: usize,
}

impl UploadFileBuilder {
    fn new(client: Client) -> Self {
        Self {
            client,
            data: None,
            content_type: None,
            chunk_size: DEFAULT_UPLOAD_CHUNK_SIZE,
        }
    }

    /// The raw file bytes to upload.
    pub fn data(mut self, data: impl Into<Vec<u8>>) -> Self {
        self.data = Some(data.into());
        self
    }

    /// The MIME type of the file — one of the spec §3 video MIME types,
    /// e.g. `"video/mp4"`.
    pub fn content_type(mut self, content_type: impl Into<String>) -> Self {
        self.content_type = Some(content_type.into());
        self
    }

    /// Override the resumable-upload chunk size in bytes (test hook). Must be
    /// a multiple of 256 KiB (262144) per the GCS resumable protocol.
    #[doc(hidden)]
    pub fn chunk_size(mut self, n: usize) -> Self {
        self.chunk_size = n;
        self
    }

    pub async fn send(self) -> Result<UploadFileOutput, Error> {
        let data = self
            .data
            .ok_or_else(|| Error::client_validation("upload_file data is required"))?;
        let content_type = self
            .content_type
            .ok_or_else(|| Error::client_validation("upload_file content_type is required"))?;
        self.client
            .send_upload_file(data, content_type, self.chunk_size)
            .await
    }
}

/// Fluent builder for `generate_image` (spec §8).
pub struct GenerateImageBuilder {
    client: Client,
    model: Option<String>,
    prompt: Option<String>,
    n: Option<u32>,
    size: Option<String>,
    quality: Option<String>,
    response_format: Option<String>,
    user: Option<String>,
}

impl GenerateImageBuilder {
    fn new(client: Client) -> Self {
        Self {
            client,
            model: None,
            prompt: None,
            n: None,
            size: None,
            quality: None,
            response_format: None,
            user: None,
        }
    }

    /// The image model id (e.g. `"gpt-image-1"`, `"dall-e-3"`, `"imagen-4"`).
    /// Required.
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// Text description of the desired image(s); 1–32000 characters. Required.
    pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
        self.prompt = Some(prompt.into());
        self
    }

    /// Number of images to generate (1–10, model-dependent).
    pub fn n(mut self, n: u32) -> Self {
        self.n = Some(n);
        self
    }

    /// Dimensions as a `"WxH"` string (e.g. `"1024x1024"`); supported values
    /// differ per model.
    pub fn size(mut self, size: impl Into<String>) -> Self {
        self.size = Some(size.into());
        self
    }

    /// Provider-specific quality tier (e.g. `"standard"`, `"hd"`).
    pub fn quality(mut self, quality: impl Into<String>) -> Self {
        self.quality = Some(quality.into());
        self
    }

    /// Output delivery: `"url"` (default, signed download link expiring after
    /// ~24 h) or `"b64_json"` (inline base64 bytes).
    pub fn response_format(mut self, response_format: impl Into<String>) -> Self {
        self.response_format = Some(response_format.into());
        self
    }

    /// End-user identifier forwarded to the provider.
    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    pub async fn send(self) -> Result<GenerateImageOutput, Error> {
        let model = self
            .model
            .filter(|m| !m.is_empty())
            .ok_or_else(|| Error::client_validation("generate_image model is required"))?;
        let prompt = self
            .prompt
            .filter(|p| !p.is_empty())
            .ok_or_else(|| Error::client_validation("generate_image prompt is required"))?;

        // Omit absent optionals entirely (never send nulls).
        let mut body = serde_json::Map::new();
        body.insert("model".into(), Value::String(model));
        body.insert("prompt".into(), Value::String(prompt));
        if let Some(n) = self.n {
            body.insert("n".into(), Value::from(n));
        }
        if let Some(size) = self.size {
            body.insert("size".into(), Value::String(size));
        }
        if let Some(quality) = self.quality {
            body.insert("quality".into(), Value::String(quality));
        }
        if let Some(response_format) = self.response_format {
            body.insert("response_format".into(), Value::String(response_format));
        }
        if let Some(user) = self.user {
            body.insert("user".into(), Value::String(user));
        }

        self.client.send_generate_image(&Value::Object(body)).await
    }
}