openai-rust2 1.7.4

pub extern crate futures_util;
use anyhow::{anyhow, Result};
use lazy_static::lazy_static;
use std::time::Duration;

lazy_static! {
    static ref DEFAULT_BASE_URL: reqwest::Url =
        reqwest::Url::parse("https://api.openai.com/v1/models").unwrap();
}

pub struct Client {
    req_client: reqwest::Client,
    key: String,
    base_url: reqwest::Url,
    timeout: Duration,
    max_retries: u32,
}

pub mod chat;
pub mod completions;
pub mod edits;
pub mod embeddings;
pub mod images;
pub mod models;

#[derive(Debug, Clone)]
pub struct ClientBuilder {
    api_key: String,
    base_url: Option<String>,
    timeout: Duration,
    max_retries: u32,
}

impl Default for ClientBuilder {
    fn default() -> Self {
        Self {
            api_key: String::new(),
            base_url: None,
            timeout: Duration::from_secs(60),
            max_retries: 3,
        }
    }
}

impl ClientBuilder {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            api_key: api_key.into(),
            ..Default::default()
        }
    }

    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
        self.base_url = Some(base_url.into());
        self
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
        self.max_retries = max_retries;
        self
    }

    pub fn build(self) -> Result<Client> {
        let req_client = reqwest::ClientBuilder::new()
            .timeout(self.timeout)
            .build()
            .map_err(|e| anyhow!("failed to build reqwest client: {}", e))?;
        let base_url = match self.base_url {
            Some(url) => reqwest::Url::parse(&url)
                .map_err(|e| anyhow!("invalid base URL '{}': {}", url, e))?,
            None => DEFAULT_BASE_URL.clone(),
        };
        Ok(Client {
            req_client,
            key: self.api_key,
            base_url,
            timeout: self.timeout,
            max_retries: self.max_retries,
        })
    }
}

impl Client {
    pub fn new(api_key: &str) -> Client {
        let timeout = Duration::from_secs(60);
        let req_client = reqwest::ClientBuilder::new()
            .timeout(timeout)
            .build()
            .unwrap();
        Client {
            req_client,
            key: api_key.to_owned(),
            base_url: DEFAULT_BASE_URL.clone(),
            timeout,
            max_retries: 3,
        }
    }

    pub fn new_with_client(api_key: &str, req_client: reqwest::Client) -> Client {
        let timeout = Duration::from_secs(60);
        Client {
            req_client,
            key: api_key.to_owned(),
            base_url: DEFAULT_BASE_URL.clone(),
            timeout,
            max_retries: 3,
        }
    }

    pub fn new_with_base_url(api_key: &str, base_url: &str) -> Client {
        let timeout = Duration::from_secs(60);
        let req_client = reqwest::ClientBuilder::new()
            .timeout(timeout)
            .build()
            .unwrap();
        let base_url = reqwest::Url::parse(base_url).unwrap();
        Client {
            req_client,
            key: api_key.to_owned(),
            base_url,
            timeout,
            max_retries: 3,
        }
    }

    pub fn new_with_client_and_base_url(
        api_key: &str,
        req_client: reqwest::Client,
        base_url: &str,
    ) -> Client {
        let timeout = Duration::from_secs(60);
        Client {
            req_client,
            key: api_key.to_owned(),
            base_url: reqwest::Url::parse(base_url).unwrap(),
            timeout,
            max_retries: 3,
        }
    }

    /// Read the response body as text and deserialize it as JSON, surfacing
    /// the HTTP status and the raw body in the error when either step fails.
    ///
    /// This is the path that `create_chat` and the other `create_*` methods
    /// use for the success branch. Routing everything through a single helper
    /// guarantees that callers never see a bare "error decoding response body"
    /// from reqwest without context, which previously made it impossible to
    /// tell whether a 200 response was actually a non-UTF-8 binary blob, a
    /// truncated stream, or some other transport-level failure.
    pub async fn read_and_parse_json<T: serde::de::DeserializeOwned>(
        res: reqwest::Response,
        error_context: &str,
    ) -> Result<T, anyhow::Error> {
        let status = res.status();
        match res.text().await {
            Ok(text) => serde_json::from_str(&text).map_err(|e| {
                anyhow!(
                    "{} failed to parse JSON response (status {}): {}. Raw body ({} bytes): {}",
                    error_context,
                    status,
                    e,
                    text.len(),
                    truncate_for_error(&text, 4096)
                )
            }),
            Err(e) => Err(anyhow!(
                "{} failed to read response body (status {}): {}",
                error_context,
                status,
                e
            )),
        }
    }

    pub async fn list_models(
        &self,
        opt_url_path: Option<String>,
    ) -> Result<Vec<models::Model>, anyhow::Error> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/models")));

        let res = self
            .req_client
            .get(url)
            .bearer_auth(&self.key)
            .send()
            .await?;

        if res.status() == 200 {
            let parsed: models::ListModelsResponse =
                Self::read_and_parse_json(res, "list_models").await?;
            Ok(parsed.data)
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "list_models failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    pub async fn create_chat(
        &self,
        args: chat::ChatArguments,
        opt_url_path: Option<String>,
    ) -> Result<chat::ChatCompletion, anyhow::Error> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/chat/completions")));

