mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
feat: add local inference provider with llama.cpp backend and HuggingFace model management (#6933)
Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: jh-block <jhugo@block.xyz> Co-authored-by: Spence <spencermartin@squareup.com> Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Douwe Osinga
@@ -854,6 +854,13 @@ enum Command {
|
||||
#[command(subcommand)]
|
||||
command: TermCommand,
|
||||
},
|
||||
/// Manage local inference models
|
||||
#[command(about = "Manage local inference models", visible_alias = "lm")]
|
||||
LocalModels {
|
||||
#[command(subcommand)]
|
||||
command: LocalModelsCommand,
|
||||
},
|
||||
|
||||
/// Generate completions for various shells
|
||||
#[command(about = "Generate the autocompletion script for the specified shell")]
|
||||
Completion {
|
||||
@@ -875,6 +882,38 @@ enum Command {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum LocalModelsCommand {
|
||||
/// Search HuggingFace for GGUF models
|
||||
#[command(about = "Search HuggingFace for GGUF models")]
|
||||
Search {
|
||||
/// Search query
|
||||
query: String,
|
||||
|
||||
/// Maximum number of results
|
||||
#[arg(short, long, default_value = "10")]
|
||||
limit: usize,
|
||||
},
|
||||
|
||||
/// Download a model from HuggingFace
|
||||
#[command(about = "Download a GGUF model (e.g. bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M)")]
|
||||
Download {
|
||||
/// Model spec in user/repo:quantization format
|
||||
spec: String,
|
||||
},
|
||||
|
||||
/// List downloaded local models
|
||||
#[command(about = "List downloaded local models")]
|
||||
List,
|
||||
|
||||
/// Delete a downloaded model
|
||||
#[command(about = "Delete a downloaded local model")]
|
||||
Delete {
|
||||
/// Model ID to delete
|
||||
id: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum TermCommand {
|
||||
/// Print shell initialization script
|
||||
@@ -964,6 +1003,7 @@ fn get_command_name(command: &Option<Command>) -> &'static str {
|
||||
Some(Command::Recipe { .. }) => "recipe",
|
||||
Some(Command::Web { .. }) => "web",
|
||||
Some(Command::Term { .. }) => "term",
|
||||
Some(Command::LocalModels { .. }) => "local-models",
|
||||
Some(Command::Completion { .. }) => "completion",
|
||||
Some(Command::ValidateExtensions { .. }) => "validate-extensions",
|
||||
None => "default_session",
|
||||
@@ -1400,6 +1440,170 @@ async fn handle_term_subcommand(command: TermCommand) -> Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_local_models_command(command: LocalModelsCommand) -> Result<()> {
|
||||
use goose::providers::local_inference::hf_models;
|
||||
use goose::providers::local_inference::local_model_registry::{
|
||||
display_name_from_repo, get_registry, model_id_from_repo, LocalModelEntry,
|
||||
};
|
||||
|
||||
match command {
|
||||
LocalModelsCommand::Search { query, limit } => {
|
||||
println!("Searching HuggingFace for '{}'...", query);
|
||||
let results = hf_models::search_gguf_models(&query, limit).await?;
|
||||
|
||||
if results.is_empty() {
|
||||
println!("No GGUF models found.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for model in &results {
|
||||
println!(
|
||||
"\n{} (by {}) — {} downloads",
|
||||
model.model_name, model.author, model.downloads
|
||||
);
|
||||
for file in &model.gguf_files {
|
||||
let size = if file.size_bytes > 0 {
|
||||
format!(
|
||||
"{:.1}GB",
|
||||
file.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
|
||||
)
|
||||
} else {
|
||||
"unknown".to_string()
|
||||
};
|
||||
println!(" {} — {}", file.quantization, size);
|
||||
}
|
||||
println!(
|
||||
" Download: goose local-models download {}:<quantization>",
|
||||
model.repo_id
|
||||
);
|
||||
}
|
||||
}
|
||||
LocalModelsCommand::Download { spec } => {
|
||||
println!("Resolving {}...", spec);
|
||||
let (repo_id, file) = hf_models::resolve_model_spec(&spec).await?;
|
||||
let model_id = model_id_from_repo(&repo_id, &file.quantization);
|
||||
let display_name = display_name_from_repo(&repo_id, &file.quantization);
|
||||
let local_path =
|
||||
goose::config::paths::Paths::in_data_dir("models").join(&file.filename);
|
||||
|
||||
println!(
|
||||
"Downloading {} ({})...",
|
||||
display_name,
|
||||
if file.size_bytes > 0 {
|
||||
format!(
|
||||
"{:.1}GB",
|
||||
file.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
|
||||
)
|
||||
} else {
|
||||
"unknown size".to_string()
|
||||
}
|
||||
);
|
||||
|
||||
// Register
|
||||
let entry = LocalModelEntry {
|
||||
id: model_id.clone(),
|
||||
display_name: display_name.clone(),
|
||||
repo_id: repo_id.clone(),
|
||||
filename: file.filename.clone(),
|
||||
quantization: file.quantization.clone(),
|
||||
local_path: local_path.clone(),
|
||||
source_url: file.download_url.clone(),
|
||||
settings: Default::default(),
|
||||
size_bytes: file.size_bytes,
|
||||
};
|
||||
|
||||
{
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?;
|
||||
registry.add_model(entry)?;
|
||||
}
|
||||
|
||||
// Download
|
||||
let manager = goose::download_manager::get_download_manager();
|
||||
manager
|
||||
.download_model(
|
||||
format!("{}-model", model_id),
|
||||
file.download_url,
|
||||
local_path,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Poll progress
|
||||
loop {
|
||||
if let Some(progress) = manager.get_progress(&format!("{}-model", model_id)) {
|
||||
match progress.status {
|
||||
goose::download_manager::DownloadStatus::Downloading => {
|
||||
print!(
|
||||
"\r {:.1}% ({:.0}MB / {:.0}MB)",
|
||||
progress.progress_percent,
|
||||
progress.bytes_downloaded as f64 / (1024.0 * 1024.0),
|
||||
progress.total_bytes as f64 / (1024.0 * 1024.0),
|
||||
);
|
||||
use std::io::Write;
|
||||
std::io::stdout().flush().ok();
|
||||
}
|
||||
goose::download_manager::DownloadStatus::Completed => {
|
||||
println!("\nDownloaded: {} (id: {})", display_name, model_id);
|
||||
break;
|
||||
}
|
||||
goose::download_manager::DownloadStatus::Failed => {
|
||||
let err = progress.error.unwrap_or_default();
|
||||
anyhow::bail!("Download failed: {}", err);
|
||||
}
|
||||
goose::download_manager::DownloadStatus::Cancelled => {
|
||||
println!("\nDownload cancelled.");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
LocalModelsCommand::List => {
|
||||
let registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?;
|
||||
let models = registry.list_models();
|
||||
|
||||
if models.is_empty() {
|
||||
println!("No local models downloaded.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("{:<40} {:<20} {:<10} Downloaded", "ID", "Name", "Quant");
|
||||
println!("{}", "-".repeat(80));
|
||||
for m in models {
|
||||
println!(
|
||||
"{:<40} {:<20} {:<10} {}",
|
||||
m.id,
|
||||
m.display_name,
|
||||
m.quantization,
|
||||
if m.is_downloaded() { "✓" } else { "✗" }
|
||||
);
|
||||
}
|
||||
}
|
||||
LocalModelsCommand::Delete { id } => {
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?;
|
||||
|
||||
if let Some(entry) = registry.get_model(&id) {
|
||||
if entry.local_path.exists() {
|
||||
std::fs::remove_file(&entry.local_path)?;
|
||||
}
|
||||
registry.remove_model(&id)?;
|
||||
println!("Deleted model: {}", id);
|
||||
} else {
|
||||
println!("Model not found: {}", id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_default_session() -> Result<()> {
|
||||
if !Config::global().exists() {
|
||||
return handle_configure().await;
|
||||
@@ -1530,6 +1734,7 @@ pub async fn cli() -> anyhow::Result<()> {
|
||||
no_auth,
|
||||
}) => crate::commands::web::handle_web(port, host, open, auth_token, no_auth).await,
|
||||
Some(Command::Term { command }) => handle_term_subcommand(command).await,
|
||||
Some(Command::LocalModels { command }) => handle_local_models_command(command).await,
|
||||
Some(Command::ValidateExtensions { file }) => {
|
||||
use goose::agents::validate_extensions::validate_bundled_extensions;
|
||||
match validate_bundled_extensions(&file) {
|
||||
|
||||
@@ -4,7 +4,7 @@ use goose::agents::ExtensionConfig;
|
||||
use goose::config::permission::PermissionLevel;
|
||||
use goose::config::ExtensionEntry;
|
||||
use goose::conversation::Conversation;
|
||||
use goose::dictation::download_manager::{DownloadProgress, DownloadStatus};
|
||||
use goose::download_manager::{DownloadProgress, DownloadStatus};
|
||||
use goose::model::ModelConfig;
|
||||
use goose::permission::permission_confirmation::{Permission, PrincipalType};
|
||||
use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType};
|
||||
@@ -424,6 +424,15 @@ derive_utoipa!(Icon as IconSchema);
|
||||
super::routes::dictation::get_download_progress,
|
||||
super::routes::dictation::cancel_download,
|
||||
super::routes::dictation::delete_model,
|
||||
super::routes::local_inference::list_local_models,
|
||||
super::routes::local_inference::search_hf_models,
|
||||
super::routes::local_inference::get_repo_files,
|
||||
super::routes::local_inference::download_hf_model,
|
||||
super::routes::local_inference::get_local_model_download_progress,
|
||||
super::routes::local_inference::cancel_local_model_download,
|
||||
super::routes::local_inference::delete_local_model,
|
||||
super::routes::local_inference::get_model_settings,
|
||||
super::routes::local_inference::update_model_settings,
|
||||
),
|
||||
components(schemas(
|
||||
super::routes::config_management::UpsertConfigQuery,
|
||||
@@ -592,6 +601,15 @@ derive_utoipa!(Icon as IconSchema);
|
||||
goose::dictation::providers::DictationProvider,
|
||||
super::routes::dictation::DictationProviderStatus,
|
||||
super::routes::dictation::WhisperModelResponse,
|
||||
super::routes::local_inference::LocalModelResponse,
|
||||
super::routes::local_inference::ModelDownloadStatus,
|
||||
super::routes::local_inference::DownloadModelRequest,
|
||||
goose::providers::local_inference::hf_models::HfModelInfo,
|
||||
goose::providers::local_inference::hf_models::HfGgufFile,
|
||||
goose::providers::local_inference::hf_models::HfQuantVariant,
|
||||
super::routes::local_inference::RepoVariantsResponse,
|
||||
goose::providers::local_inference::local_model_registry::ModelSettings,
|
||||
goose::providers::local_inference::local_model_registry::SamplingConfig,
|
||||
DownloadProgress,
|
||||
DownloadStatus,
|
||||
))
|
||||
|
||||
@@ -7,11 +7,11 @@ use axum::{
|
||||
Json, Router,
|
||||
};
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
|
||||
use goose::dictation::download_manager::{get_download_manager, DownloadProgress};
|
||||
use goose::dictation::providers::{
|
||||
is_configured, transcribe_local, transcribe_with_provider, DictationProvider, PROVIDERS,
|
||||
};
|
||||
use goose::dictation::whisper;
|
||||
use goose::download_manager::{get_download_manager, DownloadProgress};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
@@ -257,11 +257,16 @@ pub async fn download_model(Path(model_id): Path<String>) -> Result<StatusCode,
|
||||
.ok_or_else(|| ErrorResponse::bad_request("Model not found"))?;
|
||||
|
||||
let manager = get_download_manager();
|
||||
let model_id_for_config = model.id.to_string();
|
||||
manager
|
||||
.download_model(
|
||||
model.id.to_string(),
|
||||
model.url.to_string(),
|
||||
model.local_path(),
|
||||
Some(Box::new(move || {
|
||||
let _ = goose::config::Config::global()
|
||||
.set_param(whisper::LOCAL_WHISPER_MODEL_CONFIG_KEY, model_id_for_config);
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.map_err(convert_error)?;
|
||||
|
||||
@@ -0,0 +1,466 @@
|
||||
use crate::routes::errors::ErrorResponse;
|
||||
use crate::state::AppState;
|
||||
use axum::{
|
||||
extract::{Path, Query},
|
||||
http::StatusCode,
|
||||
routing::{delete, get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use goose::config::paths::Paths;
|
||||
use goose::download_manager::{get_download_manager, DownloadProgress};
|
||||
use goose::providers::local_inference::hf_models::{self, HfModelInfo, HfQuantVariant};
|
||||
use goose::providers::local_inference::{
|
||||
available_inference_memory_bytes,
|
||||
hf_models::{resolve_model_spec, HfGgufFile},
|
||||
local_model_registry::{
|
||||
display_name_from_repo, get_registry, is_featured_model, model_id_from_repo,
|
||||
LocalModelEntry, ModelDownloadStatus as RegistryDownloadStatus, ModelSettings,
|
||||
FEATURED_MODELS,
|
||||
},
|
||||
recommend_local_model,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tracing::debug;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(tag = "state")]
|
||||
pub enum ModelDownloadStatus {
|
||||
NotDownloaded,
|
||||
Downloading {
|
||||
progress_percent: f32,
|
||||
bytes_downloaded: u64,
|
||||
total_bytes: u64,
|
||||
speed_bps: Option<u64>,
|
||||
},
|
||||
Downloaded,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct LocalModelResponse {
|
||||
pub id: String,
|
||||
pub display_name: String,
|
||||
pub repo_id: String,
|
||||
pub filename: String,
|
||||
pub quantization: String,
|
||||
pub size_bytes: u64,
|
||||
pub status: ModelDownloadStatus,
|
||||
pub recommended: bool,
|
||||
pub settings: ModelSettings,
|
||||
}
|
||||
|
||||
async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
|
||||
let mut entries_to_add = Vec::new();
|
||||
|
||||
for spec in FEATURED_MODELS {
|
||||
let (repo_id, quantization) = match hf_models::parse_model_spec(spec) {
|
||||
Ok(parts) => parts,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let model_id = model_id_from_repo(&repo_id, &quantization);
|
||||
|
||||
{
|
||||
let registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
if registry.has_model(&model_id) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let hf_file = match resolve_model_spec(spec).await {
|
||||
Ok((_repo, file)) => file,
|
||||
Err(_) => {
|
||||
let filename = format!(
|
||||
"{}-{}.gguf",
|
||||
repo_id.split('/').next_back().unwrap_or("model"),
|
||||
quantization
|
||||
);
|
||||
HfGgufFile {
|
||||
filename: filename.clone(),
|
||||
size_bytes: 0,
|
||||
quantization: quantization.to_string(),
|
||||
download_url: format!(
|
||||
"https://huggingface.co/{}/resolve/main/{}",
|
||||
repo_id, filename
|
||||
),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let local_path = Paths::in_data_dir("models").join(&hf_file.filename);
|
||||
|
||||
entries_to_add.push(LocalModelEntry {
|
||||
id: model_id,
|
||||
display_name: display_name_from_repo(&repo_id, &quantization),
|
||||
repo_id,
|
||||
filename: hf_file.filename,
|
||||
quantization,
|
||||
local_path,
|
||||
source_url: hf_file.download_url,
|
||||
settings: ModelSettings::default(),
|
||||
size_bytes: hf_file.size_bytes,
|
||||
});
|
||||
}
|
||||
|
||||
if !entries_to_add.is_empty() {
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
registry.sync_with_featured(entries_to_add);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/local-inference/models",
|
||||
responses(
|
||||
(status = 200, description = "List of available local LLM models", body = Vec<LocalModelResponse>)
|
||||
)
|
||||
)]
|
||||
pub async fn list_local_models(
|
||||
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
|
||||
) -> Result<Json<Vec<LocalModelResponse>>, ErrorResponse> {
|
||||
ensure_featured_models_in_registry().await?;
|
||||
|
||||
let recommended_id = recommend_local_model(&state.inference_runtime);
|
||||
|
||||
let registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
|
||||
let mut models: Vec<LocalModelResponse> = Vec::new();
|
||||
|
||||
for entry in registry.list_models() {
|
||||
let goose_status = entry.download_status();
|
||||
|
||||
let status = match goose_status {
|
||||
RegistryDownloadStatus::NotDownloaded => ModelDownloadStatus::NotDownloaded,
|
||||
RegistryDownloadStatus::Downloading {
|
||||
progress_percent,
|
||||
bytes_downloaded,
|
||||
total_bytes,
|
||||
speed_bps,
|
||||
} => ModelDownloadStatus::Downloading {
|
||||
progress_percent,
|
||||
bytes_downloaded,
|
||||
total_bytes,
|
||||
speed_bps: Some(speed_bps),
|
||||
},
|
||||
RegistryDownloadStatus::Downloaded => ModelDownloadStatus::Downloaded,
|
||||
};
|
||||
|
||||
let size_bytes = entry.file_size();
|
||||
|
||||
models.push(LocalModelResponse {
|
||||
id: entry.id.clone(),
|
||||
display_name: entry.display_name.clone(),
|
||||
repo_id: entry.repo_id.clone(),
|
||||
filename: entry.filename.clone(),
|
||||
quantization: entry.quantization.clone(),
|
||||
size_bytes,
|
||||
status,
|
||||
recommended: recommended_id == entry.id,
|
||||
settings: entry.settings.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
models.sort_by(|a, b| {
|
||||
let a_downloaded = matches!(a.status, ModelDownloadStatus::Downloaded);
|
||||
let b_downloaded = matches!(b.status, ModelDownloadStatus::Downloaded);
|
||||
match (b_downloaded, a_downloaded) {
|
||||
(true, false) => std::cmp::Ordering::Greater,
|
||||
(false, true) => std::cmp::Ordering::Less,
|
||||
_ => a.display_name.cmp(&b.display_name),
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Json(models))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SearchQuery {
|
||||
pub q: String,
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, ToSchema)]
|
||||
pub struct RepoVariantsResponse {
|
||||
pub variants: Vec<HfQuantVariant>,
|
||||
pub recommended_index: Option<usize>,
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/local-inference/search",
|
||||
params(
|
||||
("q" = String, Query, description = "Search query"),
|
||||
("limit" = Option<usize>, Query, description = "Max results")
|
||||
),
|
||||
responses(
|
||||
(status = 200, description = "Search results", body = Vec<HfModelInfo>),
|
||||
(status = 500, description = "Search failed")
|
||||
)
|
||||
)]
|
||||
pub async fn search_hf_models(
|
||||
Query(params): Query<SearchQuery>,
|
||||
) -> Result<Json<Vec<HfModelInfo>>, ErrorResponse> {
|
||||
let limit = params.limit.unwrap_or(20).min(50);
|
||||
let results = hf_models::search_gguf_models(¶ms.q, limit)
|
||||
.await
|
||||
.map_err(|e| ErrorResponse::internal(format!("Search failed: {}", e)))?;
|
||||
Ok(Json(results))
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/local-inference/repo/{author}/{repo}/files",
|
||||
responses(
|
||||
(status = 200, description = "GGUF files in the repo", body = RepoVariantsResponse)
|
||||
)
|
||||
)]
|
||||
pub async fn get_repo_files(
|
||||
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
|
||||
Path((author, repo)): Path<(String, String)>,
|
||||
) -> Result<Json<RepoVariantsResponse>, ErrorResponse> {
|
||||
let repo_id = format!("{}/{}", author, repo);
|
||||
let variants = hf_models::get_repo_gguf_variants(&repo_id)
|
||||
.await
|
||||
.map_err(|e| ErrorResponse::internal(format!("Failed to fetch repo files: {}", e)))?;
|
||||
|
||||
let available_memory = available_inference_memory_bytes(&state.inference_runtime);
|
||||
let recommended_index = hf_models::recommend_variant(&variants, available_memory);
|
||||
|
||||
Ok(Json(RepoVariantsResponse {
|
||||
variants,
|
||||
recommended_index,
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, ToSchema)]
|
||||
pub struct DownloadModelRequest {
|
||||
/// Model spec like "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M"
|
||||
pub spec: String,
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/local-inference/download",
|
||||
request_body = DownloadModelRequest,
|
||||
responses(
|
||||
(status = 202, description = "Download started", body = String),
|
||||
(status = 400, description = "Invalid request")
|
||||
)
|
||||
)]
|
||||
pub async fn download_hf_model(
|
||||
Json(req): Json<DownloadModelRequest>,
|
||||
) -> Result<(StatusCode, Json<String>), ErrorResponse> {
|
||||
let (repo_id, quantization) = hf_models::parse_model_spec(&req.spec)
|
||||
.map_err(|e| ErrorResponse::bad_request(format!("Invalid spec format: {e}")))?;
|
||||
|
||||
let (_repo, hf_file) = resolve_model_spec(&req.spec)
|
||||
.await
|
||||
.map_err(|e| ErrorResponse::bad_request(format!("Invalid spec: {}", e)))?;
|
||||
|
||||
let model_id = model_id_from_repo(&repo_id, &quantization);
|
||||
let local_path = Paths::in_data_dir("models").join(&hf_file.filename);
|
||||
let download_url = hf_file.download_url.clone();
|
||||
|
||||
let entry = LocalModelEntry {
|
||||
id: model_id.clone(),
|
||||
display_name: display_name_from_repo(&repo_id, &quantization),
|
||||
repo_id,
|
||||
filename: hf_file.filename,
|
||||
quantization,
|
||||
local_path: local_path.clone(),
|
||||
source_url: download_url.clone(),
|
||||
settings: ModelSettings::default(),
|
||||
size_bytes: hf_file.size_bytes,
|
||||
};
|
||||
|
||||
{
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
registry
|
||||
.add_model(entry)
|
||||
.map_err(|e| ErrorResponse::internal(format!("{}", e)))?;
|
||||
}
|
||||
|
||||
let dm = get_download_manager();
|
||||
dm.download_model(
|
||||
format!("{}-model", model_id),
|
||||
download_url,
|
||||
local_path,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| ErrorResponse::internal(format!("Download failed: {}", e)))?;
|
||||
|
||||
Ok((StatusCode::ACCEPTED, Json(model_id)))
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/local-inference/models/{model_id}/download",
|
||||
responses(
|
||||
(status = 200, description = "Download progress", body = DownloadProgress),
|
||||
(status = 404, description = "No active download")
|
||||
)
|
||||
)]
|
||||
pub async fn get_local_model_download_progress(
|
||||
Path(model_id): Path<String>,
|
||||
) -> Result<Json<DownloadProgress>, ErrorResponse> {
|
||||
let download_id = format!("{}-model", model_id);
|
||||
debug!(model_id = %model_id, download_id = %download_id, "Getting download progress");
|
||||
|
||||
let manager = get_download_manager();
|
||||
|
||||
let model_progress = manager
|
||||
.get_progress(&download_id)
|
||||
.ok_or_else(|| ErrorResponse::not_found("No active download"))?;
|
||||
|
||||
Ok(Json(model_progress))
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
delete,
|
||||
path = "/local-inference/models/{model_id}/download",
|
||||
responses(
|
||||
(status = 200, description = "Download cancelled"),
|
||||
(status = 404, description = "No active download")
|
||||
)
|
||||
)]
|
||||
pub async fn cancel_local_model_download(
|
||||
Path(model_id): Path<String>,
|
||||
) -> Result<StatusCode, ErrorResponse> {
|
||||
let manager = get_download_manager();
|
||||
manager
|
||||
.cancel_download(&format!("{}-model", model_id))
|
||||
.map_err(|e| ErrorResponse::internal(format!("{}", e)))?;
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
delete,
|
||||
path = "/local-inference/models/{model_id}",
|
||||
responses(
|
||||
(status = 200, description = "Model deleted"),
|
||||
(status = 404, description = "Model not found")
|
||||
)
|
||||
)]
|
||||
pub async fn delete_local_model(Path(model_id): Path<String>) -> Result<StatusCode, ErrorResponse> {
|
||||
let local_path = {
|
||||
let registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
let entry = registry
|
||||
.get_model(&model_id)
|
||||
.ok_or_else(|| ErrorResponse::not_found("Model not found"))?;
|
||||
entry.local_path.clone()
|
||||
};
|
||||
|
||||
if local_path.exists() {
|
||||
tokio::fs::remove_file(&local_path)
|
||||
.await
|
||||
.map_err(|e| ErrorResponse::internal(format!("Failed to delete: {}", e)))?;
|
||||
}
|
||||
|
||||
// Only remove non-featured models from registry (featured ones stay as placeholders)
|
||||
if !is_featured_model(&model_id) {
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
registry
|
||||
.remove_model(&model_id)
|
||||
.map_err(|e| ErrorResponse::internal(format!("{}", e)))?;
|
||||
}
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/local-inference/models/{model_id}/settings",
|
||||
responses(
|
||||
(status = 200, description = "Model settings", body = ModelSettings),
|
||||
(status = 404, description = "Model not found")
|
||||
)
|
||||
)]
|
||||
pub async fn get_model_settings(
|
||||
Path(model_id): Path<String>,
|
||||
) -> Result<Json<ModelSettings>, ErrorResponse> {
|
||||
let registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
|
||||
if let Some(settings) = registry.get_model_settings(&model_id) {
|
||||
return Ok(Json(settings.clone()));
|
||||
}
|
||||
|
||||
Err(ErrorResponse::not_found("Model not found"))
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
put,
|
||||
path = "/local-inference/models/{model_id}/settings",
|
||||
request_body = ModelSettings,
|
||||
responses(
|
||||
(status = 200, description = "Settings updated", body = ModelSettings),
|
||||
(status = 404, description = "Model not found"),
|
||||
(status = 500, description = "Failed to save settings")
|
||||
)
|
||||
)]
|
||||
pub async fn update_model_settings(
|
||||
Path(model_id): Path<String>,
|
||||
Json(settings): Json<ModelSettings>,
|
||||
) -> Result<Json<ModelSettings>, ErrorResponse> {
|
||||
let mut registry = get_registry()
|
||||
.lock()
|
||||
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
|
||||
|
||||
registry
|
||||
.update_model_settings(&model_id, settings.clone())
|
||||
.map_err(|e| ErrorResponse::not_found(format!("{}", e)))?;
|
||||
|
||||
Ok(Json(settings))
|
||||
}
|
||||
|
||||
pub fn routes(state: Arc<AppState>) -> Router {
|
||||
goose::download_manager::cleanup_partial_downloads(&Paths::in_data_dir("models"));
|
||||
|
||||
Router::new()
|
||||
.route("/local-inference/models", get(list_local_models))
|
||||
.route("/local-inference/search", get(search_hf_models))
|
||||
.route(
|
||||
"/local-inference/repo/{author}/{repo}/files",
|
||||
get(get_repo_files),
|
||||
)
|
||||
.route("/local-inference/download", post(download_hf_model))
|
||||
.route(
|
||||
"/local-inference/models/{model_id}/download",
|
||||
get(get_local_model_download_progress),
|
||||
)
|
||||
.route(
|
||||
"/local-inference/models/{model_id}/download",
|
||||
delete(cancel_local_model_download),
|
||||
)
|
||||
.route(
|
||||
"/local-inference/models/{model_id}",
|
||||
delete(delete_local_model),
|
||||
)
|
||||
.route(
|
||||
"/local-inference/models/{model_id}/settings",
|
||||
get(get_model_settings),
|
||||
)
|
||||
.route(
|
||||
"/local-inference/models/{model_id}/settings",
|
||||
axum::routing::put(update_model_settings),
|
||||
)
|
||||
.with_state(state)
|
||||
}
|
||||
@@ -3,6 +3,7 @@ pub mod agent;
|
||||
pub mod config_management;
|
||||
pub mod dictation;
|
||||
pub mod errors;
|
||||
pub mod local_inference;
|
||||
pub mod mcp_app_proxy;
|
||||
pub mod mcp_ui_proxy;
|
||||
pub mod prompts;
|
||||
@@ -30,6 +31,7 @@ pub fn configure(state: Arc<crate::state::AppState>, secret_key: String) -> Rout
|
||||
.merge(action_required::routes(state.clone()))
|
||||
.merge(agent::routes(state.clone()))
|
||||
.merge(dictation::routes(state.clone()))
|
||||
.merge(local_inference::routes(state.clone()))
|
||||
.merge(config_management::routes(state.clone()))
|
||||
.merge(prompts::routes())
|
||||
.merge(recipe::routes(state.clone()))
|
||||
|
||||
@@ -94,6 +94,11 @@ pub fn inspect_keys(
|
||||
pub fn check_provider_configured(metadata: &ProviderMetadata, provider_type: ProviderType) -> bool {
|
||||
let config = Config::global();
|
||||
|
||||
// Special override
|
||||
if metadata.name == "local" {
|
||||
return true;
|
||||
}
|
||||
|
||||
if provider_type == ProviderType::Custom || provider_type == ProviderType::Declarative {
|
||||
if let Ok(loaded_provider) = load_provider(metadata.name.as_str()) {
|
||||
if !loaded_provider.config.requires_auth {
|
||||
|
||||
@@ -14,6 +14,7 @@ use tokio::task::JoinHandle;
|
||||
|
||||
use crate::tunnel::TunnelManager;
|
||||
use goose::agents::ExtensionLoadResult;
|
||||
use goose::providers::local_inference::InferenceRuntime;
|
||||
|
||||
type ExtensionLoadingTasks =
|
||||
Arc<Mutex<HashMap<String, Arc<Mutex<Option<JoinHandle<Vec<ExtensionLoadResult>>>>>>>>;
|
||||
@@ -26,6 +27,7 @@ pub struct AppState {
|
||||
recipe_session_tracker: Arc<Mutex<HashSet<String>>>,
|
||||
pub tunnel_manager: Arc<TunnelManager>,
|
||||
pub extension_loading_tasks: ExtensionLoadingTasks,
|
||||
pub inference_runtime: Arc<InferenceRuntime>,
|
||||
}
|
||||
|
||||
fn spawn_developer(r: tokio::io::DuplexStream, w: tokio::io::DuplexStream) {
|
||||
@@ -57,6 +59,7 @@ impl AppState {
|
||||
recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())),
|
||||
tunnel_manager,
|
||||
extension_loading_tasks: Arc::new(Mutex::new(HashMap::new())),
|
||||
inference_runtime: InferenceRuntime::get_or_init(),
|
||||
}))
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ description.workspace = true
|
||||
|
||||
[features]
|
||||
default = []
|
||||
cuda = ["candle-core/cuda", "candle-nn/cuda"]
|
||||
cuda = ["candle-core/cuda", "candle-nn/cuda", "llama-cpp-2/cuda"]
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -126,14 +126,17 @@ ignore = { workspace = true }
|
||||
which = { workspace = true }
|
||||
pctx_code_mode = "^0.2.3"
|
||||
unbinder = "0.1.7"
|
||||
llama-cpp-2 = { version = "0.1.133", features = ["sampler"] }
|
||||
encoding_rs = "0.8.35"
|
||||
|
||||
[target.'cfg(target_os = "windows")'.dependencies]
|
||||
winapi = { version = "0.3", features = ["wincred"] }
|
||||
|
||||
# Platform-specific GPU acceleration for Whisper
|
||||
# Platform-specific GPU acceleration for Whisper and local inference
|
||||
[target.'cfg(target_os = "macos")'.dependencies]
|
||||
candle-core = { version = "0.9", default-features = false, features = ["metal"] }
|
||||
candle-nn = { version = "0.9", default-features = false, features = ["metal"] }
|
||||
llama-cpp-2 = { version = "0.1.133", features = ["sampler", "metal"] }
|
||||
|
||||
[dev-dependencies]
|
||||
serial_test = { workspace = true }
|
||||
@@ -156,6 +159,7 @@ path = "examples/agent.rs"
|
||||
name = "databricks_oauth"
|
||||
path = "examples/databricks_oauth.rs"
|
||||
|
||||
|
||||
[[bin]]
|
||||
name = "build_canonical_models"
|
||||
path = "src/providers/canonical/build_canonical_models.rs"
|
||||
|
||||
@@ -236,7 +236,12 @@ async fn child_process_client(
|
||||
command.env("PATH", path);
|
||||
}
|
||||
|
||||
if let Some(dir) = working_dir {
|
||||
// Use explicitly passed working_dir, falling back to GOOSE_WORKING_DIR env var
|
||||
let effective_working_dir = working_dir
|
||||
.map(|p| p.to_path_buf())
|
||||
.or_else(|| std::env::var("GOOSE_WORKING_DIR").ok().map(PathBuf::from));
|
||||
|
||||
if let Some(ref dir) = effective_working_dir {
|
||||
if dir.exists() && dir.is_dir() {
|
||||
tracing::info!("Setting MCP process working directory: {:?}", dir);
|
||||
command.current_dir(dir);
|
||||
@@ -710,6 +715,7 @@ impl ExtensionManager {
|
||||
})?;
|
||||
let mut context = self.context.clone();
|
||||
context.extension_manager = Some(Arc::downgrade(self));
|
||||
|
||||
(def.client_factory)(context)
|
||||
}
|
||||
ExtensionConfig::InlinePython {
|
||||
@@ -1584,6 +1590,16 @@ impl ExtensionManager {
|
||||
session_id: &str,
|
||||
working_dir: &std::path::Path,
|
||||
) -> Option<String> {
|
||||
// Skip MOIM for models with small context windows to avoid consuming limited context
|
||||
const MIN_CONTEXT_FOR_MOIM: usize = 32_000;
|
||||
if let Ok(provider_guard) = self.provider.try_lock() {
|
||||
if let Some(provider) = provider_guard.as_ref() {
|
||||
if provider.get_model_config().context_limit() < MIN_CONTEXT_FOR_MOIM {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Use minute-level granularity to prevent conversation changes every second
|
||||
let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:00").to_string();
|
||||
let mut content = format!(
|
||||
|
||||
@@ -1,3 +1,2 @@
|
||||
pub mod download_manager;
|
||||
pub mod providers;
|
||||
pub mod whisper;
|
||||
|
||||
+40
-8
@@ -1,12 +1,33 @@
|
||||
use crate::dictation::whisper::LOCAL_WHISPER_MODEL_CONFIG_KEY;
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tracing::info;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
fn partial_path_for(destination: &Path) -> PathBuf {
|
||||
destination.with_extension(
|
||||
destination
|
||||
.extension()
|
||||
.map(|e| format!("{}.part", e.to_string_lossy()))
|
||||
.unwrap_or_else(|| "part".to_string()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Remove any leftover `.part` files in the given directory.
|
||||
pub fn cleanup_partial_downloads(dir: &Path) {
|
||||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().is_some_and(|e| e == "part") {
|
||||
let _ = std::fs::remove_file(&path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct DownloadProgress {
|
||||
/// Model ID being downloaded
|
||||
@@ -78,8 +99,9 @@ impl DownloadManager {
|
||||
model_id: String,
|
||||
url: String,
|
||||
destination: PathBuf,
|
||||
on_complete: Option<Box<dyn FnOnce() + Send + 'static>>,
|
||||
) -> Result<()> {
|
||||
// Initialize progress
|
||||
info!(model_id = %model_id, url = %url, destination = ?destination, "Starting model download");
|
||||
{
|
||||
let mut downloads = self
|
||||
.downloads
|
||||
@@ -115,10 +137,13 @@ impl DownloadManager {
|
||||
let downloads = self.downloads.clone();
|
||||
let model_id_clone = model_id.clone();
|
||||
|
||||
let destination_for_cleanup = destination.clone();
|
||||
|
||||
// Download in background task
|
||||
tokio::spawn(async move {
|
||||
match Self::download_file(&url, &destination, &downloads, &model_id_clone).await {
|
||||
Ok(_) => {
|
||||
info!(model_id = %model_id_clone, "Download completed successfully");
|
||||
if let Ok(mut downloads) = downloads.lock() {
|
||||
if let Some(progress) = downloads.get_mut(&model_id_clone) {
|
||||
progress.status = DownloadStatus::Completed;
|
||||
@@ -126,10 +151,15 @@ impl DownloadManager {
|
||||
}
|
||||
}
|
||||
|
||||
let _ = crate::config::Config::global()
|
||||
.set_param(LOCAL_WHISPER_MODEL_CONFIG_KEY, model_id_clone.clone());
|
||||
if let Some(callback) = on_complete {
|
||||
callback();
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Clean up partial file on failure
|
||||
let partial = partial_path_for(&destination_for_cleanup);
|
||||
let _ = tokio::fs::remove_file(&partial).await;
|
||||
|
||||
if let Ok(mut downloads) = downloads.lock() {
|
||||
if let Some(progress) = downloads.get_mut(&model_id_clone) {
|
||||
progress.status = DownloadStatus::Failed;
|
||||
@@ -166,7 +196,8 @@ impl DownloadManager {
|
||||
}
|
||||
}
|
||||
|
||||
let mut file = tokio::fs::File::create(destination).await?;
|
||||
let partial_path = partial_path_for(destination);
|
||||
let mut file = tokio::fs::File::create(&partial_path).await?;
|
||||
let mut bytes_downloaded = 0u64;
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
@@ -185,8 +216,7 @@ impl DownloadManager {
|
||||
};
|
||||
|
||||
if should_cancel {
|
||||
// Clean up partial download
|
||||
let _ = tokio::fs::remove_file(destination).await;
|
||||
let _ = tokio::fs::remove_file(&partial_path).await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -226,6 +256,8 @@ impl DownloadManager {
|
||||
}
|
||||
|
||||
file.flush().await?;
|
||||
drop(file);
|
||||
tokio::fs::rename(&partial_path, destination).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ pub mod config;
|
||||
pub mod context_mgmt;
|
||||
pub mod conversation;
|
||||
pub mod dictation;
|
||||
pub mod download_manager;
|
||||
pub mod execution;
|
||||
pub mod goose_apps;
|
||||
pub mod hints;
|
||||
|
||||
@@ -39,6 +39,14 @@ static TEMPLATE_REGISTRY: &[(&str, &str)] = &[
|
||||
"plan.md",
|
||||
"Prompt used when goose creates step-by-step plans. CLI only",
|
||||
),
|
||||
(
|
||||
"tiny_model_system.md",
|
||||
"System prompt for tiny local models using shell command emulation",
|
||||
),
|
||||
(
|
||||
"session_name.md",
|
||||
"System prompt for generating short session names from conversation history",
|
||||
),
|
||||
];
|
||||
|
||||
/// Information about a template including its content and customization status
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
Generate a short title (four words or less) that describes the topic of the user's messages. Reply with only the title, nothing else.
|
||||
|
||||
Examples:
|
||||
- "how do I reverse a list in python?" → Python list reversal
|
||||
- "what's the weather in Tokyo?" → Tokyo weather
|
||||
- "explain how transformers work in ML" → ML transformers explained
|
||||
@@ -0,0 +1,22 @@
|
||||
You are goose, an autonomous AI agent created by Block. You act on the user's
|
||||
behalf — you do not explain how to do things, you DO them directly.
|
||||
|
||||
The OS is {{os}}, the shell is {{shell}}, and the working directory is {{working_directory}}
|
||||
|
||||
When the user asks you to do something, take action immediately. Do not describe
|
||||
what you would do or give instructions — execute the commands yourself.
|
||||
|
||||
To run a shell command, start a new line with $:
|
||||
|
||||
$ ls
|
||||
|
||||
Keep your responses brief. State what you are doing, then do it. For example:
|
||||
|
||||
User: how many files are in /tmp?
|
||||
You: Let me check.
|
||||
$ ls -1 /tmp | wc -l
|
||||
|
||||
After a command runs, you will see its output. Use the output to answer the user
|
||||
or take the next step. Do not repeat commands you have already run.
|
||||
|
||||
Do not use shell commands if you already know the answer.
|
||||
@@ -17,10 +17,22 @@ use rmcp::model::Tool;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
use once_cell::sync::Lazy;
|
||||
use regex::Regex;
|
||||
use std::ops::{Add, AddAssign};
|
||||
use std::pin::Pin;
|
||||
use std::sync::LazyLock;
|
||||
use std::sync::Mutex;
|
||||
|
||||
fn strip_xml_tags(text: &str) -> String {
|
||||
static BLOCK_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?</[a-zA-Z][a-zA-Z0-9_]*>").unwrap()
|
||||
});
|
||||
static TAG_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"</?[a-zA-Z][a-zA-Z0-9_]*[^>]*>").unwrap());
|
||||
let pass1 = BLOCK_RE.replace_all(text, "");
|
||||
TAG_RE.replace_all(&pass1, "").into_owned()
|
||||
}
|
||||
|
||||
/// A global store for the current model being used, we use this as when a provider returns, it tells us the real model, not an alias
|
||||
pub static CURRENT_MODEL: Lazy<Mutex<Option<String>>> = Lazy::new(|| Mutex::new(None));
|
||||
|
||||
@@ -594,20 +606,28 @@ pub trait Provider: Send + Sync {
|
||||
messages: &Conversation,
|
||||
) -> Result<String, ProviderError> {
|
||||
let context = self.get_initial_user_messages(messages);
|
||||
let prompt = self.create_session_name_prompt(&context);
|
||||
let message = Message::user().with_text(&prompt);
|
||||
let system = crate::prompt_template::render_template(
|
||||
"session_name.md",
|
||||
&std::collections::HashMap::<String, String>::new(),
|
||||
)
|
||||
.map_err(|e| ProviderError::ContextLengthExceeded(e.to_string()))?;
|
||||
|
||||
let user_text = format!(
|
||||
"---BEGIN USER MESSAGES---\n{}\n---END USER MESSAGES---\n\nGenerate a short title for the above messages.",
|
||||
context.join("\n")
|
||||
);
|
||||
let message = Message::user().with_text(&user_text);
|
||||
let result = self
|
||||
.complete_fast(
|
||||
session_id,
|
||||
"Reply with only a description in four words or less",
|
||||
&[message],
|
||||
&[],
|
||||
)
|
||||
.complete_fast(session_id, &system, &[message], &[])
|
||||
.await?;
|
||||
|
||||
let description = result
|
||||
let raw: String = result
|
||||
.0
|
||||
.as_concat_text()
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| c.as_text())
|
||||
.collect();
|
||||
let description = strip_xml_tags(&raw)
|
||||
.split_whitespace()
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
@@ -615,21 +635,6 @@ pub trait Provider: Send + Sync {
|
||||
Ok(safe_truncate(&description, 100))
|
||||
}
|
||||
|
||||
// Generate a prompt for a session name based on the conversation history
|
||||
fn create_session_name_prompt(&self, context: &[String]) -> String {
|
||||
// Create a prompt for a concise description
|
||||
let mut prompt = "Based on the conversation so far, provide a concise description of this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the description".to_string();
|
||||
|
||||
if !context.is_empty() {
|
||||
prompt = format!(
|
||||
"Here are the first few user messages:\n{}\n\n{}",
|
||||
context.join("\n"),
|
||||
prompt
|
||||
);
|
||||
}
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Configure OAuth authentication for this provider
|
||||
///
|
||||
/// This method is called when a provider has configuration keys marked with oauth_flow = true.
|
||||
@@ -719,6 +724,42 @@ mod tests {
|
||||
use test_case::test_case;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_strip_xml_tags() {
|
||||
assert_eq!(strip_xml_tags("<think>reasoning</think>answer"), "answer");
|
||||
assert_eq!(strip_xml_tags("before<t>mid</t>after"), "beforeafter");
|
||||
assert_eq!(strip_xml_tags("<a>x</a><b>y</b>z"), "z");
|
||||
assert_eq!(strip_xml_tags("no tags here"), "no tags here");
|
||||
assert_eq!(strip_xml_tags("a < b > c"), "a < b > c");
|
||||
assert_eq!(strip_xml_tags("<think>über</think>ok"), "ok");
|
||||
assert_eq!(strip_xml_tags("<think>日本語</think>hello"), "hello");
|
||||
assert_eq!(strip_xml_tags(""), "");
|
||||
assert_eq!(strip_xml_tags("<>stuff</>"), "<>stuff</>");
|
||||
// attributes
|
||||
assert_eq!(
|
||||
strip_xml_tags(r#"<think class="deep">reasoning</think>answer"#),
|
||||
"answer"
|
||||
);
|
||||
// self-closing tags
|
||||
assert_eq!(strip_xml_tags("<br/>self closing"), "self closing");
|
||||
// orphan closing tags
|
||||
assert_eq!(strip_xml_tags("orphan </think> tag"), "orphan tag");
|
||||
// multiline content
|
||||
assert_eq!(
|
||||
strip_xml_tags("<think>\nline1\nline2\n</think>result"),
|
||||
"result"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_creation() {
|
||||
let usage = Usage::new(Some(10), Some(20), Some(30));
|
||||
assert_eq!(usage.input_tokens, Some(10));
|
||||
assert_eq!(usage.output_tokens, Some(20));
|
||||
assert_eq!(usage.total_tokens, Some(30));
|
||||
}
|
||||
|
||||
fn content_from_str(s: String) -> MessageContent {
|
||||
if let Some(img_data) = s.strip_prefix("*img:") {
|
||||
MessageContent::image(format!("http://example.com/{}", img_data), "image/png")
|
||||
@@ -798,11 +839,10 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_collect_stream_defaults_usage() {
|
||||
// Should not error when usage is missing
|
||||
let stream = create_test_stream(vec!["Hello".to_string()]);
|
||||
let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap();
|
||||
assert_eq!(content_to_strings(&msg), vec!["Hello"]);
|
||||
assert_eq!(usage.model, "unknown"); // Default usage
|
||||
assert_eq!(usage.model, "unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -16,6 +16,7 @@ use super::{
|
||||
google::GoogleProvider,
|
||||
lead_worker::LeadWorkerProvider,
|
||||
litellm::LiteLLMProvider,
|
||||
local_inference::LocalInferenceProvider,
|
||||
ollama::OllamaProvider,
|
||||
openai::OpenAiProvider,
|
||||
openrouter::OpenRouterProvider,
|
||||
@@ -47,6 +48,7 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
|
||||
registry.register::<AnthropicProvider>(true);
|
||||
registry.register::<AzureProvider>(false);
|
||||
registry.register::<BedrockProvider>(false);
|
||||
registry.register::<LocalInferenceProvider>(false);
|
||||
registry.register::<ChatGptCodexProvider>(true);
|
||||
registry.register::<ClaudeCodeProvider>(true);
|
||||
registry.register::<CodexProvider>(true);
|
||||
|
||||
@@ -0,0 +1,647 @@
|
||||
pub mod hf_models;
|
||||
mod inference_emulated_tools;
|
||||
mod inference_engine;
|
||||
mod inference_native_tools;
|
||||
pub mod local_model_registry;
|
||||
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;
|
||||
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 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};
|
||||
use rmcp::model::{Role, Tool};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
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";
|
||||
|
||||
type ModelSlot = Arc<Mutex<Option<LoadedModel>>>;
|
||||
|
||||
/// 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<HashMap<String, ModelSlot>>,
|
||||
backend: LlamaBackend,
|
||||
}
|
||||
|
||||
/// Global weak reference used to share a single `InferenceRuntime` across
|
||||
/// all providers and server routes. Only a `Weak` is stored — strong `Arc`s
|
||||
/// live in providers and `AppState`. When all strong refs drop (normal
|
||||
/// shutdown), the runtime is deallocated and the backend freed. The `Weak`
|
||||
/// left behind is inert during `__cxa_finalize`, so no ggml statics race.
|
||||
static RUNTIME: StdMutex<Weak<InferenceRuntime>> = StdMutex::new(Weak::new());
|
||||
|
||||
impl InferenceRuntime {
|
||||
pub fn get_or_init() -> Arc<Self> {
|
||||
let mut guard = RUNTIME.lock().expect("runtime lock poisoned");
|
||||
if let Some(runtime) = guard.upgrade() {
|
||||
return 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) => panic!("Failed to init llama backend: {}", e),
|
||||
};
|
||||
let runtime = Arc::new(Self {
|
||||
models: StdMutex::new(HashMap::new()),
|
||||
backend,
|
||||
});
|
||||
*guard = Arc::downgrade(&runtime);
|
||||
runtime
|
||||
}
|
||||
|
||||
pub fn backend(&self) -> &LlamaBackend {
|
||||
&self.backend
|
||||
}
|
||||
|
||||
fn get_or_create_model_slot(&self, model_id: &str) -> ModelSlot {
|
||||
let mut map = self.models.lock().expect("model cache lock poisoned");
|
||||
map.entry(model_id.to_string())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(None)))
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn other_model_slots(&self, keep_model_id: &str) -> Vec<ModelSlot> {
|
||||
let map = self.models.lock().expect("model cache lock poisoned");
|
||||
map.iter()
|
||||
.filter(|(id, _)| id.as_str() != keep_model_id)
|
||||
.map(|(_, slot)| slot.clone())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
const PROVIDER_NAME: &str = "local";
|
||||
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";
|
||||
|
||||
/// Resolve model path, context limit, and settings for a model ID from the registry.
|
||||
pub fn resolve_model_path(
|
||||
model_id: &str,
|
||||
) -> Option<(
|
||||
PathBuf,
|
||||
usize,
|
||||
crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
)> {
|
||||
use crate::providers::local_inference::local_model_registry::get_registry;
|
||||
|
||||
if let Ok(registry) = get_registry().lock() {
|
||||
if let Some(entry) = registry.get_model(model_id) {
|
||||
let ctx = entry.settings.context_size.unwrap_or(0) as usize;
|
||||
return Some((entry.local_path.clone(), ctx, entry.settings.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
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| {
|
||||
matches!(
|
||||
d.device_type,
|
||||
LlamaBackendDeviceType::Gpu
|
||||
| LlamaBackendDeviceType::IntegratedGpu
|
||||
| LlamaBackendDeviceType::Accelerator
|
||||
)
|
||||
})
|
||||
.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 {
|
||||
use local_model_registry::{get_registry, is_featured_model, FEATURED_MODELS};
|
||||
|
||||
let available_memory = available_inference_memory_bytes(runtime);
|
||||
|
||||
if let Ok(registry) = get_registry().lock() {
|
||||
let mut models: Vec<_> = registry
|
||||
.list_models()
|
||||
.iter()
|
||||
.filter(|m| is_featured_model(&m.id) && m.size_bytes > 0)
|
||||
.collect();
|
||||
models.sort_by(|a, b| b.size_bytes.cmp(&a.size_bytes));
|
||||
|
||||
// Return largest that fits in available memory
|
||||
for model in &models {
|
||||
if available_memory >= model.size_bytes {
|
||||
return model.id.clone();
|
||||
}
|
||||
}
|
||||
|
||||
// If nothing fits, return smallest
|
||||
if let Some(smallest) = models.last() {
|
||||
return smallest.id.clone();
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to first featured model
|
||||
FEATURED_MODELS[0].to_string()
|
||||
}
|
||||
|
||||
fn build_openai_messages_json(system: &str, messages: &[Message]) -> String {
|
||||
use crate::providers::formats::openai::format_messages;
|
||||
use crate::providers::utils::ImageFormat;
|
||||
|
||||
let mut arr: Vec<Value> = vec![json!({"role": "system", "content": system})];
|
||||
arr.extend(format_messages(messages, &ImageFormat::OpenAi));
|
||||
serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string())
|
||||
}
|
||||
|
||||
/// Convert a message into plain text for the emulator path's chat history.
|
||||
///
|
||||
/// This is the emulator-path counterpart of [`format_messages`] used by the native
|
||||
/// path. It reconstructs the text-based tool syntax that the emulator prompt teaches
|
||||
/// the model:
|
||||
///
|
||||
/// - `ToolRequest` with a `"command"` argument → `$ command`
|
||||
/// - `ToolRequest` with a `"code"` argument → `` ```execute\n…\n``` ``
|
||||
/// - `ToolResponse` → `Command output:\n…`
|
||||
///
|
||||
/// Only `developer__shell` and `code_execution__execute` style tool calls are
|
||||
/// recognized (by argument shape, not tool name). Tool calls from other extensions
|
||||
/// (e.g. custom MCP tools made by a native-tool-calling model earlier in the
|
||||
/// conversation) are silently dropped, since the emulator path has no syntax to
|
||||
/// represent them.
|
||||
fn extract_text_content(msg: &Message) -> String {
|
||||
let mut parts = Vec::new();
|
||||
|
||||
for content in &msg.content {
|
||||
match content {
|
||||
MessageContent::Text(text) => {
|
||||
parts.push(text.text.clone());
|
||||
}
|
||||
MessageContent::ToolRequest(req) => {
|
||||
if let Ok(call) = &req.tool_call {
|
||||
if let Some(cmd) = call
|
||||
.arguments
|
||||
.as_ref()
|
||||
.and_then(|a| a.get("command"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
parts.push(format!("$ {}", cmd));
|
||||
} else if let Some(code) = call
|
||||
.arguments
|
||||
.as_ref()
|
||||
.and_then(|a| a.get("code"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
parts.push(format!("```execute\n{}\n```", code));
|
||||
}
|
||||
}
|
||||
}
|
||||
MessageContent::ToolResponse(response) => match &response.tool_result {
|
||||
Ok(result) => {
|
||||
let mut output_parts = Vec::new();
|
||||
for content_item in &result.content {
|
||||
if let Some(text_content) = content_item.as_text() {
|
||||
output_parts.push(text_content.text.to_string());
|
||||
}
|
||||
}
|
||||
if !output_parts.is_empty() {
|
||||
parts.push(format!("Command output:\n{}", output_parts.join("\n")));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
parts.push(format!("Command error: {}", e));
|
||||
}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
parts.join("\n")
|
||||
}
|
||||
|
||||
/// Build a `ProviderUsage` and write the request log entry.
|
||||
fn finalize_usage(
|
||||
log: &mut RequestLog,
|
||||
model_name: String,
|
||||
path_label: &str,
|
||||
prompt_token_count: usize,
|
||||
output_token_count: i32,
|
||||
extra_log_fields: Option<(&str, &str)>,
|
||||
) -> ProviderUsage {
|
||||
let input_tokens = prompt_token_count as i32;
|
||||
let total_tokens = input_tokens + output_token_count;
|
||||
let usage = Usage::new(
|
||||
Some(input_tokens),
|
||||
Some(output_token_count),
|
||||
Some(total_tokens),
|
||||
);
|
||||
let mut log_json = serde_json::json!({
|
||||
"path": path_label,
|
||||
"prompt_tokens": input_tokens,
|
||||
"output_tokens": output_token_count,
|
||||
});
|
||||
if let Some((key, value)) = extra_log_fields {
|
||||
log_json[key] = serde_json::json!(value);
|
||||
}
|
||||
let _ = log.write(&log_json, Some(&usage));
|
||||
ProviderUsage::new(model_name, usage)
|
||||
}
|
||||
|
||||
type StreamSender =
|
||||
tokio::sync::mpsc::Sender<Result<(Option<Message>, Option<ProviderUsage>), ProviderError>>;
|
||||
|
||||
pub struct LocalInferenceProvider {
|
||||
runtime: Arc<InferenceRuntime>,
|
||||
model: ModelSlot,
|
||||
model_config: ModelConfig,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl LocalInferenceProvider {
|
||||
pub async fn from_env(model: ModelConfig, _extensions: Vec<ExtensionConfig>) -> Result<Self> {
|
||||
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<LoadedModel, ProviderError> {
|
||||
let (model_path, _context_limit, _) = resolve_model_path(model_id)
|
||||
.ok_or_else(|| ProviderError::ExecutionError(format!("Unknown model: {}", model_id)))?;
|
||||
|
||||
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
|
||||
))
|
||||
})?
|
||||
}
|
||||
};
|
||||
|
||||
tracing::info!("Model loaded successfully");
|
||||
|
||||
Ok(LoadedModel { model, template })
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for LocalInferenceProvider {
|
||||
type Provider = Self;
|
||||
|
||||
fn metadata() -> ProviderMetadata
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
use crate::providers::local_inference::local_model_registry::{
|
||||
get_registry, FEATURED_MODELS,
|
||||
};
|
||||
|
||||
let mut known_models: Vec<&str> = FEATURED_MODELS.to_vec();
|
||||
|
||||
// Add any registry models not already in the featured list
|
||||
let mut dynamic_models = Vec::new();
|
||||
if let Ok(registry) = get_registry().lock() {
|
||||
for entry in registry.list_models() {
|
||||
if !known_models.contains(&entry.id.as_str()) {
|
||||
dynamic_models.push(entry.id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
let dynamic_refs: Vec<&str> = dynamic_models.iter().map(|s| s.as_str()).collect();
|
||||
known_models.extend(dynamic_refs);
|
||||
|
||||
ProviderMetadata::new(
|
||||
PROVIDER_NAME,
|
||||
"Local Inference",
|
||||
"Local inference using quantized GGUF models (llama.cpp)",
|
||||
DEFAULT_MODEL,
|
||||
known_models,
|
||||
"https://github.com/utilityai/llama-cpp-rs",
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
|
||||
fn from_env(
|
||||
model: ModelConfig,
|
||||
extensions: Vec<ExtensionConfig>,
|
||||
) -> BoxFuture<'static, Result<Self::Provider>>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
Box::pin(Self::from_env(model, extensions))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for LocalInferenceProvider {
|
||||
fn get_name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
self.model_config.clone()
|
||||
}
|
||||
|
||||
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
||||
use crate::providers::local_inference::local_model_registry::get_registry;
|
||||
|
||||
let mut all_models: Vec<String> = Vec::new();
|
||||
|
||||
if let Ok(registry) = get_registry().lock() {
|
||||
for entry in registry.list_models() {
|
||||
all_models.push(entry.id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_models)
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
system: &str,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let (_model_path, model_context_limit, model_settings) =
|
||||
resolve_model_path(&model_config.model_name).ok_or_else(|| {
|
||||
ProviderError::ExecutionError(format!(
|
||||
"Model not found: {}",
|
||||
model_config.model_name
|
||||
))
|
||||
})?;
|
||||
|
||||
// Ensure model is loaded — unload any other models first to free memory.
|
||||
{
|
||||
let mut model_lock = self.model.lock().await;
|
||||
if model_lock.is_none() {
|
||||
for slot in self.runtime.other_model_slots(&model_config.model_name) {
|
||||
let mut other = slot.lock().await;
|
||||
if other.is_some() {
|
||||
tracing::info!("Unloading previous model to free memory");
|
||||
*other = None;
|
||||
}
|
||||
}
|
||||
|
||||
let model_id = model_config.model_name.clone();
|
||||
let settings_for_load = model_settings.clone();
|
||||
let runtime_for_load = self.runtime.clone();
|
||||
let loaded = tokio::task::spawn_blocking(move || {
|
||||
Self::load_model_sync(&runtime_for_load, &model_id, &settings_for_load)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))??;
|
||||
*model_lock = Some(loaded);
|
||||
}
|
||||
}
|
||||
|
||||
// Models that support native OpenAI-compatible tool-call JSON use the
|
||||
// native path (template-based tool calling with JSON output). All other
|
||||
// models use the emulator which parses `$ command` and ```execute blocks.
|
||||
// Only use emulator when there are actually tools to emulate - utility calls
|
||||
// like compaction and session naming pass empty tools and should preserve
|
||||
// their system prompts.
|
||||
let use_emulator = !model_settings.native_tool_calling && !tools.is_empty();
|
||||
let system_prompt = if use_emulator {
|
||||
load_tiny_model_prompt()
|
||||
} else {
|
||||
system.to_string()
|
||||
};
|
||||
|
||||
// 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 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 {
|
||||
Some(build_openai_messages_json(&system_prompt, messages))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let model_arc = self.model.clone();
|
||||
let runtime = self.runtime.clone();
|
||||
let model_name = model_config.model_name.clone();
|
||||
let context_limit = model_context_limit;
|
||||
let settings = model_settings;
|
||||
|
||||
let log_payload = serde_json::json!({
|
||||
"system": &system_prompt,
|
||||
"messages": messages.iter().map(|m| {
|
||||
serde_json::json!({
|
||||
"role": match m.role { Role::User => "user", Role::Assistant => "assistant" },
|
||||
"content": extract_text_content(m),
|
||||
})
|
||||
}).collect::<Vec<_>>(),
|
||||
"tools": tools.iter().map(|t| &t.name).collect::<Vec<_>>(),
|
||||
"settings": {
|
||||
"use_jinja": settings.use_jinja,
|
||||
"native_tool_calling": settings.native_tool_calling,
|
||||
"context_size": settings.context_size,
|
||||
"sampling": settings.sampling,
|
||||
},
|
||||
});
|
||||
|
||||
let mut log = RequestLog::start(&self.model_config, &log_payload)
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<
|
||||
Result<(Option<Message>, Option<ProviderUsage>), ProviderError>,
|
||||
>(32);
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// Macro to log errors before sending them through the channel
|
||||
macro_rules! send_err {
|
||||
($err:expr) => {{
|
||||
let err = $err;
|
||||
let msg = match &err {
|
||||
ProviderError::ExecutionError(s) => s.as_str(),
|
||||
ProviderError::ContextLengthExceeded(s) => s.as_str(),
|
||||
_ => "unknown error",
|
||||
};
|
||||
let _ = log.error(msg);
|
||||
let _ = tx.blocking_send(Err(err));
|
||||
return;
|
||||
}};
|
||||
}
|
||||
|
||||
let model_guard = model_arc.blocking_lock();
|
||||
let loaded = match model_guard.as_ref() {
|
||||
Some(l) => l,
|
||||
None => {
|
||||
send_err!(ProviderError::ExecutionError(
|
||||
"Model not loaded".to_string()
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let message_id = Uuid::new_v4().to_string();
|
||||
|
||||
let mut gen_ctx = GenerationContext {
|
||||
loaded,
|
||||
runtime: &runtime,
|
||||
chat_messages: &chat_messages,
|
||||
settings: &settings,
|
||||
context_limit,
|
||||
model_name,
|
||||
message_id: &message_id,
|
||||
tx: &tx,
|
||||
log: &mut log,
|
||||
};
|
||||
|
||||
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(),
|
||||
)
|
||||
};
|
||||
|
||||
if let Err(err) = result {
|
||||
let msg = match &err {
|
||||
ProviderError::ExecutionError(s) => s.as_str(),
|
||||
ProviderError::ContextLengthExceeded(s) => s.as_str(),
|
||||
_ => "unknown error",
|
||||
};
|
||||
let _ = log.error(msg);
|
||||
let _ = tx.blocking_send(Err(err));
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Box::pin(try_stream! {
|
||||
while let Some(result) = rx.recv().await {
|
||||
let item = result?;
|
||||
yield item;
|
||||
}
|
||||
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,486 @@
|
||||
use anyhow::{bail, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use utoipa::ToSchema;
|
||||
|
||||
const HF_API_BASE: &str = "https://huggingface.co/api/models";
|
||||
const HF_DOWNLOAD_BASE: &str = "https://huggingface.co";
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct HfModelInfo {
|
||||
pub repo_id: String,
|
||||
pub author: String,
|
||||
pub model_name: String,
|
||||
pub downloads: u64,
|
||||
pub gguf_files: Vec<HfGgufFile>,
|
||||
}
|
||||
|
||||
/// A single downloadable GGUF file (used internally and for downloads).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct HfGgufFile {
|
||||
pub filename: String,
|
||||
pub size_bytes: u64,
|
||||
pub quantization: String,
|
||||
pub download_url: String,
|
||||
}
|
||||
|
||||
/// A quantization variant — groups sharded files into one logical entry.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct HfQuantVariant {
|
||||
pub quantization: String,
|
||||
pub size_bytes: u64,
|
||||
pub filename: String,
|
||||
pub download_url: String,
|
||||
pub description: &'static str,
|
||||
pub quality_rank: u8,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct HfApiModel {
|
||||
id: Option<String>,
|
||||
author: Option<String>,
|
||||
downloads: Option<u64>,
|
||||
siblings: Option<Vec<HfApiSibling>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct HfApiSibling {
|
||||
rfilename: String,
|
||||
#[serde(default)]
|
||||
size: Option<u64>,
|
||||
}
|
||||
|
||||
struct QuantInfo {
|
||||
description: &'static str,
|
||||
quality_rank: u8,
|
||||
}
|
||||
|
||||
const QUANT_TABLE: &[(&str, &str, u8)] = &[
|
||||
("IQ1_S", "Extremely small, very low quality", 1),
|
||||
("IQ1_M", "Extremely small, very low quality", 2),
|
||||
("IQ2_XXS", "Very small, low quality", 3),
|
||||
("IQ2_XS", "Very small, low quality", 4),
|
||||
("IQ2_S", "Very small, low quality", 5),
|
||||
("IQ2_M", "Very small, low quality", 6),
|
||||
("Q2_K", "Small, low quality", 7),
|
||||
("Q2_K_S", "Small, low quality", 7),
|
||||
("IQ3_XXS", "Very small, moderate quality loss", 8),
|
||||
("IQ3_XS", "Small, moderate quality loss", 9),
|
||||
("IQ3_S", "Small, moderate quality loss", 9),
|
||||
("Q3_K_S", "Small, moderate quality loss", 10),
|
||||
("IQ3_M", "Small, moderate quality loss", 11),
|
||||
("Q3_K_M", "Small, balanced quality/size", 12),
|
||||
("Q3_K_L", "Medium-small, decent quality", 13),
|
||||
("IQ4_XS", "Medium, good quality", 14),
|
||||
("IQ4_NL", "Medium, good quality", 15),
|
||||
("Q4_0", "Medium, good quality", 16),
|
||||
("Q4_1", "Medium, good quality", 17),
|
||||
("Q4_K_S", "Medium, good quality/size balance", 18),
|
||||
(
|
||||
"Q4_K_M",
|
||||
"Medium, recommended balance of quality and size",
|
||||
19,
|
||||
),
|
||||
("Q5_0", "Medium-large, high quality", 20),
|
||||
("Q5_1", "Medium-large, high quality", 21),
|
||||
("Q5_K_S", "Medium-large, high quality", 22),
|
||||
("Q5_K_M", "Medium-large, very high quality", 23),
|
||||
("Q6_K", "Large, near-lossless quality", 24),
|
||||
("Q8_0", "Large, near-lossless quality", 25),
|
||||
("F16", "Full size, original quality (16-bit)", 26),
|
||||
("BF16", "Full size, original quality (bfloat16)", 27),
|
||||
("F32", "Full size, original quality (32-bit)", 28),
|
||||
(
|
||||
"MXFP4_MOE",
|
||||
"Medium, mixed-precision 4-bit for MoE models",
|
||||
18,
|
||||
),
|
||||
("TQ1_0", "Tiny, ternary quantization", 1),
|
||||
("Q2_K_XL", "Extended-layer variant", 15),
|
||||
("Q3_K_XL", "Extended-layer variant", 15),
|
||||
("Q4_K_XL", "Extended-layer variant", 15),
|
||||
("Q2_K_L", "Small, low quality (large variant)", 8),
|
||||
("Q4_K_L", "Medium, good quality (large variant)", 20),
|
||||
];
|
||||
|
||||
fn quant_info(quant: &str) -> QuantInfo {
|
||||
QUANT_TABLE
|
||||
.iter()
|
||||
.find(|(name, _, _)| *name == quant)
|
||||
.map(|(_, description, quality_rank)| QuantInfo {
|
||||
description,
|
||||
quality_rank: *quality_rank,
|
||||
})
|
||||
.unwrap_or(QuantInfo {
|
||||
description: "",
|
||||
quality_rank: 15,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn parse_quantization_from_filename(filename: &str) -> String {
|
||||
parse_quantization(filename)
|
||||
}
|
||||
|
||||
fn parse_quantization(filename: &str) -> String {
|
||||
// Strip directory prefix (e.g. "Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf")
|
||||
let basename = filename.rsplit('/').next().unwrap_or(filename);
|
||||
let stem = basename.trim_end_matches(".gguf");
|
||||
|
||||
// Strip shard suffix like "-00001-of-00004"
|
||||
let stem = if let Some(pos) = stem.rfind("-of-") {
|
||||
stem.get(..pos)
|
||||
.and_then(|s| s.rsplit_once('-').map(|(prefix, _)| prefix))
|
||||
.unwrap_or(stem)
|
||||
} else {
|
||||
stem
|
||||
};
|
||||
|
||||
// The quantization tag is typically the last hyphen-separated component
|
||||
// that looks like a quant identifier (starts with Q, IQ, F, BF, TQ, MXFP, etc.)
|
||||
// e.g. "Qwen3-Coder-Next-Q4_K_M" -> "Q4_K_M"
|
||||
// "Model-UD-IQ1_M" -> "IQ1_M"
|
||||
if let Some((_, candidate)) = stem.rsplit_once('-') {
|
||||
if looks_like_quant(candidate) {
|
||||
return candidate.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: try dot separator (e.g. "model.Q4_K_M")
|
||||
if let Some((_, candidate)) = stem.rsplit_once('.') {
|
||||
if looks_like_quant(candidate) {
|
||||
return candidate.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
"unknown".to_string()
|
||||
}
|
||||
|
||||
fn looks_like_quant(s: &str) -> bool {
|
||||
let upper = s.to_uppercase();
|
||||
upper.starts_with("Q")
|
||||
|| upper.starts_with("IQ")
|
||||
|| upper.starts_with("TQ")
|
||||
|| upper.starts_with("MXFP")
|
||||
|| upper == "F16"
|
||||
|| upper == "F32"
|
||||
|| upper == "BF16"
|
||||
}
|
||||
|
||||
fn is_shard_file(filename: &str) -> bool {
|
||||
// Matches patterns like "-00001-of-00003.gguf"
|
||||
let basename = filename.rsplit('/').next().unwrap_or(filename);
|
||||
let stem = basename.trim_end_matches(".gguf");
|
||||
if let Some(pos) = stem.rfind("-of-") {
|
||||
stem.get(..pos)
|
||||
.and_then(|before| before.rsplit('-').next())
|
||||
.map(|s| !s.is_empty() && s.chars().all(|c| c.is_ascii_digit()))
|
||||
.unwrap_or(false)
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn build_download_url(repo_id: &str, filename: &str) -> String {
|
||||
format!("{}/{}/resolve/main/{}", HF_DOWNLOAD_BASE, repo_id, filename)
|
||||
}
|
||||
|
||||
/// Collect single-file GGUFs into quantization variants (sharded files are excluded).
|
||||
fn group_into_variants(repo_id: &str, files: Vec<HfApiSibling>) -> Vec<HfQuantVariant> {
|
||||
let mut variants: Vec<HfQuantVariant> = files
|
||||
.into_iter()
|
||||
.filter(|s| {
|
||||
s.rfilename.ends_with(".gguf")
|
||||
&& !is_shard_file(&s.rfilename)
|
||||
&& parse_quantization(&s.rfilename) != "unknown"
|
||||
})
|
||||
.map(|s| {
|
||||
let quant = parse_quantization(&s.rfilename);
|
||||
let info = quant_info(&quant);
|
||||
let download_url = build_download_url(repo_id, &s.rfilename);
|
||||
HfQuantVariant {
|
||||
quantization: quant,
|
||||
size_bytes: s.size.unwrap_or(0),
|
||||
filename: s.rfilename,
|
||||
download_url,
|
||||
description: info.description,
|
||||
quality_rank: info.quality_rank,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
variants.sort_by_key(|v| v.quality_rank);
|
||||
variants
|
||||
}
|
||||
|
||||
pub async fn search_gguf_models(query: &str, limit: usize) -> Result<Vec<HfModelInfo>> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!(
|
||||
"{}?search={}&filter=gguf&sort=downloads&direction=-1&limit={}",
|
||||
HF_API_BASE, query, limit
|
||||
);
|
||||
|
||||
let response = client
|
||||
.get(&url)
|
||||
.header("User-Agent", "goose-ai-agent")
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
bail!("HuggingFace API returned status {}", response.status());
|
||||
}
|
||||
|
||||
let models: Vec<HfApiModel> = response.json().await?;
|
||||
|
||||
let results = models
|
||||
.into_iter()
|
||||
.filter_map(|m| {
|
||||
let repo_id = m.id?;
|
||||
let siblings = m.siblings.unwrap_or_default();
|
||||
|
||||
// The search endpoint may not include `siblings`; parse whatever
|
||||
// is available. Files are fetched on-demand via `get_repo_gguf_variants`.
|
||||
let gguf_files: Vec<HfGgufFile> = siblings
|
||||
.into_iter()
|
||||
.filter(|s| s.rfilename.ends_with(".gguf"))
|
||||
.map(|s| {
|
||||
let quantization = parse_quantization(&s.rfilename);
|
||||
let download_url = build_download_url(&repo_id, &s.rfilename);
|
||||
HfGgufFile {
|
||||
filename: s.rfilename,
|
||||
size_bytes: s.size.unwrap_or(0),
|
||||
quantization,
|
||||
download_url,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let author = m
|
||||
.author
|
||||
.unwrap_or_else(|| repo_id.split('/').next().unwrap_or_default().to_string());
|
||||
let model_name = repo_id
|
||||
.split('/')
|
||||
.next_back()
|
||||
.unwrap_or(&repo_id)
|
||||
.to_string();
|
||||
|
||||
Some(HfModelInfo {
|
||||
repo_id,
|
||||
author,
|
||||
model_name,
|
||||
downloads: m.downloads.unwrap_or(0),
|
||||
gguf_files,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Fetch GGUF files for a repo and return them grouped by quantization.
|
||||
pub async fn get_repo_gguf_variants(repo_id: &str) -> Result<Vec<HfQuantVariant>> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id);
|
||||
|
||||
let response = client
|
||||
.get(&url)
|
||||
.header("User-Agent", "goose-ai-agent")
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
bail!(
|
||||
"HuggingFace API returned status {} for repo {}",
|
||||
response.status(),
|
||||
repo_id
|
||||
);
|
||||
}
|
||||
|
||||
let model: HfApiModel = response.json().await?;
|
||||
let siblings = model.siblings.unwrap_or_default();
|
||||
|
||||
Ok(group_into_variants(repo_id, siblings))
|
||||
}
|
||||
|
||||
/// Fetch raw GGUF files (kept for resolve_model_spec).
|
||||
pub async fn get_repo_gguf_files(repo_id: &str) -> Result<Vec<HfGgufFile>> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id);
|
||||
|
||||
let response = client
|
||||
.get(&url)
|
||||
.header("User-Agent", "goose-ai-agent")
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
bail!(
|
||||
"HuggingFace API returned status {} for repo {}",
|
||||
response.status(),
|
||||
repo_id
|
||||
);
|
||||
}
|
||||
|
||||
let model: HfApiModel = response.json().await?;
|
||||
let siblings = model.siblings.unwrap_or_default();
|
||||
|
||||
let files = siblings
|
||||
.into_iter()
|
||||
.filter(|s| s.rfilename.ends_with(".gguf"))
|
||||
.filter(|s| !is_shard_file(&s.rfilename))
|
||||
.map(|s| {
|
||||
let quantization = parse_quantization(&s.rfilename);
|
||||
let download_url = build_download_url(repo_id, &s.rfilename);
|
||||
HfGgufFile {
|
||||
filename: s.rfilename,
|
||||
size_bytes: s.size.unwrap_or(0),
|
||||
quantization,
|
||||
download_url,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(files)
|
||||
}
|
||||
|
||||
/// Parse a model spec like "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M" into (repo_id, quantization).
|
||||
pub fn parse_model_spec(spec: &str) -> Result<(String, String)> {
|
||||
let (repo_id, quant) = spec.rsplit_once(':').ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"Invalid model spec '{}': expected format 'user/repo:quantization'",
|
||||
spec
|
||||
)
|
||||
})?;
|
||||
|
||||
if !repo_id.contains('/') {
|
||||
bail!("Invalid repo_id '{}': expected format 'user/repo'", repo_id);
|
||||
}
|
||||
|
||||
Ok((repo_id.to_string(), quant.to_string()))
|
||||
}
|
||||
|
||||
/// Resolve a model spec to a specific GGUF file from the repo.
|
||||
pub async fn resolve_model_spec(spec: &str) -> Result<(String, HfGgufFile)> {
|
||||
let (repo_id, quant) = parse_model_spec(spec)?;
|
||||
let files = get_repo_gguf_files(&repo_id).await?;
|
||||
|
||||
let file = files
|
||||
.into_iter()
|
||||
.find(|f| f.quantization.eq_ignore_ascii_case(&quant))
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"No GGUF file with quantization '{}' found in {}",
|
||||
quant,
|
||||
repo_id
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok((repo_id, file))
|
||||
}
|
||||
|
||||
/// Recommend which quantization variant to use based on available memory.
|
||||
pub fn recommend_variant(
|
||||
variants: &[HfQuantVariant],
|
||||
available_memory_bytes: u64,
|
||||
) -> Option<usize> {
|
||||
// We need ~10-20% overhead beyond model size for inference context.
|
||||
// Pick the highest-quality variant that fits.
|
||||
let usable = (available_memory_bytes as f64 * 0.85) as u64;
|
||||
|
||||
let mut best: Option<usize> = None;
|
||||
for (i, v) in variants.iter().enumerate() {
|
||||
if v.size_bytes <= usable {
|
||||
match best {
|
||||
Some(bi) if variants[bi].quality_rank < v.quality_rank => best = Some(i),
|
||||
None => best = Some(i),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_quantization() {
|
||||
assert_eq!(parse_quantization("Model-Q4_K_M.gguf"), "Q4_K_M");
|
||||
assert_eq!(parse_quantization("Model-Q8_0.gguf"), "Q8_0");
|
||||
assert_eq!(parse_quantization("Model-IQ4_NL.gguf"), "IQ4_NL");
|
||||
assert_eq!(parse_quantization("Model-F16.gguf"), "F16");
|
||||
assert_eq!(parse_quantization("random-name.gguf"), "unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_quantization_with_directory() {
|
||||
assert_eq!(
|
||||
parse_quantization("Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf"),
|
||||
"Q5_K_M"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_quantization_extended_tags() {
|
||||
assert_eq!(parse_quantization("Model-MXFP4_MOE.gguf"), "MXFP4_MOE");
|
||||
assert_eq!(parse_quantization("Model-UD-TQ1_0.gguf"), "TQ1_0");
|
||||
assert_eq!(parse_quantization("Model-Q2_K_L.gguf"), "Q2_K_L");
|
||||
assert_eq!(parse_quantization("Model-UD-Q4_K_XL.gguf"), "Q4_K_XL");
|
||||
assert_eq!(parse_quantization("Model-UD-IQ1_M.gguf"), "IQ1_M");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_shard_file() {
|
||||
assert!(is_shard_file("Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf"));
|
||||
assert!(is_shard_file("Model-BF16-00003-of-00004.gguf"));
|
||||
assert!(!is_shard_file("Model-Q4_K_M.gguf"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_model_spec() {
|
||||
let (repo, quant) =
|
||||
parse_model_spec("bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M").unwrap();
|
||||
assert_eq!(repo, "bartowski/Llama-3.2-1B-Instruct-GGUF");
|
||||
assert_eq!(quant, "Q4_K_M");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_model_spec_invalid() {
|
||||
assert!(parse_model_spec("no-colon").is_err());
|
||||
assert!(parse_model_spec("noslash:Q4_K_M").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recommend_variant() {
|
||||
let variants = vec![
|
||||
HfQuantVariant {
|
||||
quantization: "Q2_K".into(),
|
||||
size_bytes: 2_000_000_000,
|
||||
filename: "m-Q2_K.gguf".into(),
|
||||
download_url: String::new(),
|
||||
description: "Small",
|
||||
quality_rank: 7,
|
||||
},
|
||||
HfQuantVariant {
|
||||
quantization: "Q4_K_M".into(),
|
||||
size_bytes: 4_000_000_000,
|
||||
filename: "m-Q4_K_M.gguf".into(),
|
||||
download_url: String::new(),
|
||||
description: "Medium",
|
||||
quality_rank: 19,
|
||||
},
|
||||
HfQuantVariant {
|
||||
quantization: "Q8_0".into(),
|
||||
size_bytes: 8_000_000_000,
|
||||
filename: "m-Q8_0.gguf".into(),
|
||||
download_url: String::new(),
|
||||
description: "Large",
|
||||
quality_rank: 25,
|
||||
},
|
||||
];
|
||||
|
||||
assert_eq!(recommend_variant(&variants, 5_000_000_000), Some(1));
|
||||
assert_eq!(recommend_variant(&variants, 10_000_000_000), Some(2));
|
||||
assert_eq!(recommend_variant(&variants, 1_000_000_000), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,685 @@
|
||||
//! Tool call emulation for models without native tool-calling support.
|
||||
//!
|
||||
//! The model is prompted to emit shell commands as `$ command` on a new line and
|
||||
//! code blocks as `` ```execute `` fenced blocks. A streaming parser detects these
|
||||
//! patterns and converts them into tool-call messages.
|
||||
//!
|
||||
//! # Known false-positive scenarios
|
||||
//!
|
||||
//! Because detection is purely text-based, the parser can misinterpret model output:
|
||||
//!
|
||||
//! - **`$` at line start in explanatory text.** If the model writes a line starting
|
||||
//! with `$` as an example (e.g. "$ is the jQuery selector"), it will be treated as
|
||||
//! a shell command. Mid-sentence `$` (e.g. "costs $50") is safe — only `\n$` or
|
||||
//! `$` at the very start of output triggers command detection.
|
||||
//!
|
||||
//! - **`` ```execute `` in explanatory code fences.** If the model uses this exact
|
||||
//! fence tag in prose, the content will be executed. Standard `` ```js `` or
|
||||
//! `` ```python `` fences are not affected.
|
||||
//!
|
||||
//! These are inherent to text-based tool emulation. Models with native tool-calling
|
||||
//! support should use the `inference_native_tools` path instead.
|
||||
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::providers::errors::ProviderError;
|
||||
use llama_cpp_2::model::AddBos;
|
||||
use rmcp::model::{CallToolRequestParams, Tool};
|
||||
use serde_json::json;
|
||||
use std::borrow::Cow;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::inference_engine::{
|
||||
create_and_prefill_context, generation_loop, validate_and_compute_context, GenerationContext,
|
||||
TokenAction,
|
||||
};
|
||||
use super::{finalize_usage, StreamSender, CODE_EXECUTION_TOOL, SHELL_TOOL};
|
||||
|
||||
const HOLD_BACK_CODE_MODE: usize = " ```execute\n".len();
|
||||
const HOLD_BACK_SHELL_ONLY: usize = "\n$".len();
|
||||
|
||||
pub(super) fn load_tiny_model_prompt() -> String {
|
||||
use std::env;
|
||||
|
||||
let os = if cfg!(target_os = "macos") {
|
||||
"macos"
|
||||
} else if cfg!(target_os = "linux") {
|
||||
"linux"
|
||||
} else if cfg!(target_os = "windows") {
|
||||
"windows"
|
||||
} else {
|
||||
"unknown"
|
||||
};
|
||||
|
||||
let working_directory = env::current_dir()
|
||||
.map(|p| p.display().to_string())
|
||||
.unwrap_or_else(|_| "unknown".to_string());
|
||||
|
||||
let shell = env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string());
|
||||
|
||||
let context = json!({
|
||||
"os": os,
|
||||
"working_directory": working_directory,
|
||||
"shell": shell,
|
||||
});
|
||||
|
||||
crate::prompt_template::render_template("tiny_model_system.md", &context).unwrap_or_else(|e| {
|
||||
tracing::warn!("Failed to load tiny_model_system.md: {:?}", e);
|
||||
"You are Goose, an AI assistant. You can execute shell commands by starting lines with $."
|
||||
.to_string()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn build_emulator_tool_description(tools: &[Tool], code_mode_enabled: bool) -> String {
|
||||
let mut tool_desc = String::new();
|
||||
|
||||
if code_mode_enabled {
|
||||
tool_desc.push_str("\n\n# Running Code\n\n");
|
||||
tool_desc.push_str(
|
||||
"You can call tools by writing code in a ```execute block. \
|
||||
The code runs immediately — do not explain it, just run it.\n\n",
|
||||
);
|
||||
tool_desc.push_str("Example — counting files in /tmp:\n\n");
|
||||
tool_desc.push_str("```execute\nasync function run() {\n");
|
||||
tool_desc.push_str(
|
||||
" const result = await Developer.shell({ command: \"ls -1 /tmp | wc -l\" });\n",
|
||||
);
|
||||
tool_desc.push_str(" return result;\n}\n```\n\n");
|
||||
tool_desc.push_str("Rules:\n");
|
||||
tool_desc.push_str("- Code MUST define async function run() and return a result\n");
|
||||
tool_desc.push_str("- All function calls are async — use await\n");
|
||||
tool_desc.push_str("- Use ```execute for tool calls, $ for simple shell one-liners\n\n");
|
||||
tool_desc.push_str("Available functions:\n\n");
|
||||
|
||||
for tool in tools {
|
||||
if tool.name.starts_with("code_execution__") {
|
||||
continue;
|
||||
}
|
||||
let parts: Vec<&str> = tool.name.splitn(2, "__").collect();
|
||||
if parts.len() == 2 {
|
||||
let namespace = {
|
||||
let mut c = parts[0].chars();
|
||||
match c.next() {
|
||||
None => String::new(),
|
||||
Some(first) => first.to_uppercase().chain(c).collect::<String>(),
|
||||
}
|
||||
};
|
||||
let camel_name: String = parts[1]
|
||||
.split('_')
|
||||
.enumerate()
|
||||
.map(|(i, part)| {
|
||||
if i == 0 {
|
||||
part.to_string()
|
||||
} else {
|
||||
let mut c = part.chars();
|
||||
match c.next() {
|
||||
None => String::new(),
|
||||
Some(first) => first.to_uppercase().chain(c).collect(),
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let desc = tool.description.as_ref().map(|d| d.as_ref()).unwrap_or("");
|
||||
tool_desc.push_str(&format!("- {namespace}.{camel_name}(): {desc}\n"));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tool_desc.push_str("\n\n# Tools\n\nYou have access to the following tools:\n\n");
|
||||
for tool in tools {
|
||||
let desc = tool
|
||||
.description
|
||||
.as_ref()
|
||||
.map(|d| d.as_ref())
|
||||
.unwrap_or("No description");
|
||||
tool_desc.push_str(&format!("- {}: {}\n", tool.name, desc));
|
||||
}
|
||||
}
|
||||
|
||||
tool_desc
|
||||
}
|
||||
|
||||
enum EmulatorAction {
|
||||
Text(String),
|
||||
ShellCommand(String),
|
||||
ExecuteCode(String),
|
||||
}
|
||||
|
||||
enum ParserState {
|
||||
Normal,
|
||||
InCommand,
|
||||
InExecuteBlock,
|
||||
}
|
||||
|
||||
struct StreamingEmulatorParser {
|
||||
buffer: String,
|
||||
state: ParserState,
|
||||
code_mode_enabled: bool,
|
||||
}
|
||||
|
||||
impl StreamingEmulatorParser {
|
||||
fn new(code_mode_enabled: bool) -> Self {
|
||||
Self {
|
||||
buffer: String::new(),
|
||||
state: ParserState::Normal,
|
||||
code_mode_enabled,
|
||||
}
|
||||
}
|
||||
|
||||
fn process_chunk(&mut self, chunk: &str) -> Vec<EmulatorAction> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut results = Vec::new();
|
||||
|
||||
loop {
|
||||
match self.state {
|
||||
ParserState::InCommand => {
|
||||
if let Some((command_line, rest)) = self.buffer.split_once('\n') {
|
||||
if let Some(command) = command_line.strip_prefix('$') {
|
||||
let command = command.trim();
|
||||
if !command.is_empty() {
|
||||
results.push(EmulatorAction::ShellCommand(command.to_string()));
|
||||
}
|
||||
}
|
||||
self.buffer = rest.to_string();
|
||||
self.state = ParserState::Normal;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
ParserState::InExecuteBlock => {
|
||||
// Look for closing ``` to end the execute block
|
||||
if let Some(end_idx) = self.buffer.find("\n```") {
|
||||
#[allow(clippy::string_slice)]
|
||||
let code = self.buffer[..end_idx].to_string();
|
||||
// Skip past the closing ``` and any trailing newline
|
||||
#[allow(clippy::string_slice)]
|
||||
let rest = &self.buffer[end_idx + 4..];
|
||||
let rest = rest.strip_prefix('\n').unwrap_or(rest);
|
||||
self.buffer = rest.to_string();
|
||||
self.state = ParserState::Normal;
|
||||
if !code.trim().is_empty() {
|
||||
results.push(EmulatorAction::ExecuteCode(code));
|
||||
}
|
||||
} else {
|
||||
// Still accumulating code — wait for closing fence
|
||||
break;
|
||||
}
|
||||
}
|
||||
ParserState::Normal => {
|
||||
// Check for ```execute block (code mode)
|
||||
if self.code_mode_enabled {
|
||||
if let Some((before, after)) = self.buffer.split_once("```execute\n") {
|
||||
if !before.trim().is_empty() {
|
||||
results.push(EmulatorAction::Text(before.to_string()));
|
||||
}
|
||||
self.buffer = after.to_string();
|
||||
self.state = ParserState::InExecuteBlock;
|
||||
continue;
|
||||
}
|
||||
// Also handle without newline after tag (accumulating)
|
||||
if self.buffer.ends_with("```execute") {
|
||||
let before = self.buffer.trim_end_matches("```execute");
|
||||
if !before.trim().is_empty() {
|
||||
results.push(EmulatorAction::Text(before.to_string()));
|
||||
}
|
||||
self.buffer.clear();
|
||||
self.state = ParserState::InExecuteBlock;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Check for $ command
|
||||
if let Some((before_dollar, from_dollar)) = self.buffer.split_once("\n$") {
|
||||
let text = format!("{}\n", before_dollar);
|
||||
if !text.trim().is_empty() {
|
||||
results.push(EmulatorAction::Text(text));
|
||||
}
|
||||
self.buffer = format!("${}", from_dollar);
|
||||
self.state = ParserState::InCommand;
|
||||
} else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() {
|
||||
self.state = ParserState::InCommand;
|
||||
} else {
|
||||
let hold_back = if self.code_mode_enabled {
|
||||
HOLD_BACK_CODE_MODE
|
||||
} else {
|
||||
HOLD_BACK_SHELL_ONLY
|
||||
};
|
||||
let char_count = self.buffer.chars().count();
|
||||
if char_count > hold_back && !self.buffer.ends_with('\n') {
|
||||
let mut chars = self.buffer.chars();
|
||||
let emit_count = char_count - hold_back;
|
||||
let emit_text: String = chars.by_ref().take(emit_count).collect();
|
||||
let keep_text: String = chars.collect();
|
||||
if !emit_text.is_empty() {
|
||||
results.push(EmulatorAction::Text(emit_text));
|
||||
}
|
||||
self.buffer = keep_text;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> Vec<EmulatorAction> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
if !self.buffer.is_empty() {
|
||||
match self.state {
|
||||
ParserState::InCommand => {
|
||||
let command_line = self.buffer.trim();
|
||||
if let Some(command) = command_line.strip_prefix('$') {
|
||||
let command = command.trim();
|
||||
if !command.is_empty() {
|
||||
results.push(EmulatorAction::ShellCommand(command.to_string()));
|
||||
}
|
||||
} else if !command_line.is_empty() {
|
||||
results.push(EmulatorAction::Text(self.buffer.clone()));
|
||||
}
|
||||
}
|
||||
ParserState::InExecuteBlock => {
|
||||
let code = self.buffer.trim();
|
||||
if !code.is_empty() {
|
||||
results.push(EmulatorAction::ExecuteCode(code.to_string()));
|
||||
}
|
||||
}
|
||||
ParserState::Normal => {
|
||||
results.push(EmulatorAction::Text(self.buffer.clone()));
|
||||
}
|
||||
}
|
||||
self.buffer.clear();
|
||||
self.state = ParserState::Normal;
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
fn send_emulator_action(
|
||||
action: &EmulatorAction,
|
||||
message_id: &str,
|
||||
tx: &StreamSender,
|
||||
) -> Result<bool, ()> {
|
||||
match action {
|
||||
EmulatorAction::Text(text) => {
|
||||
let mut message = Message::assistant().with_text(text);
|
||||
message.id = Some(message_id.to_string());
|
||||
tx.blocking_send(Ok((Some(message), None)))
|
||||
.map_err(|_| ())?;
|
||||
Ok(false)
|
||||
}
|
||||
EmulatorAction::ShellCommand(command) => {
|
||||
let tool_id = Uuid::new_v4().to_string();
|
||||
let mut args = serde_json::Map::new();
|
||||
args.insert("command".to_string(), json!(command));
|
||||
let tool_call = CallToolRequestParams {
|
||||
meta: None,
|
||||
task: None,
|
||||
name: Cow::Borrowed(SHELL_TOOL),
|
||||
arguments: Some(args),
|
||||
};
|
||||
let mut message = Message::assistant();
|
||||
message
|
||||
.content
|
||||
.push(MessageContent::tool_request(tool_id, Ok(tool_call)));
|
||||
message.id = Some(message_id.to_string());
|
||||
tx.blocking_send(Ok((Some(message), None)))
|
||||
.map_err(|_| ())?;
|
||||
Ok(true)
|
||||
}
|
||||
EmulatorAction::ExecuteCode(code) => {
|
||||
let tool_id = Uuid::new_v4().to_string();
|
||||
let wrapped = if code.contains("async function run()") {
|
||||
code.clone()
|
||||
} else {
|
||||
format!("async function run() {{\n{}\n}}", code)
|
||||
};
|
||||
let mut args = serde_json::Map::new();
|
||||
args.insert("code".to_string(), json!(wrapped));
|
||||
let tool_call = CallToolRequestParams {
|
||||
meta: None,
|
||||
task: None,
|
||||
name: Cow::Borrowed(CODE_EXECUTION_TOOL),
|
||||
arguments: Some(args),
|
||||
};
|
||||
let mut message = Message::assistant();
|
||||
message
|
||||
.content
|
||||
.push(MessageContent::tool_request(tool_id, Ok(tool_call)));
|
||||
message.id = Some(message_id.to_string());
|
||||
tx.blocking_send(Ok((Some(message), None)))
|
||||
.map_err(|_| ())?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn generate_with_emulated_tools(
|
||||
ctx: &mut GenerationContext<'_>,
|
||||
code_mode_enabled: bool,
|
||||
) -> Result<(), ProviderError> {
|
||||
let prompt = ctx
|
||||
.loaded
|
||||
.model
|
||||
.apply_chat_template(&ctx.loaded.template, ctx.chat_messages, true)
|
||||
.map_err(|e| {
|
||||
ProviderError::ExecutionError(format!("Failed to apply chat template: {}", e))
|
||||
})?;
|
||||
|
||||
let tokens = ctx
|
||||
.loaded
|
||||
.model
|
||||
.str_to_token(&prompt, AddBos::Never)
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||
|
||||
let (prompt_token_count, effective_ctx) = validate_and_compute_context(
|
||||
ctx.loaded,
|
||||
ctx.runtime,
|
||||
tokens.len(),
|
||||
ctx.context_limit,
|
||||
ctx.settings,
|
||||
)?;
|
||||
let mut llama_ctx = create_and_prefill_context(
|
||||
ctx.loaded,
|
||||
ctx.runtime,
|
||||
&tokens,
|
||||
effective_ctx,
|
||||
ctx.settings,
|
||||
)?;
|
||||
|
||||
let message_id = ctx.message_id;
|
||||
let tx = ctx.tx;
|
||||
let mut emulator_parser = StreamingEmulatorParser::new(code_mode_enabled);
|
||||
let mut tool_call_emitted = false;
|
||||
let mut send_failed = false;
|
||||
|
||||
let output_token_count = generation_loop(
|
||||
&ctx.loaded.model,
|
||||
&mut llama_ctx,
|
||||
ctx.settings,
|
||||
prompt_token_count,
|
||||
effective_ctx,
|
||||
|piece| {
|
||||
let actions = emulator_parser.process_chunk(piece);
|
||||
for action in actions {
|
||||
match send_emulator_action(&action, message_id, tx) {
|
||||
Ok(is_tool) => {
|
||||
if is_tool {
|
||||
tool_call_emitted = true;
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
send_failed = true;
|
||||
return Ok(TokenAction::Stop);
|
||||
}
|
||||
}
|
||||
}
|
||||
if tool_call_emitted {
|
||||
Ok(TokenAction::Stop)
|
||||
} else {
|
||||
Ok(TokenAction::Continue)
|
||||
}
|
||||
},
|
||||
)?;
|
||||
|
||||
if !send_failed {
|
||||
for action in emulator_parser.flush() {
|
||||
if send_emulator_action(&action, message_id, tx).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let provider_usage = finalize_usage(
|
||||
ctx.log,
|
||||
std::mem::take(&mut ctx.model_name),
|
||||
"emulator",
|
||||
prompt_token_count,
|
||||
output_token_count,
|
||||
None,
|
||||
);
|
||||
let _ = ctx.tx.blocking_send(Ok((None, Some(provider_usage))));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Collect all actions from feeding chunks through the parser, then flushing.
|
||||
fn parse_chunks(chunks: &[&str], code_mode: bool) -> Vec<EmulatorAction> {
|
||||
let mut parser = StreamingEmulatorParser::new(code_mode);
|
||||
let mut actions = Vec::new();
|
||||
for chunk in chunks {
|
||||
actions.extend(parser.process_chunk(chunk));
|
||||
}
|
||||
actions.extend(parser.flush());
|
||||
actions
|
||||
}
|
||||
|
||||
fn parse_all(input: &str, code_mode: bool) -> Vec<EmulatorAction> {
|
||||
parse_chunks(&[input], code_mode)
|
||||
}
|
||||
|
||||
fn assert_text(action: &EmulatorAction, expected: &str) {
|
||||
match action {
|
||||
EmulatorAction::Text(t) => assert_eq!(t.trim(), expected.trim(), "text mismatch"),
|
||||
other => panic!("expected Text, got {:?}", action_label(other)),
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_shell(action: &EmulatorAction, expected: &str) {
|
||||
match action {
|
||||
EmulatorAction::ShellCommand(cmd) => {
|
||||
assert_eq!(cmd, expected, "shell command mismatch")
|
||||
}
|
||||
other => panic!("expected ShellCommand, got {:?}", action_label(other)),
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_execute(action: &EmulatorAction, expected: &str) {
|
||||
match action {
|
||||
EmulatorAction::ExecuteCode(code) => {
|
||||
assert_eq!(code.trim(), expected.trim(), "execute code mismatch")
|
||||
}
|
||||
other => panic!("expected ExecuteCode, got {:?}", action_label(other)),
|
||||
}
|
||||
}
|
||||
|
||||
fn action_label(a: &EmulatorAction) -> &'static str {
|
||||
match a {
|
||||
EmulatorAction::Text(_) => "Text",
|
||||
EmulatorAction::ShellCommand(_) => "ShellCommand",
|
||||
EmulatorAction::ExecuteCode(_) => "ExecuteCode",
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plain_text_no_tools() {
|
||||
let actions = parse_all("Hello, world!", false);
|
||||
// Hold-back may split text across actions; concatenate all text
|
||||
let all_text: String = actions
|
||||
.iter()
|
||||
.map(|a| match a {
|
||||
EmulatorAction::Text(t) => t.as_str(),
|
||||
_ => panic!("expected only Text actions"),
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(all_text.trim(), "Hello, world!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_shell_command() {
|
||||
let actions = parse_all("$ ls -la\n", false);
|
||||
assert_eq!(actions.len(), 1);
|
||||
assert_shell(&actions[0], "ls -la");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_then_shell_command() {
|
||||
let actions = parse_all("Let me check:\n$ ls -la\n", false);
|
||||
assert!(actions.len() >= 2);
|
||||
assert_text(&actions[0], "Let me check:");
|
||||
assert_shell(&actions[actions.len() - 1], "ls -la");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shell_command_at_start_of_output() {
|
||||
let actions = parse_all("$ whoami\n", false);
|
||||
assert_eq!(actions.len(), 1);
|
||||
assert_shell(&actions[0], "whoami");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shell_command_without_trailing_newline() {
|
||||
// Flush should handle unterminated command
|
||||
let actions = parse_all("$ whoami", false);
|
||||
assert_eq!(actions.len(), 1);
|
||||
assert_shell(&actions[0], "whoami");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dollar_sign_mid_sentence_is_not_command() {
|
||||
let actions = parse_all("It costs $50 per month", false);
|
||||
for action in &actions {
|
||||
assert!(
|
||||
matches!(action, EmulatorAction::Text(_)),
|
||||
"mid-sentence $ should not trigger a shell command"
|
||||
);
|
||||
}
|
||||
let all_text: String = actions
|
||||
.iter()
|
||||
.filter_map(|a| match a {
|
||||
EmulatorAction::Text(t) => Some(t.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(all_text.trim(), "It costs $50 per month");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execute_block() {
|
||||
let input = "Here's the code:\n```execute\nconsole.log('hi');\n```\n";
|
||||
let actions = parse_all(input, true);
|
||||
assert!(actions.len() >= 2);
|
||||
assert_text(&actions[0], "Here's the code:");
|
||||
assert_execute(&actions[actions.len() - 1], "console.log('hi');");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execute_block_not_detected_without_code_mode() {
|
||||
let input = "```execute\nconsole.log('hi');\n```\n";
|
||||
let actions = parse_all(input, false);
|
||||
// Should be treated as plain text
|
||||
for action in &actions {
|
||||
assert!(matches!(action, EmulatorAction::Text(_)));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dollar_split_across_chunks() {
|
||||
// The \n and $ arrive in separate chunks
|
||||
let actions = parse_chunks(&["Let me check\n", "$ ls -la\n"], false);
|
||||
let shells: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
|
||||
.collect();
|
||||
assert_eq!(shells.len(), 1);
|
||||
assert_shell(shells[0], "ls -la");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execute_fence_split_across_chunks() {
|
||||
let actions = parse_chunks(&["Here:\n```ex", "ecute\nlet x = 1;\n", "```\n"], true);
|
||||
let executes: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
|
||||
.collect();
|
||||
assert_eq!(executes.len(), 1);
|
||||
assert_execute(executes[0], "let x = 1;");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_commands_on_separate_lines() {
|
||||
// In practice, generation stops after the first tool call. But the
|
||||
// parser should detect commands separated by \n$ when fed as chunks.
|
||||
let actions = parse_chunks(&["Here:\n$ cd /tmp\n", "Done.\n$ ls\n"], false);
|
||||
let shells: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
|
||||
.collect();
|
||||
assert_eq!(shells.len(), 2);
|
||||
assert_shell(shells[0], "cd /tmp");
|
||||
assert_shell(shells[1], "ls");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn regular_code_fence_not_treated_as_execute() {
|
||||
let input = "```python\nprint('hi')\n```\n";
|
||||
let actions = parse_all(input, true);
|
||||
for action in &actions {
|
||||
assert!(
|
||||
matches!(action, EmulatorAction::Text(_)),
|
||||
"regular code fence should be text"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_command_ignored() {
|
||||
let actions = parse_all("$\n", false);
|
||||
// Empty command after $ should not produce a ShellCommand
|
||||
let shells: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
|
||||
.collect();
|
||||
assert_eq!(shells.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_by_token_streaming() {
|
||||
// Simulate LLM generating one token at a time
|
||||
let input = "$ echo hello\n";
|
||||
let chars: Vec<String> = input.chars().map(|c| c.to_string()).collect();
|
||||
let chunks: Vec<&str> = chars.iter().map(|s| s.as_str()).collect();
|
||||
let actions = parse_chunks(&chunks, false);
|
||||
let shells: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
|
||||
.collect();
|
||||
assert_eq!(shells.len(), 1);
|
||||
assert_shell(shells[0], "echo hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execute_block_with_multiline_code() {
|
||||
let input = "```execute\nasync function run() {\n const r = await Developer.shell({ command: \"ls\" });\n return r;\n}\n```\n";
|
||||
let actions = parse_all(input, true);
|
||||
let executes: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
|
||||
.collect();
|
||||
assert_eq!(executes.len(), 1);
|
||||
match executes[0] {
|
||||
EmulatorAction::ExecuteCode(code) => {
|
||||
assert!(code.contains("async function run()"));
|
||||
assert!(code.contains("Developer.shell"));
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unclosed_execute_block_flushed() {
|
||||
// Model stops generating mid-block
|
||||
let input = "```execute\nlet x = 1;";
|
||||
let actions = parse_all(input, true);
|
||||
let executes: Vec<_> = actions
|
||||
.iter()
|
||||
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
|
||||
.collect();
|
||||
assert_eq!(executes.len(), 1);
|
||||
assert_execute(executes[0], "let x = 1;");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,387 @@
|
||||
use crate::providers::errors::ProviderError;
|
||||
use crate::providers::local_inference::local_model_registry::ModelSettings;
|
||||
use crate::providers::utils::RequestLog;
|
||||
use llama_cpp_2::context::params::LlamaContextParams;
|
||||
use llama_cpp_2::llama_batch::LlamaBatch;
|
||||
use llama_cpp_2::model::{LlamaChatMessage, LlamaChatTemplate, LlamaModel};
|
||||
use llama_cpp_2::sampling::LlamaSampler;
|
||||
use std::num::NonZeroU32;
|
||||
|
||||
use super::{InferenceRuntime, StreamSender};
|
||||
|
||||
pub(super) struct GenerationContext<'a> {
|
||||
pub loaded: &'a LoadedModel,
|
||||
pub runtime: &'a InferenceRuntime,
|
||||
pub chat_messages: &'a [LlamaChatMessage],
|
||||
pub settings: &'a ModelSettings,
|
||||
pub context_limit: usize,
|
||||
pub model_name: String,
|
||||
pub message_id: &'a str,
|
||||
pub tx: &'a StreamSender,
|
||||
pub log: &'a mut RequestLog,
|
||||
}
|
||||
|
||||
pub(super) struct LoadedModel {
|
||||
pub model: LlamaModel,
|
||||
pub template: LlamaChatTemplate,
|
||||
}
|
||||
|
||||
/// Estimate the maximum context length that can fit in available accelerator/CPU
|
||||
/// memory based on the model's KV cache requirements.
|
||||
///
|
||||
/// Returns `None` if the model architecture values are unavailable.
|
||||
pub(super) fn estimate_max_context_for_memory(
|
||||
model: &LlamaModel,
|
||||
runtime: &InferenceRuntime,
|
||||
) -> Option<usize> {
|
||||
let available = super::available_inference_memory_bytes(runtime);
|
||||
if available == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Reserve memory for computation scratch buffers (attention, etc.) and other overhead.
|
||||
// The compute buffer can be 40-50% of the KV cache size for large models, so we
|
||||
// conservatively use only half the available memory for the KV cache.
|
||||
let usable = (available as f64 * 0.5) as u64;
|
||||
|
||||
let n_layer = model.n_layer() as u64;
|
||||
let n_head_kv = model.n_head_kv() as u64;
|
||||
let n_head = model.n_head() as u64;
|
||||
let n_embd = model.n_embd() as u64;
|
||||
|
||||
if n_head == 0 || n_layer == 0 || n_head_kv == 0 || n_embd == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
// For MLA (Multi-head Latent Attention) models like DeepSeek/GLM, the actual KV cache
|
||||
// dimensions differ from n_head_kv * head_dim. Read the true dimensions from GGUF metadata.
|
||||
let arch = model
|
||||
.meta_val_str("general.architecture")
|
||||
.unwrap_or_default();
|
||||
let head_dim = n_embd / n_head;
|
||||
let k_per_head = model
|
||||
.meta_val_str(&format!("{arch}.attention.key_length"))
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.unwrap_or(head_dim);
|
||||
let v_per_head = model
|
||||
.meta_val_str(&format!("{arch}.attention.value_length"))
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.unwrap_or(head_dim);
|
||||
|
||||
// Total KV dimensions across all KV heads, times n_layer, times 2 bytes (f16) per element
|
||||
let bytes_per_token = (k_per_head + v_per_head) * n_head_kv * n_layer * 2;
|
||||
|
||||
if bytes_per_token == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some((usable / bytes_per_token) as usize)
|
||||
}
|
||||
|
||||
pub(super) fn context_cap(
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
context_limit: usize,
|
||||
n_ctx_train: usize,
|
||||
memory_max_ctx: Option<usize>,
|
||||
) -> usize {
|
||||
if let Some(ctx_size) = settings.context_size {
|
||||
return ctx_size as usize;
|
||||
}
|
||||
|
||||
let limit = if context_limit > 0 {
|
||||
context_limit
|
||||
} else {
|
||||
n_ctx_train
|
||||
};
|
||||
|
||||
match memory_max_ctx {
|
||||
Some(mem_max) if mem_max < limit => {
|
||||
tracing::info!(
|
||||
"Capping context from {} to {} based on available memory",
|
||||
limit,
|
||||
mem_max,
|
||||
);
|
||||
mem_max
|
||||
}
|
||||
_ => limit,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn effective_context_size(
|
||||
prompt_token_count: usize,
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
context_limit: usize,
|
||||
n_ctx_train: usize,
|
||||
memory_max_ctx: Option<usize>,
|
||||
) -> usize {
|
||||
let limit = context_cap(settings, context_limit, n_ctx_train, memory_max_ctx);
|
||||
let min_generation_headroom = 512;
|
||||
let needed = prompt_token_count + min_generation_headroom;
|
||||
if needed > limit {
|
||||
tracing::warn!(
|
||||
"Prompt ({} tokens) + headroom exceeds context limit ({}), capping to limit",
|
||||
prompt_token_count,
|
||||
limit,
|
||||
);
|
||||
}
|
||||
needed.min(limit)
|
||||
}
|
||||
|
||||
pub(super) fn build_context_params(
|
||||
ctx_size: u32,
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
) -> LlamaContextParams {
|
||||
let mut params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(ctx_size));
|
||||
|
||||
if let Some(n_batch) = settings.n_batch {
|
||||
params = params.with_n_batch(n_batch);
|
||||
}
|
||||
if let Some(n_threads) = settings.n_threads {
|
||||
params = params.with_n_threads(n_threads);
|
||||
params = params.with_n_threads_batch(n_threads);
|
||||
}
|
||||
if let Some(flash_attn) = settings.flash_attention {
|
||||
let policy = if flash_attn { 1 } else { 0 };
|
||||
params = params.with_flash_attention_policy(policy);
|
||||
}
|
||||
|
||||
params
|
||||
}
|
||||
|
||||
pub(super) fn build_sampler(
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
) -> LlamaSampler {
|
||||
use crate::providers::local_inference::local_model_registry::SamplingConfig;
|
||||
|
||||
let has_penalties = settings.repeat_penalty != 1.0
|
||||
|| settings.frequency_penalty != 0.0
|
||||
|| settings.presence_penalty != 0.0;
|
||||
|
||||
let mut samplers: Vec<LlamaSampler> = Vec::new();
|
||||
|
||||
if has_penalties {
|
||||
samplers.push(LlamaSampler::penalties(
|
||||
settings.repeat_last_n,
|
||||
settings.repeat_penalty,
|
||||
settings.frequency_penalty,
|
||||
settings.presence_penalty,
|
||||
));
|
||||
}
|
||||
|
||||
match &settings.sampling {
|
||||
SamplingConfig::Greedy => {
|
||||
samplers.push(LlamaSampler::greedy());
|
||||
}
|
||||
SamplingConfig::Temperature {
|
||||
temperature,
|
||||
top_k,
|
||||
top_p,
|
||||
min_p,
|
||||
seed,
|
||||
} => {
|
||||
samplers.push(LlamaSampler::top_k(*top_k));
|
||||
samplers.push(LlamaSampler::top_p(*top_p, 1));
|
||||
samplers.push(LlamaSampler::min_p(*min_p, 1));
|
||||
samplers.push(LlamaSampler::temp(*temperature));
|
||||
samplers.push(LlamaSampler::dist(seed.unwrap_or(0)));
|
||||
}
|
||||
SamplingConfig::MirostatV2 { tau, eta, seed } => {
|
||||
samplers.push(LlamaSampler::mirostat_v2(seed.unwrap_or(0), *tau, *eta));
|
||||
}
|
||||
}
|
||||
|
||||
if samplers.len() == 1 {
|
||||
samplers.pop().unwrap()
|
||||
} else {
|
||||
LlamaSampler::chain_simple(samplers)
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate prompt tokens against memory limits and compute the effective
|
||||
/// context size. Returns `(prompt_token_count, effective_ctx)`.
|
||||
pub(super) fn validate_and_compute_context(
|
||||
loaded: &LoadedModel,
|
||||
runtime: &InferenceRuntime,
|
||||
prompt_token_count: usize,
|
||||
context_limit: usize,
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
) -> Result<(usize, usize), ProviderError> {
|
||||
let n_ctx_train = loaded.model.n_ctx_train() as usize;
|
||||
let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, runtime);
|
||||
let effective_ctx = effective_context_size(
|
||||
prompt_token_count,
|
||||
settings,
|
||||
context_limit,
|
||||
n_ctx_train,
|
||||
memory_max_ctx,
|
||||
);
|
||||
if let Some(mem_max) = memory_max_ctx {
|
||||
if prompt_token_count > mem_max {
|
||||
return Err(ProviderError::ContextLengthExceeded(format!(
|
||||
"Prompt ({} tokens) exceeds estimated memory capacity ({} tokens). \
|
||||
Try a smaller model or reduce conversation length.",
|
||||
prompt_token_count, mem_max,
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok((prompt_token_count, effective_ctx))
|
||||
}
|
||||
|
||||
/// Create a llama context and prefill (decode) all prompt tokens.
|
||||
pub(super) fn create_and_prefill_context<'model>(
|
||||
loaded: &'model LoadedModel,
|
||||
runtime: &InferenceRuntime,
|
||||
tokens: &[llama_cpp_2::token::LlamaToken],
|
||||
effective_ctx: usize,
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
) -> Result<llama_cpp_2::context::LlamaContext<'model>, ProviderError> {
|
||||
let ctx_params = build_context_params(effective_ctx as u32, settings);
|
||||
let mut ctx = loaded
|
||||
.model
|
||||
.new_context(runtime.backend(), ctx_params)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {}", e)))?;
|
||||
|
||||
let n_batch = ctx.n_batch() as usize;
|
||||
for chunk in tokens.chunks(n_batch) {
|
||||
let mut batch = LlamaBatch::get_one(chunk)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Failed to create batch: {}", e)))?;
|
||||
ctx.decode(&mut batch)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Prefill decode failed: {}", e)))?;
|
||||
}
|
||||
|
||||
Ok(ctx)
|
||||
}
|
||||
|
||||
/// Action to take after processing a generated token piece.
|
||||
pub(super) enum TokenAction {
|
||||
Continue,
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// Run the autoregressive generation loop. Calls `on_piece` for each non-empty
|
||||
/// token piece. The callback returns `TokenAction::Stop` to break early.
|
||||
/// Returns the total number of generated tokens.
|
||||
pub(super) fn generation_loop(
|
||||
model: &LlamaModel,
|
||||
ctx: &mut llama_cpp_2::context::LlamaContext<'_>,
|
||||
settings: &crate::providers::local_inference::local_model_registry::ModelSettings,
|
||||
prompt_token_count: usize,
|
||||
effective_ctx: usize,
|
||||
mut on_piece: impl FnMut(&str) -> Result<TokenAction, ProviderError>,
|
||||
) -> Result<i32, ProviderError> {
|
||||
let mut sampler = build_sampler(settings);
|
||||
let max_output = if let Some(max) = settings.max_output_tokens {
|
||||
effective_ctx.saturating_sub(prompt_token_count).min(max)
|
||||
} else {
|
||||
effective_ctx.saturating_sub(prompt_token_count)
|
||||
};
|
||||
let mut decoder = encoding_rs::UTF_8.new_decoder();
|
||||
let mut output_token_count: i32 = 0;
|
||||
|
||||
for _ in 0..max_output {
|
||||
let token = sampler.sample(ctx, -1);
|
||||
sampler.accept(token);
|
||||
|
||||
if model.is_eog_token(token) {
|
||||
break;
|
||||
}
|
||||
|
||||
output_token_count += 1;
|
||||
|
||||
let piece = model
|
||||
.token_to_piece(token, &mut decoder, true, None)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Failed to decode token: {}", e)))?;
|
||||
|
||||
if !piece.is_empty() && matches!(on_piece(&piece)?, TokenAction::Stop) {
|
||||
break;
|
||||
}
|
||||
|
||||
let next_tokens = [token];
|
||||
let mut next_batch = LlamaBatch::get_one(&next_tokens)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Failed to create batch: {}", e)))?;
|
||||
ctx.decode(&mut next_batch)
|
||||
.map_err(|e| ProviderError::ExecutionError(format!("Decode failed: {}", e)))?;
|
||||
}
|
||||
|
||||
Ok(output_token_count)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::providers::local_inference::local_model_registry::ModelSettings;
|
||||
|
||||
fn default_settings() -> ModelSettings {
|
||||
ModelSettings::default()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_basic() {
|
||||
assert_eq!(
|
||||
effective_context_size(100, &default_settings(), 4096, 4096, None),
|
||||
612
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_capped_by_limit() {
|
||||
assert_eq!(
|
||||
effective_context_size(100, &default_settings(), 1024, 8192, None),
|
||||
612
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_capped_by_memory() {
|
||||
assert_eq!(
|
||||
effective_context_size(100, &default_settings(), 4096, 4096, Some(800)),
|
||||
612
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_memory_smaller_than_needed() {
|
||||
assert_eq!(
|
||||
effective_context_size(600, &default_settings(), 4096, 4096, Some(700)),
|
||||
700
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_zero_limit_uses_train() {
|
||||
assert_eq!(
|
||||
effective_context_size(100, &default_settings(), 0, 2048, None),
|
||||
612
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_context_size_prompt_exceeds_all_limits() {
|
||||
assert_eq!(
|
||||
effective_context_size(5000, &default_settings(), 4096, 4096, None),
|
||||
4096
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_cap_with_settings_override() {
|
||||
let mut settings = default_settings();
|
||||
settings.context_size = Some(2048);
|
||||
assert_eq!(context_cap(&settings, 4096, 8192, Some(1024)), 2048);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_cap_without_override() {
|
||||
assert_eq!(context_cap(&default_settings(), 4096, 8192, None), 4096);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_cap_memory_limited() {
|
||||
assert_eq!(
|
||||
context_cap(&default_settings(), 4096, 8192, Some(2048)),
|
||||
2048
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
use crate::conversation::message::Message;
|
||||
use crate::providers::errors::ProviderError;
|
||||
use llama_cpp_2::model::AddBos;
|
||||
use llama_cpp_2::openai::OpenAIChatTemplateParams;
|
||||
|
||||
use super::finalize_usage;
|
||||
use super::inference_engine::{
|
||||
context_cap, create_and_prefill_context, estimate_max_context_for_memory, generation_loop,
|
||||
validate_and_compute_context, GenerationContext, TokenAction,
|
||||
};
|
||||
use super::tool_parsing::{
|
||||
extract_tool_call_messages, extract_xml_tool_call_messages, safe_stream_end,
|
||||
split_content_and_tool_calls, split_content_and_xml_tool_calls,
|
||||
};
|
||||
|
||||
pub(super) fn generate_with_native_tools(
|
||||
ctx: &mut GenerationContext<'_>,
|
||||
oai_messages_json: &Option<String>,
|
||||
full_tools_json: Option<&str>,
|
||||
compact_tools: Option<&str>,
|
||||
) -> Result<(), ProviderError> {
|
||||
let min_generation_headroom = 512;
|
||||
let n_ctx_train = ctx.loaded.model.n_ctx_train() as usize;
|
||||
let memory_max_ctx = estimate_max_context_for_memory(&ctx.loaded.model, ctx.runtime);
|
||||
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 apply_template = |tools: Option<&str>| {
|
||||
if let Some(ref messages_json) = oai_messages_json {
|
||||
let params = OpenAIChatTemplateParams {
|
||||
messages_json: messages_json.as_str(),
|
||||
tools_json: tools,
|
||||
tool_choice: None,
|
||||
json_schema: None,
|
||||
grammar: None,
|
||||
reasoning_format: None,
|
||||
chat_template_kwargs: None,
|
||||
add_generation_prompt: true,
|
||||
use_jinja: true,
|
||||
parallel_tool_calls: false,
|
||||
enable_thinking: false,
|
||||
add_bos: false,
|
||||
add_eos: false,
|
||||
parse_tool_calls: true,
|
||||
};
|
||||
ctx.loaded
|
||||
.model
|
||||
.apply_chat_template_oaicompat(&ctx.loaded.template, ¶ms)
|
||||
} else {
|
||||
ctx.loaded.model.apply_chat_template_with_tools_oaicompat(
|
||||
&ctx.loaded.template,
|
||||
ctx.chat_messages,
|
||||
tools,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let template_result = match apply_template(full_tools_json) {
|
||||
Ok(r) => {
|
||||
let token_count = ctx
|
||||
.loaded
|
||||
.model
|
||||
.str_to_token(&r.prompt, AddBos::Never)
|
||||
.map(|t| t.len())
|
||||
.unwrap_or(0);
|
||||
if token_count > token_budget {
|
||||
apply_template(compact_tools).unwrap_or(r)
|
||||
} else {
|
||||
r
|
||||
}
|
||||
}
|
||||
Err(_) => apply_template(compact_tools).map_err(|e| {
|
||||
ProviderError::ExecutionError(format!("Failed to apply chat template: {}", e))
|
||||
})?,
|
||||
};
|
||||
|
||||
let _ = ctx.log.write(
|
||||
&serde_json::json!({"applied_prompt": &template_result.prompt}),
|
||||
None,
|
||||
);
|
||||
|
||||
let tokens = ctx
|
||||
.loaded
|
||||
.model
|
||||
.str_to_token(&template_result.prompt, AddBos::Never)
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||
|
||||
let (prompt_token_count, effective_ctx) = validate_and_compute_context(
|
||||
ctx.loaded,
|
||||
ctx.runtime,
|
||||
tokens.len(),
|
||||
ctx.context_limit,
|
||||
ctx.settings,
|
||||
)?;
|
||||
let mut llama_ctx = create_and_prefill_context(
|
||||
ctx.loaded,
|
||||
ctx.runtime,
|
||||
&tokens,
|
||||
effective_ctx,
|
||||
ctx.settings,
|
||||
)?;
|
||||
|
||||
let message_id = ctx.message_id;
|
||||
let tx = ctx.tx;
|
||||
let mut generated_text = String::new();
|
||||
let mut streamed_len: usize = 0;
|
||||
|
||||
let output_token_count = generation_loop(
|
||||
&ctx.loaded.model,
|
||||
&mut llama_ctx,
|
||||
ctx.settings,
|
||||
prompt_token_count,
|
||||
effective_ctx,
|
||||
|piece| {
|
||||
generated_text.push_str(piece);
|
||||
|
||||
let has_xml_tc = split_content_and_xml_tool_calls(&generated_text).is_some();
|
||||
let (content, tc) = split_content_and_tool_calls(&generated_text);
|
||||
let stream_up_to = if tc.is_some() {
|
||||
content.len()
|
||||
} else if has_xml_tc {
|
||||
split_content_and_xml_tool_calls(&generated_text)
|
||||
.map(|(c, _)| c.len())
|
||||
.unwrap_or(0)
|
||||
} else {
|
||||
safe_stream_end(&generated_text)
|
||||
};
|
||||
if stream_up_to > streamed_len {
|
||||
#[allow(clippy::string_slice)]
|
||||
let new_text = &generated_text[streamed_len..stream_up_to];
|
||||
if !new_text.is_empty() {
|
||||
let mut msg = Message::assistant().with_text(new_text);
|
||||
msg.id = Some(message_id.to_string());
|
||||
if tx.blocking_send(Ok((Some(msg), None))).is_err() {
|
||||
return Ok(TokenAction::Stop);
|
||||
}
|
||||
}
|
||||
streamed_len = stream_up_to;
|
||||
}
|
||||
|
||||
let should_stop = template_result
|
||||
.additional_stops
|
||||
.iter()
|
||||
.any(|stop| generated_text.ends_with(stop));
|
||||
if should_stop {
|
||||
Ok(TokenAction::Stop)
|
||||
} else {
|
||||
Ok(TokenAction::Continue)
|
||||
}
|
||||
},
|
||||
)?;
|
||||
|
||||
let (content, tool_call_msgs) =
|
||||
if let Some((xml_content, xml_calls)) = split_content_and_xml_tool_calls(&generated_text) {
|
||||
let msgs = extract_xml_tool_call_messages(xml_calls, message_id);
|
||||
(xml_content, msgs)
|
||||
} else {
|
||||
let (json_content, tool_calls_json) = split_content_and_tool_calls(&generated_text);
|
||||
let msgs = tool_calls_json
|
||||
.map(|tc| extract_tool_call_messages(&tc, message_id))
|
||||
.unwrap_or_default();
|
||||
(json_content, msgs)
|
||||
};
|
||||
|
||||
if content.len() > streamed_len {
|
||||
#[allow(clippy::string_slice)]
|
||||
let remaining = &content[streamed_len..];
|
||||
if !remaining.is_empty() {
|
||||
let mut msg = Message::assistant().with_text(remaining);
|
||||
msg.id = Some(message_id.to_string());
|
||||
let _ = tx.blocking_send(Ok((Some(msg), None)));
|
||||
}
|
||||
}
|
||||
|
||||
if !tool_call_msgs.is_empty() {
|
||||
for msg in tool_call_msgs {
|
||||
let _ = tx.blocking_send(Ok((Some(msg), None)));
|
||||
}
|
||||
} else if content.is_empty() && !generated_text.is_empty() {
|
||||
let mut msg = Message::assistant().with_text(&generated_text);
|
||||
msg.id = Some(message_id.to_string());
|
||||
let _ = tx.blocking_send(Ok((Some(msg), None)));
|
||||
}
|
||||
|
||||
let provider_usage = finalize_usage(
|
||||
ctx.log,
|
||||
std::mem::take(&mut ctx.model_name),
|
||||
"native",
|
||||
prompt_token_count,
|
||||
output_token_count,
|
||||
Some(("generated_text", &generated_text)),
|
||||
);
|
||||
let _ = ctx.tx.blocking_send(Ok((None, Some(provider_usage))));
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
use crate::config::paths::Paths;
|
||||
use crate::download_manager::{get_download_manager, DownloadStatus};
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum SamplingConfig {
|
||||
Greedy,
|
||||
Temperature {
|
||||
temperature: f32,
|
||||
top_k: i32,
|
||||
top_p: f32,
|
||||
min_p: f32,
|
||||
seed: Option<u32>,
|
||||
},
|
||||
MirostatV2 {
|
||||
tau: f32,
|
||||
eta: f32,
|
||||
seed: Option<u32>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Default for SamplingConfig {
|
||||
fn default() -> Self {
|
||||
SamplingConfig::Temperature {
|
||||
temperature: 0.8,
|
||||
top_k: 40,
|
||||
top_p: 0.95,
|
||||
min_p: 0.05,
|
||||
seed: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ModelSettings {
|
||||
pub context_size: Option<u32>,
|
||||
pub max_output_tokens: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub sampling: SamplingConfig,
|
||||
#[serde(default = "default_repeat_penalty")]
|
||||
pub repeat_penalty: f32,
|
||||
#[serde(default = "default_repeat_last_n")]
|
||||
pub repeat_last_n: i32,
|
||||
#[serde(default)]
|
||||
pub frequency_penalty: f32,
|
||||
#[serde(default)]
|
||||
pub presence_penalty: f32,
|
||||
pub n_batch: Option<u32>,
|
||||
pub n_gpu_layers: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub use_mlock: bool,
|
||||
pub flash_attention: Option<bool>,
|
||||
pub n_threads: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub native_tool_calling: bool,
|
||||
#[serde(default)]
|
||||
pub use_jinja: bool,
|
||||
}
|
||||
|
||||
fn default_repeat_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repeat_last_n() -> i32 {
|
||||
64
|
||||
}
|
||||
|
||||
impl Default for ModelSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
context_size: None,
|
||||
max_output_tokens: None,
|
||||
sampling: SamplingConfig::default(),
|
||||
repeat_penalty: 1.0,
|
||||
repeat_last_n: 64,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
n_batch: None,
|
||||
n_gpu_layers: None,
|
||||
use_mlock: false,
|
||||
flash_attention: None,
|
||||
n_threads: None,
|
||||
native_tool_calling: false,
|
||||
use_jinja: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Featured models — HuggingFace specs in "author/repo-GGUF:quantization" format.
|
||||
pub const FEATURED_MODELS: &[&str] = &[
|
||||
"bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M",
|
||||
"bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M",
|
||||
"bartowski/Hermes-2-Pro-Mistral-7B-GGUF:Q4_K_M",
|
||||
"bartowski/Mistral-Small-24B-Instruct-2501-GGUF:Q4_K_M",
|
||||
];
|
||||
|
||||
/// Check if a model ID corresponds to a featured model.
|
||||
pub fn is_featured_model(model_id: &str) -> bool {
|
||||
use super::hf_models::parse_model_spec;
|
||||
FEATURED_MODELS.iter().any(|spec| {
|
||||
if let Ok((repo_id, quant)) = parse_model_spec(spec) {
|
||||
model_id_from_repo(&repo_id, &quant) == model_id
|
||||
} else {
|
||||
false
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
static REGISTRY: OnceLock<Mutex<LocalModelRegistry>> = OnceLock::new();
|
||||
|
||||
pub fn get_registry() -> &'static Mutex<LocalModelRegistry> {
|
||||
REGISTRY.get_or_init(|| {
|
||||
let registry = LocalModelRegistry::load().unwrap_or_default();
|
||||
Mutex::new(registry)
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LocalModelEntry {
|
||||
pub id: String,
|
||||
pub display_name: String,
|
||||
pub repo_id: String,
|
||||
pub filename: String,
|
||||
pub quantization: String,
|
||||
pub local_path: PathBuf,
|
||||
pub source_url: String,
|
||||
#[serde(default)]
|
||||
pub settings: ModelSettings,
|
||||
#[serde(default)]
|
||||
pub size_bytes: u64,
|
||||
}
|
||||
|
||||
impl LocalModelEntry {
|
||||
pub fn is_downloaded(&self) -> bool {
|
||||
self.local_path.exists()
|
||||
}
|
||||
|
||||
pub fn is_downloading(&self) -> bool {
|
||||
let download_id = format!("{}-model", self.id);
|
||||
let manager = get_download_manager();
|
||||
manager.get_progress(&download_id).is_some()
|
||||
}
|
||||
|
||||
pub fn download_status(&self) -> ModelDownloadStatus {
|
||||
if self.local_path.exists() {
|
||||
return ModelDownloadStatus::Downloaded;
|
||||
}
|
||||
|
||||
let download_id = format!("{}-model", self.id);
|
||||
let manager = get_download_manager();
|
||||
if let Some(progress) = manager.get_progress(&download_id) {
|
||||
return match progress.status {
|
||||
DownloadStatus::Downloading => ModelDownloadStatus::Downloading {
|
||||
progress_percent: progress.progress_percent,
|
||||
bytes_downloaded: progress.bytes_downloaded,
|
||||
total_bytes: progress.total_bytes,
|
||||
speed_bps: progress.speed_bps.unwrap_or(0),
|
||||
},
|
||||
DownloadStatus::Completed => ModelDownloadStatus::Downloaded,
|
||||
DownloadStatus::Failed | DownloadStatus::Cancelled => {
|
||||
ModelDownloadStatus::NotDownloaded
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
ModelDownloadStatus::NotDownloaded
|
||||
}
|
||||
|
||||
pub fn file_size(&self) -> u64 {
|
||||
if self.size_bytes > 0 {
|
||||
return self.size_bytes;
|
||||
}
|
||||
std::fs::metadata(&self.local_path)
|
||||
.map(|m| m.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum ModelDownloadStatus {
|
||||
NotDownloaded,
|
||||
Downloading {
|
||||
progress_percent: f32,
|
||||
bytes_downloaded: u64,
|
||||
total_bytes: u64,
|
||||
speed_bps: u64,
|
||||
},
|
||||
Downloaded,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct LocalModelRegistry {
|
||||
pub models: Vec<LocalModelEntry>,
|
||||
}
|
||||
|
||||
impl LocalModelRegistry {
|
||||
fn registry_path() -> PathBuf {
|
||||
Paths::in_data_dir("models/registry.json")
|
||||
}
|
||||
|
||||
pub fn load() -> Result<Self> {
|
||||
let path = Self::registry_path();
|
||||
if path.exists() {
|
||||
let lock_path = path.with_extension("json.lock");
|
||||
let lock_file = std::fs::File::create(&lock_path)?;
|
||||
fs2::FileExt::lock_shared(&lock_file)?;
|
||||
let contents = std::fs::read_to_string(&path)?;
|
||||
fs2::FileExt::unlock(&lock_file)?;
|
||||
let registry: LocalModelRegistry = serde_json::from_str(&contents)?;
|
||||
Ok(registry)
|
||||
} else {
|
||||
Ok(Self::default())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn save(&self) -> Result<()> {
|
||||
let path = Self::registry_path();
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
let lock_path = path.with_extension("json.lock");
|
||||
let lock_file = std::fs::File::create(&lock_path)?;
|
||||
fs2::FileExt::lock_exclusive(&lock_file)?;
|
||||
|
||||
let mut tmp = tempfile::NamedTempFile::new_in(path.parent().unwrap())?;
|
||||
let contents = serde_json::to_string_pretty(self)?;
|
||||
std::io::Write::write_all(&mut tmp, contents.as_bytes())?;
|
||||
tmp.persist(&path)?;
|
||||
|
||||
fs2::FileExt::unlock(&lock_file)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Sync registry with featured models:
|
||||
/// add any featured models that are missing, remove non-downloaded non-featured models.
|
||||
pub fn sync_with_featured(&mut self, featured_entries: Vec<LocalModelEntry>) {
|
||||
let mut changed = false;
|
||||
|
||||
for entry in featured_entries {
|
||||
if !self.models.iter().any(|m| m.id == entry.id) {
|
||||
self.models.push(entry);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
let before_len = self.models.len();
|
||||
self.models
|
||||
.retain(|m| m.is_downloaded() || m.is_downloading() || is_featured_model(&m.id));
|
||||
if self.models.len() != before_len {
|
||||
changed = true;
|
||||
}
|
||||
|
||||
if changed {
|
||||
let _ = self.save();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_model(&mut self, entry: LocalModelEntry) -> Result<()> {
|
||||
if let Some(existing) = self.models.iter_mut().find(|m| m.id == entry.id) {
|
||||
*existing = entry;
|
||||
} else {
|
||||
self.models.push(entry);
|
||||
}
|
||||
self.save()
|
||||
}
|
||||
|
||||
pub fn remove_model(&mut self, id: &str) -> Result<()> {
|
||||
self.models.retain(|m| m.id != id);
|
||||
self.save()
|
||||
}
|
||||
|
||||
pub fn get_model(&self, id: &str) -> Option<&LocalModelEntry> {
|
||||
self.models.iter().find(|m| m.id == id)
|
||||
}
|
||||
|
||||
pub fn has_model(&self, id: &str) -> bool {
|
||||
self.models.iter().any(|m| m.id == id)
|
||||
}
|
||||
|
||||
pub fn get_model_settings(&self, id: &str) -> Option<&ModelSettings> {
|
||||
self.models.iter().find(|m| m.id == id).map(|m| &m.settings)
|
||||
}
|
||||
|
||||
pub fn update_model_settings(&mut self, id: &str, settings: ModelSettings) -> Result<()> {
|
||||
let entry = self
|
||||
.models
|
||||
.iter_mut()
|
||||
.find(|m| m.id == id)
|
||||
.ok_or_else(|| anyhow::anyhow!("Model not found: {}", id))?;
|
||||
entry.settings = settings;
|
||||
self.save()
|
||||
}
|
||||
|
||||
pub fn list_models(&self) -> &[LocalModelEntry] {
|
||||
&self.models
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a unique ID for a model from its repo_id and quantization.
|
||||
pub fn model_id_from_repo(repo_id: &str, quantization: &str) -> String {
|
||||
format!("{}:{}", repo_id, quantization)
|
||||
}
|
||||
|
||||
/// Generate a display name from repo_id and quantization.
|
||||
pub fn display_name_from_repo(repo_id: &str, quantization: &str) -> String {
|
||||
let model_name = repo_id
|
||||
.split('/')
|
||||
.next_back()
|
||||
.unwrap_or(repo_id)
|
||||
.trim_end_matches("-GGUF")
|
||||
.trim_end_matches("-gguf");
|
||||
format!("{} ({})", model_name, quantization)
|
||||
}
|
||||
@@ -0,0 +1,601 @@
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use rmcp::model::{CallToolRequestParams, Tool};
|
||||
use serde_json::{json, Value};
|
||||
use std::borrow::Cow;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub(super) fn compact_tools_json(tools: &[Tool]) -> Option<String> {
|
||||
let compact: Vec<Value> = tools
|
||||
.iter()
|
||||
.map(|t| {
|
||||
json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.name,
|
||||
"description": t.description.as_ref().map(|d| d.as_ref()).unwrap_or(""),
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
serde_json::to_string(&compact).ok()
|
||||
}
|
||||
|
||||
/// Split generated text into (content, tool_calls_json).
|
||||
/// Looks for the last top-level JSON object containing `"tool_calls"`.
|
||||
/// Returns the text before it as content, and the JSON string if found.
|
||||
#[allow(clippy::string_slice)]
|
||||
pub(super) fn split_content_and_tool_calls(text: &str) -> (String, Option<String>) {
|
||||
let trimmed = text.trim_end();
|
||||
if !trimmed.ends_with('}') {
|
||||
return (text.to_string(), None);
|
||||
}
|
||||
|
||||
// Scan backwards for the matching '{' of the final '}'.
|
||||
// We only match on ASCII braces so `start` is always a char boundary.
|
||||
let bytes = trimmed.as_bytes();
|
||||
let mut depth = 0i32;
|
||||
let mut json_start = None;
|
||||
for i in (0..bytes.len()).rev() {
|
||||
match bytes[i] {
|
||||
b'}' => depth += 1,
|
||||
b'{' => {
|
||||
depth -= 1;
|
||||
if depth == 0 {
|
||||
json_start = Some(i);
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let Some(start) = json_start else {
|
||||
return (text.to_string(), None);
|
||||
};
|
||||
|
||||
let json_str = &trimmed[start..];
|
||||
let parsed: Value = match serde_json::from_str(json_str) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return (text.to_string(), None),
|
||||
};
|
||||
|
||||
if parsed
|
||||
.get("tool_calls")
|
||||
.and_then(|v| v.as_array())
|
||||
.is_none()
|
||||
{
|
||||
return (text.to_string(), None);
|
||||
}
|
||||
|
||||
let content = trimmed[..start].trim_end().to_string();
|
||||
(content, Some(json_str.to_string()))
|
||||
}
|
||||
|
||||
/// Return the byte length of text that is safe to stream.
|
||||
/// Everything before the last unmatched top-level `{` is safe — the `{` could
|
||||
/// be the start of a tool-call JSON block still being generated.
|
||||
/// If all braces are balanced the entire text is safe.
|
||||
pub(super) fn safe_stream_end(text: &str) -> usize {
|
||||
// Hold back from the start of any incomplete <tool_call> tag.
|
||||
// If we find an unmatched opening, nothing from that point should be streamed.
|
||||
let xml_hold = text.find("<tool_call>").unwrap_or(text.len());
|
||||
|
||||
let bytes = text.as_bytes();
|
||||
let mut safe_end = bytes.len();
|
||||
let mut depth = 0i32;
|
||||
for (i, &b) in bytes.iter().enumerate() {
|
||||
match b {
|
||||
b'{' => {
|
||||
if depth == 0 {
|
||||
safe_end = i;
|
||||
}
|
||||
depth += 1;
|
||||
}
|
||||
b'}' => {
|
||||
depth -= 1;
|
||||
if depth == 0 {
|
||||
safe_end = i + 1;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if depth == 0 {
|
||||
safe_end = i + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Also hold back a partial `<tool_call` prefix at the end of the text.
|
||||
// The tag is 11 chars; if the last N chars are a prefix of `<tool_call>`, hold them.
|
||||
let tag = b"<tool_call>";
|
||||
let tail_hold = {
|
||||
let mut hold = safe_end;
|
||||
let check_len = tag.len().min(bytes.len());
|
||||
for start in (safe_end.saturating_sub(check_len))..safe_end {
|
||||
let tail = &bytes[start..safe_end];
|
||||
if tag.starts_with(tail) {
|
||||
hold = start;
|
||||
break;
|
||||
}
|
||||
}
|
||||
hold
|
||||
};
|
||||
|
||||
safe_end.min(xml_hold).min(tail_hold)
|
||||
}
|
||||
|
||||
/// Extract tool call messages from a JSON object containing "tool_calls".
|
||||
/// Handles both the model's native format (name/arguments at top level)
|
||||
/// and the OpenAI format (function.name/function.arguments).
|
||||
pub(super) fn extract_tool_call_messages(tool_calls_json: &str, message_id: &str) -> Vec<Message> {
|
||||
let parsed: Value = match serde_json::from_str(tool_calls_json) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return vec![],
|
||||
};
|
||||
|
||||
let Some(tool_calls) = parsed.get("tool_calls").and_then(|v| v.as_array()) else {
|
||||
return vec![];
|
||||
};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for tc in tool_calls {
|
||||
// Try OpenAI format first: {"function": {"name": ..., "arguments": ...}, "id": ...}
|
||||
// Then model's native format: {"name": ..., "arguments": {...}, "id": ...}
|
||||
let (name, arguments) = if let Some(func) = tc.get("function") {
|
||||
let n = func.get("name").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let args_str = func
|
||||
.get("arguments")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("{}");
|
||||
let args: Option<serde_json::Map<String, Value>> = serde_json::from_str(args_str).ok();
|
||||
(n.to_string(), args)
|
||||
} else {
|
||||
let n = tc.get("name").and_then(|v| v.as_str()).unwrap_or("");
|
||||
// Arguments may be an object directly (model format) or a string (OAI format)
|
||||
let args = if let Some(obj) = tc.get("arguments").and_then(|v| v.as_object()) {
|
||||
Some(obj.clone())
|
||||
} else if let Some(s) = tc.get("arguments").and_then(|v| v.as_str()) {
|
||||
serde_json::from_str(s).ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
(n.to_string(), args)
|
||||
};
|
||||
|
||||
if name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let id = tc
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||||
|
||||
let tool_call = CallToolRequestParams {
|
||||
meta: None,
|
||||
task: None,
|
||||
name: Cow::Owned(name),
|
||||
arguments,
|
||||
};
|
||||
|
||||
let mut msg = Message::assistant();
|
||||
msg.content
|
||||
.push(MessageContent::tool_request(id, Ok(tool_call)));
|
||||
msg.id = Some(message_id.to_string());
|
||||
messages.push(msg);
|
||||
}
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
/// Parse XML-style tool calls used by models like qwen3-coder.
|
||||
/// Format:
|
||||
/// ```text
|
||||
/// <tool_call>
|
||||
/// <function=tool_name>
|
||||
/// <parameter=param1>value1</parameter>
|
||||
/// <parameter=param2>value2</parameter>
|
||||
/// </function>
|
||||
/// </tool_call>
|
||||
/// ```
|
||||
/// Returns (content_before_tool_calls, vec_of_tool_calls) or None if no XML tool calls found.
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub(super) fn split_content_and_xml_tool_calls(
|
||||
text: &str,
|
||||
) -> Option<(String, Vec<(String, serde_json::Map<String, Value>)>)> {
|
||||
let (content, first_block_and_rest) = text.split_once("<tool_call>")?;
|
||||
let content = content.trim_end().to_string();
|
||||
let mut tool_calls = Vec::new();
|
||||
|
||||
// Process the first block, then keep splitting on subsequent <tool_call> tags
|
||||
let mut remaining = first_block_and_rest;
|
||||
loop {
|
||||
// Split off the block up to </tool_call> (or take the rest if unclosed)
|
||||
let (block, after_close) = remaining
|
||||
.split_once("</tool_call>")
|
||||
.unwrap_or((remaining, ""));
|
||||
|
||||
if let Some(tool_call) = parse_single_xml_tool_call(block) {
|
||||
tool_calls.push(tool_call);
|
||||
}
|
||||
|
||||
// Try to find the next <tool_call> in what remains
|
||||
match after_close.split_once("<tool_call>") {
|
||||
Some((_between, next_remaining)) => remaining = next_remaining,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
if tool_calls.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some((content, tool_calls))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_single_xml_tool_call(block: &str) -> Option<(String, serde_json::Map<String, Value>)> {
|
||||
// Try <function=NAME><parameter=K>V</parameter>...</function> format first
|
||||
if let Some(result) = parse_xml_function_format(block) {
|
||||
return Some(result);
|
||||
}
|
||||
// Try GLM-style: TOOL_NAME<arg_key>K</arg_key><arg_value>V</arg_value>...
|
||||
parse_xml_arg_key_value_format(block)
|
||||
}
|
||||
|
||||
fn parse_xml_function_format(block: &str) -> Option<(String, serde_json::Map<String, Value>)> {
|
||||
let (_, after_func_eq) = block.split_once("<function=")?;
|
||||
let (func_name, func_body) = after_func_eq.split_once('>')?;
|
||||
let func_name = func_name.trim().to_string();
|
||||
|
||||
let mut args = serde_json::Map::new();
|
||||
let mut rest = func_body;
|
||||
|
||||
while let Some((_, after_param_eq)) = rest.split_once("<parameter=") {
|
||||
let Some((param_name, after_name_close)) = after_param_eq.split_once('>') else {
|
||||
break;
|
||||
};
|
||||
let param_name = param_name.trim().to_string();
|
||||
|
||||
let (value, after_value) = after_name_close
|
||||
.split_once("</parameter>")
|
||||
.unwrap_or((after_name_close, ""));
|
||||
let value = value.trim();
|
||||
|
||||
let json_value =
|
||||
serde_json::from_str(value).unwrap_or_else(|_| Value::String(value.to_string()));
|
||||
args.insert(param_name, json_value);
|
||||
|
||||
rest = after_value;
|
||||
}
|
||||
|
||||
Some((func_name, args))
|
||||
}
|
||||
|
||||
/// Parse GLM-style tool calls: `NAME<arg_key>K</arg_key><arg_value>V</arg_value>...`
|
||||
/// Also handles zero-argument calls like just `NAME`.
|
||||
fn parse_xml_arg_key_value_format(block: &str) -> Option<(String, serde_json::Map<String, Value>)> {
|
||||
let func_name_end = block.find("<arg_key>").unwrap_or(block.len());
|
||||
// Safe: find returns a byte offset at the start of an ASCII '<' character,
|
||||
// and block.len() is always a valid boundary.
|
||||
#[allow(clippy::string_slice)]
|
||||
let func_name = block[..func_name_end].trim().to_string();
|
||||
if func_name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut args = serde_json::Map::new();
|
||||
#[allow(clippy::string_slice)]
|
||||
let mut rest = &block[func_name_end..];
|
||||
|
||||
while let Some((_, after_key_open)) = rest.split_once("<arg_key>") {
|
||||
let Some((key, after_key_close)) = after_key_open.split_once("</arg_key>") else {
|
||||
break;
|
||||
};
|
||||
let key = key.trim().to_string();
|
||||
|
||||
let Some((_, after_val_open)) = after_key_close.split_once("<arg_value>") else {
|
||||
break;
|
||||
};
|
||||
let (value, after_val_close) = after_val_open
|
||||
.split_once("</arg_value>")
|
||||
.unwrap_or((after_val_open, ""));
|
||||
let value = value.trim();
|
||||
|
||||
let json_value =
|
||||
serde_json::from_str(value).unwrap_or_else(|_| Value::String(value.to_string()));
|
||||
args.insert(key, json_value);
|
||||
|
||||
rest = after_val_close;
|
||||
}
|
||||
|
||||
Some((func_name, args))
|
||||
}
|
||||
|
||||
pub(super) fn extract_xml_tool_call_messages(
|
||||
tool_calls: Vec<(String, serde_json::Map<String, Value>)>,
|
||||
message_id: &str,
|
||||
) -> Vec<Message> {
|
||||
tool_calls
|
||||
.into_iter()
|
||||
.map(|(name, args)| {
|
||||
let tool_call = CallToolRequestParams {
|
||||
meta: None,
|
||||
task: None,
|
||||
name: Cow::Owned(name),
|
||||
arguments: if args.is_empty() { None } else { Some(args) },
|
||||
};
|
||||
let mut msg = Message::assistant();
|
||||
msg.content.push(MessageContent::tool_request(
|
||||
Uuid::new_v4().to_string(),
|
||||
Ok(tool_call),
|
||||
));
|
||||
msg.id = Some(message_id.to_string());
|
||||
msg
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
const SHELL_TOOL: &str = "developer__shell";
|
||||
|
||||
#[test]
|
||||
fn test_parse_xml_tool_call_single() {
|
||||
let text = "I'll search for that.\n\n<tool_call>\n<function=search__files>\n<parameter=pattern>local.*inference</parameter>\n</function>\n</tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert_eq!(content, "I'll search for that.");
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, "search__files");
|
||||
assert_eq!(calls[0].1.get("pattern").unwrap(), "local.*inference");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_xml_tool_call_multiple_params() {
|
||||
let text = "<tool_call>\n<function=developer__shell>\n<parameter=command>ls -la</parameter>\n<parameter=timeout>30</parameter>\n</function>\n</tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert!(content.is_empty());
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, SHELL_TOOL);
|
||||
assert_eq!(calls[0].1.get("command").unwrap(), "ls -la");
|
||||
// 30 should be parsed as a number
|
||||
assert_eq!(calls[0].1.get("timeout").unwrap(), &json!(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_xml_tool_call_no_tool_call() {
|
||||
let text = "Just some regular text with no tool calls.";
|
||||
assert!(split_content_and_xml_tool_calls(text).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_xml_tool_call_multiple_calls() {
|
||||
let text = "Doing two things:\n<tool_call>\n<function=foo__bar>\n<parameter=x>1</parameter>\n</function>\n</tool_call>\n<tool_call>\n<function=baz__qux>\n<parameter=y>hello</parameter>\n</function>\n</tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert_eq!(content, "Doing two things:");
|
||||
assert_eq!(calls.len(), 2);
|
||||
assert_eq!(calls[0].0, "foo__bar");
|
||||
assert_eq!(calls[1].0, "baz__qux");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_xml_tool_call_multiline_value() {
|
||||
let text = "<tool_call>\n<function=developer__write_file>\n<parameter=path>test.py</parameter>\n<parameter=content>def hello():\n print(\"world\")</parameter>\n</function>\n</tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (_content, calls) = result.unwrap();
|
||||
assert_eq!(calls[0].0, "developer__write_file");
|
||||
assert_eq!(
|
||||
calls[0].1.get("content").unwrap(),
|
||||
"def hello():\n print(\"world\")"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_holds_back_tool_call_tag() {
|
||||
let text = "Some text before <tool_call>\n<function=foo>";
|
||||
let safe = safe_stream_end(text);
|
||||
assert!(safe <= text.find("<tool_call>").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_holds_back_partial_tag() {
|
||||
let text = "Some text <tool_ca";
|
||||
let safe = safe_stream_end(text);
|
||||
// Should hold back the partial tag
|
||||
assert!(safe <= text.find('<').unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_glm_style_tool_call() {
|
||||
let text = "<tool_call>developer__shell<arg_key>command</arg_key><arg_value>ls -la</arg_value></tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert!(content.is_empty());
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, SHELL_TOOL);
|
||||
assert_eq!(calls[0].1.get("command").unwrap(), "ls -la");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_glm_style_tool_call_no_args() {
|
||||
let text = "Some text\n<tool_call>load</tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert_eq!(content, "Some text");
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, "load");
|
||||
assert!(calls[0].1.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_glm_style_tool_call_multiple_args() {
|
||||
let text = "Let me check.\n<tool_call>execute<arg_key>code</arg_key><arg_value>async function run() { return 1; }</arg_value><arg_key>tool_graph</arg_key><arg_value>[{\"tool\": \"shell\"}]</arg_value></tool_call>";
|
||||
let result = split_content_and_xml_tool_calls(text);
|
||||
assert!(result.is_some());
|
||||
let (content, calls) = result.unwrap();
|
||||
assert_eq!(content, "Let me check.");
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, "execute");
|
||||
assert_eq!(
|
||||
calls[0].1.get("code").unwrap(),
|
||||
"async function run() { return 1; }"
|
||||
);
|
||||
// tool_graph should be parsed as JSON array
|
||||
assert!(calls[0].1.get("tool_graph").unwrap().is_array());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_tool_call_messages() {
|
||||
let calls = vec![(
|
||||
SHELL_TOOL.to_string(),
|
||||
serde_json::Map::from_iter(vec![("command".to_string(), json!("ls"))]),
|
||||
)];
|
||||
let msgs = extract_xml_tool_call_messages(calls, "test-id");
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].id, Some("test-id".to_string()));
|
||||
match &msgs[0].content[0] {
|
||||
MessageContent::ToolRequest(req) => {
|
||||
let call = req.tool_call.as_ref().unwrap();
|
||||
assert_eq!(&*call.name, SHELL_TOOL);
|
||||
assert_eq!(
|
||||
call.arguments.as_ref().unwrap().get("command").unwrap(),
|
||||
"ls"
|
||||
);
|
||||
}
|
||||
_ => panic!("Expected ToolRequest"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_content_and_tool_calls_with_tool() {
|
||||
let text = "Here is the result.\n{\"tool_calls\": [{\"function\": {\"name\": \"shell\", \"arguments\": \"{}\"}, \"id\": \"abc\"}]}";
|
||||
let (content, tc) = split_content_and_tool_calls(text);
|
||||
assert_eq!(content, "Here is the result.");
|
||||
assert!(tc.is_some());
|
||||
let parsed: serde_json::Value = serde_json::from_str(&tc.unwrap()).unwrap();
|
||||
assert_eq!(parsed["tool_calls"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_content_and_tool_calls_no_tool() {
|
||||
let text = "Just regular text, no JSON.";
|
||||
let (content, tc) = split_content_and_tool_calls(text);
|
||||
assert_eq!(content, text);
|
||||
assert!(tc.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_content_and_tool_calls_json_without_tool_calls_key() {
|
||||
let text = "{\"key\": \"value\"}";
|
||||
let (content, tc) = split_content_and_tool_calls(text);
|
||||
assert_eq!(content, text);
|
||||
assert!(tc.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_call_messages_openai_format() {
|
||||
let json_str = r#"{"tool_calls": [{"function": {"name": "developer__shell", "arguments": "{\"command\": \"ls\"}"}, "id": "call-1"}]}"#;
|
||||
let msgs = extract_tool_call_messages(json_str, "msg-1");
|
||||
assert_eq!(msgs.len(), 1);
|
||||
match &msgs[0].content[0] {
|
||||
MessageContent::ToolRequest(req) => {
|
||||
let call = req.tool_call.as_ref().unwrap();
|
||||
assert_eq!(&*call.name, SHELL_TOOL);
|
||||
assert_eq!(
|
||||
call.arguments.as_ref().unwrap().get("command").unwrap(),
|
||||
"ls"
|
||||
);
|
||||
}
|
||||
_ => panic!("Expected ToolRequest"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_call_messages_native_format() {
|
||||
let json_str = r#"{"tool_calls": [{"name": "developer__shell", "arguments": {"command": "ls"}, "id": "call-2"}]}"#;
|
||||
let msgs = extract_tool_call_messages(json_str, "msg-2");
|
||||
assert_eq!(msgs.len(), 1);
|
||||
match &msgs[0].content[0] {
|
||||
MessageContent::ToolRequest(req) => {
|
||||
let call = req.tool_call.as_ref().unwrap();
|
||||
assert_eq!(&*call.name, SHELL_TOOL);
|
||||
}
|
||||
_ => panic!("Expected ToolRequest"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_call_messages_invalid_json() {
|
||||
assert!(extract_tool_call_messages("not json", "msg-3").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_call_messages_empty_name_skipped() {
|
||||
let json_str = r#"{"tool_calls": [{"name": "", "arguments": {}, "id": "x"}]}"#;
|
||||
assert!(extract_tool_call_messages(json_str, "msg-4").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compact_tools_json_produces_minimal_output() {
|
||||
use rmcp::model::Tool;
|
||||
use rmcp::object;
|
||||
|
||||
let tools = vec![Tool::new(
|
||||
"developer__shell".to_string(),
|
||||
"Run shell commands".to_string(),
|
||||
object!({"type": "object", "properties": {"command": {"type": "string"}}}),
|
||||
)];
|
||||
let result = compact_tools_json(&tools);
|
||||
assert!(result.is_some());
|
||||
let parsed: Vec<serde_json::Value> = serde_json::from_str(&result.unwrap()).unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
let func = &parsed[0]["function"];
|
||||
assert_eq!(func["name"], "developer__shell");
|
||||
assert_eq!(func["description"], "Run shell commands");
|
||||
// Should not contain full parameter schemas
|
||||
assert!(func.get("parameters").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compact_tools_json_empty() {
|
||||
let result = compact_tools_json(&[]);
|
||||
assert!(result.is_some());
|
||||
let parsed: Vec<serde_json::Value> = serde_json::from_str(&result.unwrap()).unwrap();
|
||||
assert!(parsed.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_balanced_braces() {
|
||||
let text = "Result: {\"key\": \"value\"} done";
|
||||
assert_eq!(safe_stream_end(text), text.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_unbalanced_open_brace() {
|
||||
let text = "Some text {\"tool_calls\": [";
|
||||
assert_eq!(safe_stream_end(text), "Some text ".len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_empty() {
|
||||
assert_eq!(safe_stream_end(""), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_stream_end_no_braces() {
|
||||
let text = "plain text here";
|
||||
assert_eq!(safe_stream_end(text), text.len());
|
||||
}
|
||||
}
|
||||
@@ -23,6 +23,7 @@ pub mod google;
|
||||
mod init;
|
||||
pub mod lead_worker;
|
||||
pub mod litellm;
|
||||
pub mod local_inference;
|
||||
pub mod oauth;
|
||||
pub mod ollama;
|
||||
pub mod openai;
|
||||
|
||||
@@ -7,16 +7,13 @@ use super::utils::{ImageFormat, RequestLog};
|
||||
use crate::config::declarative_providers::DeclarativeProviderConfig;
|
||||
use crate::config::GooseMode;
|
||||
use crate::conversation::message::Message;
|
||||
use crate::conversation::Conversation;
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::formats::ollama::{create_request, response_to_streaming_message_ollama};
|
||||
use crate::utils::safe_truncate;
|
||||
use anyhow::{Error, Result};
|
||||
use async_stream::try_stream;
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use futures::TryStreamExt;
|
||||
use regex::Regex;
|
||||
use reqwest::Response;
|
||||
use rmcp::model::Tool;
|
||||
use serde_json::{json, Value};
|
||||
@@ -215,30 +212,6 @@ impl Provider for OllamaProvider {
|
||||
self.model.clone()
|
||||
}
|
||||
|
||||
async fn generate_session_name(
|
||||
&self,
|
||||
session_id: &str,
|
||||
messages: &Conversation,
|
||||
) -> Result<String, ProviderError> {
|
||||
let context = self.get_initial_user_messages(messages);
|
||||
let message = Message::user().with_text(self.create_session_name_prompt(&context));
|
||||
let model_config = self.get_model_config();
|
||||
let result = self
|
||||
.complete(
|
||||
&model_config,
|
||||
session_id,
|
||||
"You are a title generator. Output only the requested title of 4 words or less, with no additional text, reasoning, or explanations.",
|
||||
&[message],
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut description = result.0.as_concat_text();
|
||||
description = Self::filter_reasoning_tokens(&description);
|
||||
|
||||
Ok(safe_truncate(&description, 100))
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
model_config: &ModelConfig,
|
||||
@@ -318,46 +291,6 @@ impl Provider for OllamaProvider {
|
||||
}
|
||||
}
|
||||
|
||||
impl OllamaProvider {
|
||||
fn filter_reasoning_tokens(text: &str) -> String {
|
||||
let mut filtered = text.to_string();
|
||||
|
||||
let reasoning_patterns = [
|
||||
r"<think>.*?</think>",
|
||||
r"<thinking>.*?</thinking>",
|
||||
r"Let me think.*?\n",
|
||||
r"I need to.*?\n",
|
||||
r"First, I.*?\n",
|
||||
r"Okay, .*?\n",
|
||||
r"So, .*?\n",
|
||||
r"Well, .*?\n",
|
||||
r"Hmm, .*?\n",
|
||||
r"Actually, .*?\n",
|
||||
r"Based on.*?I think",
|
||||
r"Looking at.*?I would say",
|
||||
];
|
||||
|
||||
for pattern in reasoning_patterns {
|
||||
if let Ok(re) = Regex::new(pattern) {
|
||||
filtered = re.replace_all(&filtered, "").to_string();
|
||||
}
|
||||
}
|
||||
filtered = filtered
|
||||
.replace("<think>", "")
|
||||
.replace("</think>", "")
|
||||
.replace("<thinking>", "")
|
||||
.replace("</thinking>", "");
|
||||
filtered = filtered
|
||||
.lines()
|
||||
.map(|line| line.trim())
|
||||
.filter(|line| !line.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
|
||||
filtered
|
||||
}
|
||||
}
|
||||
|
||||
/// Ollama-specific streaming handler with XML tool call fallback.
|
||||
/// Uses the Ollama format module which buffers text when XML tool calls are detected,
|
||||
/// preventing duplicate content from being emitted to the UI.
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
//! Integration tests for LocalInferenceProvider.
|
||||
//!
|
||||
//! These tests require a downloaded GGUF model and are ignored by default.
|
||||
//! Run with: cargo test -p goose --test local_inference_integration -- --ignored
|
||||
|
||||
use futures::StreamExt;
|
||||
use goose::conversation::message::Message;
|
||||
use goose::model::ModelConfig;
|
||||
use goose::providers::create;
|
||||
use std::time::Instant;
|
||||
|
||||
const TEST_MODEL: &str = "llama-3.2-1b";
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_local_inference_stream_produces_output() {
|
||||
let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config");
|
||||
let provider = create("local", model_config.clone(), Vec::new())
|
||||
.await
|
||||
.expect("provider creation should succeed");
|
||||
|
||||
let system = "You are a helpful assistant. Be brief.";
|
||||
let messages = vec![Message::user().with_text("Say hello.")];
|
||||
|
||||
let mut stream = provider
|
||||
.stream(&model_config, "test-session", system, &messages, &[])
|
||||
.await
|
||||
.expect("stream should start");
|
||||
|
||||
let mut got_text = false;
|
||||
let mut got_usage = false;
|
||||
|
||||
while let Some(result) = stream.next().await {
|
||||
let (msg, usage) = result.expect("stream item should be Ok");
|
||||
if msg.is_some() {
|
||||
got_text = true;
|
||||
}
|
||||
if let Some(u) = usage {
|
||||
got_usage = true;
|
||||
let usage_inner = u.usage;
|
||||
assert!(
|
||||
usage_inner.input_tokens.unwrap_or(0) > 0,
|
||||
"should have input tokens"
|
||||
);
|
||||
assert!(
|
||||
usage_inner.output_tokens.unwrap_or(0) > 0,
|
||||
"should have output tokens"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_text, "stream should produce at least one text message");
|
||||
assert!(got_usage, "stream should produce usage info");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_local_inference_cold_and_warm_performance() {
|
||||
let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config");
|
||||
let provider = create("local", model_config.clone(), Vec::new())
|
||||
.await
|
||||
.expect("provider creation should succeed");
|
||||
|
||||
// Cold start (includes model loading)
|
||||
let messages = vec![Message::user().with_text("what is the capital of Moldova?")];
|
||||
let start = Instant::now();
|
||||
let (response, _usage) = provider
|
||||
.complete(&model_config, "test-session", "", &messages, &[])
|
||||
.await
|
||||
.expect("cold completion should succeed");
|
||||
let cold_elapsed = start.elapsed();
|
||||
|
||||
let text = response.as_concat_text();
|
||||
assert!(!text.is_empty(), "cold start should produce a response");
|
||||
println!(
|
||||
"Cold start: {cold_elapsed:.2?}, response length: {}",
|
||||
text.len()
|
||||
);
|
||||
|
||||
// Warm run (model already loaded)
|
||||
let messages2 = vec![Message::user().with_text("what is the capital of France?")];
|
||||
let start2 = Instant::now();
|
||||
let (response2, _usage2) = provider
|
||||
.complete(&model_config, "test-session", "", &messages2, &[])
|
||||
.await
|
||||
.expect("warm completion should succeed");
|
||||
let warm_elapsed = start2.elapsed();
|
||||
|
||||
let text2 = response2.as_concat_text();
|
||||
assert!(!text2.is_empty(), "warm run should produce a response");
|
||||
println!(
|
||||
"Warm run: {warm_elapsed:.2?}, response length: {}",
|
||||
text2.len()
|
||||
);
|
||||
assert!(
|
||||
warm_elapsed < cold_elapsed,
|
||||
"warm run ({warm_elapsed:.2?}) should be faster than cold start ({cold_elapsed:.2?})"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_local_inference_large_prompt() {
|
||||
let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config");
|
||||
let provider = create("local", model_config.clone(), Vec::new())
|
||||
.await
|
||||
.expect("provider creation should succeed");
|
||||
|
||||
// Build a large prompt (~3500 tokens) to exercise prefill performance
|
||||
let padding = "You are Goose, a highly capable AI assistant.\n".repeat(80);
|
||||
let prompt = format!("{padding}\nNow answer this: what is the capital of Moldova?");
|
||||
let messages = vec![Message::user().with_text(&prompt)];
|
||||
|
||||
let start = Instant::now();
|
||||
let (response, _usage) = provider
|
||||
.complete(&model_config, "test-session", "", &messages, &[])
|
||||
.await
|
||||
.expect("large prompt completion should succeed");
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
let text = response.as_concat_text();
|
||||
assert!(!text.is_empty(), "large prompt should produce a response");
|
||||
println!(
|
||||
"Large prompt: {elapsed:.2?}, prompt ~{} chars, response length: {}",
|
||||
prompt.len(),
|
||||
text.len()
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user