audacity-sdk 0.1.0

//! 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`].
    pub fn new(config: &Config) -> Self {
        let http = reqwest::Client::builder()
            .timeout(config.timeout)
            .build()
            .expect("failed to build reqwest client");
        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()?;
        Ok(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);
        let (resp, start) = self.request_with_retries(&body, "application/json").await?;

        let latency = start.elapsed().as_millis() as u64;
        let raw: Value = resp
            .json()
            .await
            .map_err(|e| Error::Sdk(format!("JSON decode error: {e}")))?;
        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 (resp, start) = self
            .request_with_retries(&body, "text/event-stream")
            .await?;

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

    /// POST the request body with the spec §4 retry policy (attempts, jittered
    /// backoff, `Retry-After`), returning the 200 response + request start time.
    /// Retryability is decided solely by [`Error::is_retryable`].
    async fn request_with_retries(
        &self,
        body: &Value,
        accept: &str,
    ) -> Result<(reqwest::Response, Instant), 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 start = Instant::now();
            let result = 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)
                .send()
                .await;

            let err = match result {
                Ok(resp) if resp.status().as_u16() == 200 => return Ok((resp, start)),
                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) => Error::Sdk(e.to_string()),
            };

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

/// 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::Sdk("model_id is required".into()));
        }
        if self.messages.is_empty() {
            return Err(Error::Sdk("messages must not be empty".into()));
        }
        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
    }
}