        let mut attempt = 0;
        loop {
            let res = self
                .req_client
                .post(url.clone())
                .bearer_auth(&self.key)
                .json(&args)
                .send()
                .await?;

            let status = res.status();
            if status.is_success() {
                // Use the helper so that body read failures (reqwest's
                // "error decoding response body" from a connection reset or
                // truncated stream) are treated as transient and retried
                // instead of bubbling up as a bare `Kind::Decode` error.
                match Self::read_and_parse_json::<chat::ChatCompletion>(res, "create_chat").await {
                    Ok(parsed) => return Ok(parsed),
                    Err(e) => {
                        if attempt >= self.max_retries {
                            return Err(e);
                        }
                        let backoff = backoff_for_attempt(attempt);
                        tokio::time::sleep(Duration::from_secs(backoff)).await;
                        attempt += 1;
                        continue;
                    }
                }
            }

            let should_retry = status == 429
                || status.as_u16() >= 500
                || matches!(
                    status.as_u16(),
                    502 | 503 | 504
                );

            if !should_retry || attempt >= self.max_retries {
                let body_text = res.text().await.unwrap_or_default();
                return Err(anyhow!(
                    "create_chat failed after {} attempts: status={} body={}",
                    attempt + 1,
                    status,
                    truncate_for_error(&body_text, 4096)
                ));
            }

            let retry_after = res
                .headers()
                .get("retry-after")
                .and_then(|v| v.to_str().ok())
                .and_then(|s| s.parse::<u64>().ok())
                .unwrap_or_else(|| 2u64.saturating_pow(attempt));

            // Consume the body (best-effort) so the connection can be reused
            // before we sleep. Errors here are non-fatal.
            let _ = res.text().await;

            tokio::time::sleep(Duration::from_secs(retry_after)).await;
            attempt += 1;
        }
    }

    pub async fn create_chat_stream(
        &self,
        args: chat::ChatArguments,
        opt_url_path: Option<String>,
    ) -> Result<chat::stream::ChatCompletionChunkStream> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/chat/completions")));

        let mut args = args;
        args.stream = Some(true);

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            Ok(chat::stream::ChatCompletionChunkStream::new(Box::pin(
                res.bytes_stream(),
            )))
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_chat_stream failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    pub async fn create_completion(
        &self,
        args: completions::CompletionArguments,
        opt_url_path: Option<String>,
    ) -> Result<completions::CompletionResponse> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/completions")));

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            Self::read_and_parse_json(res, "create_completion").await
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_completion failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    pub async fn create_embeddings(
        &self,
        args: embeddings::EmbeddingsArguments,
        opt_url_path: Option<String>,
    ) -> Result<embeddings::EmbeddingsResponse> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/embeddings")));

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            Self::read_and_parse_json(res, "create_embeddings").await
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_embeddings failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    pub async fn create_image_old(
        &self,
        args: images::ImageArguments,
        opt_url_path: Option<String>,
    ) -> Result<Vec<String>> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/images/generations")));

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            let parsed: images::ImageResponse =
                Self::read_and_parse_json(res, "create_image_old").await?;
            Ok(parsed
                .data
                .iter()
                .map(|o| match o {
                    images::ImageObject::Url(s) => s.to_string(),
                    images::ImageObject::Base64JSON(s) => s.to_string(),
                })
                .collect())
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_image_old failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    pub async fn create_image(
        &self,
        args: images::ImageArguments,
        opt_url_path: Option<String>,
    ) -> Result<Vec<String>> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/images/generations")));

        let image_args = images::ImageArguments {
            prompt: args.prompt,
            model: Some("gpt-image-1".to_string()),
            n: Some(1),
            size: Some("1024x1024".to_string()),
            quality: Some("auto".to_string()), // valid quality values are 'low', 'medium', 'high' and 'auto'
            //TODO: Make this an enum parameter to create_image
            user: None,
        };

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&image_args)
            .send()
            .await?;

        if res.status() == 200 {
            let parsed: images::ImageResponse =
                Self::read_and_parse_json(res, "create_image").await?;
            Ok(parsed
                .data
                .iter()
                .map(|o| match o {
                    images::ImageObject::Url(s) => s.to_string(),
                    images::ImageObject::Base64JSON(s) => s.to_string(),
                })
                .collect())
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_image failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    /// Create a response using xAI's Responses API with agentic tool calling.
    ///
    /// This method calls the `/v1/responses` endpoint which supports server-side
    /// tools like web_search, x_search, code_execution, and more.
    ///
    /// # Arguments
    /// * `args` - The ResponsesArguments containing model, input messages, and tools
    /// * `opt_url_path` - Optional URL path override (defaults to `/v1/responses`)
    ///
    /// # Example
    /// ```rust,no_run
    /// use openai_rust2::chat::{ResponsesArguments, ResponsesMessage, GrokTool};
    /// use openai_rust2::Client;
    ///
    /// async fn example() -> anyhow::Result<()> {
    ///     let client = Client::new_with_base_url("your-api-key", "https://api.x.ai/v1");
    ///     let args = ResponsesArguments::new(
    ///         "grok-4-1-fast-reasoning",
    ///         vec![ResponsesMessage {
    ///             role: "user".to_string(),
    ///             content: "What is the current Bitcoin price?".to_string(),
    ///         }],
    ///     ).with_tools(vec![GrokTool::web_search()]);
    ///
    ///     let response = client.create_responses(args, None).await?;
    ///     println!("{}", response.get_text_content());
    ///     Ok(())
    /// }
    /// ```
    pub async fn create_responses(
        &self,
        args: chat::ResponsesArguments,
        opt_url_path: Option<String>,
    ) -> Result<chat::ResponsesCompletion, anyhow::Error> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/responses")));

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            Self::read_and_parse_json(res, "create_responses").await
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_responses failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }

    /// Create a response using OpenAI's Responses API with agentic tool calling.
    ///
    /// This method calls the `/v1/responses` endpoint which supports server-side
    /// tools like web_search, file_search, and code_interpreter.
    ///
    /// Supported models: gpt-5, gpt-4o, and other models with tool support.
    ///
    /// # Arguments
    /// * `args` - The OpenAIResponsesArguments containing model, input messages, and tools
    /// * `opt_url_path` - Optional URL path override (defaults to `/v1/responses`)
    ///
    /// # Example
    /// ```rust,no_run
    /// use openai_rust2::chat::{OpenAIResponsesArguments, ResponsesMessage, OpenAITool};
    /// use openai_rust2::Client;
    ///
    /// async fn example() -> anyhow::Result<()> {
    ///     let client = Client::new("your-openai-api-key");
    ///     let args = OpenAIResponsesArguments::new(
    ///         "gpt-5",
    ///         vec![ResponsesMessage {
    ///             role: "user".to_string(),
    ///             content: "What are the latest developments in AI?".to_string(),
    ///         }],
    ///     ).with_tools(vec![OpenAITool::web_search()]);
    ///
    ///     let response = client.create_openai_responses(args, None).await?;
    ///     println!("{}", response.get_text_content());
    ///     Ok(())
    /// }
    /// ```
    pub async fn create_openai_responses(
        &self,
        args: chat::OpenAIResponsesArguments,
        opt_url_path: Option<String>,
    ) -> Result<chat::ResponsesCompletion, anyhow::Error> {
        let mut url = self.base_url.clone();
        url.set_path(&opt_url_path.unwrap_or_else(|| String::from("/v1/responses")));

        let res = self
            .req_client
            .post(url)
            .bearer_auth(&self.key)
            .json(&args)
            .send()
            .await?;

        if res.status() == 200 {
            Self::read_and_parse_json(res, "create_openai_responses").await
        } else {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            Err(anyhow!(
                "create_openai_responses failed: status={} body={}",
                status,
                truncate_for_error(&body, 4096)
            ))
        }
    }
}

