code_repo_wiki/generate/
embed.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use anyhow::{Context, Result};
4
5use crate::analysis::feature::Embedder;
6use crate::config::schema::EmbedSection;
7
8/// Embedding 引擎：将代码/文档文本转为向量表示。
9///
10/// 调用 OpenAI 兼容的嵌入 API（支持 text-embedding-3-small 等模型）。
11pub struct EmbeddingEngine {
12    client: reqwest::Client,
13    config: EmbedSection,
14    call_count: AtomicUsize,
15    /// 全局 tokio Runtime 句柄（同步 Embedder 实现经其驱动 async 请求）
16    rt: tokio::runtime::Handle,
17}
18
19impl EmbeddingEngine {
20    /// 从配置创建 Embedding 引擎。
21    ///
22    /// 优先使用 `api_key` 字段，其次从环境变量读取。
23    /// `rt` 传入全局 Runtime 句柄（语义索引与特征聚类共用）。
24    pub fn new(config: &EmbedSection, rt: tokio::runtime::Handle) -> Result<Self> {
25        let client = reqwest::Client::builder()
26            .timeout(std::time::Duration::from_secs(60))
27            .build()
28            .context("创建 Embedding HTTP 客户端失败")?;
29        Ok(Self {
30            client,
31            config: config.clone(),
32            call_count: AtomicUsize::new(0),
33            rt,
34        })
35    }
36
37    /// 解析 API Key，优先级：api_key > 环境变量 > 报错
38    fn resolve_api_key(&self) -> Result<String> {
39        self.config
40            .api_key
41            .clone()
42            .or_else(|| std::env::var(&self.config.api_key_env).ok())
43            .context(format!(
44                "Embedding API Key 未设置（api_key 为空且环境变量 {} 未定义）",
45                self.config.api_key_env
46            ))
47    }
48
49    /// 获取 API base URL
50    fn resolve_base_url(&self) -> String {
51        self.config
52            .base_url
53            .clone()
54            .unwrap_or_else(|| "https://api.openai.com/v1".to_string())
55    }
56
57    /// 批量嵌入：将多个文本转为向量。
58    ///
59    /// 按 `batch_size` 分批发往 API，自动合并结果。
60    pub async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
61        let api_key = self.resolve_api_key()?;
62        let url = format!("{}/embeddings", self.resolve_base_url());
63
64        let mut all_embeddings = Vec::with_capacity(texts.len());
65
66        for chunk in texts.chunks(crate::config::schema::EMBED_BATCH_SIZE) {
67            let body = serde_json::json!({
68                "model": self.config.model,
69                "input": chunk,
70            });
71
72            // N16：embedding 请求接入统一重试骨架（与 LLM 通道一致：429/5xx/
73            // 超时/连接失败按指数退避重试，其余 4xx 立即失败）。每轮重试重建
74            // 请求（闭包捕获 body/url/key 的引用）。
75            let resp = crate::generate::llm::retry_with_backoff(
76                crate::generate::llm::MAX_RETRIES,
77                || {
78                    let body = &body;
79                    let url = &url;
80                    let api_key = &api_key;
81                    let client = &self.client;
82                    async move {
83                        client
84                            .post(url)
85                            .bearer_auth(api_key)
86                            .json(body)
87                            .send()
88                            .await
89                    }
90                },
91            )
92            .await
93            .with_context(|| "Embedding API 请求失败")?;
94
95            if !resp.status().is_success() {
96                let status = resp.status();
97                let text = resp.text().await.unwrap_or_default();
98                anyhow::bail!("Embedding API 返回错误 ({}): {}", status, text);
99            }
100
101            self.call_count.fetch_add(1, Ordering::Relaxed);
102
103            let data: serde_json::Value = resp
104                .json()
105                .await
106                .context("解析 Embedding API 响应 JSON 失败")?;
107
108            let embeddings = data["data"]
109                .as_array()
110                .context("Embedding 响应缺少 data 字段")?
111                .iter()
112                .map(|item| {
113                    let arr = item["embedding"]
114                        .as_array()
115                        .context("嵌入向量缺失")?;
116                    // B6：元素必须全为数字——filter_map 静默丢弃非数字元素会
117                    // 让向量降维而不报错（同批一致变短时维度校验也捕获不到），
118                    // 模型输出异常必须显式失败而非产出残缺向量
119                    arr.iter()
120                        .map(|v| {
121                            v.as_f64()
122                                .map(|f| f as f32)
123                                .with_context(|| "嵌入向量包含非数字元素（模型输出异常，拒绝静默丢弃）")
124                        })
125                        .collect::<Result<Vec<f32>>>()
126                })
127                .collect::<Result<Vec<_>>>()?;
128
129            // N5 修复：响应校验——data 条数必须与请求批次一致，且同批
130            // 向量维度必须一致。此前只校验"字段存在"，条数不足时
131            // 索引错位（下游 zip 静默丢弃多余/缺失）、维度不一致时
132            // 向量库维度校验失败但错误发生在数据已被吞之后。
133            if embeddings.len() != chunk.len() {
134                anyhow::bail!(
135                    "Embedding 响应数量不匹配：请求 {} 条，返回 {} 条",
136                    chunk.len(),
137                    embeddings.len()
138                );
139            }
140            if let Some(first) = embeddings.first() {
141                let dim = first.len();
142                if let Some(bad) = embeddings.iter().find(|v| v.len() != dim) {
143                    anyhow::bail!(
144                        "Embedding 响应维度不一致：{} 维与 {} 维并存",
145                        dim,
146                        bad.len()
147                    );
148                }
149            }
150
151            all_embeddings.extend(embeddings);
152        }
153
154        Ok(all_embeddings)
155    }
156
157    /// 单文本嵌入。
158    pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
159        let results = self.embed_batch(&[text.to_string()]).await?;
160        results.into_iter().next().context("Embedding 返回空结果")
161    }
162
163    /// 余弦相似度（-1~1）。
164    pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
165        if a.len() != b.len() || a.is_empty() {
166            return 0.0;
167        }
168        let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
169        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
170        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
171        if norm_a == 0.0 || norm_b == 0.0 {
172            return 0.0;
173        }
174        (dot / (norm_a * norm_b)).clamp(-1.0, 1.0)
175    }
176
177    /// 已完成的 API 调用次数。
178    pub fn call_count(&self) -> usize {
179        self.call_count.load(Ordering::Relaxed)
180    }
181}
182
183/// 特征聚类用的 Embedder 实现（analysis::feature::Embedder）。
184///
185/// 同步方法经内部持有的 tokio Handle 驱动 async 请求。
186impl Embedder for EmbeddingEngine {
187    fn embed(&self, text: &str) -> Result<Vec<f32>> {
188        self.rt.block_on(EmbeddingEngine::embed(self, text))
189    }
190
191    fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
192        self.rt
193            .block_on(EmbeddingEngine::embed_batch(self, texts))
194    }
195
196    fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> f64 {
197        EmbeddingEngine::cosine_similarity(a, b) as f64
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn test_cosine_similarity_identical() {
207        let a = vec![1.0, 0.0, 0.0];
208        let sim = EmbeddingEngine::cosine_similarity(&a, &a);
209        assert!((sim - 1.0).abs() < 1e-6);
210    }
211
212    #[test]
213    fn test_cosine_similarity_orthogonal() {
214        let a = vec![1.0, 0.0];
215        let b = vec![0.0, 1.0];
216        let sim = EmbeddingEngine::cosine_similarity(&a, &b);
217        assert!((sim - 0.0).abs() < 1e-6);
218    }
219
220    #[test]
221    fn test_cosine_similarity_opposite() {
222        let a = vec![1.0, 0.0];
223        let b = vec![-1.0, 0.0];
224        let sim = EmbeddingEngine::cosine_similarity(&a, &b);
225        assert!((sim + 1.0).abs() < 1e-6);
226    }
227
228    #[test]
229    fn test_cosine_similarity_zero_vector() {
230        let a = vec![0.0, 0.0];
231        let b = vec![1.0, 0.0];
232        let sim = EmbeddingEngine::cosine_similarity(&a, &b);
233        assert!((sim - 0.0).abs() < 1e-6);
234    }
235
236    #[test]
237    fn test_cosine_similarity_empty() {
238        let sim = EmbeddingEngine::cosine_similarity(&[], &[]);
239        assert!((sim - 0.0).abs() < 1e-6);
240    }
241
242    #[test]
243    fn test_cosine_similarity_mismatched_length() {
244        let a = vec![1.0, 0.0];
245        let b = vec![1.0];
246        let sim = EmbeddingEngine::cosine_similarity(&a, &b);
247        assert!((sim - 0.0).abs() < 1e-6);
248    }
249}
code_repo_wiki/generate/embed.rs

code_repo_wiki/generate/
embed.rs
