lc_embeddings/
openai_compat.rs

1// lc-embeddings/src/openai_compat.rs
2//! OpenAI 兼容协议 embedding 客户端公共基类（P1-5）。
3//!
4//! DeepSeek 与 Qwen 走同一套 OpenAI `/embeddings` 协议（同样的请求体、同样的
5//! `data[index]` 对齐、同样的 Bearer 认证），二者源码几乎逐行重复。本模块抽出
6//! 通用实现，DeepSeek/Qwen 仅通过 [`CompatSpec`] 配置 URL / 模型 / 维度 / 批量大小。
7
8use crate::{EmbeddingError, Embeddings};
9use async_trait::async_trait;
10use serde::Deserialize;
11
12/// 访问 provider 配置字段的抽象——DeepSeek/Qwen 的 config 结构体字段名相同。
13pub trait CompatConfigAccess {
14    /// 返回 API key
15    fn api_key(&self) -> &str;
16    /// 返回 Base URL
17    fn base_url(&self) -> &str;
18    /// 返回模型名
19    fn model(&self) -> &str;
20}
21
22/// OpenAI 兼容协议 provider 的静态规格。
23///
24/// 实现该 trait 即可获得 `OpenAICompatEmbeddings` 提供的完整 embedding 能力，
25/// 是接入新增 OpenAI 兼容 provider 的扩展点。
26pub trait CompatSpec: CompatConfigAccess + Sized + Default {
27    /// 环境变量名：API key（用于构造期错误信息，P1-3）。
28    fn api_key_env() -> &'static str;
29    /// 单次请求的批量上限。
30    fn batch_size() -> usize;
31    /// 给定模型的向量维度；未知模型必须报错（P1-2），不得回落默认值撒谎。
32    fn dimension_for(model: &str) -> Result<usize, EmbeddingError>;
33    /// 从环境变量构造 config（复用各 config 已实现的 from_env_result）。
34    fn from_env_result() -> Result<Self, EmbeddingError>;
35}
36
37/// 通用 OpenAI 兼容 embedding 客户端（P1-5）。
38///
39/// DeepSeek/Qwen 等走 OpenAI `/embeddings` 协议的 provider 通过
40/// [`CompatSpec`] 配置规格复用本实现；`C` 即各自的 config 类型。
41pub struct OpenAICompatEmbeddings<C: CompatConfigAccess + CompatSpec> {
42    config: C,
43    client: reqwest::Client,
44    dimension: usize,
45}
46
47impl<C: CompatConfigAccess + CompatSpec> std::fmt::Debug for OpenAICompatEmbeddings<C> {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct("OpenAICompatEmbeddings")
50            .field("model", &self.config.model())
51            .field("dimension", &self.dimension)
52            .finish()
53    }
54}
55
56impl<C: CompatConfigAccess + CompatSpec> OpenAICompatEmbeddings<C> {
57    /// 构造时 fail fast（P1-3）：API key 为空立即报错，而不是拖到发请求才 401；
58    /// 同时校验模型维度已知（P1-2）。
59    pub fn new(config: C) -> Result<Self, EmbeddingError> {
60        if config.api_key().trim().is_empty() {
61            return Err(EmbeddingError::Config(format!(
62                "{} is empty",
63                C::api_key_env()
64            )));
65        }
66        let dimension = C::dimension_for(config.model())?;
67        Ok(Self {
68            config,
69            client: reqwest::Client::new(),
70            dimension,
71        })
72    }
73
74    /// Creates from environment variables, returning a Result.
75    pub fn from_env_result() -> Result<Self, EmbeddingError> {
76        let config = C::from_env_result()?;
77        Self::new(config)
78    }
79}
80
81#[async_trait]
82impl<C: CompatConfigAccess + CompatSpec + Send + Sync> Embeddings for OpenAICompatEmbeddings<C> {
83    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
84        if text.trim().is_empty() {
85            return Err(EmbeddingError::EmptyInput);
86        }
87
88        let url = format!("{}/embeddings", self.config.base_url());
89
90        let body = serde_json::json!({
91            "model": self.config.model(),
92            "input": text,
93        });
94
95        // P2-5: 429/5xx 指数退避重试。
96        let response = crate::retry::post_json_with_retry(
97            &self.client,
98            &url,
99            self.config.api_key(),
100            &body,
101            &crate::retry::DEFAULT_RETRY,
102        )
103        .await
104        .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
105
106        let status = response.status();
107        if !status.is_success() {
108            // P1-4: 读失败的错误体也要报错，不能 unwrap_or_default() 吞掉。
109            let error_text = response.text().await.map_err(|e| {
110                EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
111            })?;
112            return Err(EmbeddingError::ApiError(format!(
113                "HTTP {}: {}",
114                status, error_text
115            )));
116        }
117
118        let embedding_response: EmbeddingResponse = response
119            .json()
120            .await
121            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
122
123        let mut embedding = embedding_response
124            .data
125            .first()
126            .ok_or_else(|| EmbeddingError::ApiError("No embedding data in response".to_string()))?
127            .embedding
128            .clone();
129        // P2-8: 统一 L2 归一化,保证单位长度。
130        crate::l2_normalize(&mut embedding);
131        Ok(embedding)
132    }
133
134    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
135        // P1-1: 空切片不是错误（无事可做），含空/全空白文本才报错。
136        if texts.is_empty() {
137            return Ok(Vec::new());
138        }
139        if texts.iter().any(|t| t.trim().is_empty()) {
140            return Err(EmbeddingError::EmptyInput);
141        }
142
143        let url = format!("{}/embeddings", self.config.base_url());
144        let batch_size = C::batch_size().max(1);
145        // P0-1: 用 Option 槽位逐项收集,拒绝静默空向量。某 chunk 少返回/错位会
146        // 留下 None 槽位并在收尾时报错,而不是产出零向量被下游当成"不相似"。
147        let mut all_results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
148        let mut offset = 0;
149
150        for chunk in texts.chunks(batch_size) {
151            let body = serde_json::json!({
152                "model": self.config.model(),
153                "input": chunk,
154            });
155
156            // P2-5: 429/5xx 指数退避重试。
157            let response = crate::retry::post_json_with_retry(
158                &self.client,
159                &url,
160                self.config.api_key(),
161                &body,
162                &crate::retry::DEFAULT_RETRY,
163            )
164            .await
165            .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
166
167            let status = response.status();
168            if !status.is_success() {
169                // P1-4: 读失败的错误体也要报错，不能 unwrap_or_default() 吞掉。
170                let error_text = response.text().await.map_err(|e| {
171                    EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
172                })?;
173                return Err(EmbeddingError::ApiError(format!(
174                    "HTTP {}: {}",
175                    status, error_text
176                )));
177            }
178
179            let embedding_response: EmbeddingResponse = response
180                .json()
181                .await
182                .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
183
184            for item in embedding_response.data {
185                let global_index = offset + item.index as usize;
186                if global_index >= all_results.len() {
187                    // 服务端 index 超出请求范围 = 批次错位,直接报错。
188                    return Err(EmbeddingError::BatchMismatch {
189                        expected: all_results.len(),
190                        actual: global_index + 1,
191                    });
192                }
193                all_results[global_index] = Some(item.embedding);
194            }
195            offset += chunk.len();
196        }
197
198        // 展开为 Result:任一槽位空缺即显式报错,而非留下零向量;并统一 L2 归一化(P2-8)。
199        all_results
200            .into_iter()
201            .map(|opt| {
202                let mut v = opt.ok_or(EmbeddingError::EmptyVectorInBatch)?;
203                crate::l2_normalize(&mut v);
204                Ok(v)
205            })
206            .collect()
207    }
208
209    fn dimension(&self) -> usize {
210        self.dimension
211    }
212
213    fn model_name(&self) -> &str {
214        self.config.model()
215    }
216}
217
218/// OpenAI 兼容协议的 embedding 响应体（DeepSeek/Qwen 共用）。
219#[derive(Debug, Deserialize)]
220struct EmbeddingResponse {
221    data: Vec<EmbeddingData>,
222}
223
224#[derive(Debug, Deserialize)]
225struct EmbeddingData {
226    embedding: Vec<f32>,
227    index: i32,
228}
lc_embeddings/openai_compat.rs

lc_embeddings/
openai_compat.rs