/// Exponential backoff for `attempt` (0-based): 1, 2, 4, 8 … seconds,
/// capped at 30s so a runaway retry loop can never sleep for a minute
/// per attempt. Used for the body-decode retry path inside `create_chat`.
fn backoff_for_attempt(attempt: u32) -> u64 {
    (2u64.saturating_pow(attempt)).min(30)
}

/// Truncate a string for inclusion in error messages, marking the cut point
/// when bytes were dropped so logs don't get spammed by 10MB error bodies.
fn truncate_for_error(text: &str, max_bytes: usize) -> String {
    if text.len() <= max_bytes {
        text.to_string()
    } else {
        let mut cut = max_bytes;
        while !text.is_char_boundary(cut) && cut > 0 {
            cut -= 1;
        }
        format!("{}…[truncated, total {} bytes]", &text[..cut], text.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn truncate_for_error_keeps_short_strings() {
        let s = "hello".to_string();
        assert_eq!(truncate_for_error(&s, 10), "hello");
    }

    #[test]
    fn truncate_for_error_marks_cut_point() {
        let s = "a".repeat(100);
        let out = truncate_for_error(&s, 10);
        assert!(out.starts_with(&"a".repeat(10)));
        assert!(out.contains("truncated"));
        assert!(out.contains("100 bytes"));
    }

    #[test]
    fn truncate_for_error_respects_char_boundaries() {
        // "á" is 2 bytes in UTF-8; cutting at byte 1 would split it.
        let s = "ááááá";
        let out = truncate_for_error(&s, 3);
        // We should have cut down to a valid char boundary, never producing
        // a panic when downstream code tries to re-serialize the string.
        assert!(out.is_char_boundary(out.len() - 0));
    }

    #[test]
    fn backoff_grows_then_caps() {
        assert_eq!(backoff_for_attempt(0), 1);
        assert_eq!(backoff_for_attempt(1), 2);
        assert_eq!(backoff_for_attempt(2), 4);
        assert_eq!(backoff_for_attempt(3), 8);
        assert_eq!(backoff_for_attempt(10), 30); // capped
    }
}