diff --git a/crates/goose/src/providers/local_inference.rs b/crates/goose/src/providers/local_inference.rs index 75540ef6a8..970a35311f 100644 --- a/crates/goose/src/providers/local_inference.rs +++ b/crates/goose/src/providers/local_inference.rs @@ -1,19 +1,10 @@ +mod backend; pub mod hf_models; -mod inference_emulated_tools; -mod inference_engine; -mod inference_native_tools; +mod llamacpp; pub mod local_model_registry; pub(crate) mod multimodal; mod tool_parsing; -use inference_emulated_tools::{ - build_emulator_tool_description, generate_with_emulated_tools, load_tiny_model_prompt, -}; -use inference_engine::GenerationContext; -use inference_engine::LoadedModel; -use inference_native_tools::generate_with_native_tools; -use tool_parsing::compact_tools_json; - use crate::config::ExtensionConfig; use crate::conversation::message::{Message, MessageContent}; use crate::model::ModelConfig; @@ -21,18 +12,14 @@ use crate::providers::base::{ MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage, }; use crate::providers::errors::ProviderError; -use crate::providers::formats::openai::format_tools; use crate::providers::utils::RequestLog; use anyhow::Result; use async_stream::try_stream; use async_trait::async_trait; +use backend::{BackendLoadedModel, LocalInferenceBackend}; use futures::future::BoxFuture; -use llama_cpp_2::llama_backend::LlamaBackend; -use llama_cpp_2::model::params::LlamaModelParams; -use llama_cpp_2::model::{LlamaChatMessage, LlamaChatTemplate, LlamaModel}; -use llama_cpp_2::{list_llama_ggml_backend_devices, LlamaBackendDeviceType, LogOptions}; -use multimodal::ExtractedImage; -use rmcp::model::{Role, Tool}; +use llamacpp::{LlamaCppBackend, LLAMACPP_BACKEND_ID}; +use rmcp::model::Tool; use serde_json::{json, Value}; use std::collections::HashMap; use std::path::PathBuf; @@ -40,18 +27,26 @@ use std::sync::{Arc, Mutex as StdMutex, Weak}; use tokio::sync::Mutex; use uuid::Uuid; -const SHELL_TOOL: &str = "developer__shell"; -const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript"; +type ModelSlot = Arc>>>; -type ModelSlot = Arc>>; +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct ModelCacheKey { + backend_id: &'static str, + model_id: String, +} + +impl ModelCacheKey { + fn new(backend_id: &'static str, model_id: impl Into) -> Self { + Self { + backend_id, + model_id: model_id.into(), + } + } +} -/// Owns the llama backend and all cached models. Field order matters: -/// `models` is declared before `backend` so Rust drops all loaded models -/// (and their Metal/GPU resources) before the backend calls -/// `llama_backend_free()`, avoiding the ggml-metal assertion on shutdown. pub struct InferenceRuntime { - models: StdMutex>, - backend: LlamaBackend, + models: StdMutex>, + backends: HashMap<&'static str, Arc>, } /// Global weak reference used to share a single `InferenceRuntime` across @@ -67,49 +62,47 @@ impl InferenceRuntime { if let Some(runtime) = guard.upgrade() { return Ok(runtime); } - // Safety invariant: the Weak::upgrade() check and LlamaBackend::init() - // both execute inside this same mutex guard, so there is no window where - // another thread could drop the Arc and re-enter concurrently. - // BackendAlreadyInitialized therefore means LlamaBackend::drop() did not - // reset the C library's init flag — a llama-cpp-rs bug, not a race. - let backend = match LlamaBackend::init() { - Ok(b) => b, - Err(llama_cpp_2::LlamaCppError::BackendAlreadyInitialized) => { - unreachable!( - "LlamaBackend already initialized but Weak was dead; \ - the mutex guard prevents concurrent re-init" - ) - } - Err(e) => { - tracing::error!(error = %e, "failed to initialize local inference runtime"); - return Err(anyhow::anyhow!("Failed to init llama backend: {}", e)); - } - }; - llama_cpp_2::send_logs_to_tracing(LogOptions::default()); - log_inference_backend_devices(); + let llamacpp_backend: Arc = Arc::new(LlamaCppBackend::new()?); + let mut backends = HashMap::new(); + backends.insert(LLAMACPP_BACKEND_ID, llamacpp_backend); let runtime = Arc::new(Self { models: StdMutex::new(HashMap::new()), - backend, + backends, }); *guard = Arc::downgrade(&runtime); Ok(runtime) } - pub fn backend(&self) -> &LlamaBackend { - &self.backend + fn default_backend(&self) -> &dyn LocalInferenceBackend { + self.backends + .get(LLAMACPP_BACKEND_ID) + .expect("default local inference backend registered") + .as_ref() } - fn get_or_create_model_slot(&self, model_id: &str) -> ModelSlot { + fn backend_for_model( + &self, + _resolved: &ResolvedModelPaths, + ) -> Result, ProviderError> { + self.backends + .get(LLAMACPP_BACKEND_ID) + .cloned() + .ok_or_else(|| { + ProviderError::ExecutionError("Local inference backend unavailable".to_string()) + }) + } + + fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlot { let mut map = self.models.lock().expect("model cache lock poisoned"); - map.entry(model_id.to_string()) + map.entry(key) .or_insert_with(|| Arc::new(Mutex::new(None))) .clone() } - fn other_model_slots(&self, keep_model_id: &str) -> Vec { + fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec { let map = self.models.lock().expect("model cache lock poisoned"); map.iter() - .filter(|(id, _)| id.as_str() != keep_model_id) + .filter(|(key, _)| *key != keep_key) .map(|(_, slot)| slot.clone()) .collect() } @@ -120,7 +113,8 @@ const DEFAULT_MODEL: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M"; pub const LOCAL_LLM_MODEL_CONFIG_KEY: &str = "LOCAL_LLM_MODEL"; -pub struct ResolvedModelPaths { +#[derive(Clone)] +pub(super) struct ResolvedModelPaths { pub model_path: PathBuf, pub context_limit: usize, pub settings: crate::providers::local_inference::local_model_registry::ModelSettings, @@ -128,7 +122,7 @@ pub struct ResolvedModelPaths { } /// Resolve model path, context limit, settings, and mmproj path for a model ID from the registry. -pub fn resolve_model_path(model_id: &str) -> Option { +fn resolve_model_path(model_id: &str) -> Option { use crate::providers::local_inference::local_model_registry::{ default_settings_for_model, get_registry, }; @@ -157,69 +151,8 @@ pub fn resolve_model_path(model_id: &str) -> Option { None } -fn is_accelerator_device(device_type: LlamaBackendDeviceType) -> bool { - matches!( - device_type, - LlamaBackendDeviceType::Gpu - | LlamaBackendDeviceType::IntegratedGpu - | LlamaBackendDeviceType::Accelerator - ) -} - -fn is_non_cpu_device(device_type: LlamaBackendDeviceType) -> bool { - !matches!(device_type, LlamaBackendDeviceType::Cpu) -} - -fn log_inference_backend_devices() { - let devices = list_llama_ggml_backend_devices(); - let non_cpu_devices: Vec<_> = devices - .iter() - .filter(|device| is_non_cpu_device(device.device_type)) - .collect(); - - if non_cpu_devices.is_empty() { - tracing::info!( - device_count = devices.len(), - "No non-CPU llama.cpp backend devices detected for local inference" - ); - return; - } - - for device in non_cpu_devices { - tracing::info!( - index = device.index, - backend = %device.backend, - name = %device.name, - description = %device.description, - device_type = ?device.device_type, - memory_total_bytes = device.memory_total as u64, - memory_free_bytes = device.memory_free as u64, - "Non-CPU llama.cpp backend device detected for local inference" - ); - } -} - pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { - let _ = &runtime.backend; - let devices = list_llama_ggml_backend_devices(); - - let accel_memory = devices - .iter() - .filter(|d| is_accelerator_device(d.device_type)) - .map(|d| d.memory_free as u64) - .max() - .unwrap_or(0); - - if accel_memory > 0 { - accel_memory - } else { - devices - .iter() - .filter(|d| d.device_type == LlamaBackendDeviceType::Cpu) - .map(|d| d.memory_free as u64) - .max() - .unwrap_or(0) - } + runtime.default_backend().available_memory_bytes() } pub fn recommend_local_model(runtime: &InferenceRuntime) -> String { @@ -389,7 +322,6 @@ type StreamSender = pub struct LocalInferenceProvider { runtime: Arc, - model: ModelSlot, model_config: ModelConfig, name: String, } @@ -397,102 +329,12 @@ pub struct LocalInferenceProvider { impl LocalInferenceProvider { pub async fn from_env(model: ModelConfig, _extensions: Vec) -> Result { let runtime = InferenceRuntime::get_or_init()?; - let model_slot = runtime.get_or_create_model_slot(&model.model_name); Ok(Self { runtime, - model: model_slot, model_config: model, name: PROVIDER_NAME.to_string(), }) } - - fn load_model_sync( - runtime: &InferenceRuntime, - model_id: &str, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, - ) -> Result { - let resolved = resolve_model_path(model_id) - .ok_or_else(|| ProviderError::ExecutionError(format!("Unknown model: {}", model_id)))?; - let model_path = resolved.model_path; - - if !model_path.exists() { - return Err(ProviderError::ExecutionError(format!( - "Model not downloaded: {}. Please download it from Settings > Local Inference.", - model_id - ))); - } - - tracing::info!("Loading {} from: {}", model_id, model_path.display()); - - let backend = runtime.backend(); - - let mut params = LlamaModelParams::default(); - if let Some(n_gpu_layers) = settings.n_gpu_layers { - params = params.with_n_gpu_layers(n_gpu_layers); - } - if settings.use_mlock { - params = params.with_use_mlock(true); - } - let model = LlamaModel::load_from_file(backend, &model_path, ¶ms) - .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; - - let template = match model.chat_template(None) { - Ok(t) => t, - Err(_) => { - tracing::warn!("Model has no embedded chat template, falling back to chatml"); - LlamaChatTemplate::new("chatml").map_err(|e| { - ProviderError::ExecutionError(format!( - "Failed to create fallback chat template: {}", - e - )) - })? - } - }; - - let mtmd_ctx = Self::init_mtmd_context(&model, &resolved.mmproj_path, settings); - - tracing::info!(model_id = model_id, "Model loaded successfully"); - - Ok(LoadedModel { - model, - template, - mtmd_ctx, - }) - } - - fn init_mtmd_context( - model: &LlamaModel, - mmproj_path: &Option, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, - ) -> Option { - use llama_cpp_2::mtmd::{MtmdContext, MtmdContextParams}; - - let mmproj_path = mmproj_path.as_ref().filter(|p| p.exists())?; - - let params = MtmdContextParams { - use_gpu: true, - n_threads: settings - .n_threads - .unwrap_or_else(|| MtmdContextParams::default().n_threads), - ..MtmdContextParams::default() - }; - - match MtmdContext::init_from_file(mmproj_path.to_str().unwrap_or_default(), model, ¶ms) - { - Ok(ctx) => { - tracing::info!( - vision = ctx.support_vision(), - audio = ctx.support_audio(), - "Multimodal context initialized" - ); - Some(ctx) - } - Err(e) => { - tracing::warn!(error = %e, "Failed to init multimodal context"); - None - } - } - } } impl ProviderDef for LocalInferenceProvider { @@ -577,14 +419,17 @@ impl Provider for LocalInferenceProvider { let resolved = resolve_model_path(&model_config.model_name).ok_or_else(|| { ProviderError::ExecutionError(format!("Model not found: {}", model_config.model_name)) })?; + let backend = self.runtime.backend_for_model(&resolved)?; let model_context_limit = resolved.context_limit; - let model_settings = resolved.settings; + let model_settings = resolved.settings.clone(); + let cache_key = ModelCacheKey::new(backend.id(), model_config.model_name.clone()); + let model_slot = self.runtime.get_or_create_model_slot(cache_key.clone()); // Ensure model is loaded — unload any other models first to free memory. { - let mut model_lock = self.model.lock().await; + let mut model_lock = model_slot.lock().await; if model_lock.is_none() { - for slot in self.runtime.other_model_slots(&model_config.model_name) { + for slot in self.runtime.other_model_slots(&cache_key) { let mut other = slot.lock().await; if other.is_some() { tracing::info!("Unloading previous model to free memory"); @@ -593,10 +438,11 @@ impl Provider for LocalInferenceProvider { } let model_id = model_config.model_name.clone(); + let resolved_for_load = resolved.clone(); let settings_for_load = model_settings.clone(); - let runtime_for_load = self.runtime.clone(); + let backend_for_load = backend.clone(); let loaded = tokio::task::spawn_blocking(move || { - Self::load_model_sync(&runtime_for_load, &model_id, &settings_for_load) + backend_for_load.load_model(&model_id, &resolved_for_load, &settings_for_load) }) .await .map_err(|e| ProviderError::ExecutionError(e.to_string()))??; @@ -612,98 +458,20 @@ impl Provider for LocalInferenceProvider { model_settings.enable_thinking = false; } - // Use the model's native_tool_calling setting to decide the path. - // Featured models have this set explicitly; user-added models default to false. - let native_tool_calling = model_settings.native_tool_calling; - let use_emulator = !native_tool_calling && !tools.is_empty(); - let system_prompt = if use_emulator { - load_tiny_model_prompt() - } else { - system.to_string() - }; - - // Extract images for vision-capable models, replacing them with markers. - // For non-vision models, leave messages unchanged (existing strip logic handles them). - let has_vision = resolved.mmproj_path.is_some(); - let marker = llama_cpp_2::mtmd::mtmd_default_marker(); - let (images, vision_messages): (Vec, Option>) = if has_vision { - let (imgs, msgs) = multimodal::extract_images_from_messages(messages, marker); - (imgs, Some(msgs)) - } else { - (Vec::new(), None) - }; - let effective_messages: &[Message] = vision_messages.as_deref().unwrap_or(messages); - - // Build chat messages for the template - let mut chat_messages = - vec![ - LlamaChatMessage::new("system".to_string(), system_prompt.clone()).map_err( - |e| { - ProviderError::ExecutionError(format!( - "Failed to create system message: {}", - e - )) - }, - )?, - ]; - - let code_mode_enabled = tools.iter().any(|t| t.name == CODE_EXECUTION_TOOL); - - if use_emulator && !tools.is_empty() { - let tool_desc = build_emulator_tool_description(tools, code_mode_enabled); - chat_messages = vec![LlamaChatMessage::new( - "system".to_string(), - format!("{}{}", system_prompt, tool_desc), - ) - .map_err(|e| { - ProviderError::ExecutionError(format!("Failed to create system message: {}", e)) - })?]; - } - - for msg in effective_messages { - let role = match msg.role { - Role::User => "user", - Role::Assistant => "assistant", - }; - let content = extract_text_content(msg); - if !content.trim().is_empty() { - chat_messages.push(LlamaChatMessage::new(role.to_string(), content).map_err( - |e| ProviderError::ExecutionError(format!("Failed to create message: {}", e)), - )?); - } - } - - let (full_tools_json, compact_tools) = if !use_emulator && !tools.is_empty() { - let full = format_tools(tools) - .ok() - .and_then(|spec| serde_json::to_string(&spec).ok()); - let compact = compact_tools_json(tools); - (full, compact) - } else { - (None, None) - }; - - let oai_messages_json = if model_settings.use_jinja || native_tool_calling { - Some(build_openai_messages_json( - &system_prompt, - effective_messages, - )) - } else { - None - }; - - let model_arc = self.model.clone(); - let runtime = self.runtime.clone(); + let model_arc = model_slot.clone(); + let backend = backend.clone(); let model_name = model_config.model_name.clone(); let context_limit = model_context_limit; let settings = model_settings; - let mmproj_path = resolved.mmproj_path.clone(); - + let resolved_model = resolved.clone(); + let system = system.to_string(); + let messages = messages.to_vec(); + let tools = tools.to_vec(); let log_payload = serde_json::json!({ - "system": &system_prompt, + "system": &system, "messages": messages.iter().map(|m| { serde_json::json!({ - "role": match m.role { Role::User => "user", Role::Assistant => "assistant" }, + "role": match m.role { rmcp::model::Role::User => "user", rmcp::model::Role::Assistant => "assistant" }, "content": extract_text_content(m), }) }).collect::>(), @@ -749,41 +517,22 @@ impl Provider for LocalInferenceProvider { } }; - // Lazily initialize the multimodal context if the vision encoder - // was downloaded after the model was loaded. - if !images.is_empty() && loaded.mtmd_ctx.is_none() { - loaded.mtmd_ctx = LocalInferenceProvider::init_mtmd_context( - &loaded.model, - &mmproj_path, - &settings, - ); - } - let message_id = Uuid::new_v4().to_string(); - let mut gen_ctx = GenerationContext { - loaded, - runtime: &runtime, - chat_messages: &chat_messages, + let request = backend::LocalGenerationRequest { + model_name, + system: &system, + messages: &messages, + tools: &tools, settings: &settings, context_limit, - model_name, + resolved_model: &resolved_model, message_id: &message_id, tx: &tx, log: &mut log, - images: &images, }; - let result = if use_emulator { - generate_with_emulated_tools(&mut gen_ctx, code_mode_enabled) - } else { - generate_with_native_tools( - &mut gen_ctx, - &oai_messages_json, - full_tools_json.as_deref(), - compact_tools.as_deref(), - ) - }; + let result = backend.generate(loaded.as_mut(), request); if let Err(err) = result { let msg = match &err { diff --git a/crates/goose/src/providers/local_inference/backend.rs b/crates/goose/src/providers/local_inference/backend.rs new file mode 100644 index 0000000000..0968aafcf1 --- /dev/null +++ b/crates/goose/src/providers/local_inference/backend.rs @@ -0,0 +1,45 @@ +use rmcp::model::Tool; +use std::any::Any; + +use crate::conversation::message::Message; +use crate::providers::errors::ProviderError; +use crate::providers::local_inference::local_model_registry::ModelSettings; +use crate::providers::utils::RequestLog; + +use super::{ResolvedModelPaths, StreamSender}; + +pub(super) trait BackendLoadedModel: Send { + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +pub(super) struct LocalGenerationRequest<'a> { + pub model_name: String, + pub system: &'a str, + pub messages: &'a [Message], + pub tools: &'a [Tool], + pub settings: &'a ModelSettings, + pub context_limit: usize, + pub resolved_model: &'a ResolvedModelPaths, + pub message_id: &'a str, + pub tx: &'a StreamSender, + pub log: &'a mut RequestLog, +} + +pub(super) trait LocalInferenceBackend: Send + Sync { + fn id(&self) -> &'static str; + + fn load_model( + &self, + model_id: &str, + resolved: &ResolvedModelPaths, + settings: &ModelSettings, + ) -> Result, ProviderError>; + + fn generate( + &self, + loaded: &mut dyn BackendLoadedModel, + request: LocalGenerationRequest<'_>, + ) -> Result<(), ProviderError>; + + fn available_memory_bytes(&self) -> u64; +} diff --git a/crates/goose/src/providers/local_inference/inference_emulated_tools.rs b/crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs similarity index 98% rename from crates/goose/src/providers/local_inference/inference_emulated_tools.rs rename to crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs index 90732de690..1f18612a68 100644 --- a/crates/goose/src/providers/local_inference/inference_emulated_tools.rs +++ b/crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs @@ -28,11 +28,14 @@ use serde_json::json; use std::borrow::Cow; use uuid::Uuid; +use super::super::{finalize_usage, StreamSender}; use super::inference_engine::{ create_and_prefill_context, create_and_prefill_multimodal, generation_loop, validate_and_compute_context, GenerationContext, TokenAction, }; -use super::{finalize_usage, StreamSender, CODE_EXECUTION_TOOL, SHELL_TOOL}; + +const SHELL_TOOL: &str = "developer__shell"; +const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript"; const HOLD_BACK_CODE_MODE: usize = " ```execute_typescript\n".len(); const HOLD_BACK_SHELL_ONLY: usize = "\n$".len(); @@ -373,7 +376,7 @@ pub(super) fn generate_with_emulated_tools( let (mut llama_ctx, prompt_token_count, effective_ctx) = if !ctx.images.is_empty() { create_and_prefill_multimodal( ctx.loaded, - ctx.runtime, + ctx.backend, &prompt, ctx.images, ctx.context_limit, @@ -387,13 +390,13 @@ pub(super) fn generate_with_emulated_tools( .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; let (ptc, ectx) = validate_and_compute_context( ctx.loaded, - ctx.runtime, + ctx.backend, tokens.len(), ctx.context_limit, ctx.settings, )?; let lctx = - create_and_prefill_context(ctx.loaded, ctx.runtime, &tokens, ectx, ctx.settings)?; + create_and_prefill_context(ctx.loaded, ctx.backend, &tokens, ectx, ctx.settings)?; (lctx, ptc, ectx) }; diff --git a/crates/goose/src/providers/local_inference/inference_engine.rs b/crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs similarity index 96% rename from crates/goose/src/providers/local_inference/inference_engine.rs rename to crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs index f9ae886b7a..86d2988c05 100644 --- a/crates/goose/src/providers/local_inference/inference_engine.rs +++ b/crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs @@ -1,4 +1,5 @@ use crate::providers::errors::ProviderError; +use crate::providers::local_inference::backend::LocalInferenceBackend; use crate::providers::local_inference::local_model_registry::ModelSettings; use crate::providers::local_inference::multimodal::ExtractedImage; use crate::providers::utils::RequestLog; @@ -9,11 +10,12 @@ use llama_cpp_2::mtmd::{MtmdBitmap, MtmdContext, MtmdInputText}; use llama_cpp_2::sampling::LlamaSampler; use std::num::NonZeroU32; -use super::{InferenceRuntime, StreamSender}; +use super::super::StreamSender; +use super::LlamaCppBackend; pub(super) struct GenerationContext<'a> { pub loaded: &'a LoadedModel, - pub runtime: &'a InferenceRuntime, + pub backend: &'a LlamaCppBackend, pub chat_messages: &'a [LlamaChatMessage], pub settings: &'a ModelSettings, pub context_limit: usize, @@ -37,10 +39,10 @@ pub(super) struct LoadedModel { /// Returns `None` if the model architecture values are unavailable. pub(super) fn estimate_max_context_for_memory( model: &LlamaModel, - runtime: &InferenceRuntime, + backend: &LlamaCppBackend, mmproj_overhead_bytes: u64, ) -> Option { - let raw_available = super::available_inference_memory_bytes(runtime); + let raw_available = backend.available_memory_bytes(); if raw_available == 0 { return None; } @@ -210,7 +212,7 @@ pub(super) fn build_sampler( /// context size. Returns `(prompt_token_count, effective_ctx)`. pub(super) fn validate_and_compute_context( loaded: &LoadedModel, - runtime: &InferenceRuntime, + backend: &LlamaCppBackend, prompt_token_count: usize, context_limit: usize, settings: &crate::providers::local_inference::local_model_registry::ModelSettings, @@ -221,7 +223,7 @@ pub(super) fn validate_and_compute_context( } else { 0 }; - let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, runtime, mmproj_overhead); + let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, backend, mmproj_overhead); let effective_ctx = effective_context_size( prompt_token_count, settings, @@ -251,7 +253,7 @@ pub(super) fn validate_and_compute_context( /// Create a llama context and prefill (decode) all prompt tokens. pub(super) fn create_and_prefill_context<'model>( loaded: &'model LoadedModel, - runtime: &InferenceRuntime, + backend: &LlamaCppBackend, tokens: &[llama_cpp_2::token::LlamaToken], effective_ctx: usize, settings: &crate::providers::local_inference::local_model_registry::ModelSettings, @@ -259,7 +261,7 @@ pub(super) fn create_and_prefill_context<'model>( let ctx_params = build_context_params(effective_ctx as u32, settings); let mut ctx = loaded .model - .new_context(runtime.backend(), ctx_params) + .new_context(backend.llama_backend(), ctx_params) .map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {}", e)))?; let n_batch = ctx.n_batch() as usize; @@ -279,7 +281,7 @@ pub(super) fn create_and_prefill_context<'model>( /// and the effective context size. pub(super) fn create_and_prefill_multimodal<'model>( loaded: &'model LoadedModel, - runtime: &InferenceRuntime, + backend: &LlamaCppBackend, prompt_text: &str, images: &[ExtractedImage], context_limit: usize, @@ -316,7 +318,7 @@ pub(super) fn create_and_prefill_multimodal<'model>( let n_ctx_train = loaded.model.n_ctx_train() as usize; let mmproj_overhead = settings.mmproj_size_bytes; - let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, runtime, mmproj_overhead); + let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, backend, mmproj_overhead); let effective_ctx = effective_context_size( prompt_token_count, settings, @@ -336,7 +338,7 @@ pub(super) fn create_and_prefill_multimodal<'model>( let ctx_params = build_context_params(effective_ctx as u32, settings); let llama_ctx = loaded .model - .new_context(runtime.backend(), ctx_params) + .new_context(backend.llama_backend(), ctx_params) .map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {e}")))?; let n_batch = llama_ctx.n_batch() as i32; diff --git a/crates/goose/src/providers/local_inference/inference_native_tools.rs b/crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs similarity index 98% rename from crates/goose/src/providers/local_inference/inference_native_tools.rs rename to crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs index e89c50c4a0..662b802445 100644 --- a/crates/goose/src/providers/local_inference/inference_native_tools.rs +++ b/crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs @@ -7,7 +7,7 @@ use serde_json::Value; use std::borrow::Cow; use uuid::Uuid; -use super::finalize_usage; +use super::super::finalize_usage; use super::inference_engine::{ context_cap, create_and_prefill_context, create_and_prefill_multimodal, estimate_max_context_for_memory, generation_loop, validate_and_compute_context, @@ -28,7 +28,7 @@ pub(super) fn generate_with_native_tools( 0 }; let memory_max_ctx = - estimate_max_context_for_memory(&ctx.loaded.model, ctx.runtime, mmproj_overhead); + estimate_max_context_for_memory(&ctx.loaded.model, ctx.backend, mmproj_overhead); let cap = context_cap(ctx.settings, ctx.context_limit, n_ctx_train, memory_max_ctx); let token_budget = cap.saturating_sub(min_generation_headroom); @@ -97,7 +97,7 @@ pub(super) fn generate_with_native_tools( let (mut llama_ctx, prompt_token_count, effective_ctx) = if !ctx.images.is_empty() { create_and_prefill_multimodal( ctx.loaded, - ctx.runtime, + ctx.backend, &template_result.prompt, ctx.images, ctx.context_limit, @@ -111,13 +111,13 @@ pub(super) fn generate_with_native_tools( .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; let (ptc, ectx) = validate_and_compute_context( ctx.loaded, - ctx.runtime, + ctx.backend, tokens.len(), ctx.context_limit, ctx.settings, )?; let lctx = - create_and_prefill_context(ctx.loaded, ctx.runtime, &tokens, ectx, ctx.settings)?; + create_and_prefill_context(ctx.loaded, ctx.backend, &tokens, ectx, ctx.settings)?; (lctx, ptc, ectx) }; diff --git a/crates/goose/src/providers/local_inference/llamacpp/mod.rs b/crates/goose/src/providers/local_inference/llamacpp/mod.rs new file mode 100644 index 0000000000..0b4141ad72 --- /dev/null +++ b/crates/goose/src/providers/local_inference/llamacpp/mod.rs @@ -0,0 +1,355 @@ +mod inference_emulated_tools; +mod inference_engine; +mod inference_native_tools; + +use std::any::Any; +use std::path::PathBuf; + +use anyhow::Result; +use llama_cpp_2::llama_backend::LlamaBackend; +use llama_cpp_2::model::params::LlamaModelParams; +use llama_cpp_2::model::{LlamaChatMessage, LlamaChatTemplate, LlamaModel}; +use llama_cpp_2::{list_llama_ggml_backend_devices, LlamaBackendDeviceType, LogOptions}; +use rmcp::model::Role; + +use self::inference_emulated_tools::{ + build_emulator_tool_description, generate_with_emulated_tools, load_tiny_model_prompt, +}; +use self::inference_engine::{GenerationContext, LoadedModel}; +use self::inference_native_tools::generate_with_native_tools; +use crate::providers::errors::ProviderError; +use crate::providers::formats::openai::format_tools; +use crate::providers::local_inference::backend::{ + BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend, +}; +use crate::providers::local_inference::multimodal::ExtractedImage; +use crate::providers::local_inference::tool_parsing::compact_tools_json; +use crate::providers::local_inference::{ + build_openai_messages_json, extract_text_content, ResolvedModelPaths, +}; + +pub(super) const LLAMACPP_BACKEND_ID: &str = "llamacpp"; + +const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript"; + +pub(super) struct LlamaCppBackend { + backend: LlamaBackend, +} + +impl LlamaCppBackend { + pub(super) fn new() -> Result { + let backend = match LlamaBackend::init() { + Ok(backend) => backend, + Err(llama_cpp_2::LlamaCppError::BackendAlreadyInitialized) => { + unreachable!( + "LlamaBackend already initialized but Weak was dead; \ + the runtime mutex prevents concurrent re-init" + ) + } + Err(e) => { + tracing::error!(error = %e, "failed to initialize local inference runtime"); + return Err(anyhow::anyhow!("Failed to init llama backend: {}", e)); + } + }; + + llama_cpp_2::send_logs_to_tracing(LogOptions::default()); + log_inference_backend_devices(); + + Ok(Self { backend }) + } + + pub(super) fn llama_backend(&self) -> &LlamaBackend { + &self.backend + } + + fn init_mtmd_context( + model: &LlamaModel, + mmproj_path: &Option, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + ) -> Option { + use llama_cpp_2::mtmd::{MtmdContext, MtmdContextParams}; + + let mmproj_path = mmproj_path.as_ref().filter(|p| p.exists())?; + + let params = MtmdContextParams { + use_gpu: true, + n_threads: settings + .n_threads + .unwrap_or_else(|| MtmdContextParams::default().n_threads), + ..MtmdContextParams::default() + }; + + match MtmdContext::init_from_file(mmproj_path.to_str().unwrap_or_default(), model, ¶ms) + { + Ok(ctx) => { + tracing::info!( + vision = ctx.support_vision(), + audio = ctx.support_audio(), + "Multimodal context initialized" + ); + Some(ctx) + } + Err(e) => { + tracing::warn!(error = %e, "Failed to init multimodal context"); + None + } + } + } +} + +impl LocalInferenceBackend for LlamaCppBackend { + fn id(&self) -> &'static str { + LLAMACPP_BACKEND_ID + } + + fn load_model( + &self, + model_id: &str, + resolved: &ResolvedModelPaths, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + ) -> Result, ProviderError> { + let model_path = &resolved.model_path; + + if !model_path.exists() { + return Err(ProviderError::ExecutionError(format!( + "Model not downloaded: {}. Please download it from Settings > Local Inference.", + model_id + ))); + } + + tracing::info!( + backend = self.id(), + "Loading {} from: {}", + model_id, + model_path.display() + ); + + let mut params = LlamaModelParams::default(); + if let Some(n_gpu_layers) = settings.n_gpu_layers { + params = params.with_n_gpu_layers(n_gpu_layers); + } + if settings.use_mlock { + params = params.with_use_mlock(true); + } + let model = LlamaModel::load_from_file(&self.backend, model_path, ¶ms) + .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; + + let template = match model.chat_template(None) { + Ok(t) => t, + Err(_) => { + tracing::warn!("Model has no embedded chat template, falling back to chatml"); + LlamaChatTemplate::new("chatml").map_err(|e| { + ProviderError::ExecutionError(format!( + "Failed to create fallback chat template: {}", + e + )) + })? + } + }; + + let mtmd_ctx = Self::init_mtmd_context(&model, &resolved.mmproj_path, settings); + + tracing::info!( + backend = self.id(), + model_id = model_id, + "Model loaded successfully" + ); + + Ok(Box::new(LoadedModel { + model, + template, + mtmd_ctx, + })) + } + + fn generate( + &self, + loaded: &mut dyn BackendLoadedModel, + request: LocalGenerationRequest<'_>, + ) -> Result<(), ProviderError> { + let loaded = loaded + .as_any_mut() + .downcast_mut::() + .ok_or_else(|| { + ProviderError::ExecutionError("Loaded model backend mismatch".to_string()) + })?; + + let native_tool_calling = request.settings.native_tool_calling; + let use_emulator = !native_tool_calling && !request.tools.is_empty(); + let system_prompt = if use_emulator { + load_tiny_model_prompt() + } else { + request.system.to_string() + }; + + let has_vision = request.resolved_model.mmproj_path.is_some(); + let marker = llama_cpp_2::mtmd::mtmd_default_marker(); + let (images, vision_messages): (Vec, Option>) = if has_vision { + let (imgs, msgs) = + super::multimodal::extract_images_from_messages(request.messages, marker); + (imgs, Some(msgs)) + } else { + (Vec::new(), None) + }; + let effective_messages = vision_messages.as_deref().unwrap_or(request.messages); + + let mut chat_messages = + vec![ + LlamaChatMessage::new("system".to_string(), system_prompt.clone()).map_err( + |e| { + ProviderError::ExecutionError(format!( + "Failed to create system message: {}", + e + )) + }, + )?, + ]; + + let code_mode_enabled = request.tools.iter().any(|t| t.name == CODE_EXECUTION_TOOL); + + if use_emulator && !request.tools.is_empty() { + let tool_desc = build_emulator_tool_description(request.tools, code_mode_enabled); + chat_messages = vec![LlamaChatMessage::new( + "system".to_string(), + format!("{}{}", system_prompt, tool_desc), + ) + .map_err(|e| { + ProviderError::ExecutionError(format!("Failed to create system message: {}", e)) + })?]; + } + + for msg in effective_messages { + let role = match msg.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + let content = extract_text_content(msg); + if !content.trim().is_empty() { + chat_messages.push(LlamaChatMessage::new(role.to_string(), content).map_err( + |e| ProviderError::ExecutionError(format!("Failed to create message: {}", e)), + )?); + } + } + + let (full_tools_json, compact_tools) = if !use_emulator && !request.tools.is_empty() { + let full = format_tools(request.tools) + .ok() + .and_then(|spec| serde_json::to_string(&spec).ok()); + let compact = compact_tools_json(request.tools); + (full, compact) + } else { + (None, None) + }; + + let oai_messages_json = if request.settings.use_jinja || native_tool_calling { + Some(build_openai_messages_json( + &system_prompt, + effective_messages, + )) + } else { + None + }; + + if !images.is_empty() && loaded.mtmd_ctx.is_none() { + loaded.mtmd_ctx = Self::init_mtmd_context( + &loaded.model, + &request.resolved_model.mmproj_path, + request.settings, + ); + } + + let mut gen_ctx = GenerationContext { + loaded, + backend: self, + chat_messages: &chat_messages, + settings: request.settings, + context_limit: request.context_limit, + model_name: request.model_name, + message_id: request.message_id, + tx: request.tx, + log: request.log, + images: &images, + }; + + if use_emulator { + generate_with_emulated_tools(&mut gen_ctx, code_mode_enabled) + } else { + generate_with_native_tools( + &mut gen_ctx, + &oai_messages_json, + full_tools_json.as_deref(), + compact_tools.as_deref(), + ) + } + } + + fn available_memory_bytes(&self) -> u64 { + let devices = list_llama_ggml_backend_devices(); + + let accel_memory = devices + .iter() + .filter(|d| is_accelerator_device(d.device_type)) + .map(|d| d.memory_free as u64) + .max() + .unwrap_or(0); + + if accel_memory > 0 { + accel_memory + } else { + devices + .iter() + .filter(|d| d.device_type == LlamaBackendDeviceType::Cpu) + .map(|d| d.memory_free as u64) + .max() + .unwrap_or(0) + } + } +} + +impl BackendLoadedModel for LoadedModel { + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +fn is_accelerator_device(device_type: LlamaBackendDeviceType) -> bool { + matches!( + device_type, + LlamaBackendDeviceType::Gpu + | LlamaBackendDeviceType::IntegratedGpu + | LlamaBackendDeviceType::Accelerator + ) +} + +fn is_non_cpu_device(device_type: LlamaBackendDeviceType) -> bool { + !matches!(device_type, LlamaBackendDeviceType::Cpu) +} + +fn log_inference_backend_devices() { + let devices = list_llama_ggml_backend_devices(); + let non_cpu_devices: Vec<_> = devices + .iter() + .filter(|device| is_non_cpu_device(device.device_type)) + .collect(); + + if non_cpu_devices.is_empty() { + tracing::info!( + device_count = devices.len(), + "No non-CPU llama.cpp backend devices detected for local inference" + ); + return; + } + + for device in non_cpu_devices { + tracing::info!( + index = device.index, + backend = %device.backend, + name = %device.name, + description = %device.description, + device_type = ?device.device_type, + memory_total_bytes = device.memory_total as u64, + memory_free_bytes = device.memory_free as u64, + "Non-CPU llama.cpp backend device detected for local inference" + ); + } +}