fix: improved storage manager, prep for s3

ingestion-pipeline/src/utils/file_text_extraction.rs

@@ -1,63 +1,200 @@
+use anyhow::anyhow;
 use common::{
     error::AppError,
-    storage::{db::SurrealDbClient, store, types::file_info::FileInfo},
+    storage::{db::SurrealDbClient, store::StorageManager, types::file_info::FileInfo},
     utils::config::AppConfig,
 };
+use std::{
+    env,
+    io::{Error as IoError, ErrorKind},
+    path::{Path, PathBuf},
+};
+use uuid::Uuid;
 
 use super::{
     audio_transcription::transcribe_audio_file, image_parsing::extract_text_from_image,
     pdf_ingestion::extract_pdf_content,
 };
 
+struct TempPathGuard {
+    path: PathBuf,
+}
+
+impl TempPathGuard {
+    fn as_path(&self) -> &Path {
+        &self.path
+    }
+}
+
+impl Drop for TempPathGuard {
+    fn drop(&mut self) {
+        let _ = std::fs::remove_file(&self.path);
+    }
+}
+
+async fn materialize_temp_file(
+    bytes: &[u8],
+    extension: Option<&str>,
+) -> Result<TempPathGuard, AppError> {
+    let mut path = env::temp_dir();
+    let mut file_name = format!("minne-ingest-{}", Uuid::new_v4());
+
+    if let Some(ext) = extension {
+        if !ext.is_empty() {
+            file_name.push('.');
+            file_name.push_str(ext);
+        }
+    }
+
+    path.push(file_name);
+
+    tokio::fs::write(&path, bytes).await?;
+
+    Ok(TempPathGuard { path })
+}
+
+async fn resolve_existing_local_path(storage: &StorageManager, location: &str) -> Option<PathBuf> {
+    let path = storage.resolve_local_path(location)?;
+    match tokio::fs::metadata(&path).await {
+        Ok(_) => Some(path),
+        Err(_) => None,
+    }
+}
+
+fn infer_extension(file_info: &FileInfo) -> Option<String> {
+    Path::new(&file_info.path)
+        .extension()
+        .and_then(|ext| ext.to_str())
+        .map(|ext| ext.to_string())
+}
+
 pub async fn extract_text_from_file(
     file_info: &FileInfo,
     db_client: &SurrealDbClient,
     openai_client: &async_openai::Client<async_openai::config::OpenAIConfig>,
     config: &AppConfig,
+    storage: &StorageManager,
 ) -> Result<String, AppError> {
-    let base_path = store::resolve_base_dir(config);
-    let absolute_path = base_path.join(&file_info.path);
+    let file_bytes = storage
+        .get(&file_info.path)
+        .await
+        .map_err(|e| AppError::from(anyhow!(e)))?;
+    let local_path = resolve_existing_local_path(storage, &file_info.path).await;
 
     match file_info.mime_type.as_str() {
         "text/plain" | "text/markdown" | "application/octet-stream" | "text/x-rust" => {
-            let content = tokio::fs::read_to_string(&absolute_path).await?;
+            let content = String::from_utf8(file_bytes.to_vec())
+                .map_err(|err| AppError::Io(IoError::new(ErrorKind::InvalidData, err)))?;
             Ok(content)
         }
         "application/pdf" => {
-            extract_pdf_content(
-                &absolute_path,
+            if let Some(path) = local_path.as_ref() {
+                return extract_pdf_content(
+                    path,
+                    db_client,
+                    openai_client,
+                    &config.pdf_ingest_mode,
+                )
+                .await;
+            }
+
+            let temp_guard = materialize_temp_file(file_bytes.as_ref(), Some("pdf")).await?;
+            let result = extract_pdf_content(
+                temp_guard.as_path(),
                 db_client,
                 openai_client,
                 &config.pdf_ingest_mode,
             )
-            .await
+            .await;
+            drop(temp_guard);
+            result
         }
         "image/png" | "image/jpeg" => {
-            let path_str = absolute_path
-                .to_str()
-                .ok_or_else(|| {
-                    AppError::Processing(format!(
-                        "Encountered a non-UTF8 path while reading image {}",
-                        file_info.id
-                    ))
-                })?
-                .to_string();
-            let content = extract_text_from_image(&path_str, db_client, openai_client).await?;
+            let content =
+                extract_text_from_image(file_bytes.as_ref(), db_client, openai_client).await?;
             Ok(content)
         }
         "audio/mpeg" | "audio/mp3" | "audio/wav" | "audio/x-wav" | "audio/webm" | "audio/mp4"
         | "audio/ogg" | "audio/flac" => {
-            let path_str = absolute_path
-                .to_str()
-                .ok_or_else(|| {
+            if let Some(path) = local_path.as_ref() {
+                let path_str = path.to_str().ok_or_else(|| {
                     AppError::Processing(format!(
                         "Encountered a non-UTF8 path while reading audio {}",
                         file_info.id
                     ))
-                })?
-                .to_string();
-            transcribe_audio_file(&path_str, db_client, openai_client).await
+                })?;
+                return transcribe_audio_file(path_str, db_client, openai_client).await;
+            }
+
+            let extension = infer_extension(file_info);
+            let temp_guard =
+                materialize_temp_file(file_bytes.as_ref(), extension.as_deref()).await?;
+            let path_str = temp_guard.as_path().to_str().ok_or_else(|| {
+                AppError::Processing(format!(
+                    "Encountered a non-UTF8 path while reading audio {}",
+                    file_info.id
+                ))
+            })?;
+            let result = transcribe_audio_file(path_str, db_client, openai_client).await;
+            drop(temp_guard);
+            result
         }
         _ => Err(AppError::NotFound(file_info.mime_type.clone())),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use async_openai::{config::OpenAIConfig, Client};
+    use bytes::Bytes;
+    use chrono::Utc;
+    use common::{
+        storage::{db::SurrealDbClient, store::StorageManager},
+        utils::config::{AppConfig, StorageKind},
+    };
+
+    #[tokio::test]
+    async fn extracts_text_using_memory_storage_backend() {
+        let mut config = AppConfig::default();
+        config.storage = StorageKind::Memory;
+
+        let storage = StorageManager::new(&config)
+            .await
+            .expect("create storage manager");
+
+        let location = "user/test/file.txt";
+        let contents = b"hello from memory storage";
+
+        storage
+            .put(location, Bytes::from(contents.as_slice().to_vec()))
+            .await
+            .expect("write object");
+
+        let now = Utc::now();
+        let file_info = FileInfo {
+            id: "file".into(),
+            created_at: now,
+            updated_at: now,
+            sha256: "sha256".into(),
+            path: location.to_string(),
+            file_name: "file.txt".into(),
+            mime_type: "text/plain".into(),
+            user_id: "user".into(),
+        };
+
+        let namespace = "test_ns";
+        let database = &Uuid::new_v4().to_string();
+        let db = SurrealDbClient::memory(namespace, database)
+            .await
+            .expect("create surreal memory");
+
+        let openai_client = Client::with_config(OpenAIConfig::default());
+
+        let text = extract_text_from_file(&file_info, &db, &openai_client, &config, &storage)
+            .await
+            .expect("extract text");
+
+        assert_eq!(text, String::from_utf8_lossy(contents));
+    }
+}

ingestion-pipeline/src/utils/image_parsing.rs

@@ -10,14 +10,13 @@ use common::{
 };
 
 pub async fn extract_text_from_image(
-    path: &str,
+    image_bytes: &[u8],
     db: &SurrealDbClient,
     client: &async_openai::Client<async_openai::config::OpenAIConfig>,
 ) -> Result<String, AppError> {
     let system_settings = SystemSettings::get_current(db).await?;
-    let image_bytes = tokio::fs::read(&path).await?;
 
-    let base64_image = STANDARD.encode(&image_bytes);
+    let base64_image = STANDARD.encode(image_bytes);
 
     let image_url = format!("data:image/png;base64,{base64_image}");
 

ingestion-pipeline/src/utils/url_text_retrieval.rs

@@ -3,8 +3,7 @@ use axum_typed_multipart::{FieldData, FieldMetadata};
 use chrono::Utc;
 use common::{
     error::AppError,
-    storage::{db::SurrealDbClient, types::file_info::FileInfo},
-    utils::config::AppConfig,
+    storage::{db::SurrealDbClient, store::StorageManager, types::file_info::FileInfo},
 };
 use dom_smoothie::{Article, Readability, TextMode};
 use headless_chrome::Browser;
@@ -19,7 +18,7 @@ pub async fn extract_text_from_url(
     url: &str,
     db: &SurrealDbClient,
     user_id: &str,
-    config: &AppConfig,
+    storage: &StorageManager,
 ) -> Result<(Article, FileInfo), AppError> {
     info!("Fetching URL: {}", url);
     let now = Instant::now();
@@ -81,7 +80,7 @@ pub async fn extract_text_from_url(
         metadata,
     };
 
-    let file_info = FileInfo::new(field_data, db, user_id, config).await?;
+    let file_info = FileInfo::new_with_storage(field_data, db, user_id, storage).await?;
 
     let config = dom_smoothie::Config {
         text_mode: TextMode::Markdown,