Implement embed for inference (#44)

## Problem

The `embed` endpoint was not implemented.

## Solution

This PR implements the `embed` endpoint for inference.

## Type of Change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan

All new and existing test cases pass.

Author: emily-emily

pinecone_sdk/src/pinecone/inference.rs

@@ -0,0 +1,123 @@
+use openapi::apis::inference_api;
+use openapi::models::EmbedRequest;
+
+use crate::pinecone::PineconeClient;
+use crate::utils::errors::PineconeError;
+
+pub use openapi::models::{EmbedRequestParameters, EmbeddingsList};
+
+impl PineconeClient {
+    /// Generate embeddings for input data.
+    ///
+    /// # Arguments
+    /// * `model: &str` - The model to use for embedding.
+    /// * `parameters: Option<EmbedRequestParameters>` - Model-specific parameters.
+    /// * `inputs: &Vec<&str>` - The input data to embed.
+    pub async fn embed(
+        &self,
+        model: &str,
+        parameters: Option<EmbedRequestParameters>,
+        inputs: &Vec<&str>,
+    ) -> Result<EmbeddingsList, PineconeError> {
+        let request = EmbedRequest {
+            model: model.to_string(),
+            parameters: parameters.map(|x| Box::new(x)),
+            inputs: inputs
+                .iter()
+                .map(|&x| openapi::models::EmbedRequestInputsInner {
+                    text: Some(x.to_string()),
+                })
+                .collect(),
+        };
+
+        let res = inference_api::embed(&self.openapi_config, Some(request))
+            .await
+            .map_err(|e| PineconeError::from(e))?;
+
+        Ok(res)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use httpmock::prelude::*;
+    use tokio;
+
+    #[tokio::test]
+    async fn test_embed() -> Result<(), PineconeError> {
+        let server = MockServer::start();
+
+        let mock = server.mock(|when, then| {
+            when.method(POST).path("/embed");
+            then.status(200)
+                .header("content-type", "application/json")
+                .body(
+                    r#"
+                    {
+                        "model": "multilingual-e5-large",
+                        "data": [
+                          {"values": [0.01849365234375, -0.003767013549804688, -0.037261962890625, 0.0222930908203125]}
+                        ],
+                        "usage": {"total_tokens": 1632}
+                    }
+                    "#,
+                );
+        });
+
+        let client = PineconeClient::new(None, Some(server.base_url().as_str()), None, None)?;
+        let response = client
+            .embed("multilingual-e5-large", None, &vec!["Hello, world!"])
+            .await
+            .expect("Failed to embed");
+
+        mock.assert();
+
+        assert_eq!(response.model.unwrap(), "multilingual-e5-large");
+        assert_eq!(response.data.unwrap().len(), 1);
+        assert_eq!(response.usage.unwrap().total_tokens, Some(1632));
+
+        Ok(())
+    }
+
+    async fn test_embed_invalid_arguments() -> Result<(), PineconeError> {
+        let server = MockServer::start();
+
+        let mock = server.mock(|when, then| {
+            when.method(POST).path("/embed");
+            then.status(400)
+                .header("content-type", "application/json")
+                .body(
+                    r#"
+                    {
+                        "error": {
+                          "code": "INVALID_ARGUMENT",
+                          "message": "Invalid parameter value input_type='bad-parameter' for model 'multilingual-e5-large', must be one of [query, passage]"
+                        },
+                        "status": 400
+                      }
+                    "#,
+                );
+        });
+
+        let client = PineconeClient::new(None, None, None, None).unwrap();
+
+        let parameters = EmbedRequestParameters {
+            input_type: Some("bad-parameter".to_string()),
+            truncate: Some("bad-parameter".to_string()),
+        };
+
+        let _ = client
+            .embed(
+                "multilingual-e5-large",
+                Some(parameters),
+                &vec!["Hello, world!"],
+            )
+            .await
+            .expect_err("Expected to fail embedding with invalid arguments");
+
+        mock.assert();
+
+        Ok(())
+    }
+}

pinecone_sdk/src/pinecone/mod.rs

@@ -15,6 +15,9 @@ pub mod control;
 /// Data plane module.
 pub mod data;
 
+/// Inference module.
+pub mod inference;
+
 /// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK.
 #[derive(Debug, Clone)]
 pub struct PineconeClient {

pinecone_sdk/src/utils/errors.rs

@@ -131,6 +131,12 @@ pub enum PineconeError {
         /// Error status
         status: tonic::Status,
     },
+
+    /// InferenceError: Failed to perform an inference operation.
+    InferenceError {
+        /// Error status
+        status: tonic::Status,
+    },
 }
 
 // Implement the conversion from OpenApiError to PineconeError for CreateIndexError.
@@ -198,67 +204,70 @@ impl std::fmt::Display for PineconeError {
                     "Unknown response error: status: {}, message: {}",
                     status, message
                 )
-            },
+            }
             PineconeError::ResourceAlreadyExistsError { source } => {
                 write!(f, "Resource already exists error: {}", source)
-            },
+            }
             PineconeError::UnprocessableEntityError { source } => {
                 write!(f, "Unprocessable entity error: {}", source)
-            },
+            }
             PineconeError::PendingCollectionError { source } => {
                 write!(f, "Pending collection error: {}", source)
-            },
+            }
             PineconeError::InternalServerError { source } => {
                 write!(f, "Internal server error: {}", source)
-            },
+            }
             PineconeError::ReqwestError { source } => {
                 write!(f, "Reqwest error: {}", source.to_string())
-            },
+            }
             PineconeError::SerdeError { source } => {
                 write!(f, "Serde error: {}", source.to_string())
-            },
+            }
             PineconeError::IoError { message } => {
                 write!(f, "IO error: {}", message)
-            },
+            }
             PineconeError::BadRequestError { source } => {
                 write!(f, "Bad request error: {}", source)
-            },
+            }
             PineconeError::UnauthorizedError { source } => {
                 write!(f, "Unauthorized error: status: {}", source)
-            },
+            }
             PineconeError::PodQuotaExceededError { source } => {
                 write!(f, "Pod quota exceeded error: {}", source)
-            },
+            }
             PineconeError::CollectionsQuotaExceededError { source } => {
                 write!(f, "Collections quota exceeded error: {}", source)
-            },
+            }
             PineconeError::InvalidCloudError { source } => {
                 write!(f, "Invalid cloud error: status: {}", source)
-            },
+            }
             PineconeError::InvalidRegionError { source } => {
                 write!(f, "Invalid region error: {}", source)
-            },
+            }
             PineconeError::CollectionNotFoundError { source } => {
                 write!(f, "Collection not found error: {}", source)
-            },
+            }
             PineconeError::IndexNotFoundError { source } => {
                 write!(f, "Index not found error: status: {}", source)
-            },
+            }
             PineconeError::APIKeyMissingError { message } => {
                 write!(f, "API key missing error: {}", message)
-            },
+            }
             PineconeError::InvalidHeadersError { message } => {
                 write!(f, "Invalid headers error: {}", message)
-            },
+            }
             PineconeError::TimeoutError { message } => {
                 write!(f, "Timeout error: {}", message)
-            },
+            }
             PineconeError::ConnectionError { source } => {
                 write!(f, "Connection error: {}", source)
             }
             PineconeError::DataPlaneError { status } => {
                 write!(f, "Data plane error: {}", status)
             }
