openai-interface 0.14.0

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

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

use crate::errors::OapiError;
use crate::rest::{Auth, RequestOptions};

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

/// Applies the per-request authentication and extra headers to a request
/// under construction. Headers in [`RequestOptions::extra_headers`] replace
/// previously set headers of the same name.
pub(crate) fn apply_options(
    builder: reqwest::RequestBuilder,
    options: &RequestOptions,
) -> reqwest::RequestBuilder {
    let builder = match &options.auth {
        Auth::Bearer(token) => builder.bearer_auth(token),
        Auth::None => builder,
    };
    builder.headers(options.extra_headers.clone())
}

/// Sends a multipart form as a POST request and returns the raw response
/// body. Shared by the endpoints that upload files (`files`, `images`,
/// `audio`, `uploads`).
pub(crate) async fn post_multipart_json(
    client: &reqwest::Client,
    url: String,
    form: reqwest::multipart::Form,
    options: &RequestOptions,
) -> Result<String, OapiError> {
    let response = apply_options(
        client.post(url).header("Accept", "application/json"),
        options,
    )
    .multipart(form)
    .send()
    .await?;

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

/// Appends the request's `extra_body_map` entries to a multipart form as
/// text fields: strings verbatim, any other JSON value in its serialized
/// form. Multipart endpoints have no JSON body, so without this the
/// flattened map would silently never reach the wire.
pub(crate) fn append_extra_body_map(
    mut form: reqwest::multipart::Form,
    extra_body_map: &Option<serde_json::Map<String, serde_json::Value>>,
) -> reqwest::multipart::Form {
    if let Some(map) = extra_body_map {
        for (key, value) in map {
            let text = match value {
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string(),
            };
            form = form.text(key.clone(), text);
        }
    }
    form
}

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,
        options: &RequestOptions,
    ) -> impl Future<Output = Result<String, OapiError>> + Send + Sync {
        async move {
            if self.is_streaming() {
                return Err(OapiError::NonStreamingViolation);
            }

            let response = apply_options(
                client
                    .post(self.build_url(base_url)?)
                    .header("Accept", "application/json"),
                options,
            )
            .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,
        options: &RequestOptions,
    ) -> impl Future<Output = Result<Self::Response, OapiError>> + Send + Sync {
        async move {
            let text = self.get_response_string(client, url, options).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 + 'static;

    /// Sends a streaming POST request and returns the raw SSE data items as
    /// strings.
    ///
    /// The returned stream yields every SSE data payload verbatim, including
    /// the final `"[DONE]"` sentinel; use
    /// [`Self::get_stream_response`] when you want the sentinel consumed and
    /// the items deserialized. SSE framing errors surface as
    /// [`OapiError::SseParseError`] items.
    ///
    /// 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, RequestOptions};
    ///
    /// 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>> {
    /// // Needs the `ferritls` cargo feature; drop this line if you install
    /// // your own rustls crypto provider (see `openai_interface::rest`).
    /// # #[cfg(feature = "ferritls")]
    /// openai_interface::rest::install_crypto_provider().ok();
    ///
    /// let request = RequestBody {
    ///     messages: vec![
    ///         Message::system("This is a request of test purpose. Reply briefly"),
    ///         Message::user("What's your name?"),
    ///     ],
    ///     model: DEEPSEEK_MODEL.to_string(),
    ///     stream: Some(true),
    ///     ..Default::default()
    /// };
    ///
    /// let options = RequestOptions::bearer(DEEPSEEK_API_KEY);
    /// let mut response = request
    ///     .get_stream_response_string(&default_client(), DEEPSEEK_CHAT_URL, &options)
    ///     .await?;
    ///
    /// while let Some(chunk) = response.next().await {
    ///     println!("{}", chunk?);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    fn get_stream_response_string(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        options: &RequestOptions,
    ) -> impl Future<
        Output = Result<
            impl Stream<Item = Result<String, OapiError>> + Send + Unpin + 'static,
            OapiError,
        >,
    > + Send
    + Sync {
        async move {
            if !self.is_streaming() {
                return Err(OapiError::StreamingViolation);
            }

            let response = apply_options(
                client
                    .post(self.build_url(base_url)?)
                    .header("Accept", "text/event-stream"),
                options,
            )
            .json(self)
            .send()
            .await?;

            // Parse the body as an SSE stream. Framing errors become stream
            // items rather than terminating 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}"))),
                });

            Ok(stream)
        }
    }

    /// Sends a streaming POST request and deserializes every SSE data item
    /// into [`Self::Response`].
    ///
    /// The stream ends after the `data: [DONE]` sentinel. Deserialization
    /// failures surface as [`OapiError::DeserializationError`] items without
    /// ending the stream; wrap the stream with
    /// [`crate::rest::skip_deserialization_errors`] to drop them instead.
    fn get_stream_response(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        options: &RequestOptions,
    ) -> impl Future<
        Output = Result<
            impl Stream<Item = Result<Self::Response, OapiError>> + Send + Unpin + 'static,
            OapiError,
        >,
    > + Send
    + Sync {
        async move {
            let stream = self
                .get_stream_response_string(client, base_url, options)
                .await?;

            let parsed_stream = stream
                // The `data: [DONE]` sentinel is not JSON; end the stream
                // there. Everything else — including SSE framing errors —
                // flows through to the consumer.
                // `future::ready` keeps the combinators `Unpin` (async blocks
                // are not).
                .take_while(|result| {
                    std::future::ready(!matches!(result, Ok(data) if data == "[DONE]"))
                })
                .and_then(|data| std::future::ready(Self::Response::from_str(&data)));

            Ok(parsed_stream)
        }
    }
}

