Classify terminal input with a shared local language model
This commit is contained in:
@@ -219,6 +219,7 @@ pub struct BlocklistAIInputModel {
|
||||
terminal_surface_id: EntityId,
|
||||
|
||||
autodetect_abort_handle: Option<AbortHandle>,
|
||||
autodetection_generation: u64,
|
||||
model: Arc<FairMutex<TerminalModel>>,
|
||||
}
|
||||
|
||||
@@ -405,6 +406,7 @@ impl BlocklistAIInputModel {
|
||||
last_explicit_input_type_set_at: None,
|
||||
was_lock_set_with_empty_buffer: false,
|
||||
autodetect_abort_handle: None,
|
||||
autodetection_generation: 0,
|
||||
model,
|
||||
}
|
||||
}
|
||||
@@ -683,6 +685,7 @@ impl BlocklistAIInputModel {
|
||||
|
||||
/// Aborts any in progress work for autodetection.
|
||||
pub fn abort_in_progress_detection(&mut self) {
|
||||
self.autodetection_generation = self.autodetection_generation.wrapping_add(1);
|
||||
if let Some(handle) = self.autodetect_abort_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
@@ -786,9 +789,35 @@ impl BlocklistAIInputModel {
|
||||
};
|
||||
|
||||
let classifier = InputClassifierModel::as_ref(ctx).classifier();
|
||||
let generation = self.autodetection_generation;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
let use_local_model = AISettings::as_ref(ctx).should_use_local_model();
|
||||
let handle = ctx
|
||||
.spawn(
|
||||
async move {
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
if use_local_model {
|
||||
// A cancelled debounce never starts inference. The local worker
|
||||
// also cancels obsolete requests already queued or generating.
|
||||
warpui::r#async::Timer::after(std::time::Duration::from_millis(300)).await;
|
||||
let result = crate::local_inference::classify(
|
||||
buffer_cloned.to_string(),
|
||||
history_entries.as_deref().unwrap_or_default(),
|
||||
is_agent_follow_up,
|
||||
)
|
||||
.await;
|
||||
return match result {
|
||||
Some(input_type) => (
|
||||
input_type,
|
||||
InputClassifierDecisionSource::LocalLanguageModel.into(),
|
||||
),
|
||||
None => (
|
||||
current_input_type,
|
||||
InputClassifierDecisionSource::InputClassifierFallbackCurrentInput
|
||||
.into(),
|
||||
),
|
||||
};
|
||||
}
|
||||
// First check if the token is a natural language word, if current input type is AI
|
||||
if matches!(current_input_type, InputType::AI)
|
||||
&& is_one_off_natural_language_word(first_token_str.to_lowercase().as_str())
|
||||
@@ -852,7 +881,9 @@ impl BlocklistAIInputModel {
|
||||
// If the autodetect abort handle is none, it means we aborted autodetection.
|
||||
// It's possible that the future already completed before we aborted, and then we reach this callback after abort.
|
||||
// In this case, don't set the input type.
|
||||
if me.autodetect_abort_handle.is_none() {
|
||||
if me.autodetect_abort_handle.is_none()
|
||||
|| me.autodetection_generation != generation
|
||||
{
|
||||
return;
|
||||
}
|
||||
me.set_input_config_internal(
|
||||
|
||||
@@ -70,6 +70,8 @@ mod interval_timer;
|
||||
mod linear;
|
||||
#[cfg(feature = "local_fs")]
|
||||
mod local_control;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
mod local_inference;
|
||||
mod local_object_repository;
|
||||
#[cfg(any(target_os = "macos", target_os = "windows"))]
|
||||
mod login_item;
|
||||
|
||||
+148
-3
@@ -1,3 +1,148 @@
|
||||
// Local inference engine module — currently unused.
|
||||
// The local_inference crate provides the engine and task definitions;
|
||||
// this module is reserved for future app-layer integration.
|
||||
//! One lazy local classifier worker shared by all terminal panes.
|
||||
use std::sync::{mpsc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::channel::oneshot;
|
||||
use input_classifier::InputType;
|
||||
use instant::Instant;
|
||||
use local_inference::{
|
||||
CancellationToken, Device, InferenceEngine, InputCategory, InputClassificationInput,
|
||||
InputClassificationTask,
|
||||
};
|
||||
|
||||
struct Request {
|
||||
input: InputClassificationInput,
|
||||
cancel: CancellationToken,
|
||||
reply: oneshot::Sender<Option<InputType>>,
|
||||
}
|
||||
|
||||
struct CancelOnDrop(CancellationToken);
|
||||
|
||||
impl Drop for CancelOnDrop {
|
||||
fn drop(&mut self) {
|
||||
self.0.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
static WORKER: OnceLock<Option<mpsc::SyncSender<Request>>> = OnceLock::new();
|
||||
|
||||
fn start_worker() -> Option<mpsc::SyncSender<Request>> {
|
||||
// Bound queued work across panes. Superseded requests are skipped by the worker.
|
||||
let (sender, receiver) = mpsc::sync_channel::<Request>(16);
|
||||
let result = std::thread::Builder::new()
|
||||
.name("local-input-classifier".to_string())
|
||||
.spawn(move || {
|
||||
let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
|
||||
Ok(runtime) => runtime,
|
||||
Err(error) => {
|
||||
log::warn!("Could not initialize local classifier runtime: {error}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let mut engine = None;
|
||||
let mut retry_after = None;
|
||||
let mut cached: Option<(InputClassificationInput, InputType)> = None;
|
||||
for request in receiver {
|
||||
if request.cancel.is_cancelled() || request.reply.is_canceled() {
|
||||
continue;
|
||||
}
|
||||
if let Some((input, result)) = &cached {
|
||||
if input == &request.input {
|
||||
log::info!("Local input classification: category={result:?}, cached=true");
|
||||
let _ = request.reply.send(Some(*result));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
if engine.is_none() {
|
||||
if retry_after.is_some_and(|deadline| Instant::now() < deadline) {
|
||||
let _ = request.reply.send(None);
|
||||
continue;
|
||||
}
|
||||
match runtime.block_on(InferenceEngine::new(Device::best_available())) {
|
||||
Ok(loaded) => engine = Some(loaded),
|
||||
Err(error) => {
|
||||
log::warn!("Local input classifier unavailable; preserving input mode: {error:#}");
|
||||
retry_after = Some(Instant::now() + Duration::from_secs(60));
|
||||
let _ = request.reply.send(None);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
if request.cancel.is_cancelled() || request.reply.is_canceled() {
|
||||
continue;
|
||||
}
|
||||
let started = Instant::now();
|
||||
let result = runtime.block_on(InputClassificationTask.run_cancellable(
|
||||
engine.as_ref().expect("engine was initialized"),
|
||||
request.input.clone(),
|
||||
&request.cancel,
|
||||
));
|
||||
let classification = match result {
|
||||
Ok(result) => {
|
||||
let input_type = match result.category {
|
||||
InputCategory::Shell => InputType::Shell,
|
||||
InputCategory::AgentPrompt => InputType::AI,
|
||||
};
|
||||
log::info!("Local input classification: category={input_type:?}, label_probability={:.3}, elapsed_ms={}", result.confidence, started.elapsed().as_millis());
|
||||
cached = Some((request.input, input_type));
|
||||
Some(input_type)
|
||||
}
|
||||
Err(error) => {
|
||||
if !request.cancel.is_cancelled() {
|
||||
log::warn!("Local classifier returned no decision; preserving input mode: {error:#}");
|
||||
}
|
||||
None
|
||||
}
|
||||
};
|
||||
let _ = request.reply.send(classification);
|
||||
}
|
||||
});
|
||||
match result {
|
||||
Ok(_) => Some(sender),
|
||||
Err(error) => {
|
||||
log::warn!("Could not start local classifier: {error}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn classify(
|
||||
user_input: String,
|
||||
recent_commands: &[String],
|
||||
is_follow_up: bool,
|
||||
) -> Option<InputType> {
|
||||
// Do not classify a truncated script or document as though it were the full input.
|
||||
if user_input.trim().is_empty() || user_input.len() > 2048 {
|
||||
return None;
|
||||
}
|
||||
let sender = WORKER.get_or_init(start_worker).as_ref()?;
|
||||
let cancel = CancellationToken::new();
|
||||
let _cancel_on_drop = CancelOnDrop(cancel.clone());
|
||||
let (reply, receiver) = oneshot::channel();
|
||||
sender
|
||||
.try_send(Request {
|
||||
input: InputClassificationInput {
|
||||
user_input,
|
||||
recent_commands: recent_commands
|
||||
.iter()
|
||||
.rev()
|
||||
.take(3)
|
||||
.map(|command| command.chars().take(120).collect())
|
||||
.collect(),
|
||||
is_follow_up,
|
||||
},
|
||||
cancel,
|
||||
reply,
|
||||
})
|
||||
.ok()?;
|
||||
let timeout = galaxyui::r#async::Timer::after(Duration::from_secs(2));
|
||||
match futures::future::select(receiver, Box::pin(timeout)).await {
|
||||
futures::future::Either::Left((result, _)) => result.ok().flatten(),
|
||||
futures::future::Either::Right(_) => {
|
||||
log::info!(
|
||||
"Local input classification timed out after 2 seconds; preserving input mode"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1235,9 +1235,8 @@ define_settings_group!(AISettings, settings: [
|
||||
toml_path: "agents.warp_agent.input.ai_command_denylist",
|
||||
description: "Commands to exclude from AI natural language autodetection.",
|
||||
}
|
||||
// Whether to use the local AI model for input classification and suggestions
|
||||
// instead of making Bedrock API calls. When enabled, SmolLM2-135M handles
|
||||
// autodetection, prompt suggestions, and code banners locally.
|
||||
// Use SmolLM2-135M for input classification instead of legacy heuristics.
|
||||
// The model is downloaded once and inference stays on the device.
|
||||
use_local_model: UseLocalModel {
|
||||
type: bool,
|
||||
default: true,
|
||||
@@ -1245,7 +1244,7 @@ define_settings_group!(AISettings, settings: [
|
||||
sync_to_cloud: SyncToCloud::Never,
|
||||
private: false,
|
||||
toml_path: "agents.warp_agent.input.use_local_model",
|
||||
description: "Use a locally downloaded AI model for command recognition and suggestions instead of Bedrock.",
|
||||
description: "Use a locally downloaded AI model to distinguish commands from agent messages instead of legacy heuristics.",
|
||||
},
|
||||
// This field should not be referenced directly to lookup intelligent autosuggestion enablement
|
||||
// -- use the `is_intelligent_autosuggestions_enabled()` getter.
|
||||
|
||||
Reference in New Issue
Block a user