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

lc_embeddings/
openai_compat.rs
