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    /// Delete by IDs.
117    pub async fn delete(&self, ids: &[String]) -> Result<(), String> {
118        let url = format!("{}/vectors/delete", self.host);
119        let body = serde_json::json!({ "ids": ids });
120        let resp = self
121            .client
122            .post(&url)
123            .header("Api-Key", &self.api_key)
124            .json(&body)
125            .send()
126            .await
127            .map_err(|e| e.to_string())?;
128        if !resp.status().is_success() {
129            return Err(format!("Pinecone delete error: {}", resp.status()));
130        }
131        Ok(())
132    }
133}
134
135#[async_trait]
136impl VectorStore for PineconeStore {
137    async fn add_documents(
138        &self,
139        documents: Vec<Document>,
140        embeddings: Vec<Vec<f32>>,
141    ) -> Result<Vec<String>, VectorStoreError> {
142        let ids: Vec<String> = documents
143            .iter()
144            .map(|d| {
145                d.id.clone()
146                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
147            })
148            .collect();
149
150        // Build upsert body with pre-computed embeddings
151        let body = Self::build_upsert_body(&documents, &embeddings);
152        let url = format!("{}/vectors/upsert", self.host);
153        let resp = self
154            .client
155            .post(&url)
156            .header("Api-Key", &self.api_key)
157            .json(&body)
158            .send()
159            .await
160            .map_err(|e| {
161                VectorStoreError::StorageError(format!("Pinecone upsert failed: {}", e))
162            })?;
163
164        if !resp.status().is_success() {
165            return Err(VectorStoreError::StorageError(format!(
166                "Pinecone upsert HTTP error: {}",
167                resp.status()
168            )));
169        }
170
171        Ok(ids)
172    }
173
174    async fn similarity_search(
175        &self,
176        query_embedding: &[f32],
177        k: usize,
178    ) -> Result<Vec<SearchResult>, VectorStoreError> {
179        let body = Self::build_query_body(query_embedding, k);
180        let url = format!("{}/query", self.host);
181        let resp = self
182            .client
183            .post(&url)
184            .header("Api-Key", &self.api_key)
185            .json(&body)
186            .send()
187            .await
188            .map_err(|e| VectorStoreError::StorageError(format!("Pinecone query failed: {}", e)))?;
189
190        if !resp.status().is_success() {
191            return Err(VectorStoreError::StorageError(format!(
192                "Pinecone query HTTP error: {}",
193                resp.status()
194            )));
195        }
196
197        let query_resp: QueryResponse = resp.json().await.map_err(|e| {
198            VectorStoreError::StorageError(format!("Pinecone query parse error: {}", e))
199        })?;
200
201        let results = query_resp
202            .matches
203            .into_iter()
204            .map(|m| {
205                let content = m
206                    .metadata
207                    .as_ref()
208                    .and_then(|md| md.get("content").cloned())
209                    .unwrap_or_default();
210                let doc = Document {
211                    content,
212                    metadata: m.metadata.unwrap_or_default(),
213                    id: Some(m.id.clone()),
214                };
215                SearchResult {
216                    document: doc,
217                    score: m.score as f32,
218                }
219            })
220            .collect();
221
222        Ok(results)
223    }
224
225    async fn get_document(&self, _id: &str) -> Result<Option<Document>, VectorStoreError> {
226        // Pinecone HTTP API doesn't support direct fetch by ID in the basic plan.
227        // Use similarity_search with the ID as metadata filter instead.
228        Err(VectorStoreError::StorageError(
229            "Pinecone does not support direct document fetch by ID via HTTP API".to_string(),
230        ))
231    }
232
233    async fn get_embedding(&self, _id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
234        Err(VectorStoreError::StorageError(
235            "Pinecone does not support direct embedding fetch by ID via HTTP API".to_string(),
236        ))
237    }
238
239    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
240        self.delete(&[id.to_string()])
241            .await
242            .map_err(VectorStoreError::StorageError)
243    }
244
245    async fn count(&self) -> usize {
246        // Pinecone HTTP API doesn't expose a simple count endpoint in the basic API.
247        // Return 0 as a placeholder; use the Pinecone dashboard for accurate counts.
248        0
249    }
250
251    async fn clear(&self) -> Result<(), VectorStoreError> {
252        Err(VectorStoreError::StorageError(
253            "Pinecone does not support clearing all vectors via HTTP API. Delete by namespace or IDs instead.".to_string()
254        ))
255    }
256}
257
258#[derive(Deserialize)]
259struct QueryResponse {
260    matches: Vec<QueryMatch>,
261}
262
263#[derive(Deserialize)]
264struct QueryMatch {
265    id: String,
266    #[allow(dead_code)]
267    score: f64,
268    metadata: Option<HashMap<String, String>>,
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    fn doc(id: &str, content: &str) -> Document {
276        Document {
277            content: content.to_string(),
278            metadata: HashMap::new(),
279            id: Some(id.to_string()),
280        }
281    }
282
283    #[test]
284    fn test_build_upsert_body() {
285        let docs = vec![doc("1", "hello"), doc("2", "world")];
286        let vectors = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
287        let body = PineconeStore::build_upsert_body(&docs, &vectors);
288        let vectors_arr = body.get("vectors").unwrap().as_array().unwrap();
289        assert_eq!(vectors_arr.len(), 2);
290        assert_eq!(vectors_arr[0]["id"], "1");
291        assert_eq!(vectors_arr[0]["values"][0], 1.0);
292    }
293
294    #[test]
295    fn test_build_upsert_body_generates_id_if_missing() {
296        let mut d = doc("", "x");
297        d.id = None;
298        let body = PineconeStore::build_upsert_body(&[d], &[vec![0.1]]);
299        let id = body["vectors"][0]["id"].as_str().unwrap();
300        assert!(!id.is_empty());
301    }
302
303    #[test]
304    fn test_build_query_body() {
305        let body = PineconeStore::build_query_body(&[1.0, 2.0, 3.0], 5);
306        assert_eq!(body["topK"], 5);
307        assert_eq!(body["includeMetadata"], true);
308        assert_eq!(body["vector"][2], 3.0);
309    }
310
311    #[test]
312    fn test_new() {
313        let store = PineconeStore::new("key", "https://index.svc.env.pinecone.io");
314        assert_eq!(store.host, "https://index.svc.env.pinecone.io");
315    }
316}
lc_vector_stores/pinecone.rs

lc_vector_stores/
pinecone.rs
