refactor local inference around backends (#9137)
Signed-off-by: jh-block <jhugo@block.xyz>
This commit is contained in:
@@ -1,19 +1,10 @@
|
|||||||
|
mod backend;
|
||||||
pub mod hf_models;
|
pub mod hf_models;
|
||||||
mod inference_emulated_tools;
|
mod llamacpp;
|
||||||
mod inference_engine;
|
|
||||||
mod inference_native_tools;
|
|
||||||
pub mod local_model_registry;
|
pub mod local_model_registry;
|
||||||
pub(crate) mod multimodal;
|
pub(crate) mod multimodal;
|
||||||
mod tool_parsing;
|
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::config::ExtensionConfig;
|
||||||
use crate::conversation::message::{Message, MessageContent};
|
use crate::conversation::message::{Message, MessageContent};
|
||||||
use crate::model::ModelConfig;
|
use crate::model::ModelConfig;
|
||||||
@@ -21,18 +12,14 @@ use crate::providers::base::{
|
|||||||
MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage,
|
MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage,
|
||||||
};
|
};
|
||||||
use crate::providers::errors::ProviderError;
|
use crate::providers::errors::ProviderError;
|
||||||
use crate::providers::formats::openai::format_tools;
|
|
||||||
use crate::providers::utils::RequestLog;
|
use crate::providers::utils::RequestLog;
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use async_stream::try_stream;
|
use async_stream::try_stream;
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use backend::{BackendLoadedModel, LocalInferenceBackend};
|
||||||
use futures::future::BoxFuture;
|
use futures::future::BoxFuture;
|
||||||
use llama_cpp_2::llama_backend::LlamaBackend;
|
use llamacpp::{LlamaCppBackend, LLAMACPP_BACKEND_ID};
|
||||||
use llama_cpp_2::model::params::LlamaModelParams;
|
use rmcp::model::Tool;
|
||||||
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 serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
@@ -40,18 +27,26 @@ use std::sync::{Arc, Mutex as StdMutex, Weak};
|
|||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
const SHELL_TOOL: &str = "developer__shell";
|
type ModelSlot = Arc<Mutex<Option<Box<dyn BackendLoadedModel>>>>;
|
||||||
const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript";
|
|
||||||
|
|
||||||
type ModelSlot = Arc<Mutex<Option<LoadedModel>>>;
|
#[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<String>) -> 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 {
|
pub struct InferenceRuntime {
|
||||||
models: StdMutex<HashMap<String, ModelSlot>>,
|
models: StdMutex<HashMap<ModelCacheKey, ModelSlot>>,
|
||||||
backend: LlamaBackend,
|
backends: HashMap<&'static str, Arc<dyn LocalInferenceBackend>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Global weak reference used to share a single `InferenceRuntime` across
|
/// Global weak reference used to share a single `InferenceRuntime` across
|
||||||
@@ -67,49 +62,47 @@ impl InferenceRuntime {
|
|||||||
if let Some(runtime) = guard.upgrade() {
|
if let Some(runtime) = guard.upgrade() {
|
||||||
return Ok(runtime);
|
return Ok(runtime);
|
||||||
}
|
}
|
||||||
// Safety invariant: the Weak::upgrade() check and LlamaBackend::init()
|
let llamacpp_backend: Arc<dyn LocalInferenceBackend> = Arc::new(LlamaCppBackend::new()?);
|
||||||
// both execute inside this same mutex guard, so there is no window where
|
let mut backends = HashMap::new();
|
||||||
// another thread could drop the Arc and re-enter concurrently.
|
backends.insert(LLAMACPP_BACKEND_ID, llamacpp_backend);
|
||||||
// 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 runtime = Arc::new(Self {
|
let runtime = Arc::new(Self {
|
||||||
models: StdMutex::new(HashMap::new()),
|
models: StdMutex::new(HashMap::new()),
|
||||||
backend,
|
backends,
|
||||||
});
|
});
|
||||||
*guard = Arc::downgrade(&runtime);
|
*guard = Arc::downgrade(&runtime);
|
||||||
Ok(runtime)
|
Ok(runtime)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn backend(&self) -> &LlamaBackend {
|
fn default_backend(&self) -> &dyn LocalInferenceBackend {
|
||||||
&self.backend
|
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<Arc<dyn LocalInferenceBackend>, 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");
|
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)))
|
.or_insert_with(|| Arc::new(Mutex::new(None)))
|
||||||
.clone()
|
.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn other_model_slots(&self, keep_model_id: &str) -> Vec<ModelSlot> {
|
fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec<ModelSlot> {
|
||||||
let map = self.models.lock().expect("model cache lock poisoned");
|
let map = self.models.lock().expect("model cache lock poisoned");
|
||||||
map.iter()
|
map.iter()
|
||||||
.filter(|(id, _)| id.as_str() != keep_model_id)
|
.filter(|(key, _)| *key != keep_key)
|
||||||
.map(|(_, slot)| slot.clone())
|
.map(|(_, slot)| slot.clone())
|
||||||
.collect()
|
.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 const LOCAL_LLM_MODEL_CONFIG_KEY: &str = "LOCAL_LLM_MODEL";
|
||||||
|
|
||||||
pub struct ResolvedModelPaths {
|
#[derive(Clone)]
|
||||||
|
pub(super) struct ResolvedModelPaths {
|
||||||
pub model_path: PathBuf,
|
pub model_path: PathBuf,
|
||||||
pub context_limit: usize,
|
pub context_limit: usize,
|
||||||
pub settings: crate::providers::local_inference::local_model_registry::ModelSettings,
|
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.
|
/// 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<ResolvedModelPaths> {
|
fn resolve_model_path(model_id: &str) -> Option<ResolvedModelPaths> {
|
||||||
use crate::providers::local_inference::local_model_registry::{
|
use crate::providers::local_inference::local_model_registry::{
|
||||||
default_settings_for_model, get_registry,
|
default_settings_for_model, get_registry,
|
||||||
};
|
};
|
||||||
@@ -157,69 +151,8 @@ pub fn resolve_model_path(model_id: &str) -> Option<ResolvedModelPaths> {
|
|||||||
None
|
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 {
|
pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 {
|
||||||
let _ = &runtime.backend;
|
runtime.default_backend().available_memory_bytes()
|
||||||
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)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn recommend_local_model(runtime: &InferenceRuntime) -> String {
|
pub fn recommend_local_model(runtime: &InferenceRuntime) -> String {
|
||||||
@@ -389,7 +322,6 @@ type StreamSender =
|
|||||||
|
|
||||||
pub struct LocalInferenceProvider {
|
pub struct LocalInferenceProvider {
|
||||||
runtime: Arc<InferenceRuntime>,
|
runtime: Arc<InferenceRuntime>,
|
||||||
model: ModelSlot,
|
|
||||||
model_config: ModelConfig,
|
model_config: ModelConfig,
|
||||||
name: String,
|
name: String,
|
||||||
}
|
}
|
||||||
@@ -397,102 +329,12 @@ pub struct LocalInferenceProvider {
|
|||||||
impl LocalInferenceProvider {
|
impl LocalInferenceProvider {
|
||||||
pub async fn from_env(model: ModelConfig, _extensions: Vec<ExtensionConfig>) -> Result<Self> {
|
pub async fn from_env(model: ModelConfig, _extensions: Vec<ExtensionConfig>) -> Result<Self> {
|
||||||
let runtime = InferenceRuntime::get_or_init()?;
|
let runtime = InferenceRuntime::get_or_init()?;
|
||||||
let model_slot = runtime.get_or_create_model_slot(&model.model_name);
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
runtime,
|
runtime,
|
||||||
model: model_slot,
|
|
||||||
model_config: model,
|
model_config: model,
|
||||||
name: PROVIDER_NAME.to_string(),
|
name: PROVIDER_NAME.to_string(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn load_model_sync(
|
|
||||||
runtime: &InferenceRuntime,
|
|
||||||
model_id: &str,
|
|
||||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
|
||||||
) -> Result<LoadedModel, ProviderError> {
|
|
||||||
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<PathBuf>,
|
|
||||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
|
||||||
) -> Option<llama_cpp_2::mtmd::MtmdContext> {
|
|
||||||
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 {
|
impl ProviderDef for LocalInferenceProvider {
|
||||||
@@ -577,14 +419,17 @@ impl Provider for LocalInferenceProvider {
|
|||||||
let resolved = resolve_model_path(&model_config.model_name).ok_or_else(|| {
|
let resolved = resolve_model_path(&model_config.model_name).ok_or_else(|| {
|
||||||
ProviderError::ExecutionError(format!("Model not found: {}", model_config.model_name))
|
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_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.
|
// 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() {
|
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;
|
let mut other = slot.lock().await;
|
||||||
if other.is_some() {
|
if other.is_some() {
|
||||||
tracing::info!("Unloading previous model to free memory");
|
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 model_id = model_config.model_name.clone();
|
||||||
|
let resolved_for_load = resolved.clone();
|
||||||
let settings_for_load = model_settings.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 || {
|
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
|
.await
|
||||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))??;
|
.map_err(|e| ProviderError::ExecutionError(e.to_string()))??;
|
||||||
@@ -612,98 +458,20 @@ impl Provider for LocalInferenceProvider {
|
|||||||
model_settings.enable_thinking = false;
|
model_settings.enable_thinking = false;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Use the model's native_tool_calling setting to decide the path.
|
let model_arc = model_slot.clone();
|
||||||
// Featured models have this set explicitly; user-added models default to false.
|
let backend = backend.clone();
|
||||||
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<ExtractedImage>, Option<Vec<Message>>) = 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_name = model_config.model_name.clone();
|
let model_name = model_config.model_name.clone();
|
||||||
let context_limit = model_context_limit;
|
let context_limit = model_context_limit;
|
||||||
let settings = model_settings;
|
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!({
|
let log_payload = serde_json::json!({
|
||||||
"system": &system_prompt,
|
"system": &system,
|
||||||
"messages": messages.iter().map(|m| {
|
"messages": messages.iter().map(|m| {
|
||||||
serde_json::json!({
|
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),
|
"content": extract_text_content(m),
|
||||||
})
|
})
|
||||||
}).collect::<Vec<_>>(),
|
}).collect::<Vec<_>>(),
|
||||||
@@ -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 message_id = Uuid::new_v4().to_string();
|
||||||
|
|
||||||
let mut gen_ctx = GenerationContext {
|
let request = backend::LocalGenerationRequest {
|
||||||
loaded,
|
model_name,
|
||||||
runtime: &runtime,
|
system: &system,
|
||||||
chat_messages: &chat_messages,
|
messages: &messages,
|
||||||
|
tools: &tools,
|
||||||
settings: &settings,
|
settings: &settings,
|
||||||
context_limit,
|
context_limit,
|
||||||
model_name,
|
resolved_model: &resolved_model,
|
||||||
message_id: &message_id,
|
message_id: &message_id,
|
||||||
tx: &tx,
|
tx: &tx,
|
||||||
log: &mut log,
|
log: &mut log,
|
||||||
images: &images,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let result = if use_emulator {
|
let result = backend.generate(loaded.as_mut(), request);
|
||||||
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(),
|
|
||||||
)
|
|
||||||
};
|
|
||||||
|
|
||||||
if let Err(err) = result {
|
if let Err(err) = result {
|
||||||
let msg = match &err {
|
let msg = match &err {
|
||||||
|
|||||||
@@ -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<Box<dyn BackendLoadedModel>, ProviderError>;
|
||||||
|
|
||||||
|
fn generate(
|
||||||
|
&self,
|
||||||
|
loaded: &mut dyn BackendLoadedModel,
|
||||||
|
request: LocalGenerationRequest<'_>,
|
||||||
|
) -> Result<(), ProviderError>;
|
||||||
|
|
||||||
|
fn available_memory_bytes(&self) -> u64;
|
||||||
|
}
|
||||||
+7
-4
@@ -28,11 +28,14 @@ use serde_json::json;
|
|||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use super::super::{finalize_usage, StreamSender};
|
||||||
use super::inference_engine::{
|
use super::inference_engine::{
|
||||||
create_and_prefill_context, create_and_prefill_multimodal, generation_loop,
|
create_and_prefill_context, create_and_prefill_multimodal, generation_loop,
|
||||||
validate_and_compute_context, GenerationContext, TokenAction,
|
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_CODE_MODE: usize = " ```execute_typescript\n".len();
|
||||||
const HOLD_BACK_SHELL_ONLY: usize = "\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() {
|
let (mut llama_ctx, prompt_token_count, effective_ctx) = if !ctx.images.is_empty() {
|
||||||
create_and_prefill_multimodal(
|
create_and_prefill_multimodal(
|
||||||
ctx.loaded,
|
ctx.loaded,
|
||||||
ctx.runtime,
|
ctx.backend,
|
||||||
&prompt,
|
&prompt,
|
||||||
ctx.images,
|
ctx.images,
|
||||||
ctx.context_limit,
|
ctx.context_limit,
|
||||||
@@ -387,13 +390,13 @@ pub(super) fn generate_with_emulated_tools(
|
|||||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||||
let (ptc, ectx) = validate_and_compute_context(
|
let (ptc, ectx) = validate_and_compute_context(
|
||||||
ctx.loaded,
|
ctx.loaded,
|
||||||
ctx.runtime,
|
ctx.backend,
|
||||||
tokens.len(),
|
tokens.len(),
|
||||||
ctx.context_limit,
|
ctx.context_limit,
|
||||||
ctx.settings,
|
ctx.settings,
|
||||||
)?;
|
)?;
|
||||||
let lctx =
|
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)
|
(lctx, ptc, ectx)
|
||||||
};
|
};
|
||||||
|
|
||||||
+13
-11
@@ -1,4 +1,5 @@
|
|||||||
use crate::providers::errors::ProviderError;
|
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::local_model_registry::ModelSettings;
|
||||||
use crate::providers::local_inference::multimodal::ExtractedImage;
|
use crate::providers::local_inference::multimodal::ExtractedImage;
|
||||||
use crate::providers::utils::RequestLog;
|
use crate::providers::utils::RequestLog;
|
||||||
@@ -9,11 +10,12 @@ use llama_cpp_2::mtmd::{MtmdBitmap, MtmdContext, MtmdInputText};
|
|||||||
use llama_cpp_2::sampling::LlamaSampler;
|
use llama_cpp_2::sampling::LlamaSampler;
|
||||||
use std::num::NonZeroU32;
|
use std::num::NonZeroU32;
|
||||||
|
|
||||||
use super::{InferenceRuntime, StreamSender};
|
use super::super::StreamSender;
|
||||||
|
use super::LlamaCppBackend;
|
||||||
|
|
||||||
pub(super) struct GenerationContext<'a> {
|
pub(super) struct GenerationContext<'a> {
|
||||||
pub loaded: &'a LoadedModel,
|
pub loaded: &'a LoadedModel,
|
||||||
pub runtime: &'a InferenceRuntime,
|
pub backend: &'a LlamaCppBackend,
|
||||||
pub chat_messages: &'a [LlamaChatMessage],
|
pub chat_messages: &'a [LlamaChatMessage],
|
||||||
pub settings: &'a ModelSettings,
|
pub settings: &'a ModelSettings,
|
||||||
pub context_limit: usize,
|
pub context_limit: usize,
|
||||||
@@ -37,10 +39,10 @@ pub(super) struct LoadedModel {
|
|||||||
/// Returns `None` if the model architecture values are unavailable.
|
/// Returns `None` if the model architecture values are unavailable.
|
||||||
pub(super) fn estimate_max_context_for_memory(
|
pub(super) fn estimate_max_context_for_memory(
|
||||||
model: &LlamaModel,
|
model: &LlamaModel,
|
||||||
runtime: &InferenceRuntime,
|
backend: &LlamaCppBackend,
|
||||||
mmproj_overhead_bytes: u64,
|
mmproj_overhead_bytes: u64,
|
||||||
) -> Option<usize> {
|
) -> Option<usize> {
|
||||||
let raw_available = super::available_inference_memory_bytes(runtime);
|
let raw_available = backend.available_memory_bytes();
|
||||||
if raw_available == 0 {
|
if raw_available == 0 {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -210,7 +212,7 @@ pub(super) fn build_sampler(
|
|||||||
/// context size. Returns `(prompt_token_count, effective_ctx)`.
|
/// context size. Returns `(prompt_token_count, effective_ctx)`.
|
||||||
pub(super) fn validate_and_compute_context(
|
pub(super) fn validate_and_compute_context(
|
||||||
loaded: &LoadedModel,
|
loaded: &LoadedModel,
|
||||||
runtime: &InferenceRuntime,
|
backend: &LlamaCppBackend,
|
||||||
prompt_token_count: usize,
|
prompt_token_count: usize,
|
||||||
context_limit: usize,
|
context_limit: usize,
|
||||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||||
@@ -221,7 +223,7 @@ pub(super) fn validate_and_compute_context(
|
|||||||
} else {
|
} else {
|
||||||
0
|
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(
|
let effective_ctx = effective_context_size(
|
||||||
prompt_token_count,
|
prompt_token_count,
|
||||||
settings,
|
settings,
|
||||||
@@ -251,7 +253,7 @@ pub(super) fn validate_and_compute_context(
|
|||||||
/// Create a llama context and prefill (decode) all prompt tokens.
|
/// Create a llama context and prefill (decode) all prompt tokens.
|
||||||
pub(super) fn create_and_prefill_context<'model>(
|
pub(super) fn create_and_prefill_context<'model>(
|
||||||
loaded: &'model LoadedModel,
|
loaded: &'model LoadedModel,
|
||||||
runtime: &InferenceRuntime,
|
backend: &LlamaCppBackend,
|
||||||
tokens: &[llama_cpp_2::token::LlamaToken],
|
tokens: &[llama_cpp_2::token::LlamaToken],
|
||||||
effective_ctx: usize,
|
effective_ctx: usize,
|
||||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
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 ctx_params = build_context_params(effective_ctx as u32, settings);
|
||||||
let mut ctx = loaded
|
let mut ctx = loaded
|
||||||
.model
|
.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)))?;
|
.map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {}", e)))?;
|
||||||
|
|
||||||
let n_batch = ctx.n_batch() as usize;
|
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.
|
/// and the effective context size.
|
||||||
pub(super) fn create_and_prefill_multimodal<'model>(
|
pub(super) fn create_and_prefill_multimodal<'model>(
|
||||||
loaded: &'model LoadedModel,
|
loaded: &'model LoadedModel,
|
||||||
runtime: &InferenceRuntime,
|
backend: &LlamaCppBackend,
|
||||||
prompt_text: &str,
|
prompt_text: &str,
|
||||||
images: &[ExtractedImage],
|
images: &[ExtractedImage],
|
||||||
context_limit: usize,
|
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 n_ctx_train = loaded.model.n_ctx_train() as usize;
|
||||||
let mmproj_overhead = settings.mmproj_size_bytes;
|
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(
|
let effective_ctx = effective_context_size(
|
||||||
prompt_token_count,
|
prompt_token_count,
|
||||||
settings,
|
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 ctx_params = build_context_params(effective_ctx as u32, settings);
|
||||||
let llama_ctx = loaded
|
let llama_ctx = loaded
|
||||||
.model
|
.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}")))?;
|
.map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {e}")))?;
|
||||||
|
|
||||||
let n_batch = llama_ctx.n_batch() as i32;
|
let n_batch = llama_ctx.n_batch() as i32;
|
||||||
+5
-5
@@ -7,7 +7,7 @@ use serde_json::Value;
|
|||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use super::finalize_usage;
|
use super::super::finalize_usage;
|
||||||
use super::inference_engine::{
|
use super::inference_engine::{
|
||||||
context_cap, create_and_prefill_context, create_and_prefill_multimodal,
|
context_cap, create_and_prefill_context, create_and_prefill_multimodal,
|
||||||
estimate_max_context_for_memory, generation_loop, validate_and_compute_context,
|
estimate_max_context_for_memory, generation_loop, validate_and_compute_context,
|
||||||
@@ -28,7 +28,7 @@ pub(super) fn generate_with_native_tools(
|
|||||||
0
|
0
|
||||||
};
|
};
|
||||||
let memory_max_ctx =
|
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 cap = context_cap(ctx.settings, ctx.context_limit, n_ctx_train, memory_max_ctx);
|
||||||
let token_budget = cap.saturating_sub(min_generation_headroom);
|
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() {
|
let (mut llama_ctx, prompt_token_count, effective_ctx) = if !ctx.images.is_empty() {
|
||||||
create_and_prefill_multimodal(
|
create_and_prefill_multimodal(
|
||||||
ctx.loaded,
|
ctx.loaded,
|
||||||
ctx.runtime,
|
ctx.backend,
|
||||||
&template_result.prompt,
|
&template_result.prompt,
|
||||||
ctx.images,
|
ctx.images,
|
||||||
ctx.context_limit,
|
ctx.context_limit,
|
||||||
@@ -111,13 +111,13 @@ pub(super) fn generate_with_native_tools(
|
|||||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||||
let (ptc, ectx) = validate_and_compute_context(
|
let (ptc, ectx) = validate_and_compute_context(
|
||||||
ctx.loaded,
|
ctx.loaded,
|
||||||
ctx.runtime,
|
ctx.backend,
|
||||||
tokens.len(),
|
tokens.len(),
|
||||||
ctx.context_limit,
|
ctx.context_limit,
|
||||||
ctx.settings,
|
ctx.settings,
|
||||||
)?;
|
)?;
|
||||||
let lctx =
|
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)
|
(lctx, ptc, ectx)
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -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<Self> {
|
||||||
|
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<PathBuf>,
|
||||||
|
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||||
|
) -> Option<llama_cpp_2::mtmd::MtmdContext> {
|
||||||
|
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<Box<dyn BackendLoadedModel>, 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::<LoadedModel>()
|
||||||
|
.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<ExtractedImage>, Option<Vec<_>>) = 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"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user