openai_interface/embeddings/
mod.rs1pub 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 #[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}