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::{
9    Document, Embeddings, FilterOp, MetadataFilter, SearchResult, VectorStore, VectorStoreError,
10};
11
12/// Pinecone vector store client
13pub struct PineconeStore {
14    api_key: String,
15    host: String,
16    client: reqwest::Client,
17}
18
19impl PineconeStore {
20    /// Create a Pinecone client.
21    ///
22    /// `host` format: `https://{index-name}.svc.{environment}.pinecone.io`
23    pub fn new(api_key: impl Into<String>, host: impl Into<String>) -> Self {
24        Self {
25            api_key: api_key.into(),
26            host: host.into(),
27            client: reqwest::Client::new(),
28        }
29    }
30
31    /// Build upsert request body (pure function, convenient for testing).
32    pub fn build_upsert_body(docs: &[Document], vectors: &[Vec<f32>]) -> serde_json::Value {
33        let vectors_json: Vec<serde_json::Value> = docs
34            .iter()
35            .zip(vectors.iter())
36            .map(|(doc, vec)| {
37                serde_json::json!({
38                    "id": doc.id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
39                    "values": vec,
40                    "metadata": doc.metadata,
41                })
42            })
43            .collect();
44        serde_json::json!({ "vectors": vectors_json })
45    }
46
47    /// Build query request body (pure function, convenient for testing).
48    pub fn build_query_body(query_vec: &[f32], top_k: usize) -> serde_json::Value {
49        serde_json::json!({
50            "vector": query_vec,
51            "topK": top_k,
52            "includeMetadata": true,
53        })
54    }
55
56    /// Build query request body with a metadata filter (pure function, convenient for testing).
57    ///
58    /// S3: extends [`build_query_body`](Self::build_query_body) with a Pinecone `filter` field,
59    /// where [`filter_to_pinecone`] translates [`MetadataFilter`] → Pinecone query syntax
60    /// (field name → `$op` value, combinations via `$and`/`$or`).
61    pub fn build_query_body_filtered(
62        query_vec: &[f32],
63        top_k: usize,
64        filter: &MetadataFilter,
65    ) -> serde_json::Value {
66        let mut body = Self::build_query_body(query_vec, top_k);
67        body["filter"] = filter_to_pinecone(filter);
68        body
69    }
70
71    /// Upsert documents (auto-embed).
72    pub async fn upsert(
73        &self,
74        docs: &[Document],
75        embeddings: &dyn Embeddings,
76    ) -> Result<(), VectorStoreError> {
77        let texts: Vec<&str> = docs.iter().map(|d| d.content.as_str()).collect();
78        let vectors = embeddings
79            .embed_documents(&texts)
80            .await
81            .map_err(|e| VectorStoreError::EmbeddingError(e.to_string()))?;
82        let body = Self::build_upsert_body(docs, &vectors);
83        let url = format!("{}/vectors/upsert", self.host);
84        let resp = self
85            .client
86            .post(&url)
87            .header("Api-Key", &self.api_key)
88            .json(&body)
89            .send()
90            .await
91            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
92        if !resp.status().is_success() {
93            return Err(VectorStoreError::ConnectionError(format!(
94                "Pinecone upsert error: {}",
95                resp.status()
96            )));
97        }
98        Ok(())
99    }
100
101    /// Query similar documents.
102    pub async fn query(
103        &self,
104        query_vec: Vec<f32>,
105        top_k: usize,
106    ) -> Result<Vec<Document>, VectorStoreError> {
107        let body = Self::build_query_body(&query_vec, top_k);
108        let url = format!("{}/query", self.host);
109        let resp = self
110            .client
111            .post(&url)
112            .header("Api-Key", &self.api_key)
113            .json(&body)
114            .send()
115            .await
116            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
117        if !resp.status().is_success() {
118            return Err(VectorStoreError::ConnectionError(format!(
119                "Pinecone query error: {}",
120                resp.status()
121            )));
122        }
123        let query_resp: QueryResponse = resp
124            .json()
125            .await
126            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
127        let result = query_resp
128            .matches
129            .into_iter()
130            .map(|m| {
131                let content = m
132                    .metadata
133                    .as_ref()
134                    .and_then(|md| md.get("content").and_then(|v| v.as_str()))
135                    .unwrap_or_default()
136                    .to_string();
137                Document {
138                    content,
139                    metadata: m.metadata.unwrap_or_default(),
140                    id: Some(m.id),
141                }
142            })
143            .collect();
144        Ok(result)
145    }
146
147    /// Reads index statistics (the only reliable source for a true count).
148    ///
149    /// The Pinecone REST API provides `describe_index_stats`, returning `totalVectorCount`.
150    pub async fn describe_index_stats(&self) -> Result<PineconeIndexStats, VectorStoreError> {
151        let url = format!("{}/describe_index_stats", self.host);
152        let resp = self
153            .client
154            .post(&url)
155            .header("Api-Key", &self.api_key)
156            .send()
157            .await
158            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
159        if !resp.status().is_success() {
160            return Err(VectorStoreError::ConnectionError(format!(
161                "Pinecone describe_index_stats error: {}",
162                resp.status()
163            )));
164        }
165        resp.json()
166            .await
167            .map_err(|e| VectorStoreError::StorageError(e.to_string()))
168    }
169
170    /// Delete by IDs.
171    pub async fn delete(&self, ids: &[String]) -> Result<(), VectorStoreError> {
172        let url = format!("{}/vectors/delete", self.host);
173        let body = serde_json::json!({ "ids": ids });
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| VectorStoreError::ConnectionError(e.to_string()))?;
182        if !resp.status().is_success() {
183            return Err(VectorStoreError::ConnectionError(format!(
184                "Pinecone delete error: {}",
185                resp.status()
186            )));
187        }
188        Ok(())
189    }
190
191    /// POSTs to `/query` and parses the result (shared by plain and filtered retrieval).
192    async fn query_impl(
193        &self,
194        body: serde_json::Value,
195    ) -> Result<Vec<SearchResult>, VectorStoreError> {
196        let url = format!("{}/query", self.host);
197        let resp = self
198            .client
199            .post(&url)
200            .header("Api-Key", &self.api_key)
201            .json(&body)
202            .send()
203            .await
204            .map_err(|e| VectorStoreError::StorageError(format!("Pinecone query failed: {}", e)))?;
205
206        if !resp.status().is_success() {
207            return Err(VectorStoreError::StorageError(format!(
208                "Pinecone query HTTP error: {}",
209                resp.status()
210            )));
211        }
212
213        let query_resp: QueryResponse = resp.json().await.map_err(|e| {
214            VectorStoreError::StorageError(format!("Pinecone query parse error: {}", e))
215        })?;
216
217        Ok(query_resp
218            .matches
219            .into_iter()
220            .map(|m| {
221                let content = m
222                    .metadata
223                    .as_ref()
224                    .and_then(|md| md.get("content").and_then(|v| v.as_str()))
225                    .unwrap_or_default()
226                    .to_string();
227                let doc = Document {
228                    content,
229                    metadata: m.metadata.unwrap_or_default(),
230                    id: Some(m.id.clone()),
231                };
232                SearchResult {
233                    document: doc,
234                    score: m.score as f32,
235                }
236            })
237            .collect())
238    }
239}
240
241/// S3: translates [`MetadataFilter`] → Pinecone `filter` syntax.
242///
243/// A single-field condition becomes `{ key: { "$op": value } }` (Pinecone supports
244/// `$eq $ne $gt $gte $lt $lte $in $nin`); AND/OR combinations become
245/// `{ "$and": [...] }` / `{ "$or": [...] }`. The types map one-to-one with [`FilterOp`],
246/// with no inexpressible construct.
247pub fn filter_to_pinecone(filter: &MetadataFilter) -> serde_json::Value {
248    fn op_str(op: FilterOp) -> &'static str {
249        match op {
250            FilterOp::Eq => "$eq",
251            FilterOp::Ne => "$ne",
252            FilterOp::Gt => "$gt",
253            FilterOp::Gte => "$gte",
254            FilterOp::Lt => "$lt",
255            FilterOp::Lte => "$lte",
256            FilterOp::In => "$in",
257            FilterOp::Nin => "$nin",
258        }
259    }
260    match filter {
261        MetadataFilter::Field { key, op, value } => {
262            serde_json::json!({ key.clone(): { op_str(*op): value.clone() } })
263        }
264        MetadataFilter::And(filters) => {
265            let items: Vec<serde_json::Value> = filters.iter().map(filter_to_pinecone).collect();
266            serde_json::json!({ "$and": items })
267        }
268        MetadataFilter::Or(filters) => {
269            let items: Vec<serde_json::Value> = filters.iter().map(filter_to_pinecone).collect();
270            serde_json::json!({ "$or": items })
271        }
272    }
273}
274
275#[async_trait]
276impl VectorStore for PineconeStore {
277    async fn add_documents(
278        &self,
279        documents: Vec<Document>,
280        embeddings: Vec<Vec<f32>>,
281    ) -> Result<Vec<String>, VectorStoreError> {
282        let ids: Vec<String> = documents
283            .iter()
284            .map(|d| {
285                d.id.clone()
286                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
287            })
288            .collect();
289
290        // Build upsert body with pre-computed embeddings
291        let body = Self::build_upsert_body(&documents, &embeddings);
292        let url = format!("{}/vectors/upsert", self.host);
293        let resp = self
294            .client
295            .post(&url)
296            .header("Api-Key", &self.api_key)
297            .json(&body)
298            .send()
299            .await
300            .map_err(|e| {
301                VectorStoreError::StorageError(format!("Pinecone upsert failed: {}", e))
302            })?;
303
304        if !resp.status().is_success() {
305            return Err(VectorStoreError::StorageError(format!(
306                "Pinecone upsert HTTP error: {}",
307                resp.status()
308            )));
309        }
310
311        Ok(ids)
312    }
313
314    async fn similarity_search(
315        &self,
316        query_embedding: &[f32],
317        k: usize,
318    ) -> Result<Vec<SearchResult>, VectorStoreError> {
319        let body = Self::build_query_body(query_embedding, k);
320        self.query_impl(body).await
321    }
322
323    /// S3: similarity search with metadata filtering — filtering is delegated to the server (native Pinecone filter syntax).
324    async fn similarity_search_with_filter(
325        &self,
326        query_embedding: &[f32],
327        k: usize,
328        filter: Option<&MetadataFilter>,
329    ) -> Result<Vec<SearchResult>, VectorStoreError> {
330        let body = match filter {
331            Some(f) => Self::build_query_body_filtered(query_embedding, k, f),
332            None => Self::build_query_body(query_embedding, k),
333        };
334        self.query_impl(body).await
335    }
336
337    async fn get_document(&self, _id: &str) -> Result<Option<Document>, VectorStoreError> {
338        // Pinecone HTTP API doesn't support direct fetch by ID in the basic plan.
339        // Use similarity_search with the ID as metadata filter instead.
340        Err(VectorStoreError::StorageError(
341            "Pinecone does not support direct document fetch by ID via HTTP API".to_string(),
342        ))
343    }
344
345    async fn get_embedding(&self, _id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
346        Err(VectorStoreError::StorageError(
347            "Pinecone does not support direct embedding fetch by ID via HTTP API".to_string(),
348        ))
349    }
350
351    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
352        self.delete(&[id.to_string()]).await
353    }
354
355    async fn count(&self) -> usize {
356        // Q4: real implementation — reads totalVectorCount via describe_index_stats, no longer hardcoded to 0.
357        // the trait signature returns usize; on network failure it degrades to 0 and logs (no error raised).
358        match self.describe_index_stats().await {
359            Ok(stats) => stats.total_vector_count,
360            Err(e) => {
361                log::warn!("Pinecone count failed, treating as 0: {}", e);
362                0
363            }
364        }
365    }
366
367    async fn clear(&self) -> Result<(), VectorStoreError> {
368        Err(VectorStoreError::StorageError(
369            "Pinecone does not support clearing all vectors via HTTP API. Delete by namespace or IDs instead.".to_string()
370        ))
371    }
372}
373
374#[derive(Deserialize)]
375struct QueryResponse {
376    matches: Vec<QueryMatch>,
377}
378
379#[derive(Deserialize)]
380struct QueryMatch {
381    id: String,
382    score: f64,
383    metadata: Option<HashMap<String, serde_json::Value>>,
384}
385
386/// Response of Pinecone `describe_index_stats` (only the fields we care about).
387///
388/// Q4: returned by [`PineconeStore::describe_index_stats`], letting callers read the true
389/// total vector count; it is also the data source for [`VectorStore::count`].
390#[derive(Deserialize)]
391pub struct PineconeIndexStats {
392    /// Total number of vectors in the index
393    #[serde(default)]
394    pub total_vector_count: usize,
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400
401    fn doc(id: &str, content: &str) -> Document {
402        Document {
403            content: content.to_string(),
404            metadata: HashMap::new(),
405            id: Some(id.to_string()),
406        }
407    }
408
409    #[test]
410    fn test_build_upsert_body() {
411        let docs = vec![doc("1", "hello"), doc("2", "world")];
412        let vectors = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
413        let body = PineconeStore::build_upsert_body(&docs, &vectors);
414        let vectors_arr = body.get("vectors").unwrap().as_array().unwrap();
415        assert_eq!(vectors_arr.len(), 2);
416        assert_eq!(vectors_arr[0]["id"], "1");
417        assert_eq!(vectors_arr[0]["values"][0], 1.0);
418    }
419
420    #[test]
421    fn test_build_upsert_body_generates_id_if_missing() {
422        let mut d = doc("", "x");
423        d.id = None;
424        let body = PineconeStore::build_upsert_body(&[d], &[vec![0.1]]);
425        let id = body["vectors"][0]["id"].as_str().unwrap();
426        assert!(!id.is_empty());
427    }
428
429    #[test]
430    fn test_build_query_body() {
431        let body = PineconeStore::build_query_body(&[1.0, 2.0, 3.0], 5);
432        assert_eq!(body["topK"], 5);
433        assert_eq!(body["includeMetadata"], true);
434        assert_eq!(body["vector"][2], 3.0);
435    }
436
437    #[test]
438    fn test_new() {
439        let store = PineconeStore::new("key", "https://index.svc.env.pinecone.io");
440        assert_eq!(store.host, "https://index.svc.env.pinecone.io");
441    }
442
443    /// S3: single-field condition → Pinecone `{ key: { "$op": value } }`.
444    #[test]
445    fn test_filter_to_pinecone_field_ops() {
446        assert_eq!(
447            filter_to_pinecone(&MetadataFilter::field("lang", FilterOp::Eq, "rust")),
448            serde_json::json!({ "lang": { "$eq": "rust" } })
449        );
450        assert_eq!(
451            filter_to_pinecone(&MetadataFilter::field("year", FilterOp::Gte, 2020)),
452            serde_json::json!({ "year": { "$gte": 2020 } })
453        );
454        assert_eq!(
455            filter_to_pinecone(&MetadataFilter::field("tags", FilterOp::Nin, vec!["blog"])),
456            serde_json::json!({ "tags": { "$nin": ["blog"] } })
457        );
458    }
459
460    /// S3: AND/OR combination → nested `$and`/`$or`.
461    #[test]
462    fn test_filter_to_pinecone_and_or() {
463        let f = MetadataFilter::and(vec![
464            MetadataFilter::field("lang", FilterOp::Eq, "rust"),
465            MetadataFilter::or(vec![
466                MetadataFilter::field("year", FilterOp::Gte, 2020),
467                MetadataFilter::field("tags", FilterOp::In, vec!["ml"]),
468            ]),
469        ]);
470        assert_eq!(
471            filter_to_pinecone(&f),
472            serde_json::json!({
473                "$and": [
474                    { "lang": { "$eq": "rust" } },
475                    { "$or": [
476                        { "year": { "$gte": 2020 } },
477                        { "tags": { "$in": ["ml"] } }
478                    ]}
479                ]
480            })
481        );
482    }
483
484    /// S3: the filtered query body = the plain query body + a filter field.
485    #[test]
486    fn test_build_query_body_filtered() {
487        let f = MetadataFilter::field("lang", FilterOp::Eq, "rust");
488        let body = PineconeStore::build_query_body_filtered(&[1.0, 2.0], 5, &f);
489        assert_eq!(body["topK"], 5);
490        assert_eq!(body["includeMetadata"], true);
491        assert_eq!(
492            body["filter"],
493            serde_json::json!({ "lang": { "$eq": "rust" } })
494        );
495    }
496}
lc_vector_stores/pinecone.rs

lc_vector_stores/
pinecone.rs
