use crate::ai_manager::AIUserService; use crate::local_ai::controller::LocalAISetting; use flowy_ai_pub::cloud::LocalAIConfig; use flowy_error::{ErrorCode, FlowyError, FlowyResult}; use lib_infra::async_trait::async_trait; use crate::entities::LackOfAIResourcePB; use crate::local_ai::watch::{ is_plugin_ready, ollama_plugin_command_available, ollama_plugin_path, }; #[cfg(target_os = "macos")] use crate::local_ai::watch::{watch_offline_app, WatchContext}; use crate::notification::{ chat_notification_builder, ChatNotification, APPFLOWY_AI_NOTIFICATION_KEY, }; use appflowy_local_ai::ollama_plugin::OllamaPluginConfig; use lib_infra::util::{get_operating_system, OperatingSystem}; use reqwest::Client; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tracing::{error, info, instrument, trace}; #[async_trait] pub trait LLMResourceService: Send + Sync + 'static { /// Get local ai configuration from remote server async fn fetch_local_ai_config(&self) -> Result; fn store_setting(&self, setting: LocalAISetting) -> Result<(), anyhow::Error>; fn retrieve_setting(&self) -> Option; } const LLM_MODEL_DIR: &str = "models"; #[derive(Debug, Clone)] #[allow(dead_code)] pub enum WatchDiskEvent { Create, Remove, } pub enum PendingResource { PluginExecutableNotReady, OllamaServerNotReady, MissingModel(String), } impl PendingResource { pub fn desc(self) -> String { match self { PendingResource::PluginExecutableNotReady => "The Local AI app was not installed correctly. Please follow the instructions to install the Local AI application".to_string(), PendingResource::OllamaServerNotReady => "Ollama is not ready. Please follow the instructions to install Ollama".to_string(), PendingResource::MissingModel(model) => format!("Cannot find the model: {}. Please use the ollama pull command to install the model", model), } } } pub struct LocalAIResourceController { user_service: Arc, resource_service: Arc, resource_notify: tokio::sync::broadcast::Sender<()>, #[cfg(target_os = "macos")] #[allow(dead_code)] app_disk_watch: Option, app_state_sender: tokio::sync::broadcast::Sender, } impl LocalAIResourceController { pub fn new( user_service: Arc, resource_service: impl LLMResourceService, ) -> Self { let (resource_notify, _) = tokio::sync::broadcast::channel(1); let (app_state_sender, _) = tokio::sync::broadcast::channel(1); #[cfg(target_os = "macos")] let mut offline_app_disk_watch: Option = None; #[cfg(target_os = "macos")] { match watch_offline_app() { Ok((new_watcher, mut rx)) => { let sender = app_state_sender.clone(); tokio::spawn(async move { while let Some(event) = rx.recv().await { if let Err(err) = sender.send(event) { error!("[LLM Resource] Failed to send offline app state: {:?}", err); } } }); offline_app_disk_watch = Some(new_watcher); }, Err(err) => { error!("[LLM Resource] Failed to watch offline app path: {:?}", err); }, } } Self { user_service, resource_service: Arc::new(resource_service), #[cfg(target_os = "macos")] app_disk_watch: offline_app_disk_watch, app_state_sender, resource_notify, } } pub fn subscribe_resource_notify(&self) -> tokio::sync::broadcast::Receiver<()> { self.resource_notify.subscribe() } pub fn subscribe_app_state(&self) -> tokio::sync::broadcast::Receiver { self.app_state_sender.subscribe() } /// Returns true when all resources are downloaded and ready to use. pub async fn is_resource_ready(&self) -> bool { let sys = get_operating_system(); if !sys.is_desktop() { return false; } match self.calculate_pending_resources().await { Ok(res) => res.is_empty(), Err(_) => false, } } pub fn is_plugin_ready(&self) -> bool { ollama_plugin_path().exists() || ollama_plugin_command_available() } pub async fn get_plugin_download_link(&self) -> FlowyResult { let ai_config = self.get_local_ai_configuration().await?; Ok(ai_config.plugin.url) } /// Retrieves model information and updates the current model settings. pub fn get_llm_setting(&self) -> LocalAISetting { self.resource_service.retrieve_setting().unwrap_or_default() } #[instrument(level = "info", skip_all, err)] pub async fn set_llm_setting(&self, setting: LocalAISetting) -> FlowyResult<()> { self.resource_service.store_setting(setting)?; let mut resources = self.calculate_pending_resources().await?; if let Some(resource) = resources.pop() { chat_notification_builder( APPFLOWY_AI_NOTIFICATION_KEY, ChatNotification::LocalAIResourceUpdated, ) .payload(LackOfAIResourcePB { resource_desc: resource.desc(), }) .send(); } Ok(()) } pub async fn get_lack_of_resource(&self) -> Option { let mut resources = self.calculate_pending_resources().await.ok()?; resources.pop().map(|r| r.desc()) } /// Returns true when all resources are downloaded and ready to use. pub async fn calculate_pending_resources(&self) -> FlowyResult> { let mut resources = vec![]; let app_path = ollama_plugin_path(); if !app_path.exists() && !ollama_plugin_command_available() { trace!("[LLM Resource] offline app not found: {:?}", app_path); resources.push(PendingResource::PluginExecutableNotReady); return Ok(resources); } let setting = self.get_llm_setting(); let client = Client::builder().timeout(Duration::from_secs(5)).build()?; match client.get(&setting.ollama_server_url).send().await { Ok(resp) if resp.status().is_success() => { info!( "[LLM Resource] Ollama server is running at {}", setting.ollama_server_url ); }, _ => { info!( "[LLM Resource] Ollama server is not responding at {}", setting.ollama_server_url ); resources.push(PendingResource::OllamaServerNotReady); return Ok(resources); }, } let required_models = vec![ setting.chat_model_name, setting.embedding_model_name, // Add any additional required models here. ]; match tokio::process::Command::new("ollama") .arg("list") .output() .await { Ok(output) if output.status.success() => { let stdout = String::from_utf8_lossy(&output.stdout); for model in &required_models { if !stdout.contains(model.as_str()) { log::trace!( "[LLM Resource] required model '{}' not found in ollama list", model ); resources.push(PendingResource::MissingModel(model.clone())); return Ok(resources); } } }, Ok(output) => { error!( "[LLM Resource] 'ollama list' command failed with status: {:?}", output.status ); resources.push(PendingResource::OllamaServerNotReady); }, Err(e) => { error!("[LLM Resource] failed to execute 'ollama list': {:?}", e); resources.push(PendingResource::OllamaServerNotReady); }, } Ok(resources) } #[instrument(level = "info", skip_all)] pub async fn get_plugin_config(&self, rag_enabled: bool) -> FlowyResult { if !self.is_resource_ready().await { return Err(FlowyError::local_ai().with_context("Local AI resources are not ready")); } let llm_setting = self.get_llm_setting(); let bin_path = match get_operating_system() { OperatingSystem::MacOS | OperatingSystem::Windows => { if !is_plugin_ready() { return Err(FlowyError::new( ErrorCode::AppFlowyLAINotReady, "AppFlowyLAI not found", )); } ollama_plugin_path() }, _ => { return Err( FlowyError::local_ai_unavailable() .with_context("Local AI not available on current platform"), ); }, }; let mut config = OllamaPluginConfig::new( bin_path, "ollama_ai_plugin".to_string(), llm_setting.chat_model_name.clone(), llm_setting.embedding_model_name.clone(), Some(llm_setting.ollama_server_url.clone()), )?; if rag_enabled { let resource_dir = self.resource_dir()?; let persist_directory = resource_dir.join("vectorstore"); if !persist_directory.exists() { std::fs::create_dir_all(&persist_directory)?; } config.set_rag_enabled(&persist_directory)?; } if cfg!(debug_assertions) { config = config.with_verbose(true); } trace!("[AI Chat] config: {:?}", config); Ok(config) } /// Fetches the local AI configuration from the resource service. async fn get_local_ai_configuration(&self) -> FlowyResult { self .resource_service .fetch_local_ai_config() .await .map_err(|err| { error!("[LLM Resource] Failed to fetch local ai config: {:?}", err); FlowyError::local_ai() .with_context("Can't retrieve model info. Please try again later".to_string()) }) } pub(crate) fn user_model_folder(&self) -> FlowyResult { self.resource_dir().map(|dir| dir.join(LLM_MODEL_DIR)) } pub(crate) fn resource_dir(&self) -> FlowyResult { let user_data_dir = self.user_service.application_root_dir()?; Ok(user_data_dir.join("ai")) } } #[allow(dead_code)] fn bytes_to_readable_format(bytes: u64) -> String { const BYTES_IN_GIGABYTE: u64 = 1024 * 1024 * 1024; const BYTES_IN_MEGABYTE: u64 = 1024 * 1024; if bytes >= BYTES_IN_GIGABYTE { let gigabytes = (bytes as f64) / (BYTES_IN_GIGABYTE as f64); format!("{:.1} GB", gigabytes) } else { let megabytes = (bytes as f64) / (BYTES_IN_MEGABYTE as f64); format!("{:.2} MB", megabytes) } }