Skip to main content

openai_interface/embeddings/
mod.rs

1//! Get a vector representation of a given input that can be easily consumed
2//! by machine learning models and algorithms.
3//!
4//! Tested against the Alibaba Cloud Model Studio (Qwen, `text-embedding-v4`)
5//! OpenAI-compatible `/embeddings` endpoint. Note that, unlike OpenAI, Qwen
6//! currently only supports `encoding_format: "float"`.
7
8pub mod request;
9pub mod response;
10
11#[cfg(test)]
12mod tests {
13    use crate::rest::{default_client, post::PostNoStream};
14
15    use super::request::{EmbeddingInput, EmbeddingRequest};
16
17    const QWEN_BASE_URL: &str = "https://dashscope.aliyuncs.com/compatible-mode/v1";
18    const QWEN_EMBEDDING_MODEL: &str = "text-embedding-v4";
19
20    fn qwen_api_key() -> Option<String> {
21        std::env::var("QWEN_API_KEY")
22            .ok()
23            .map(|key| key.trim().to_string())
24            .filter(|key| !key.is_empty())
25    }
26
27    /// Qwen documents the OpenAI-compatible `/embeddings` endpoint; run a
28    /// live request when an API key is available.
29    #[tokio::test]
30    async fn test_qwen_embedding() -> Result<(), anyhow::Error> {
31        let Some(api_key) = qwen_api_key() else {
32            println!("Skipping: set QWEN_API_KEY to run this test");
33            return Ok(());
34        };
35
36        let request = EmbeddingRequest {
37            input: EmbeddingInput::String(
38                "衣服的质量杠杠的,很漂亮,不枉我等了这么久啊".to_string(),
39            ),
40            model: QWEN_EMBEDDING_MODEL.to_string(),
41            encoding_format: Some(super::request::EncodingFormat::Float),
42            ..Default::default()
43        };
44
45        let response = request
46            .get_response(&default_client(), QWEN_BASE_URL, &api_key)
47            .await?;
48        assert_eq!(response.model, QWEN_EMBEDDING_MODEL);
49        assert_eq!(response.data.len(), 1);
50        let super::response::EmbeddingVector::Floats(vector) = &response.data[0].embedding else {
51            anyhow::bail!("expected float embedding vector");
52        };
53        assert!(!vector.is_empty());
54        println!(
55            "Embedding dims: {}, usage: {} prompt tokens",
56            vector.len(),
57            response.usage.prompt_tokens
58        );
59        Ok(())
60    }
61}