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

lc_vector_stores/
pinecone.rs
