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};
use crate::mapping::{build_request_body, parse_response};
use crate::stream::{spawn_stream_driver, ConverseStreamOutputHandle};
use crate::types::{
    ConverseOutput, InferenceConfiguration, Message, SystemContentBlock, ToolConfiguration,
};
use crate::Error;

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

/// 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())
    }

    // ── 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))
    }

    /// Retried POST for `Converse` — resolves to the fully-read 200 body.
    async fn post_converse(&self, body: &Value) -> Result<String, Error> {
        match self.request_with_retries(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> {
        match self.request_with_retries(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, body: &Value, mode: RequestMode) -> Result<Attempted, Error> {
        let url = format!("{}/v1/chat/completions", self.config.base_url);
        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),
}

/// 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 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 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
    }
}