pinecone_sdk/pinecone/
mod.rs

1use crate::openapi::apis::configuration::ApiKey;
2use crate::openapi::apis::configuration::Configuration;
3use crate::utils::errors::PineconeError;
4use crate::utils::user_agent::get_user_agent;
5use crate::version::API_VERSION;
6use serde_json;
7use std::collections::HashMap;
8
9/// The `PINECONE_API_VERSION_KEY` is the key for the Pinecone API version header.
10pub const PINECONE_API_VERSION_KEY: &str = "X-Pinecone-Api-Version";
11
12/// Control plane module.
13pub mod control;
14
15/// Data plane module.
16pub mod data;
17
18/// Inference module.
19pub mod inference;
20
21/// The `PineconeClientConfig` struct takes in the parameters to configure the Pinecone client.
22#[derive(Default)]
23pub struct PineconeClientConfig {
24    /// Pinecone API key
25    pub api_key: Option<String>,
26    /// The Pinecone controller host
27    pub control_plane_host: Option<String>,
28    /// Additional headers to be included in all requests
29    pub additional_headers: Option<HashMap<String, String>>,
30    /// The source tag
31    pub source_tag: Option<String>,
32}
33
34impl PineconeClientConfig {
35    /// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK.
36    /// It is used to create, delete, and manage your indexes and collections.
37    /// This function constructs a `PineconeClient` struct using the provided configuration.
38    ///
39    /// ### Arguments
40    /// * `api_key: Option<&str>` - The API key used for authentication.
41    /// * `control_plane_host: Option<&str>` - The Pinecone controller host. Default is `https://api.pinecone.io`.
42    /// * `additional_headers: Option<HashMap<String, String>>` - Additional headers to be included in all requests. Expects a HashMap. If no api version header is provided, it will be added.
43    /// * `source_tag: Option<&str>` - A tag to identify the source of the request.
44    ///
45    /// ### Return
46    /// * `Result<PineconeClient, PineconeError>`
47    ///
48    /// ### Configuration with environment variables
49    /// If arguments are not provided, the SDK will attempt to read the following environment variables:
50    /// - `PINECONE_API_KEY`: The API key used for authentication. If not passed as an argument, it will be read from the environment variable.
51    /// - `PINECONE_CONTROLLER_HOST`: The Pinecone controller host. Default is `https://api.pinecone.io`.
52    /// - `PINECONE_ADDITIONAL_HEADERS`: Additional headers to be included in all requests. Expects JSON.
53    ///
54    /// ### Example
55    /// ```no_run
56    /// use pinecone_sdk::pinecone::{PineconeClient, PineconeClientConfig};
57    ///
58    /// // Create a Pinecone client with the API key and controller host.
59    ///
60    /// let config = PineconeClientConfig {
61    ///     api_key: Some("INSERT_API_KEY".to_string()),
62    ///     control_plane_host: Some("INSERT_CONTROLLER_HOST".to_string()),
63    ///     ..Default::default()
64    /// };
65    /// let pinecone: PineconeClient = config.client().expect("Failed to create Pinecone instance");
66    /// ```
67    pub fn client(self) -> Result<PineconeClient, PineconeError> {
68        // get api key
69        let api_key = match self.api_key {
70            Some(key) => key.to_string(),
71            None => match std::env::var("PINECONE_API_KEY") {
72                Ok(key) => key,
73                Err(_) => {
74                    let message =
75                        "API key is not provided as an argument nor as an environment variable";
76                    return Err(PineconeError::APIKeyMissingError {
77                        message: message.to_string(),
78                    });
79                }
80            },
81        };
82
83        let env_controller = std::env::var("PINECONE_CONTROLLER_HOST")
84            .unwrap_or("https://api.pinecone.io".to_string());
85        let controller_host = &self.control_plane_host.clone().unwrap_or(env_controller);
86
87        // get user agent
88        let user_agent = get_user_agent(self.source_tag.as_ref().map(|s| s.as_str()));
89
90        // get additional headers
91        let mut additional_headers =
92            self.additional_headers
93                .unwrap_or(match std::env::var("PINECONE_ADDITIONAL_HEADERS") {
94                    Ok(headers) => match serde_json::from_str(&headers) {
95                        Ok(headers) => headers,
96                        Err(_) => {
97                            let message = "Provided headers are not valid. Expects JSON.";
98                            return Err(PineconeError::InvalidHeadersError {
99                                message: message.to_string(),
100                            });
101                        }
102                    },
103                    Err(_) => HashMap::new(),
104                });
105
106        // add X-Pinecone-Api-Version header if not present
107        // case insensitive
108        if !additional_headers
109            .keys()
110            .any(|k| k.eq_ignore_ascii_case(PINECONE_API_VERSION_KEY))
111        {
112            add_api_version_header(&mut additional_headers);
113        }
114
115        // create reqwest headers
116        let headers: reqwest::header::HeaderMap =
117            (&additional_headers)
118                .try_into()
119                .map_err(|_| PineconeError::InvalidHeadersError {
120                    message: "Provided headers are not valid".to_string(),
121                })?;
122
123        // create reqwest client with headers
124        let client = reqwest::Client::builder()
125            .default_headers(headers)
126            .build()
127            .map_err(|e| PineconeError::ReqwestError { source: e.into() })?;
128
129        let openapi_config = Configuration {
130            base_path: controller_host.to_string(),
131            user_agent: Some(user_agent.to_string()),
132            api_key: Some(ApiKey {
133                prefix: None,
134                key: api_key.clone(),
135            }),
136            client,
137            ..Default::default()
138        };
139
140        // return Pinecone client
141        Ok(PineconeClient {
142            api_key,
143            controller_url: controller_host.to_string(),
144            additional_headers,
145            source_tag: self.source_tag,
146            user_agent: Some(user_agent),
147            openapi_config,
148        })
149    }
150}
151
152/// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK.
153#[derive(Debug, Clone)]
154pub struct PineconeClient {
155    /// Pinecone API key
156    api_key: String,
157    /// The Pinecone controller host
158    controller_url: String,
159    /// Additional headers to be included in all requests
160    additional_headers: HashMap<String, String>,
161    /// The source tag
162    source_tag: Option<String>,
163    /// The user agent
164    user_agent: Option<String>,
165    /// Configuration used for OpenAPI endpoint calls
166    openapi_config: Configuration,
167}
168
169/// Helper function to add the API version header to the headers.
170fn add_api_version_header(headers: &mut HashMap<String, String>) {
171    headers.insert(
172        PINECONE_API_VERSION_KEY.to_string(),
173        API_VERSION.to_string(),
174    );
175}
176
177impl TryFrom<PineconeClientConfig> for PineconeClient {
178    type Error = PineconeError;
179
180    fn try_from(config: PineconeClientConfig) -> Result<Self, Self::Error> {
181        config.client()
182    }
183}
184
185/// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK.
186/// It is used to create, delete, and manage your indexes and collections.
187/// This function constructs a `PineconeClient` struct by attempting to read in environment variables for the required parameters.
188///
189/// ### Return
190/// * `Result<PineconeClient, PineconeError>`
191///
192/// ### Configuration with environment variables
193/// If arguments are not provided, the SDK will attempt to read the following environment variables:
194/// - `PINECONE_API_KEY`: The API key used for authentication. If not passed as an argument, it will be read from the environment variable.
195/// - `PINECONE_CONTROLLER_HOST`: The Pinecone controller host. Default is `https://api.pinecone.io`.
196/// - `PINECONE_ADDITIONAL_HEADERS`: Additional headers to be included in all requests. Expects JSON.
197///
198/// ### Example
199/// ```no_run
200/// use pinecone_sdk::pinecone::PineconeClient;
201///
202/// // Create a Pinecone client with the API key and controller host read from environment variables.
203/// let pinecone: PineconeClient = pinecone_sdk::pinecone::default_client().expect("Failed to create Pinecone instance");
204/// ```
205pub fn default_client() -> Result<PineconeClient, PineconeError> {
206    PineconeClientConfig::default().client()
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use tokio;
213
214    fn empty_headers_with_api_version() -> HashMap<String, String> {
215        let mut headers = HashMap::new();
216        add_api_version_header(&mut headers);
217        headers
218    }
219
220    #[tokio::test]
221    async fn test_arg_api_key() -> Result<(), PineconeError> {
222        let mock_api_key = "mock-arg-api-key";
223        let mock_controller_host = "mock-arg-controller-host";
224
225        let config = PineconeClientConfig {
226            api_key: Some(mock_api_key.to_string()),
227            control_plane_host: Some(mock_controller_host.to_string()),
228            additional_headers: Some(HashMap::new()),
229            source_tag: None,
230        };
231
232        let pinecone = config
233            .client()
234            .expect("Expected to successfully create Pinecone instance");
235
236        assert_eq!(pinecone.api_key, mock_api_key);
237        assert_eq!(pinecone.controller_url, mock_controller_host);
238        assert_eq!(
239            pinecone.additional_headers,
240            empty_headers_with_api_version()
241        );
242        assert_eq!(pinecone.source_tag, None);
243        assert_eq!(
244            pinecone.user_agent,
245            Some("lang=rust; pinecone-rust-client=0.1.0".to_string())
246        );
247
248        Ok(())
249    }
250
251    #[tokio::test]
252    async fn test_env_api_key() -> Result<(), PineconeError> {
253        let mock_api_key = "mock-env-api-key";
254        let mock_controller_host = "mock-arg-controller-host";
255
256        temp_env::with_var("PINECONE_API_KEY", Some(mock_api_key), || {
257            let config = PineconeClientConfig {
258                control_plane_host: Some(mock_controller_host.to_string()),
259                additional_headers: Some(HashMap::new()),
260                ..Default::default()
261            };
262            let pinecone = config
263                .client()
264                .expect("Expected to successfully create Pinecone instance");
265
266            assert_eq!(pinecone.api_key, mock_api_key);
267            assert_eq!(pinecone.controller_url, mock_controller_host);
268            assert_eq!(
269                pinecone.additional_headers,
270                empty_headers_with_api_version()
271            );
272            assert_eq!(pinecone.source_tag, None);
273            assert_eq!(
274                pinecone.user_agent,
275                Some("lang=rust; pinecone-rust-client=0.1.0".to_string())
276            );
277        });
278
279        Ok(())
280    }
281
282    #[tokio::test]
283    async fn test_no_api_key() -> Result<(), PineconeError> {
284        let mock_controller_host = "mock-arg-controller-host";
285
286        temp_env::with_var_unset("PINECONE_API_KEY", || {
287            let config = PineconeClientConfig {
288                control_plane_host: Some(mock_controller_host.to_string()),
289                additional_headers: Some(HashMap::new()),
290                ..Default::default()
291            };
292            let pinecone = config
293                .client()
294                .expect_err("Expected to fail creating Pinecone instance due to missing API key");
295
296            assert!(matches!(pinecone, PineconeError::APIKeyMissingError { .. }));
297        });
298
299        Ok(())
300    }
301
302    #[tokio::test]
303    async fn test_arg_host() -> Result<(), PineconeError> {
304        let mock_api_key = "mock-arg-api-key";
305        let mock_controller_host = "mock-arg-controller-host";
306        let config = PineconeClientConfig {
307            api_key: Some(mock_api_key.to_string()),
308            control_plane_host: Some(mock_controller_host.to_string()),
309            additional_headers: Some(HashMap::new()),
310            source_tag: None,
311        };
312        let pinecone = config
313            .client()
314            .expect("Expected to successfully create Pinecone instance");
315
316        assert_eq!(pinecone.controller_url, mock_controller_host);
317
318        Ok(())
319    }
320
321    #[tokio::test]
322    async fn test_env_host() -> Result<(), PineconeError> {
323        let mock_api_key = "mock-arg-api-key";
324        let mock_controller_host = "mock-env-controller-host";
325
326        temp_env::with_var(
327            "PINECONE_CONTROLLER_HOST",
328            Some(mock_controller_host),
329            || {
330                let config = PineconeClientConfig {
331                    api_key: Some(mock_api_key.to_string()),
332                    additional_headers: Some(HashMap::new()),
333                    ..Default::default()
334                };
335
336                let pinecone = config
337                    .client()
338                    .expect("Expected to successfully create Pinecone instance with env host");
339
340                assert_eq!(pinecone.controller_url, mock_controller_host);
341            },
342        );
343
344        Ok(())
345    }
346
347    #[tokio::test]
348    async fn test_default_host() -> Result<(), PineconeError> {
349        let mock_api_key = "mock-arg-api-key";
350
351        temp_env::with_var_unset("PINECONE_CONTROLLER_HOST", || {
352            let config = PineconeClientConfig {
353                api_key: Some(mock_api_key.to_string()),
354                additional_headers: Some(HashMap::new()),
355                ..Default::default()
356            };
357
358            let pinecone = config.client().expect(
359                "Expected to successfully create Pinecone instance with default controller host",
360            );
361
362            assert_eq!(
363                pinecone.controller_url,
364                "https://api.pinecone.io".to_string()
365            );
366        });
367
368        Ok(())
369    }
370
371    #[tokio::test]
372    async fn test_arg_headers() -> Result<(), PineconeError> {
373        let mock_api_key = "mock-arg-api-key";
374        let mock_controller_host = "mock-arg-controller-host";
375        let mock_headers = HashMap::from([
376            ("argheader1".to_string(), "value1".to_string()),
377            ("argheader2".to_string(), "value2".to_string()),
378        ]);
379
380        let config = PineconeClientConfig {
381            api_key: Some(mock_api_key.to_string()),
382            control_plane_host: Some(mock_controller_host.to_string()),
383            additional_headers: Some(mock_headers.clone()),
384            source_tag: None,
385        };
386        let pinecone = config
387            .client()
388            .expect("Expected to successfully create Pinecone instance");
389
390        let expected_headers = {
391            let mut headers = mock_headers.clone();
392            add_api_version_header(&mut headers);
393            headers
394        };
395
396        assert_eq!(pinecone.additional_headers, expected_headers);
397
398        Ok(())
399    }
400
401    #[tokio::test]
402    async fn test_env_headers() -> Result<(), PineconeError> {
403        let mock_api_key = "mock-arg-api-key";
404        let mock_controller_host = "mock-arg-controller-host";
405        let mock_headers = HashMap::from([
406            ("envheader1".to_string(), "value1".to_string()),
407            ("envheader2".to_string(), "value2".to_string()),
408        ]);
409
410        temp_env::with_var(
411            "PINECONE_ADDITIONAL_HEADERS",
412            Some(serde_json::to_string(&mock_headers).unwrap().as_str()),
413            || {
414                let config = PineconeClientConfig {
415                    api_key: Some(mock_api_key.to_string()),
416                    control_plane_host: Some(mock_controller_host.to_string()),
417                    additional_headers: None,
418                    source_tag: None,
419                };
420
421                let pinecone = config
422                    .client()
423                    .expect("Expected to successfully create Pinecone instance with env headers");
424
425                let expected_headers = {
426                    let mut headers = mock_headers.clone();
427                    add_api_version_header(&mut headers);
428                    headers
429                };
430
431                assert_eq!(pinecone.additional_headers, expected_headers);
432            },
433        );
434
435        Ok(())
436    }
437
438    #[tokio::test]
439    async fn test_invalid_env_headers() -> Result<(), PineconeError> {
440        let mock_api_key = "mock-arg-api-key";
441        let mock_controller_host = "mock-arg-controller-host";
442
443        temp_env::with_var("PINECONE_ADDITIONAL_HEADERS", Some("invalid-json"), || {
444            let config = PineconeClientConfig {
445                api_key: Some(mock_api_key.to_string()),
446                control_plane_host: Some(mock_controller_host.to_string()),
447                additional_headers: None,
448                source_tag: None,
449            };
450            let pinecone = config
451                .client()
452                .expect_err("Expected to fail creating Pinecone instance due to invalid headers");
453
454            assert!(matches!(
455                pinecone,
456                PineconeError::InvalidHeadersError { .. }
457            ));
458        });
459
460        Ok(())
461    }
462
463    #[tokio::test]
464    async fn test_default_headers() -> Result<(), PineconeError> {
465        let mock_api_key = "mock-arg-api-key";
466        let mock_controller_host = "mock-arg-controller-host";
467
468        temp_env::with_var_unset("PINECONE_ADDITIONAL_HEADERS", || {
469            let config = PineconeClientConfig {
470                api_key: Some(mock_api_key.to_string()),
471                control_plane_host: Some(mock_controller_host.to_string()),
472                additional_headers: None,
473                source_tag: None,
474            };
475
476            let pinecone = config
477                .client()
478                .expect("Expected to successfully create Pinecone instance");
479
480            assert_eq!(
481                pinecone.additional_headers,
482                empty_headers_with_api_version()
483            );
484        });
485
486        Ok(())
487    }
488
489    #[tokio::test]
490    async fn test_headers_no_api_version() -> Result<(), PineconeError> {
491        let mock_api_key = "mock-arg-api-key";
492        let mock_controller_host = "mock-arg-controller-host";
493
494        temp_env::with_var_unset("PINECONE_ADDITIONAL_HEADERS", || {
495            let headers = HashMap::from([
496                ("HEADER1".to_string(), "value1".to_string()),
497                ("HEADER2".to_string(), "value2".to_string()),
498            ]);
499
500            let config = PineconeClientConfig {
501                api_key: Some(mock_api_key.to_string()),
502                control_plane_host: Some(mock_controller_host.to_string()),
503                additional_headers: Some(headers.clone()),
504                source_tag: None,
505            };
506
507            let pinecone = config
508                .client()
509                .expect("Expected to successfully create Pinecone instance");
510
511            // expect headers, except with the added API version header
512            let mut expected_headers = headers.clone();
513            expected_headers.insert(
514                PINECONE_API_VERSION_KEY.to_string(),
515                API_VERSION.to_string(),
516            );
517
518            assert_eq!(pinecone.additional_headers, expected_headers);
519        });
520
521        Ok(())
522    }
523
524    #[tokio::test]
525    async fn test_headers_api_version() -> Result<(), PineconeError> {
526        let mock_api_key = "mock-arg-api-key";
527        let mock_controller_host = "mock-arg-controller-host";
528
529        temp_env::with_var_unset("PINECONE_ADDITIONAL_HEADERS", || {
530            let headers = HashMap::from([
531                ("HEADER1".to_string(), "value1".to_string()),
532                ("HEADER2".to_string(), "value2".to_string()),
533                (
534                    PINECONE_API_VERSION_KEY.to_string(),
535                    "mock-api-version".to_string(),
536                ),
537            ]);
538
539            let config = PineconeClientConfig {
540                api_key: Some(mock_api_key.to_string()),
541                control_plane_host: Some(mock_controller_host.to_string()),
542                additional_headers: Some(headers.clone()),
543                source_tag: None,
544            };
545
546            let pinecone = config
547                .client()
548                .expect("Expected to successfully create Pinecone instance");
549
550            assert_eq!(pinecone.additional_headers, headers);
551        });
552
553        Ok(())
554    }
555
556    #[tokio::test]
557    async fn test_headers_api_version_different_casing() -> Result<(), PineconeError> {
558        let mock_api_key = "mock-arg-api-key";
559        let mock_controller_host = "mock-arg-controller-host";
560
561        temp_env::with_var_unset("PINECONE_ADDITIONAL_HEADERS", || {
562            let headers = HashMap::from([
563                ("HEADER1".to_string(), "value1".to_string()),
564                ("HEADER2".to_string(), "value2".to_string()),
565                (
566                    "X-pineCONE-api-version".to_string(),
567                    "mock-api-version".to_string(),
568                ),
569            ]);
570
571            let config = PineconeClientConfig {
572                api_key: Some(mock_api_key.to_string()),
573                control_plane_host: Some(mock_controller_host.to_string()),
574                additional_headers: Some(headers.clone()),
575                source_tag: None,
576            };
577
578            let pinecone = config
579                .client()
580                .expect("Expected to successfully create Pinecone instance");
581
582            assert_eq!(pinecone.additional_headers, headers);
583        });
584
585        Ok(())
586    }
587
588    #[tokio::test]
589    async fn test_arg_overrides_env() -> Result<(), PineconeError> {
590        let mock_arg_api_key = "mock-arg-api-key";
591        let mock_arg_controller_host = "mock-arg-controller-host";
592        let mock_arg_headers = HashMap::from([
593            ("argheader1".to_string(), "value1".to_string()),
594            ("argheader2".to_string(), "value2".to_string()),
595        ]);
596        let mock_env_api_key = "mock-env-api-key";
597        let mock_env_controller_host = "mock-env-controller-host";
598        let mock_env_headers = HashMap::from([
599            ("envheader1".to_string(), "value1".to_string()),
600            ("envheader2".to_string(), "value2".to_string()),
601        ]);
602
603        temp_env::with_vars(
604            [
605                ("PINECONE_API_KEY", Some(mock_env_api_key)),
606                ("PINECONE_CONTROLLER_HOST", Some(mock_env_controller_host)),
607                (
608                    "PINECONE_ADDITIONAL_HEADERS",
609                    Some(serde_json::to_string(&mock_env_headers).unwrap().as_str()),
610                ),
611            ],
612            || {
613                let config = PineconeClientConfig {
614                    api_key: Some(mock_arg_api_key.to_string()),
615                    control_plane_host: Some(mock_arg_controller_host.to_string()),
616                    additional_headers: Some(mock_arg_headers.clone()),
617                    source_tag: None,
618                };
619
620                let pinecone = config
621                    .client()
622                    .expect("Expected to successfully create Pinecone instance");
623
624                let expected_headers = {
625                    let mut headers = mock_arg_headers.clone();
626                    add_api_version_header(&mut headers);
627                    headers
628                };
629
630                assert_eq!(pinecone.api_key, mock_arg_api_key);
631                assert_eq!(pinecone.controller_url, mock_arg_controller_host);
632                assert_eq!(pinecone.additional_headers, expected_headers);
633            },
634        );
635
636        Ok(())
637    }
638}
pinecone_sdk/pinecone/mod.rs

pinecone_sdk/pinecone/
mod.rs