+            PineconeError::InferenceError { status } => {
+                write!(f, "Inference error: {}", status)
+            }
         }
     }
 }
@@ -290,6 +299,7 @@ impl std::error::Error for PineconeError {
             PineconeError::TimeoutError { message: _ } => None,
             PineconeError::ConnectionError { source } => Some(source.as_ref()),
             PineconeError::DataPlaneError { status } => Some(status),
+            PineconeError::InferenceError { status } => Some(status),
         }
     }
 }

pinecone_sdk/tests/integration_test_inference.rs

@@ -0,0 +1,60 @@
+use pinecone_sdk::pinecone::inference::EmbedRequestParameters;
+use pinecone_sdk::pinecone::PineconeClient;
+use pinecone_sdk::utils::errors::PineconeError;
+
+#[tokio::test]
+async fn test_embed() -> Result<(), PineconeError> {
+    let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+
+    let parameters = EmbedRequestParameters {
+        input_type: Some("query".to_string()),
+        truncate: Some("END".to_string()),
+    };
+
+    let response = pinecone
+        .embed(
+            "multilingual-e5-large",
+            Some(parameters),
+            &vec!["Hello, world!"],
+        )
+        .await
+        .expect("Failed to embed");
+
+    assert_eq!(response.model.unwrap(), "multilingual-e5-large");
+    assert_eq!(response.data.unwrap().len(), 1);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_embed_invalid_model() -> Result<(), PineconeError> {
+    let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+
+    let _ = pinecone
+        .embed("invalid-model", None, &vec!["Hello, world!"])
+        .await
+        .expect_err("Expected to fail embedding with invalid model");
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_embed_invalid_parameters() -> Result<(), PineconeError> {
+    let pinecone = PineconeClient::new(None, None, None, None).unwrap();
+
+    let parameters = EmbedRequestParameters {
+        input_type: Some("bad-parameter".to_string()),
+        truncate: Some("bad-parameter".to_string()),
+    };
+
+    let _ = pinecone
+        .embed(
+            "multilingual-e5-large",
+            Some(parameters),
+            &vec!["Hello, world!"],
+        )
+        .await
+        .expect_err("Expected to fail embedding with invalid model parameters");
+
+    Ok(())
+}