openai-interface 0.8.1

//! Get a vector representation of a given input that can be easily consumed
//! by machine learning models and algorithms.
//!
//! Tested against the Alibaba Cloud Model Studio (Qwen, `text-embedding-v4`)
//! OpenAI-compatible `/embeddings` endpoint. Note that, unlike OpenAI, Qwen
//! currently only supports `encoding_format: "float"`.

pub mod request;
pub mod response;

#[cfg(test)]
mod tests {
    use crate::rest::{default_client, post::PostNoStream};

    use super::request::{EmbeddingInput, EmbeddingRequest};

    const QWEN_BASE_URL: &str = "https://dashscope.aliyuncs.com/compatible-mode/v1";
    const QWEN_EMBEDDING_MODEL: &str = "text-embedding-v4";

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

    /// Qwen documents the OpenAI-compatible `/embeddings` endpoint; run a
    /// live request when an API key is available.
    #[tokio::test]
    async fn test_qwen_embedding() -> Result<(), anyhow::Error> {
        let Some(api_key) = qwen_api_key() else {
            println!("Skipping: set QWEN_API_KEY to run this test");
            return Ok(());
        };

        let request = EmbeddingRequest {
            input: EmbeddingInput::String(
                "衣服的质量杠杠的，很漂亮，不枉我等了这么久啊".to_string(),
            ),
            model: QWEN_EMBEDDING_MODEL.to_string(),
            encoding_format: Some(super::request::EncodingFormat::Float),
            ..Default::default()
        };

        let response = request
            .get_response(&default_client(), QWEN_BASE_URL, &api_key)
            .await?;
        assert_eq!(response.model, QWEN_EMBEDDING_MODEL);
        assert_eq!(response.data.len(), 1);
        let super::response::EmbeddingVector::Floats(vector) = &response.data[0].embedding else {
            anyhow::bail!("expected float embedding vector");
        };
        assert!(!vector.is_empty());
        println!(
            "Embedding dims: {}, usage: {} prompt tokens",
            vector.len(),
            response.usage.prompt_tokens
        );
        Ok(())
    }
}