lc_vector_stores/
pinecone.rs

1//! Pinecone vector store (HTTP API)
2
3use std::collections::HashMap;
4
5use async_trait::async_trait;
6use serde::Deserialize;
7
8use crate::{Document, Embeddings, SearchResult, VectorStore, VectorStoreError};
9
10/// Pinecone vector store client
11pub struct PineconeStore {
12    api_key: String,
13    host: String,
14    client: reqwest::Client,
15}
16
17impl PineconeStore {
18    /// Create a Pinecone client.
19    ///
20    /// `host` format: `https://{index-name}.svc.{environment}.pinecone.io`
21    pub fn new(api_key: impl Into<String>, host: impl Into<String>) -> Self {
22        Self {
23            api_key: api_key.into(),
24            host: host.into(),
25            client: reqwest::Client::new(),
26        }
27    }
28
29    /// Build upsert request body (pure function, convenient for testing).
30    pub fn build_upsert_body(docs: &[Document], vectors: &[Vec<f32>]) -> serde_json::Value {
31        let vectors_json: Vec<serde_json::Value> = docs
32            .iter()
33            .zip(vectors.iter())
34            .map(|(doc, vec)| {
35                serde_json::json!({
36                    "id": doc.id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
37                    "values": vec,
38                    "metadata": doc.metadata,
39                })
40            })
41            .collect();
42        serde_json::json!({ "vectors": vectors_json })
43    }
44
45    /// Build query request body (pure function, convenient for testing).
46    pub fn build_query_body(query_vec: &[f32], top_k: usize) -> serde_json::Value {
47        serde_json::json!({
48            "vector": query_vec,
49            "topK": top_k,
50            "includeMetadata": true,
51        })
52    }
53
54    /// Upsert documents (auto-embed).
55    pub async fn upsert(
56        &self,
57        docs: &[Document],
58        embeddings: &dyn Embeddings,
59    ) -> Result<(), String> {
60        let texts: Vec<&str> = docs.iter().map(|d| d.content.as_str()).collect();
61        let vectors = embeddings
62            .embed_documents(&texts)
63            .await
64            .map_err(|e| e.to_string())?;
65        let body = Self::build_upsert_body(docs, &vectors);
66        let url = format!("{}/vectors/upsert", self.host);
67        let resp = self
68            .client
69            .post(&url)
70            .header("Api-Key", &self.api_key)
71            .json(&body)
72            .send()
73            .await
74            .map_err(|e| e.to_string())?;
75        if !resp.status().is_success() {
76            return Err(format!("Pinecone upsert error: {}", resp.status()));
77        }
78        Ok(())
79    }
80
81    /// Query similar documents.
82    pub async fn query(&self, query_vec: Vec<f32>, top_k: usize) -> Result<Vec<Document>, String> {
83        let body = Self::build_query_body(&query_vec, top_k);
84        let url = format!("{}/query", self.host);
85        let resp = self
86            .client
87            .post(&url)
88            .header("Api-Key", &self.api_key)
89            .json(&body)
90            .send()
91            .await
92            .map_err(|e| e.to_string())?;
93        if !resp.status().is_success() {
94            return Err(format!("Pinecone query error: {}", resp.status()));
95        }
96        let query_resp: QueryResponse = resp.json().await.map_err(|e| e.to_string())?;
97        let result = query_resp
98            .matches
99            .into_iter()
100            .map(|m| {
101                let content = m
102                    .metadata
103                    .as_ref()
104                    .and_then(|md| md.get("content").cloned())
105                    .unwrap_or_default();
106                Document {
107                    content,
108                    metadata: m.metadata.unwrap_or_default(),
109                    id: Some(m.id),
110                }
111            })
112            .collect();
113        Ok(result)
114    }
115
116    /// 读取索引统计(真实 count 的唯一可靠来源)。
117    ///
118    /// Pinecone REST 提供 `describe_index_stats`,返回 `totalVectorCount`。
119    pub async fn describe_index_stats(&self) -> Result<PineconeIndexStats, String> {
120        let url = format!("{}/describe_index_stats", self.host);
121        let resp = self
122            .client
123            .post(&url)
124            .header("Api-Key", &self.api_key)
125            .send()
126            .await
127            .map_err(|e| e.to_string())?;
128        if !resp.status().is_success() {
129            return Err(format!(
130                "Pinecone describe_index_stats error: {}",
131                resp.status()
132            ));
133        }
134        resp.json().await.map_err(|e| e.to_string())
135    }
136
137    /// Delete by IDs.
138    pub async fn delete(&self, ids: &[String]) -> Result<(), String> {
139        let url = format!("{}/vectors/delete", self.host);
140        let body = serde_json::json!({ "ids": ids });
141        let resp = self
142            .client
143            .post(&url)
144            .header("Api-Key", &self.api_key)
145            .json(&body)
146            .send()
147            .await
148            .map_err(|e| e.to_string())?;
149        if !resp.status().is_success() {
150            return Err(format!("Pinecone delete error: {}", resp.status()));
151        }
152        Ok(())
153    }
154}
155
156#[async_trait]
157impl VectorStore for PineconeStore {
158    async fn add_documents(
159        &self,
160        documents: Vec<Document>,
161        embeddings: Vec<Vec<f32>>,
162    ) -> Result<Vec<String>, VectorStoreError> {
163        let ids: Vec<String> = documents
164            .iter()
165            .map(|d| {
166                d.id.clone()
167                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
168            })
169            .collect();
170
171        // Build upsert body with pre-computed embeddings
172        let body = Self::build_upsert_body(&documents, &embeddings);
173        let url = format!("{}/vectors/upsert", self.host);
174        let resp = self
175            .client
176            .post(&url)
177            .header("Api-Key", &self.api_key)
178            .json(&body)
179            .send()
180            .await
181            .map_err(|e| {
182                VectorStoreError::StorageError(format!("Pinecone upsert failed: {}", e))
183            })?;
184
185        if !resp.status().is_success() {
186            return Err(VectorStoreError::StorageError(format!(
187                "Pinecone upsert HTTP error: {}",
188                resp.status()
189            )));
190        }
191
192        Ok(ids)
193    }
194
195    async fn similarity_search(
196        &self,
197        query_embedding: &[f32],
198        k: usize,
199    ) -> Result<Vec<SearchResult>, VectorStoreError> {
200        let body = Self::build_query_body(query_embedding, k);
201        let url = format!("{}/query", self.host);
202        let resp = self
203            .client
204            .post(&url)
205            .header("Api-Key", &self.api_key)
206            .json(&body)
207            .send()
208            .await
209            .map_err(|e| VectorStoreError::StorageError(format!("Pinecone query failed: {}", e)))?;
210
211        if !resp.status().is_success() {
212            return Err(VectorStoreError::StorageError(format!(
213                "Pinecone query HTTP error: {}",
214                resp.status()
215            )));
216        }
217
218        let query_resp: QueryResponse = resp.json().await.map_err(|e| {
219            VectorStoreError::StorageError(format!("Pinecone query parse error: {}", e))
220        })?;
221
222        let results = query_resp
223            .matches
224            .into_iter()
225            .map(|m| {
226                let content = m
227                    .metadata
228                    .as_ref()
229                    .and_then(|md| md.get("content").cloned())
230                    .unwrap_or_default();
231                let doc = Document {
232                    content,
233                    metadata: m.metadata.unwrap_or_default(),
234                    id: Some(m.id.clone()),
235                };
236                SearchResult {
237                    document: doc,
238                    score: m.score as f32,
239                }
240            })
241            .collect();
242
243        Ok(results)
244    }
245
246    async fn get_document(&self, _id: &str) -> Result<Option<Document>, VectorStoreError> {
247        // Pinecone HTTP API doesn't support direct fetch by ID in the basic plan.
248        // Use similarity_search with the ID as metadata filter instead.
249        Err(VectorStoreError::StorageError(
250            "Pinecone does not support direct document fetch by ID via HTTP API".to_string(),
251        ))
252    }
253
254    async fn get_embedding(&self, _id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
255        Err(VectorStoreError::StorageError(
256            "Pinecone does not support direct embedding fetch by ID via HTTP API".to_string(),
257        ))
258    }
259
260    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
261        self.delete(&[id.to_string()])
262            .await
263            .map_err(VectorStoreError::StorageError)
264    }
265
266    async fn count(&self) -> usize {
267        // Q4: 真实现 —— 通过 describe_index_stats 读取 totalVectorCount,不再写死 0。
268        // trait 签名返回 usize,网络失败时退化为 0 并记录日志(不抛错)。
269        match self.describe_index_stats().await {
270            Ok(stats) => stats.total_vector_count,
271            Err(e) => {
272                log::warn!("Pinecone count 失败,按 0 处理: {}", e);
273                0
274            }
275        }
276    }
277
278    async fn clear(&self) -> Result<(), VectorStoreError> {
279        Err(VectorStoreError::StorageError(
280            "Pinecone does not support clearing all vectors via HTTP API. Delete by namespace or IDs instead.".to_string()
281        ))
282    }
283}
284
285#[derive(Deserialize)]
286struct QueryResponse {
287    matches: Vec<QueryMatch>,
288}
289
290#[derive(Deserialize)]
291struct QueryMatch {
292    id: String,
293    #[allow(dead_code)]
294    score: f64,
295    metadata: Option<HashMap<String, String>>,
296}
297
298/// Pinecone `describe_index_stats` 响应(只需我们关心的字段)。
299///
300/// Q4: 提供给 [`PineconeStore::describe_index_stats`],供调用方读取真实向量总数,
301/// 也是 [`VectorStore::count`](crate::VectorStore::count) 的数据来源。
302#[derive(Deserialize)]
303pub struct PineconeIndexStats {
304    #[serde(default)]
305    pub total_vector_count: usize,
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    fn doc(id: &str, content: &str) -> Document {
313        Document {
314            content: content.to_string(),
315            metadata: HashMap::new(),
316            id: Some(id.to_string()),
317        }
318    }
319
320    #[test]
321    fn test_build_upsert_body() {
322        let docs = vec![doc("1", "hello"), doc("2", "world")];
323        let vectors = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
324        let body = PineconeStore::build_upsert_body(&docs, &vectors);
325        let vectors_arr = body.get("vectors").unwrap().as_array().unwrap();
326        assert_eq!(vectors_arr.len(), 2);
327        assert_eq!(vectors_arr[0]["id"], "1");
328        assert_eq!(vectors_arr[0]["values"][0], 1.0);
329    }
330
331    #[test]
332    fn test_build_upsert_body_generates_id_if_missing() {
333        let mut d = doc("", "x");
334        d.id = None;
335        let body = PineconeStore::build_upsert_body(&[d], &[vec![0.1]]);
336        let id = body["vectors"][0]["id"].as_str().unwrap();
337        assert!(!id.is_empty());
338    }
339
340    #[test]
341    fn test_build_query_body() {
342        let body = PineconeStore::build_query_body(&[1.0, 2.0, 3.0], 5);
343        assert_eq!(body["topK"], 5);
344        assert_eq!(body["includeMetadata"], true);
345        assert_eq!(body["vector"][2], 3.0);
346    }
347
348    #[test]
349    fn test_new() {
350        let store = PineconeStore::new("key", "https://index.svc.env.pinecone.io");
351        assert_eq!(store.host, "https://index.svc.env.pinecone.io");
352    }
353}
lc_vector_stores/pinecone.rs

lc_vector_stores/
pinecone.rs