/// 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,
        options: &RequestOptions,
    ) -> impl Future<Output = Result<Vec<u8>, OapiError>> + Send + Sync {
        async move {
            let response = apply_options(
                client
                    .post(self.build_url(base_url)?)
                    .header("Accept", "application/octet-stream"),
                options,
            )
            .json(self)
            .send()
            .await?;

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

    /// Sends a POST request with a JSON body and streams the raw response
    /// body as byte chunks.
    ///
    /// Unlike [`Self::get_response_bytes`], the body is not buffered: each
    /// item is one transport-level chunk. This matches the official SDK's
    /// streaming binary responses (e.g. `POST /audio/speech` with
    /// `stream_format`, where the chunks carry either raw audio or SSE
    /// frames). Parsing any framing inside the byte stream is up to the
    /// caller.
    ///
    /// The `client` is supplied by the caller; see [`crate::rest::default_client`]
    /// for a sensible default.
    fn get_stream_response_bytes(
        &self,
        client: &reqwest::Client,
        base_url: &str,
        options: &RequestOptions,
    ) -> impl Future<
        Output = Result<
            impl Stream<Item = Result<Vec<u8>, OapiError>> + Send + Unpin + 'static,
            OapiError,
        >,
    > + Send
    + Sync {
        async move {
            let response = apply_options(
                client
                    .post(self.build_url(base_url)?)
                    .header("Accept", "application/octet-stream"),
                options,
            )
            .json(self)
            .send()
            .await?;

            let response = crate::rest::check_status(response).await?;

            Ok(response.bytes_stream().map(|chunk| {
                chunk
                    .map(|bytes| bytes.to_vec())
                    .map_err(OapiError::SendError)
            }))
        }
    }
}

#[cfg(test)]
mod test {
    use futures_util::StreamExt;
    use serde::Deserialize;
    use serde_json::json;
    use wiremock::matchers::{method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::*;
    use crate::chat::create::response::streaming::ChatCompletionChunk;
    use crate::rest::{RequestOptions, skip_deserialization_errors};

    /// Builds the default client and pins the SSE body used by the stream
    /// tests: one good chunk, one chunk that is not valid JSON, one good
    /// chunk, then the `[DONE]` sentinel.
    fn sse_body() -> String {
        const CHUNK_A: &str = r#"{"id":"1","choices":[{"index":0,"delta":{"content":"a"},"finish_reason":null}],"created":1,"model":"m","object":"chat.completion.chunk"}"#;
        const CHUNK_B: &str = r#"{"id":"1","choices":[{"index":0,"delta":{"content":"b"},"finish_reason":null}],"created":1,"model":"m","object":"chat.completion.chunk"}"#;
        format!("data: {CHUNK_A}\n\ndata: oops\n\ndata: {CHUNK_B}\n\ndata: [DONE]\n\n")
    }

    fn chunk_content(chunk: &ChatCompletionChunk) -> String {
        chunk.choices[0].delta.content.clone().unwrap_or_default()
    }

    #[derive(Serialize)]
    struct TestJsonRequest;

    impl Post for TestJsonRequest {
        fn is_streaming(&self) -> bool {
            false
        }
        fn build_url(&self, base_url: &str) -> Result<String, OapiError> {
            Ok(format!("{}/test", base_url.trim_end_matches('/')))
        }
    }

    impl PostNoStream for TestJsonRequest {
        type Response = TestResponse;
    }

    #[derive(Debug, Deserialize)]
    struct TestResponse {
        #[allow(dead_code)]
        id: String,
    }

    crate::impl_from_str!(TestResponse);

    #[derive(Serialize)]
    struct TestStreamRequest;

    impl Post for TestStreamRequest {
        fn is_streaming(&self) -> bool {
            true
        }
        fn build_url(&self, base_url: &str) -> Result<String, OapiError> {
            Ok(format!("{}/test", base_url.trim_end_matches('/')))
        }
    }

    impl PostStream for TestStreamRequest {
        type Response = ChatCompletionChunk;
    }

    #[tokio::test]
    async fn sends_bearer_auth_and_extra_headers() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/test"))
            .respond_with(
                ResponseTemplate::new(200).set_body_string(json!({"id": "x"}).to_string()),
            )
            .mount(&server)
            .await;

        let options = RequestOptions::bearer("sk-test")
            .with_header("OpenAI-Organization", "org-1")
            .unwrap();
        TestJsonRequest
            .get_response_string(&crate::rest::default_client(), &server.uri(), &options)
            .await
            .expect("request must succeed");

        let requests = server.received_requests().await.expect("recorded requests");
        assert_eq!(requests.len(), 1);
        let headers = &requests[0].headers;
        assert_eq!(headers.get("authorization").unwrap(), "Bearer sk-test");
        assert_eq!(headers.get("openai-organization").unwrap(), "org-1");
    }

    #[tokio::test]
    async fn header_only_auth_sends_no_authorization() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/test"))
            .respond_with(
                ResponseTemplate::new(200).set_body_string(json!({"id": "x"}).to_string()),
            )
            .mount(&server)
            .await;

        // Azure-style: credentials travel in `api-key`, not `Authorization`.
        let options = RequestOptions::new()
            .with_header("api-key", "azure-key")
            .unwrap();
        TestJsonRequest
            .get_response_string(&crate::rest::default_client(), &server.uri(), &options)
            .await
            .expect("request must succeed");

        let requests = server.received_requests().await.expect("recorded requests");
        assert_eq!(requests.len(), 1);
        let headers = &requests[0].headers;
        assert!(headers.get("authorization").is_none());
        assert_eq!(headers.get("api-key").unwrap(), "azure-key");
    }

    #[tokio::test]
    async fn multipart_upload_uses_build_url_and_options() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/files"))
            .respond_with(
                ResponseTemplate::new(200).set_body_string(json!({"id": "x"}).to_string()),
            )
            .mount(&server)
            .await;

        let path = std::env::temp_dir().join("openai_interface_multipart_test.txt");
        std::fs::write(&path, b"hello").expect("test file must be writable");

        let mut extra = serde_json::Map::new();
        extra.insert(
            "vendor_extension".to_string(),
            serde_json::json!({"depth": 3}),
        );

        let request = crate::files::create::request::CreateFileRequest {
            file: path,
            purpose: crate::files::create::request::FilePurpose::Batch,
            expires_after: None,
            extra_body_map: Some(extra),
        };
        let options = RequestOptions::new()
            .with_header("api-key", "azure-key")
            .unwrap();
        request
            .get_response_string(&crate::rest::default_client(), &server.uri(), &options)
            .await
            .expect("upload must succeed");

        let requests = server.received_requests().await.expect("recorded requests");
        assert_eq!(requests.len(), 1);
        // The extra body map is forwarded as multipart text fields.
        let body = String::from_utf8_lossy(&requests[0].body);
        assert!(
            body.contains("vendor_extension") && body.contains("\"depth\":3"),
            "extra body map must reach the multipart body: {body}"
        );
        let headers = &requests[0].headers;
        // The multipart overrides must route through `build_url`, so the
        // caller passes a bare base URL.
        assert_eq!(headers.get("api-key").unwrap(), "azure-key");
        assert_eq!(requests[0].url.path(), "/files");
    }

    #[tokio::test]
    async fn raw_stream_yields_every_data_item_including_done() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(ResponseTemplate::new(200).set_body_string(sse_body()))
            .mount(&server)
            .await;

        let stream = TestStreamRequest
            .get_stream_response_string(
                &crate::rest::default_client(),
                &server.uri(),
                &RequestOptions::bearer("sk-test"),
            )
            .await
            .expect("stream must start");

        let items: Vec<Result<String, OapiError>> = stream.collect().await;
        let data: Vec<&str> = items
            .iter()
            .map(|item| item.as_ref().expect("raw items must not fail").as_str())
            .collect();
        assert_eq!(data.len(), 4);
        assert_eq!(data[3], "[DONE]");
    }

    #[tokio::test]
    async fn parsed_stream_surfaces_bad_chunks_without_ending() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(ResponseTemplate::new(200).set_body_string(sse_body()))
            .mount(&server)
            .await;

        let stream = TestStreamRequest
            .get_stream_response(
                &crate::rest::default_client(),
                &server.uri(),
                &RequestOptions::bearer("sk-test"),
            )
            .await
            .expect("stream must start");

        let items: Vec<Result<ChatCompletionChunk, OapiError>> = stream.collect().await;
        assert_eq!(items.len(), 3, "good, bad, good; [DONE] ends the stream");
        assert_eq!(chunk_content(items[0].as_ref().unwrap()), "a");
        assert!(matches!(items[1], Err(OapiError::DeserializationError(_))));
        assert_eq!(chunk_content(items[2].as_ref().unwrap()), "b");
    }

    #[tokio::test]
    async fn skip_deserialization_errors_drops_bad_chunks() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(ResponseTemplate::new(200).set_body_string(sse_body()))
            .mount(&server)
            .await;

        let stream = TestStreamRequest
            .get_stream_response(
                &crate::rest::default_client(),
                &server.uri(),
                &RequestOptions::bearer("sk-test"),
            )
            .await
            .expect("stream must start");
        let mut stream = skip_deserialization_errors(stream);

        let mut contents = Vec::new();
        while let Some(item) = stream.next().await {
            contents.push(chunk_content(&item.expect("errors must be skipped")));
        }
        assert_eq!(contents, ["a", "b"]);
    }
}