openai_interface/rest/

post.rs

use std::future::Future;

use futures_util::{TryStreamExt, stream::BoxStream};
use serde::Serialize;

use crate::errors::RequestError;

pub trait Post {
    fn is_streaming(&self) -> bool;
}

pub trait NoStream: Post + Serialize + Sync + Send {
    /// Sends a POST request to the specified URL with the provided api-key.
    fn get_response(
        &self,
        url: &str,
        key: &str,
    ) -> impl Future<Output = Result<String, RequestError>> + Send + Sync {
        async move {
            if self.is_streaming() {
                return Err(RequestError::StreamingViolation);
            }

            let client = reqwest::Client::new();
            let response = client
                .post(url)
                .headers({
                    let mut headers = reqwest::header::HeaderMap::new();
                    headers.insert("Content-Type", "application/json".parse().unwrap());
                    headers.insert("Accept", "application/json".parse().unwrap());
                    headers
                })
                .bearer_auth(key)
                .json(self)
                .send()
                .await
                .map_err(|e| {
                    RequestError::SendError(format!("Failed to send request: {:#?}", e))
                })?;

            if response.status() != reqwest::StatusCode::OK {
                return Err(crate::errors::RequestError::ResponseStatus(
                    response.status().as_u16(),
                )
                .into());
            }

            let text = response.text().await.map_err(|e| {
                RequestError::ResponseError(format!("Failed to get response text: {:#?}", e))
            })?;

            Ok(text)
        }
    }
}

pub trait Stream: Post + Serialize + Sync + Send {
    /// Sends a streaming POST request to the specified URL with the provided api-key.
    ///
    /// # Example
    ///
    /// ```rust
    /// use std::sync::LazyLock;
    /// use futures_util::StreamExt;
    /// use openai_interface::chat::request::{Message, RequestBody};
    /// use openai_interface::rest::post::Stream;
    ///
    /// const DEEPSEEK_API_KEY: LazyLock<&str> =
    ///     LazyLock::new(|| include_str!("../.././keys/deepseek_domestic_key").trim());
    /// const DEEPSEEK_CHAT_URL: &'static str = "https://api.deepseek.com/chat/completions";
    /// const DEEPSEEK_MODEL: &'static str = "deepseek-chat";
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     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: true,
    ///         ..Default::default()
    ///     };
    ///
    ///     let mut response = request
    ///         .get_stream_response(DEEPSEEK_CHAT_URL, *DEEPSEEK_API_KEY)
    ///         .await
    ///         .unwrap();
    ///
    ///     while let Some(chunk) = response.next().await {
    ///         println!("{}", chunk.unwrap());
    ///     }
    /// }
    /// ```
    fn get_stream_response(
        &self,
        url: &str,
        api_key: &str,
    ) -> impl Future<
        Output = Result<BoxStream<'static, Result<String, anyhow::Error>>, anyhow::Error>,
    > + Send
    + Sync {
        async move {
            if !self.is_streaming() {
                return Err(anyhow::Error::from(RequestError::StreamingViolation));
            }

            let client = reqwest::Client::new();

            let response = client
                .post(url)
                .headers({
                    let mut headers = reqwest::header::HeaderMap::new();
                    headers.insert("Content-Type", "application/json".parse().unwrap());
                    headers.insert("Accept", "application/json".parse().unwrap());
                    headers
                })
                .bearer_auth(api_key)
                .json(self)
                .send()
                .await
                .map_err(|e| anyhow::anyhow!("Failed to send request: {}", e))?;

            if !response.status().is_success() {
                return Err(RequestError::ResponseStatus(response.status().as_u16()).into());
            }

            // The following code is generated by Qwen3-Coder-480B
            // 使用 unfold 来维护状态，正确处理 SSE 格式
            let stream = futures_util::stream::unfold(
                (response.bytes_stream(), String::new()),
                |(mut stream, mut buffer)| async move {
                    loop {
                        // 检查缓冲区中是否有完整的事件（以 \n\n 结尾）
                        while let Some(end_pos) = buffer.find("\n\n") {
                            let event = buffer[..end_pos].to_string();
                            buffer.drain(..end_pos + 2); // 移除事件和分隔符

                            // 处理 SSE 事件
                            if event.starts_with("data: ") {
                                let data = event[6..].to_string(); // 移除 "data: " 前缀
                                // 检查是否是 [DONE] 事件
                                if data == "[DONE]" {
                                    // 对于 [DONE] 事件，返回它以便调用者可以处理
                                    return Some((Ok(data), (stream, buffer)));
                                } else {
                                    // 对于其他数据事件，返回数据
                                    return Some((Ok(data), (stream, buffer)));
                                }
                            } else if event == "[DONE]" {
                                // 直接的 [DONE] 事件（不带 data: 前缀）
                                return Some((Ok("[DONE]".to_string()), (stream, buffer)));
                            } else {
                                // 其他类型的事件（如注释），忽略
                                continue;
                            }
                        }

                        // 从流中获取更多数据
                        match stream.try_next().await {
                            Ok(Some(bytes)) => {
                                if let Ok(s) = std::str::from_utf8(&bytes) {
                                    buffer.push_str(s);
                                    // 继续循环检查是否有完整事件
                                    continue;
                                } else {
                                    return Some((
                                        Err(RequestError::SseParseError(
                                            "Invalid UTF-8 in stream".to_string(),
                                        )
                                        .into()),
                                        (stream, buffer),
                                    ));
                                }
                            }
                            Ok(None) => {
                                // 流结束，检查是否有剩余数据
                                if !buffer.is_empty() {
                                    // 处理可能的不完整事件
                                    if buffer.starts_with("data: ") {
                                        let data = buffer[6..].to_string();
                                        // 检查是否是 [DONE] 事件
                                        if data == "[DONE]" {
                                            return Some((Ok(data), (stream, String::new())));
                                        } else {
                                            return Some((Ok(data), (stream, String::new())));
                                        }
                                    } else if buffer == "[DONE]" {
                                        // 直接的 [DONE] 事件
                                        return Some((
                                            Ok("[DONE]".to_string()),
                                            (stream, String::new()),
                                        ));
                                    }
                                }
                                return None; // 流结束
                            }
                            Err(e) => {
                                return Some((
                                    Err(anyhow::anyhow!("Stream error: {}", e).into()),
                                    (stream, buffer),
                                ));
                            }
                        }
                    }
                },
            );

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