diff --git a/examples/s3_example.py b/examples/s3_example.py new file mode 100644 index 00000000..d54aadeb --- /dev/null +++ b/examples/s3_example.py @@ -0,0 +1,36 @@ +""" +Example: Fetching files from AWS S3 and embedding them + +This example shows how to: +1. Create an S3Client with AWS credentials +2. Fetch files from S3 buckets +3. Embed the downloaded files +""" + +import embed_anything +from embed_anything import S3Client, EmbeddingModel, WhichModel, TextEmbedConfig +import os + +# Step 1: Create S3Client with explicit credentials +print("Example 1: Create S3Client with credentials") +s3_client = S3Client( + access_key_id="", + secret_access_key="", + region="" +) +print(f"S3Client created: {s3_client}") + +file = s3_client.get_file_from_s3(bucket_name="embed-anything", key="test.txt").save_file() +print(f"File: {file}") + +# Step 2: Embed the file +embedder = EmbeddingModel.from_pretrained_hf( + WhichModel.Jina, model_id="jinaai/jina-embeddings-v2-small-en" +) +embeddings = embedder.embed_file(file, config=TextEmbedConfig(chunk_size=1000, batch_size=32, buffer_size=64, splitting_strategy="sentence")) +print(f"Embeddings: {embeddings}") + +# Step 3: Embed a directory +directory = "bench" +embeddings = embedder.embed_directory(directory, config=TextEmbedConfig(chunk_size=1000, batch_size=32, buffer_size=64, splitting_strategy="sentence")) +print(f"Embeddings: {embeddings}") \ No newline at end of file diff --git a/python/src/lib.rs b/python/src/lib.rs index ef4fb4c4..45b3a5b4 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,5 +1,6 @@ pub mod config; pub mod models; +pub mod s3_client; use embed_anything::embeddings::embed::{TextEmbedder, VisionEmbedder}; use embed_anything::{ self, @@ -900,6 +901,8 @@ fn _embed_anything(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/python/src/s3_client.rs b/python/src/s3_client.rs new file mode 100644 index 00000000..61895124 --- /dev/null +++ b/python/src/s3_client.rs @@ -0,0 +1,128 @@ +use embed_anything::s3_loader::{S3Client as RustS3Client, S3File as RustS3File}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use tokio::runtime::Builder; + +/// Represents a file fetched from S3 +#[pyclass] +pub struct S3File { + pub inner: RustS3File, +} + +#[pymethods] +impl S3File { + /// Save the file to the local filesystem + /// + /// Args: + /// file_path: Optional local file path. If not provided, uses the S3 key basename in current directory + /// + /// Returns: + /// str: The file path where the file was saved + #[pyo3(signature = (file_path=None))] + pub fn save_file(&self, file_path: Option) -> PyResult { + self.inner + .save_file(file_path.as_deref()) + .map_err(|e| PyValueError::new_err(e.to_string())) + } + + /// Get the S3 key + #[getter] + pub fn key(&self) -> String { + self.inner.key().to_string() + } + + /// Get the file contents as bytes + #[getter] + pub fn bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { + PyBytes::new(py, self.inner.as_bytes()) + } + + /// Get the size of the file in bytes + #[getter] + pub fn size(&self) -> usize { + self.inner.as_bytes().len() + } + + fn __repr__(&self) -> String { + format!("S3File({} bytes)", self.size()) + } + + fn __len__(&self) -> usize { + self.size() + } +} + +/// S3Client for fetching files from AWS S3 buckets +#[pyclass] +pub struct S3Client { + pub inner: RustS3Client, +} + +#[pymethods] +impl S3Client { + /// Create a new S3Client with AWS credentials + /// + /// Args: + /// access_key_id: AWS access key ID + /// secret_access_key: AWS secret access key + /// region: AWS region (e.g., "us-east-1") + #[new] + pub fn new( + access_key_id: String, + secret_access_key: String, + region: String, + ) -> Self { + S3Client { + inner: RustS3Client::new( + access_key_id, + secret_access_key, + region, + ), + } + } + + /// Create an S3Client using environment variables + /// + /// Reads credentials from: + /// - AWS_ACCESS_KEY_ID + /// - AWS_SECRET_ACCESS_KEY + /// - AWS_REGION (optional, defaults to "us-east-1") + #[staticmethod] + pub fn from_env() -> PyResult { + let access_key = std::env::var("AWS_ACCESS_KEY_ID") + .map_err(|_| PyValueError::new_err("AWS_ACCESS_KEY_ID environment variable not set"))?; + let secret_key = std::env::var("AWS_SECRET_ACCESS_KEY").map_err(|_| { + PyValueError::new_err("AWS_SECRET_ACCESS_KEY environment variable not set") + })?; + let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()); + + Ok(S3Client { + inner: RustS3Client::new(access_key, secret_key, region), + }) + } + + /// Fetch a file from S3 and return it as an S3File + /// + /// Args: + /// bucket_name: The S3 bucket name + /// key: The S3 object key (file path) + /// + /// Returns: + /// S3File: The file object with methods to save or access bytes + pub fn get_file_from_s3(&self, bucket_name: String, key: String) -> PyResult { + let rt = Builder::new_multi_thread().enable_all().build().unwrap(); + rt.block_on(async { + let file = self + .inner + .get_file_from_s3(&bucket_name, &key) + .await + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(S3File { inner: file }) + }) + } + + fn __repr__(&self) -> String { + format!("S3Client(region='{}')", self.inner.region) + } +} diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ffcae48d..eec50877 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -72,6 +72,9 @@ model2vec-rs = "0.1.1" # Logging log = "0.4" +aws-sdk-s3 = {version = "1.110.0", features = ["rt-tokio"]} +aws-config = {version = "1.8.10", features = ["behavior-version-latest"]} +aws-credential-types = "1.2.8" [dev-dependencies] tempdir = "0.3.7" diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 79132d2a..82490b64 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -77,6 +77,7 @@ pub mod models; #[cfg(feature = "ort")] pub mod reranker; pub mod text_loader; +pub mod s3_loader; use anyhow::{Error, Result}; use config::{ImageEmbedConfig, TextEmbedConfig}; diff --git a/rust/src/s3_loader.rs b/rust/src/s3_loader.rs new file mode 100644 index 00000000..c156f07e --- /dev/null +++ b/rust/src/s3_loader.rs @@ -0,0 +1,253 @@ +//! S3 file loading utilities. +//! +//! Provides functionality for fetching files from AWS S3 buckets +//! for embedding processing. + +use anyhow::{Context, Result}; +use aws_config::BehaviorVersion; +use aws_credential_types::Credentials; +use aws_sdk_s3::Client; + +/// Represents a file fetched from S3 +#[derive(Debug, Clone)] +pub struct S3File { + pub bytes: Vec, + pub key: String, +} + +impl S3File { + /// Create a new S3File from bytes and S3 key + pub fn new(bytes: Vec, key: String) -> Self { + Self { bytes, key } + } + + /// Save the file to the local filesystem + /// + /// # Arguments + /// + /// * `file_path` - Optional local file path. If None, uses the S3 key basename in current directory + /// + /// # Returns + /// + /// A `Result` containing the file path on success + pub fn save_file(&self, file_path: Option<&str>) -> Result { + let path = match file_path { + Some(p) => p.to_string(), + None => { + // Extract filename from S3 key (last part after /) + let filename = self + .key + .split('/') + .last() + .unwrap_or(&self.key) + .to_string(); + filename + } + }; + + std::fs::write(&path, &self.bytes) + .context(format!("Failed to write file to: {}", path))?; + Ok(path) + } + + /// Get the bytes as a slice + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + /// Convert to Vec + pub fn into_bytes(self) -> Vec { + self.bytes + } + + /// Get the S3 key + pub fn key(&self) -> &str { + &self.key + } +} + +/// S3 client for fetching files from AWS S3 buckets +#[derive(Debug, Clone)] +pub struct S3Client { + pub access_key_id: String, + pub secret_access_key: String, + pub region: String, +} + +impl S3Client { + /// Create a new S3Client with the specified credentials and region + /// + /// # Arguments + /// + /// * `access_key_id` - AWS access key ID + /// * `secret_access_key` - AWS secret access key + /// * `region` - AWS region (e.g., "us-east-1") + pub fn new(access_key_id: String, secret_access_key: String, region: String) -> Self { + Self { + access_key_id, + secret_access_key, + region, + } + } + + /// Fetch a file from S3 and return it as an S3File + /// + /// # Arguments + /// + /// * `bucket_name` - The S3 bucket name + /// * `key` - The S3 object key (file path) + /// + /// # Returns + /// + /// A `Result` containing an `S3File` on success + /// + /// # Example + /// + /// ```no_run + /// # use embed_anything::s3_loader::S3Client; + /// # async fn example() -> anyhow::Result<()> { + /// let client = S3Client::new("key".to_string(), "secret".to_string(), "us-east-1".to_string()); + /// let file = client.get_file_from_s3("my-bucket", "file.txt").await?; + /// file.save_file("/tmp/downloaded.txt")?; + /// # Ok(()) + /// # } + /// ``` + pub async fn get_file_from_s3(&self, bucket_name: &str, key: &str) -> Result { + let credentials = Credentials::new( + &self.access_key_id, + &self.secret_access_key, + None, + None, + "static", + ); + + let config = aws_config::defaults(BehaviorVersion::latest()) + .credentials_provider(credentials) + .region(aws_config::Region::new(self.region.clone())) + .load() + .await; + + let client = Client::new(&config); + + let response = client + .get_object() + .bucket(bucket_name) + .key(key) + .send() + .await + .context(format!( + "Failed to fetch object from S3: bucket={}, key={}", + bucket_name, key + ))?; + + let body = response.body.collect().await?; + Ok(S3File::new(body.into_bytes().to_vec(), key.to_string())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + #[ignore] // Requires valid AWS credentials + async fn test_get_file_from_s3() { + // To run this test: + // 1. Set valid AWS credentials + // 2. Create a test bucket and upload a test file + // 3. Run: cargo test --lib s3_loader -- --ignored + + let client = S3Client::new( + "YOUR_ACCESS_KEY".to_string(), + "YOUR_SECRET_KEY".to_string(), + "us-east-1".to_string(), + ); + let file = client + .get_file_from_s3("your-bucket", "test.txt") + .await + .unwrap(); + println!("{}", String::from_utf8_lossy(file.as_bytes())); + } + + #[tokio::test] + #[ignore] // Requires valid AWS credentials + async fn test_download_and_save() { + use tempdir::TempDir; + + let client = S3Client::new( + "YOUR_ACCESS_KEY".to_string(), + "YOUR_SECRET_KEY".to_string(), + "us-east-1".to_string(), + ); + + let file = client + .get_file_from_s3("your-bucket", "test.txt") + .await + .unwrap(); + + let temp_dir = TempDir::new("test").unwrap(); + let path = temp_dir.path().join("downloaded.txt"); + let saved_path = file.save_file(Some(path.to_str().unwrap())).unwrap(); + + assert_eq!(saved_path, path.to_str().unwrap()); + assert!(std::path::Path::new(&saved_path).exists()); + } + + #[test] + fn test_s3file_creation() { + let bytes = vec![1, 2, 3, 4, 5]; + let file = S3File::new(bytes.clone(), "test/file.txt".to_string()); + assert_eq!(file.as_bytes(), &bytes); + assert_eq!(file.key(), "test/file.txt"); + assert_eq!(file.clone().into_bytes(), bytes); + } + + #[test] + fn test_s3file_save_with_default_name() { + use tempdir::TempDir; + + let temp_dir = TempDir::new("test").unwrap(); + std::env::set_current_dir(temp_dir.path()).unwrap(); + + let bytes = b"test content".to_vec(); + let file = S3File::new(bytes.clone(), "path/to/myfile.txt".to_string()); + + // Save without specifying path - should use basename + let saved_path = file.save_file(None).unwrap(); + assert_eq!(saved_path, "myfile.txt"); + assert!(std::path::Path::new(&saved_path).exists()); + + let content = std::fs::read(&saved_path).unwrap(); + assert_eq!(content, bytes); + } + + #[test] + fn test_s3file_save_with_custom_path() { + use tempdir::TempDir; + + let bytes = b"test content".to_vec(); + let file = S3File::new(bytes.clone(), "original.txt".to_string()); + + let temp_dir = TempDir::new("test").unwrap(); + let path = temp_dir.path().join("custom.txt"); + + let saved_path = file.save_file(Some(path.to_str().unwrap())).unwrap(); + assert_eq!(saved_path, path.to_str().unwrap()); + + let content = std::fs::read(&saved_path).unwrap(); + assert_eq!(content, bytes); + } + + #[test] + fn test_s3_client_creation() { + let client = S3Client::new( + "test_key".to_string(), + "test_secret".to_string(), + "us-west-2".to_string(), + ); + assert_eq!(client.access_key_id, "test_key"); + assert_eq!(client.secret_access_key, "test_secret"); + assert_eq!(client.region, "us-west-2"); + } + +}