Refactored PDFium initialization (#440)

runtime/src/app_window.rs

@@ -14,7 +14,7 @@ use crate::api_token::APIToken;
 use crate::dotnet::stop_dotnet_server;
 use crate::environment::{is_prod, CONFIG_DIRECTORY, DATA_DIRECTORY};
 use crate::log::switch_to_file_logging;
-use crate::metadata::META_DATA;
+use crate::pdfium::PDFIUM_LIB_PATH;
 
 /// The Tauri main window.
 static MAIN_WINDOW: Lazy<Mutex<Option<Window>>> = Lazy::new(|| Mutex::new(None));
@@ -39,8 +39,7 @@ pub fn start_tauri() {
 
             info!(Source = "Bootloader Tauri"; "Reconfigure the file logger to use the app data directory {data_path:?}");
             switch_to_file_logging(data_path).map_err(|e| error!("Failed to switch logging to file: {e}")).unwrap();
-
-            deploy_pdfium(app.path_resolver());
+            set_pdfium_path(app.path_resolver());
             Ok(())
         })
         .plugin(tauri_plugin_window_state::Builder::default().build())
@@ -323,48 +322,15 @@ pub struct FileSelectionResponse {
     selected_file_path: String,
 }
 
-fn deploy_pdfium(path_resolver: PathResolver) {
-    info!(Source = "Bootloader Tauri"; "Deploy PDFium from the resources...");
-    let working_directory = std::env::current_dir().unwrap();
-    let pdfium_target_path = working_directory;
-    let metadata = &META_DATA;
-    let metadata = metadata.lock().unwrap();
-    let arch = metadata.clone().unwrap().architecture;
-    let pdfium_filename = match arch.as_str() {
-        "linux-x64" => Some("libpdfium.so"),
-        "linux-arm64" => Some("libpdfium.so"),
-
-        "win-x64" => Some("pdfium.dll"),
-        "win-arm64" => Some("pdfium.dll"),
-
-        "osx-x64" => Some("libpdfium.dylib"),
-        "osx-arm64" => Some("libpdfium.dylib"),
-
-        _ => None,
-    };
-
-    if pdfium_filename.is_none() {
-        error!(Source = "Bootloader Tauri"; "Failed to find the PDFium library for the current platform.");
-        return;
-    }
-
-    let pdfium_filename = pdfium_filename.unwrap();
-    let pdfium_relative_source_path = String::from("resources/libraries/") + pdfium_filename;
+fn set_pdfium_path(path_resolver: PathResolver) {
+    let pdfium_relative_source_path = String::from("resources/libraries/");
     let pdfium_source_path = path_resolver.resolve_resource(pdfium_relative_source_path);
     if pdfium_source_path.is_none() {
-        error!(Source = "Bootloader Tauri"; "Failed to find the PDFium library for the current platform.");
+        error!(Source = "Bootloader Tauri"; "Failed to set the PDFium library path.");
         return;
     }
 
     let pdfium_source_path = pdfium_source_path.unwrap();
-    let pdfium_target_path = pdfium_target_path.join(pdfium_filename);
-
-    info!(Source = "Bootloader Tauri"; "Detected platform: {arch:?}, expected PDFium filename: {pdfium_filename:?}, source path: {pdfium_source_path:?}, target path: {pdfium_target_path:?}");
-
-    if let Err(e) = std::fs::copy(pdfium_source_path, pdfium_target_path) {
-        error!(Source = "Bootloader Tauri"; "Failed to copy the PDFium library for the current platform: {e}");
-        return;
-    }
-
-    info!(Source = "Bootloader Tauri"; "Successfully deployed PDFium.");
+    let pdfium_source_path = pdfium_source_path.to_str().unwrap().to_string();
+    *PDFIUM_LIB_PATH.lock().unwrap() = Some(pdfium_source_path.clone());
 }

runtime/src/file_data.rs

@@ -16,6 +16,7 @@ use rocket::tokio::select;
 use rocket::serde::Serialize;
 use rocket::get;
 use crate::api_token::APIToken;
+use crate::pdfium::PdfiumInit;
 
 #[derive(Debug, Serialize)]
 pub struct Chunk {
@@ -148,7 +149,7 @@ async fn stream_text_file(file_path: &str) -> Result<ChunkStream> {
 
 #[get("/retrieval/fs/read/pdf?<file_path>")]
 pub fn read_pdf(_token: APIToken, file_path: String) -> String {
-    let pdfium = Pdfium::default();
+    let pdfium = Pdfium::ai_studio_init();
     let doc = match pdfium.load_pdf_from_file(&file_path, None) {
         Ok(document) => document,
         Err(e) => return e.to_string(),
@@ -175,7 +176,7 @@ async fn stream_pdf(file_path: &str) -> Result<ChunkStream> {
     let (tx, rx) = mpsc::channel(10);
 
     tokio::task::spawn_blocking(move || {
-        let pdfium = Pdfium::default();
+        let pdfium = Pdfium::ai_studio_init();
         let doc = match pdfium.load_pdf_from_file(&path, None) {
             Ok(document) => document,
             Err(e) => {

runtime/src/lib.rs

@@ -11,3 +11,4 @@ pub mod runtime_api;
 pub mod certificate;
 pub mod file_data;
 pub mod metadata;
+pub mod pdfium;

runtime/src/pdfium.rs

@@ -0,0 +1,26 @@
+use std::sync::Mutex;
+use once_cell::sync::Lazy;
+use pdfium_render::prelude::Pdfium;
+
+pub static PDFIUM_LIB_PATH: Lazy<Mutex<Option<String>>> = Lazy::new(|| Mutex::new(None));
+
+pub trait PdfiumInit {
+    fn ai_studio_init() -> Pdfium;
+}
+
+impl PdfiumInit for Pdfium {
+
+    /// Initializes the PDFium library for AI Studio.
+    fn ai_studio_init() -> Pdfium {
+        let lib_path = PDFIUM_LIB_PATH.lock().unwrap();
+        if let Some(path) = lib_path.as_ref() {
+            return Pdfium::new(
+                Pdfium::bind_to_library(Pdfium::pdfium_platform_library_name_at_path(path))
+                    .or_else(|_| Pdfium::bind_to_system_library())
+                    .unwrap(),
+            );
+        }
+
+        Pdfium::new(Pdfium::bind_to_system_library().unwrap())
+    }
+}