openai-interface 0.6.0

//! Given a prompt, the model will return one or more predicted completions,
//! and can also return the probabilities of alternative tokens at each position.
//! Compared to the `chat` API, this one does not provide the ability to have
//! multiple rounds of conversation. This API is getting deprecated in favor of the
//! `chat` API.

pub mod request;
pub mod response;

#[cfg(test)]
mod tests {
    use futures_util::StreamExt;

    use crate::rest::{
        default_client,
        post::{PostNoStream, PostStream},
    };

    use super::*;

    const QWEN_MODEL: &str = "qwen-coder-turbo";
    const QWEN_URL: &str = "https://dashscope.aliyuncs.com/compatible-mode/v1";

    fn qwen_api_key() -> Option<String> {
        std::env::var("QWEN_API_KEY")
            .ok()
            .map(|key| key.trim().to_string())
            .filter(|key| !key.is_empty())
    }

    fn request_body() -> request::CompletionRequest {
        request::CompletionRequest {
            model: QWEN_MODEL.to_string(),
            prompt: request::Prompt::PromptString(
                "\"桂棹兮兰桨，击空明兮溯流光\" 出自哪里？".to_string(),
            ),
            max_tokens: Some(500),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn test_qwen_completion_no_stream() -> Result<(), anyhow::Error> {
        let Some(api_key) = qwen_api_key() else {
            println!("Skipping: set QWEN_API_KEY to run this test");
            return Ok(());
        };

        let completion = request_body()
            .get_response(&default_client(), QWEN_URL, &api_key)
            .await?;
        let text = &completion.choices[0].text;
        println!("Completion no-stream: {}", text);
        Ok(())
    }

    #[tokio::test]
    async fn test_qwen_completion_stream() -> Result<(), anyhow::Error> {
        let Some(api_key) = qwen_api_key() else {
            println!("Skipping: set QWEN_API_KEY to run this test");
            return Ok(());
        };

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

        let mut stream = request_body
            .get_stream_response(&default_client(), QWEN_URL, &api_key)
            .await?;

        while let Some(chunk) = stream.next().await {
            match chunk {
                Ok(completion) => {
                    println!(
                        "Completion stream chunk: {}; finish reason: {:?}",
                        completion.choices[0].text, completion.choices[0].finish_reason
                    );
                }
                Err(e) => {
                    eprintln!("Error receiving chunk: {:?}", e);
                    break;
                }
            }
        }

        Ok(())
    }
}