openai-interface 0.6.0

use std::{future::Future, str::FromStr};

use eventsource_stream::Eventsource;
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
use serde::{Serialize, de::DeserializeOwned};

use crate::errors::OapiError;

pub trait Post {
    fn is_streaming(&self) -> bool;
    /// Builds the URL for the request.
    ///
    /// `base_url` should be like <https://api.openai.com/v1>
    fn build_url(&self, base_url: &str) -> Result<String, OapiError>;
}

pub trait PostNoStream: Post + Serialize + Sync + Send {
    type Response: DeserializeOwned + FromStr<Err = OapiError> + Send + Sync;

    /// Sends a POST request and returns the raw response body.
    ///
    /// The `client` is supplied by the caller; see [`crate::rest::default_client`]
    /// for a sensible default.
    fn get_response_string(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        key: &str,
    ) -> impl Future<Output = Result<String, OapiError>> + Send + Sync {
        async move {
            if self.is_streaming() {
                return Err(OapiError::NonStreamingViolation);
            }

            let response = client
                .post(self.build_url(base_url)?)
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .bearer_auth(key)
                .json(self)
                .send()
                .await?;

            crate::rest::response_text_checked(response).await
        }
    }

    /// Sends a POST request and deserializes the response.
    fn get_response(
        &self,
        client: &reqwest::Client,
        url: &str,
        key: &str,
    ) -> impl Future<Output = Result<Self::Response, OapiError>> + Send + Sync {
        async move {
            let text = self.get_response_string(client, url, key).await?;
            let result = Self::Response::from_str(&text)?;
            Ok(result)
        }
    }
}

pub trait PostStream: Post + Serialize + Sync + Send {
    type Response: DeserializeOwned + FromStr<Err = OapiError> + Send + Sync;

    /// Sends a streaming POST request and returns the raw SSE data items as
    /// strings.
    ///
    /// The returned stream stops after the `data: [DONE]` sentinel. Errors
    /// terminate the stream.
    ///
    /// The `client` is supplied by the caller; see [`crate::rest::default_client`]
    /// for a sensible default.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use std::sync::LazyLock;
    /// use futures_util::StreamExt;
    /// use openai_interface::chat::create::request::{Message, RequestBody};
    /// use openai_interface::rest::{default_client, post::PostStream};
    ///
    /// const DEEPSEEK_API_KEY: &str = "YOUR_API_KEY";
    /// const DEEPSEEK_CHAT_URL: &'static str = "https://api.deepseek.com";
    /// const DEEPSEEK_MODEL: &'static str = "deepseek-chat";
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let request = RequestBody {
    ///     messages: vec![
    ///         Message::System {
    ///             content: "This is a request of test purpose. Reply briefly".to_string(),
    ///             name: None,
    ///         },
    ///         Message::User {
    ///             content: "What's your name?".to_string(),
    ///             name: None,
    ///         },
    ///     ],
    ///     model: DEEPSEEK_MODEL.to_string(),
    ///     stream: Some(true),
    ///     ..Default::default()
    /// };
    ///
    /// let mut response = request
    ///     .get_stream_response_string(&default_client(), DEEPSEEK_CHAT_URL, DEEPSEEK_API_KEY)
    ///     .await?;
    ///
    /// while let Some(chunk) = response.next().await {
    ///     println!("{}", chunk?);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    fn get_stream_response_string(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        api_key: &str,
    ) -> impl Future<Output = Result<BoxStream<'static, Result<String, OapiError>>, OapiError>>
    + Send
    + Sync {
        async move {
            if !self.is_streaming() {
                return Err(OapiError::StreamingViolation);
            }

            let response = client
                .post(self.build_url(base_url)?)
                .header("Content-Type", "application/json")
                .header("Accept", "text/event-stream")
                .bearer_auth(api_key)
                .json(self)
                .send()
                .await?;

            // Parse the body as an SSE stream. Errors end the stream.
            let stream = crate::rest::check_status(response)
                .await?
                .bytes_stream()
                .eventsource()
                .map(|event| match event {
                    Ok(event) => Ok(event.data),
                    Err(e) => Err(OapiError::SseParseError(format!("SSE parse error: {e}"))),
                })
                .boxed();

            Ok(stream)
        }
    }

    /// Sends a streaming POST request and deserializes every SSE data item
    /// into [`Self::Response`].
    ///
    /// The stream ends after the `data: [DONE]` sentinel or on the first
    /// error, whichever comes first.
    fn get_stream_response(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        api_key: &str,
    ) -> impl Future<
        Output = Result<BoxStream<'static, Result<Self::Response, OapiError>>, OapiError>,
    > + Send
    + Sync {
        async move {
            let stream = self
                .get_stream_response_string(client, base_url, api_key)
                .await?;

            let parsed_stream = stream
                // The `data: [DONE]` sentinel is not JSON; end the stream there.
                .take_while(|result| {
                    let should_continue = matches!(result, Ok(data) if data != "[DONE]");
                    async move { should_continue }
                })
                .and_then(|data| async move { Self::Response::from_str(&data) });

            Ok(Box::pin(parsed_stream) as BoxStream<'static, _>)
        }
    }
}

/// Trait for POST requests whose response body is binary rather than JSON.
///
/// Endpoints such as `POST /audio/speech` return the raw audio bytes instead
/// of a JSON object, so the body is returned as raw bytes.
pub trait PostBinary: Post + Serialize + Sync + Send {
    /// Sends a POST request with a JSON body and returns the raw response
    /// body bytes.
    ///
    /// The `client` is supplied by the caller; see [`crate::rest::default_client`]
    /// for a sensible default.
    fn get_response_bytes(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        api_key: &str,
    ) -> impl Future<Output = Result<Vec<u8>, OapiError>> + Send + Sync {
        async move {
            let response = client
                .post(self.build_url(base_url)?)
                .header("Content-Type", "application/json")
                .header("Accept", "application/octet-stream")
                .bearer_auth(api_key)
                .json(self)
                .send()
                .await?;

            crate::rest::response_bytes_checked(response).await
        }
    }
}