public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Initial public release of Warp.

Browse Source 
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
David Stern
2026-04-28 08:43:33 -05:00
commit 0dbd3d567a
4982 changed files with 1431549 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/input_classifier/Cargo.toml
+56
View File
@@ -0,0 +1,56 @@
[package]
name = "input_classifier"
authors = ["Warp Team <dev@warp.dev>"]
edition = "2024"
publish.workspace = true
license.workspace = true
[[bin]]
name = "evaluate"
path = "src/bin/evaluate.rs"
[features]
onnx_candle = [
"onnx",
"dep:candle-core",
"dep:candle-nn",
"dep:candle-onnx",
"dep:prost",
]
onnx_ort = [
"onnx",
"dep:ort",
"dep:ndarray",
"dep:parking_lot",
]
fasttext = ["dep:fasttext", "dep:rust-embed", "dep:tempfile"]
onnx = ["dep:rust-embed", "dep:tokenizers"]
[dependencies]
anyhow.workspace = true
async-trait.workspace = true
clap.workspace = true
futures.workspace = true
candle-core = { version = "0.9.2", optional = true }
candle-nn = { version = "0.9.2", optional = true }
candle-onnx = { version = "0.9.2", optional = true }
cfg-if.workspace = true
itertools.workspace = true
fasttext = { version = "0.7.8", optional = true }
futures-lite.workspace = true
instant.workspace = true
lazy_static.workspace = true
log.workspace = true
natural_language_detection.workspace = true
ndarray = { version = "0.16.1", optional = true }
ort = { version = "2.0.0-rc.10", optional = true }
parking_lot = { workspace = true, optional = true }
prost = { version = "0.14.3", optional = true }
rust-embed = { workspace = true, optional = true }
serde.workspace = true
smol_str.workspace = true
tempfile = { workspace = true, optional = true }
tokenizers = { version = "0.21.4", optional = true }
warp_completer.workspace = true
[dev-dependencies]
crates/input_classifier/models/fasttext/cmd_lang_classifier_v4.bin
BIN
View File
Binary file not shown.
crates/input_classifier/models/onnx/bert_tiny.onnx
BIN
View File
Binary file not shown.
crates/input_classifier/models/onnx/bert_tiny_tokenizer.json
+30686
View File
File diff suppressed because it is too large Load Diff
crates/input_classifier/src/bin/evaluate.rs
+410
View File
@@ -0,0 +1,410 @@
use std::{fs, process};
use clap::{Parser, Subcommand};
use futures::executor::block_on;
use input_classifier::{
ClassificationResult, Context, HeuristicClassifier, InputClassifier, InputType,
test_utils::CompletionContext,
};
/// Convert HSL to RGB values (0-255 range)
fn hsl_to_rgb(h: f32, s: f32, l: f32) -> (u8, u8, u8) {
let c = (1.0 - (2.0 * l - 1.0).abs()) * s;
let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs());
let m = l - c / 2.0;
let (r_prime, g_prime, b_prime) = if h < 60.0 {
(c, x, 0.0)
} else if h < 120.0 {
(x, c, 0.0)
} else if h < 180.0 {
(0.0, c, x)
} else if h < 240.0 {
(0.0, x, c)
} else if h < 300.0 {
(x, 0.0, c)
} else {
(c, 0.0, x)
};
let r = ((r_prime + m) * 255.0) as u8;
let g = ((g_prime + m) * 255.0) as u8;
let b = ((b_prime + m) * 255.0) as u8;
(r, g, b)
}
/// Generate ANSI color code for smooth mode using HSL saturation scaling
fn get_smooth_confidence_color(is_correct: bool, confidence: f32) -> String {
// Map confidence (0.5 to 1.0) to saturation (0.0 to 1.0)
// Confidence below 0.5 gets 0 saturation (gray), above 0.5 scales linearly
let saturation = if confidence <= 0.5 {
0.0
} else {
(confidence - 0.5) * 2.0
};
// Use different hues for correct vs incorrect
let hue = if is_correct { 120.0 } else { 0.0 }; // Green for correct, Red for incorrect
let lightness = 0.5; // Medium lightness
let (r, g, b) = hsl_to_rgb(hue, saturation, lightness);
format!("\x1b[38;2;{r};{g};{b}m")
}
/// Generate ANSI color code for binary mode using simple green/red with dim for low confidence
fn get_binary_confidence_color(is_correct: bool, is_low_confidence: bool) -> String {
if is_correct {
if is_low_confidence {
"\x1b[32m\x1b[2m".to_string() // Green + dim for correct but low confidence
} else {
"\x1b[32m".to_string() // Green for correct
}
} else if is_low_confidence {
"\x1b[31m\x1b[2m".to_string() // Red + dim for incorrect but low confidence
} else {
"\x1b[31m".to_string() // Red for incorrect
}
}
use warp_completer::{ParsedTokensSnapshot, util::parse_current_commands_and_tokens};
#[cfg(feature = "fasttext")]
use input_classifier::FasttextClassifier;
#[cfg(feature = "onnx")]
use input_classifier::{OnnxClassifier, OnnxModel};
#[derive(Parser)]
struct InputSource {
/// Input string to classify (use --file to read from file instead)
#[arg(group = "input_source")]
input: Option<String>,
/// Read input from file instead of command line argument
#[arg(long, group = "input_source")]
file: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq)]
enum ConfidenceMode {
Binary(f32), // Binary mode with confidence threshold
Smooth, // Smooth saturation scaling
}
#[derive(Parser)]
#[command(name = "evaluate")]
#[command(about = "Test input classifier implementations")]
struct Args {
/// Use heuristic classifier
#[arg(long)]
heuristic: bool,
/// Use fasttext classifier
#[cfg(feature = "fasttext")]
#[arg(long)]
fasttext: bool,
/// Use ONNX classifier
#[cfg(feature = "onnx")]
#[arg(long)]
onnx: bool,
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
enum Command {
/// Classify a single input string
Classify {
#[command(flatten)]
input_source: InputSource,
},
/// Verify classification by testing all prefixes of input string
Verify {
expected: String,
#[command(flatten)]
input_source: InputSource,
/// Confidence threshold for binary mode visualization. If specified, uses binary coloring with the given threshold (e.g., 0.9). If not specified, uses smooth saturation scaling.
#[arg(long)]
confident: Option<f32>,
},
}
/// Create classifiers based on CLI flags
fn create_classifiers(args: &Args) -> Vec<(&'static str, Box<dyn InputClassifier>)> {
let mut classifiers: Vec<(&'static str, Box<dyn InputClassifier>)> = Vec::new();
// Default to all available classifiers if none specified
let fasttext_specified = {
#[cfg(feature = "fasttext")]
{
args.fasttext
}
#[cfg(not(feature = "fasttext"))]
{
false
}
};
let onnx_specified = {
#[cfg(feature = "onnx")]
{
args.onnx
}
#[cfg(not(feature = "onnx"))]
{
false
}
};
let use_all = !args.heuristic && !fasttext_specified && !onnx_specified;
if args.heuristic || use_all {
classifiers.push(("heuristic", Box::new(HeuristicClassifier)));
}
#[cfg(feature = "fasttext")]
if args.fasttext || use_all {
match FasttextClassifier::new() {
Ok(classifier) => {
classifiers.push(("fasttext", Box::new(classifier)));
}
Err(e) => {
eprintln!("Warning: Failed to initialize FastText classifier: {e}");
}
}
}
#[cfg(feature = "onnx")]
if args.onnx || use_all {
match OnnxClassifier::new(OnnxModel::BertTiny) {
Ok(classifier) => {
classifiers.push(("onnx", Box::new(classifier)));
}
Err(e) => {
eprintln!("Warning: Failed to initialize ONNX classifier: {e}");
}
}
}
classifiers
}
/// Resolve input from either direct string or file
fn resolve_input_source(input: Option<String>, file: Option<String>) -> anyhow::Result<String> {
match (input, file) {
(Some(input_str), None) => Ok(input_str),
(None, Some(file_path)) => {
let content = fs::read_to_string(file_path.clone())
.map_err(|e| anyhow::anyhow!("Failed to read file '{}': {}", file_path, e))?;
// Trim trailing newline if present
Ok(content.trim_end().to_string())
}
(Some(_), Some(_)) => Err(anyhow::anyhow!("Cannot specify both input string and file")),
(None, None) => Err(anyhow::anyhow!(
"Must specify either input string or --file"
)),
}
}
/// Parse input string into ParsedTokensSnapshot
async fn parse_input(input: &str) -> anyhow::Result<ParsedTokensSnapshot> {
let completion_context = CompletionContext::new();
let snapshot = parse_current_commands_and_tokens(input.to_string(), &completion_context).await;
Ok(snapshot)
}
/// Handle classify command
async fn handle_classify(
input: &str,
classifiers: &[(&str, Box<dyn InputClassifier>)],
) -> anyhow::Result<()> {
let parsed_input = parse_input(input).await?;
let context = Context {
current_input_type: InputType::Shell,
is_agent_follow_up: false,
};
println!("Input: \"{input}\"");
println!("Classifications:");
for (name, classifier) in classifiers {
match classifier
.classify_input(parsed_input.clone(), &context)
.await
{
Ok(result) => {
let predicted_type = result.to_input_type();
println!(
" {}: {} (p_shell: {:.3}, p_ai: {:.3}, confidence: {:.3}, opacity: {:.3})",
name,
predicted_type,
result.p_shell(),
result.p_ai(),
result.confidence(),
opacity(&result)
);
}
Err(_) => {
// Fallback to detect_input_type if classify_input fails
let result = classifier
.detect_input_type(parsed_input.clone(), &context)
.await;
println!(" {name}: {result} (probabilities unavailable)");
}
}
}
Ok(())
}
/// Handle verify command (tests all prefixes)
async fn handle_verify(
input: &str,
expected: InputType,
classifiers: &[(&str, Box<dyn InputClassifier>)],
confidence_mode: ConfidenceMode,
) -> anyhow::Result<()> {
println!("Input: \"{input}\"");
println!("Expected: {expected}");
println!("Verification Results:");
// Generate all prefixes (1 character to full string)
let prefixes: Vec<String> = (1..=input.len()).map(|i| input[..i].to_string()).collect();
println!("Testing {} prefixes...", prefixes.len());
for (name, classifier) in classifiers {
let mut correct_count = 0;
let total_count = prefixes.len();
let mut classification_results: Vec<(bool, Option<ClassificationResult>)> = Vec::new();
for prefix in &prefixes {
let parsed_input = parse_input(prefix).await?;
let context = Context {
current_input_type: InputType::Shell,
is_agent_follow_up: false,
};
// Use classify_input to get probabilities
let classification_result = classifier
.classify_input(parsed_input.clone(), &context)
.await;
let (is_correct, result_opt) = match classification_result {
Ok(result) => {
let predicted_type = result.to_input_type();
let is_correct = predicted_type == expected;
if is_correct {
correct_count += 1;
}
(is_correct, Some(result))
}
Err(_) => {
// Fallback to detect_input_type if classify_input fails
let result = classifier
.detect_input_type(parsed_input.clone(), &context)
.await;
let is_correct = result == expected;
if is_correct {
correct_count += 1;
}
(is_correct, None)
}
};
classification_results.push((is_correct, result_opt));
}
let percentage = (correct_count as f64 / total_count as f64) * 100.0;
println!(" {name}: {correct_count}/{total_count} correct ({percentage:.1}%)");
print!(" Visual: ");
for (i, ch) in input.chars().enumerate() {
let (is_correct, classification_result) = &classification_results[i];
// Determine styling based on correctness and confidence
if let Some(result) = classification_result {
let color_code = match confidence_mode {
ConfidenceMode::Smooth => {
let confidence = result.confidence();
get_smooth_confidence_color(*is_correct, confidence)
}
ConfidenceMode::Binary(threshold) => {
let confidence = result.confidence();
let is_low_confidence = confidence < threshold;
get_binary_confidence_color(*is_correct, is_low_confidence)
}
};
print!("{color_code}{ch}\x1b[0m");
} else {
// Fallback to basic colors when no probability data is available
if *is_correct {
print!("\x1b[32m{ch}\x1b[0m"); // Green for correct
} else {
print!("\x1b[31m{ch}\x1b[0m"); // Red for incorrect
}
}
}
println!(); // New line after the colored string
}
Ok(())
}
fn main() -> anyhow::Result<()> {
let args = Args::parse();
let classifiers = create_classifiers(&args);
if classifiers.is_empty() {
eprintln!("Error: No classifiers available or selected");
process::exit(1);
}
let result: anyhow::Result<()> = block_on(async {
match &args.command {
Command::Classify { input_source } => {
let input_str =
resolve_input_source(input_source.input.clone(), input_source.file.clone())?;
handle_classify(&input_str, &classifiers).await
}
Command::Verify {
expected,
input_source,
confident,
} => {
let input_str =
resolve_input_source(input_source.input.clone(), input_source.file.clone())?;
let expected_type = expected
.parse::<InputType>()
.map_err(|e| anyhow::anyhow!("{}", e))?;
let confidence_mode = match confident {
Some(threshold) => {
if *threshold < 0.0 || *threshold > 1.0 {
return Err(anyhow::anyhow!(
"Confidence threshold must be between 0.0 and 1.0, got: {}",
threshold
));
}
ConfidenceMode::Binary(*threshold)
}
None => ConfidenceMode::Smooth,
};
handle_verify(&input_str, expected_type, &classifiers, confidence_mode).await
}
}
});
if let Err(e) = result {
eprintln!("Error: {e}");
process::exit(1);
}
Ok(())
}
/// Returns opacity value (0.0 to 1.0) scaled from confidence, where 0.5 confidence = 0.0 opacity and 1.0 confidence = 1.0 opacity
fn opacity(result: &ClassificationResult) -> f32 {
let conf = result.confidence();
if conf <= 0.5 { 0.0 } else { (conf - 0.5) * 2.0 }
}
crates/input_classifier/src/fasttext/mod.rs
+123
View File
@@ -0,0 +1,123 @@
use std::io::Write as _;
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use fasttext::FastText;
use rust_embed::RustEmbed;
use tempfile::NamedTempFile;
use crate::{
ClassificationResult, Context, InputClassifier, InputType,
parser::parse_query_into_tokens,
util::{is_likely_shell_command, is_one_off_natural_language_word},
};
#[derive(Clone, Copy, RustEmbed)]
#[folder = "models/fasttext"]
struct Models;
pub struct FasttextClassifier {
classifier: FastText,
}
impl FasttextClassifier {
pub fn new() -> Result<Self> {
Ok(Self {
classifier: Self::load_classifier()?,
})
}
fn load_classifier() -> Result<FastText> {
let model_bytes = Models::get("cmd_lang_classifier_v4.bin")
.ok_or_else(|| anyhow!("Model file not found"))?
.data;
let mut temp_file = NamedTempFile::new()?;
temp_file.write_all(model_bytes.as_ref())?;
let model_path = temp_file.path();
let mut classifier = FastText::new();
classifier
.load_model(
model_path
.to_str()
.ok_or_else(|| anyhow!("Invalid model path"))?,
)
.map_err(|_| anyhow!("Failed to load fasttext classifier"))?;
log::info!("Successfully loaded fasttext classifier");
Ok(classifier)
}
}
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl InputClassifier for FasttextClassifier {
async fn detect_input_type(
&self,
input: warp_completer::ParsedTokensSnapshot,
context: &Context,
) -> InputType {
let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
let total_word_token_count = word_tokens.len();
if total_word_token_count == 1 {
if is_one_off_natural_language_word(&word_tokens[0].to_lowercase()) {
return InputType::AI;
}
// Prevent flickering for short input
return context.current_input_type;
}
if is_likely_shell_command(&input, total_word_token_count).await {
return InputType::Shell;
}
self.classify_input(input, context)
.await
.map(|result| result.to_input_type())
.unwrap_or(context.current_input_type)
}
async fn classify_input(
&self,
input: warp_completer::ParsedTokensSnapshot,
context: &Context,
) -> anyhow::Result<ClassificationResult> {
if let Ok(classification_result) =
classify_input_with_fasttext(&self.classifier, input.buffer_text.as_str())
{
return Ok(classification_result);
}
super::HeuristicClassifier
.classify_input(input, context)
.await
}
}
/// Classify the current input text with the FastText classifier
fn classify_input_with_fasttext(
classifier: &FastText,
input: &str,
) -> anyhow::Result<ClassificationResult> {
anyhow::ensure!(!input.trim().is_empty(), "cannot classify empty input");
let predictions = classifier
.predict(input, 2, 0.0)
.map_err(|err| anyhow!("Failed to classify input: {err}"))?;
let mut classification_result = ClassificationResult {
p_shell: 0.0,
p_ai: 0.0,
};
for prediction in predictions {
if prediction.label.contains("terminal_command") {
classification_result.p_shell = prediction.prob;
} else if prediction.label.contains("natural_language") {
classification_result.p_ai = prediction.prob;
}
}
Ok(classification_result)
}
crates/input_classifier/src/heuristic_classifier/mod.rs
+147
View File
@@ -0,0 +1,147 @@
use std::borrow::Cow;
use async_trait::async_trait;
use itertools::Itertools as _;
use natural_language_detection::natural_language_words_score;
use warp_completer::ParsedTokensSnapshot;
use crate::{
ClassificationResult, Context, InputClassifier, InputType,
parser::parse_query_into_tokens,
util::{
is_installed_binary, is_likely_shell_command, is_one_off_natural_language_word_or_prefix,
},
};
/// Minimum number of tokens users' input should have before kicking off input detection
/// to switch from AI input to command input.
/// This could be tuned.
const MINIMUM_COMMAND_DETECTION_TOKEN_LENGTH: u8 = 2;
/// Minimum number of tokens users' input should have before kicking off input detection
/// to switch from command input to AI input.
/// This could be tuned.
const MINIMUM_NATURAL_LANGUAGE_DETECTION_TOKEN_LENGTH: u8 = 2;
/// The percentage of input tokens that can be recognized as a natural language word before
/// we consider the input as natural language. This could be tuned.
const DETECT_AS_NATURAL_LANGUAGE_THRESHOLD: f32 = 0.6;
/// Threshold for the case when we have a low number of input tokens and require a higher
/// confidence level. This could be tuned.
const DETECT_AS_NATURAL_LANGUAGE_LOW_TOKEN_THRESHOLD: f32 = 0.8;
const END_TOKEN_COMPLETE_KEYS: &[char] = &[' ', '?', '!', '.', '"', ','];
/// A classifier that uses simple heuristics to determine the type of input.
pub struct HeuristicClassifier;
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl InputClassifier for HeuristicClassifier {
async fn detect_input_type(&self, input: ParsedTokensSnapshot, context: &Context) -> InputType {
let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
let total_word_token_count = word_tokens.len();
if total_word_token_count == 1
&& is_one_off_natural_language_word_or_prefix(&word_tokens[0].to_lowercase())
{
return InputType::AI;
}
if is_likely_shell_command(&input, total_word_token_count).await {
return InputType::Shell;
}
self.classify_input(input, context)
.await
.map(|result| result.to_input_type())
.unwrap_or(context.current_input_type)
}
async fn classify_input(
&self,
input: warp_completer::ParsedTokensSnapshot,
context: &Context,
) -> anyhow::Result<super::ClassificationResult> {
let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
// Try autodetecting both including and not including the last token,
// since we aren't sure if the user is done typing. If either case is
// detected as AI input, set to AI input.
let result = natural_language_detection_heuristic(
input.clone(),
word_tokens.clone(),
context.current_input_type,
false,
)
.await;
if matches!(result.to_input_type(), InputType::AI) {
return Ok(result);
}
Ok(natural_language_detection_heuristic(
input,
word_tokens,
context.current_input_type,
true,
)
.await)
}
}
/// Given some input text and current input type, return what type of input we think it is
/// using a heuristic.
async fn natural_language_detection_heuristic(
input: ParsedTokensSnapshot,
word_tokens: Vec<String>,
current_input_type: InputType,
include_last_token: bool,
) -> ClassificationResult {
let word_tokens_count = word_tokens.len();
let min_token_length = if matches!(current_input_type, InputType::AI) {
MINIMUM_COMMAND_DETECTION_TOKEN_LENGTH
} else {
MINIMUM_NATURAL_LANGUAGE_DETECTION_TOKEN_LENGTH
};
if min_token_length > word_tokens_count as u8 {
return ClassificationResult::pure_shell();
}
let mut word_tokens = word_tokens.into_iter().map(Cow::Owned).collect_vec();
// If the last token is not complete AND we are configured to not always include the last token, do not consider it
// in the natural language classifier. When the total word tokens have length less than 3, we also shouldn't pop the last
// token as this could cause misclassification on any top command we didn't parse.
let last_token_is_complete = input.buffer_text.ends_with(END_TOKEN_COMPLETE_KEYS);
if !include_last_token && !last_token_is_complete && word_tokens.len() > 2 {
word_tokens.pop();
}
let updated_word_token_count = word_tokens.len();
let likely_english_token_count =
natural_language_words_score(word_tokens, is_installed_binary(&input));
// When token count is lower than 3, we should make sure all tokens
// are matching the target classification category.
let threshold = if updated_word_token_count <= 3 {
1.0
} else if updated_word_token_count <= 4 {
DETECT_AS_NATURAL_LANGUAGE_LOW_TOKEN_THRESHOLD
} else {
DETECT_AS_NATURAL_LANGUAGE_THRESHOLD
};
if likely_english_token_count >= (updated_word_token_count as f32 * threshold) as usize {
return ClassificationResult::pure_ai();
}
ClassificationResult::pure_shell()
}
#[cfg(test)]
#[path = "mod_tests.rs"]
mod tests;
crates/input_classifier/src/heuristic_classifier/mod_tests.rs
+95
View File
@@ -0,0 +1,95 @@
use warp_completer::util::parse_current_commands_and_tokens;
use crate::{Context, test_utils::CompletionContext};
use super::*;
async fn mock_parsed_input_token(buffer_text: String) -> ParsedTokensSnapshot {
let completion_context = CompletionContext::new();
parse_current_commands_and_tokens(buffer_text, &completion_context).await
}
#[test]
fn test_input_detection() {
futures::executor::block_on(async move {
let classifier = HeuristicClassifier;
let mut context = Context {
current_input_type: InputType::AI,
is_agent_follow_up: false,
};
let token = mock_parsed_input_token("cargo --version".to_string()).await;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::Shell
);
// We have to override the first token description here given the mocked completion
// parser will parse the first token always as commands.
//
// Mock the case where cargo is not installed. We should still parse this as Shell input.
let mut token = mock_parsed_input_token("cargo --version".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::Shell
);
let mut token = mock_parsed_input_token("rvm install 3.3".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::Shell
);
// Short queries with NL should be parsed as AI input when already in AI input.
let mut token = mock_parsed_input_token("Explain this".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token.clone(), &context).await,
InputType::AI
);
context.current_input_type = InputType::Shell;
// Typing "fix this" after an error block is a common use case.
let mut token = mock_parsed_input_token("fix this".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::AI,
);
// Short queries with punctuation should be parsed as AI input.
let token = mock_parsed_input_token("What went wrong?".to_string()).await;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::AI
);
// Short queries with contractions should be parsed as AI input.
let mut token = mock_parsed_input_token("What's the reason".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::AI
);
// Short queries with quotations should be parsed as AI input.
let mut token =
mock_parsed_input_token("The message is \"utils::future ... ok\"".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::AI
);
// String tokens with special shell syntax should not be treated as negative NL signal.
let mut token = mock_parsed_input_token("The type is \"<>\"".to_string()).await;
token.parsed_tokens[0].token_description = None;
assert_eq!(
classifier.detect_input_type(token, &context).await,
InputType::AI
);
});
}
crates/input_classifier/src/input_type.rs
+39
View File
@@ -0,0 +1,39 @@
use serde::{Deserialize, Serialize};
use std::str::FromStr;
/// The type of input the user has provided.
#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum InputType {
/// The user input is a shell command.
#[default]
Shell,
/// The user input is a natural language query to AI.
AI,
}
impl InputType {
pub fn is_ai(&self) -> bool {
matches!(self, InputType::AI)
}
}
impl FromStr for InputType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"shell" => Ok(InputType::Shell),
"ai" => Ok(InputType::AI),
_ => Err(format!("Invalid input type: {s}. Must be 'shell' or 'ai'")),
}
}
}
impl std::fmt::Display for InputType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
InputType::Shell => write!(f, "Shell"),
InputType::AI => write!(f, "AI"),
}
}
}
crates/input_classifier/src/lib.rs
+89
View File
@@ -0,0 +1,89 @@
#[cfg(feature = "fasttext")]
mod fasttext;
mod heuristic_classifier;
mod input_type;
#[cfg(feature = "onnx")]
mod onnx;
mod parser;
pub mod test_utils;
pub mod util;
use async_trait::async_trait;
#[cfg(feature = "fasttext")]
pub use fasttext::FasttextClassifier;
pub use heuristic_classifier::HeuristicClassifier;
pub use input_type::InputType;
#[cfg(feature = "onnx")]
pub use onnx::{Model as OnnxModel, OnnxClassifier};
/// An input classifier, which can take some parsed user input and determine
/// what type of input it is.
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
pub trait InputClassifier: 'static + Send + Sync {
async fn detect_input_type(
&self,
input: warp_completer::ParsedTokensSnapshot,
context: &Context,
) -> InputType;
async fn classify_input(
&self,
input: warp_completer::ParsedTokensSnapshot,
context: &Context,
) -> anyhow::Result<ClassificationResult>;
}
/// The result of running inference on some user input.
pub struct ClassificationResult {
/// The probability that the input is a shell command.
p_shell: f32,
/// The probability that the input is a natural language query to AI.
p_ai: f32,
}
impl ClassificationResult {
fn pure_ai() -> Self {
Self {
p_shell: 0.0,
p_ai: 1.0,
}
}
fn pure_shell() -> Self {
Self {
p_shell: 1.0,
p_ai: 0.0,
}
}
pub fn p_shell(&self) -> f32 {
self.p_shell
}
pub fn p_ai(&self) -> f32 {
self.p_ai
}
/// Returns the confidence score (0.0 to 1.0) as the maximum of the two probabilities
pub fn confidence(&self) -> f32 {
self.p_shell.max(self.p_ai)
}
pub fn to_input_type(&self) -> InputType {
if self.p_shell > self.p_ai {
InputType::Shell
} else {
InputType::AI
}
}
}
/// Context for the classifier.
pub struct Context {
/// The current input type.
pub current_input_type: InputType,
/// Whether or not the input is a follow-up to an agent query.
pub is_agent_follow_up: bool,
}
crates/input_classifier/src/onnx/candle.rs
+101
View File
@@ -0,0 +1,101 @@
use std::collections::HashMap;
use anyhow::{Context as _, Result, ensure};
use candle_core::{IndexOp as _, Tensor};
use candle_onnx::onnx::ModelProto;
use prost::Message as _;
use tokenizers::Tokenizer;
use warp_completer::ParsedTokensSnapshot;
use super::ClassificationResult;
use super::Model;
pub struct InferenceRunner {
model: ModelProto,
tokenizer: Tokenizer,
}
impl InferenceRunner {
pub fn new(model: Model) -> Result<Self> {
Ok(Self {
model: Self::load_model(model)?,
tokenizer: Self::load_tokenizer(model)?,
})
}
fn load_model(model: Model) -> Result<ModelProto> {
let model_bytes = model.bytes().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
})?;
let model = ModelProto::decode(model_bytes.as_ref())?;
Ok(model)
}
fn load_tokenizer(model: Model) -> Result<Tokenizer> {
let tokenizer_bytes = model.tokenizer_bytes().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "Tokenizer file not found")
})?;
let tokenizer = Tokenizer::from_bytes(tokenizer_bytes).map_err(|e| anyhow::anyhow!(e))?;
Ok(tokenizer)
}
}
impl super::InferenceRunner for InferenceRunner {
fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
// Encode the input text into tokens.
let encoding = self
.tokenizer
.encode_fast(input.buffer_text.as_str(), true)
.map_err(|e| anyhow::anyhow!(e))?;
// For now, we'll do all inference on the CPU.
let device = candle_core::Device::Cpu;
let input_ids = Tensor::new(
encoding
.get_ids()
.iter()
.map(|&x| x as i64)
.collect::<Vec<_>>()
.as_slice(),
&device,
)
.context("failed to build input ids tensor")?;
let attention_mask = Tensor::new(
encoding
.get_attention_mask()
.iter()
.map(|&x| x as i64)
.collect::<Vec<_>>()
.as_slice(),
&device,
)
.context("failed to build attention mask tensor")?;
// Run inference.
let outputs = candle_onnx::simple_eval(
&self.model,
HashMap::from([
("input_ids".to_string(), input_ids.unsqueeze(0)?),
("attention_mask".to_string(), attention_mask.unsqueeze(0)?),
]),
)
.context("error evaluating the model")?;
let logits = outputs.get("logits").context("failed to get logits")?;
let probabilities = candle_nn::ops::softmax_last_dim(logits)
.context("failed to compute softmax")?
.i(0)
.context("failed to get first dimension")?
.to_vec1::<f32>()
.context("failed to convert softmax output to vec")?;
ensure!(probabilities.len() == 2, "expected 2 probabilities");
Ok(ClassificationResult {
p_ai: probabilities[0],
p_shell: probabilities[1],
})
}
}
crates/input_classifier/src/onnx/mod.rs
+198
View File
@@ -0,0 +1,198 @@
#[cfg(feature = "onnx_candle")]
mod candle;
#[cfg(feature = "onnx_ort")]
mod ort;
use std::borrow::Cow;
use anyhow::Result;
use async_trait::async_trait;
use rust_embed::RustEmbed;
use warp_completer::ParsedTokensSnapshot;
use crate::{
ClassificationResult, Context, InputClassifier, InputType,
parser::parse_query_into_tokens,
util::{
is_likely_shell_command, is_one_off_natural_language_word, is_one_off_shell_command_keyword,
},
};
#[derive(Clone, Copy, RustEmbed)]
#[folder = "models/onnx"]
struct Models;
#[derive(Copy, Clone)]
pub enum Model {
BertTiny,
}
impl Model {
fn bytes(&self) -> Option<Cow<'static, [u8]>> {
Models::get(self.model_path()).map(|file| file.data)
}
fn tokenizer_bytes(&self) -> Option<Cow<'static, [u8]>> {
Models::get(self.tokenizer_path()).map(|file| file.data)
}
fn model_path(&self) -> &'static str {
match self {
Model::BertTiny => "bert_tiny.onnx",
}
}
fn tokenizer_path(&self) -> &'static str {
match self {
Model::BertTiny => "bert_tiny_tokenizer.json",
}
}
}
pub struct OnnxClassifier {
inference_runner: Box<dyn InferenceRunner>,
has_panicked: HasPanicked,
}
impl OnnxClassifier {
pub fn new(_model: Model) -> Result<Self> {
#[cfg(feature = "onnx_candle")]
match candle::InferenceRunner::new(_model).map(Box::new) {
Ok(inference_runner) => {
return Ok(Self {
inference_runner,
has_panicked: HasPanicked::new(),
});
}
Err(err) => log::warn!("Failed to initialize candle inference runner: {err:#}"),
}
#[cfg(feature = "onnx_ort")]
match ort::InferenceRunner::new(_model).map(Box::new) {
Ok(inference_runner) => {
return Ok(Self {
inference_runner,
has_panicked: HasPanicked::new(),
});
}
Err(err) => log::warn!("Failed to initialize ort inference runner: {err:#}"),
}
Err(anyhow::anyhow!("No onnx inference engine enabled"))
}
}
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl InputClassifier for OnnxClassifier {
async fn detect_input_type(&self, input: ParsedTokensSnapshot, context: &Context) -> InputType {
let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
let total_word_token_count = word_tokens.len();
// Start by applying some simple heuristics before running the full classifier.
if let Some(first_word) = word_tokens.first() {
let first_word = first_word.to_lowercase();
// If the input is a single word and the word is one of a specific set of words, classify it as AI
if word_tokens.len() == 1 && is_one_off_natural_language_word(&first_word) {
return InputType::AI;
}
// If the first token is one of a specific set of shell command keywords (e.g.: echo or sudo),
// we should classify it as shell.
if is_one_off_shell_command_keyword(&first_word) {
return InputType::Shell;
}
}
if is_likely_shell_command(&input, total_word_token_count).await {
return InputType::Shell;
}
// Otherwise, defer all decision-making to the model.
self.classify_input(input, context)
.await
.map(|result| result.to_input_type())
.unwrap_or(context.current_input_type)
}
async fn classify_input(
&self,
input: warp_completer::ParsedTokensSnapshot,
_context: &Context,
) -> anyhow::Result<ClassificationResult> {
// If we ever panicked while running inference, we should fall back to the heuristic classifier.
if self.has_panicked.has_panicked() {
return crate::heuristic_classifier::HeuristicClassifier
.classify_input(input, _context)
.await;
}
// Given that we only can get here if we have never panicked, we don't have to
// worry about attempting to use an inference runner that is in an invalid state
// due to recovering after catching a panic unwind.
let inference_runner = std::panic::AssertUnwindSafe(&self.inference_runner);
let input_ref = &input;
match std::panic::catch_unwind(move || {
let start = instant::Instant::now();
let result = inference_runner.run_inference(input_ref);
let duration = start.elapsed();
let duration_ms = duration.as_secs_f32() * 1000.0;
match result {
Ok(result) => {
log::debug!(
"Inference took {duration_ms:.2} ms; p_shell: {:.5}, p_ai: {:.5}",
result.p_shell,
result.p_ai
);
Ok(result)
}
Err(e) => {
log::error!("Failed to run inference (took {duration_ms:.2} ms): {e:#}");
Err(e)
}
}
}) {
Ok(result) => result,
Err(_) => {
log::error!(
"Caught panic while running inference; falling back to heuristic classifier."
);
self.has_panicked.on_panic();
crate::heuristic_classifier::HeuristicClassifier
.classify_input(input, _context)
.await
}
}
}
}
trait InferenceRunner: 'static + Send + Sync {
fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult>;
}
/// A simple structure that we can use to track whether the ONNX classifier has panicked.
struct HasPanicked {
inner: std::sync::Once,
}
impl HasPanicked {
fn new() -> Self {
Self {
inner: std::sync::Once::new(),
}
}
fn on_panic(&self) {
// Mark the classifier as having panicked.
self.inner.call_once(|| {});
}
fn has_panicked(&self) -> bool {
// Return true if the classifier has panicked.
self.inner.is_completed()
}
}
crates/input_classifier/src/onnx/ort.rs
+91
View File
@@ -0,0 +1,91 @@
use anyhow::{Result, ensure};
use itertools::Itertools as _;
use ort::{
execution_providers::CPUExecutionProvider, session::Session, tensor::ArrayExtensions as _,
value::Value,
};
use parking_lot::Mutex;
use tokenizers::Tokenizer;
use warp_completer::ParsedTokensSnapshot;
use super::ClassificationResult;
use super::Model;
pub struct InferenceRunner {
session: Mutex<Session>,
tokenizer: Tokenizer,
}
impl InferenceRunner {
pub fn new(model: Model) -> Result<Self> {
Ok(Self {
session: Self::init_session(model)?.into(),
tokenizer: Self::load_tokenizer(model)?,
})
}
fn init_session(model: Model) -> Result<Session> {
let model_bytes = model.bytes().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
})?;
let session = Session::builder()?
// For now, we'll do all inference on the CPU.
.with_execution_providers([CPUExecutionProvider::default().build()])?
.commit_from_memory(model_bytes.as_ref())?;
Ok(session)
}
fn load_tokenizer(model: Model) -> Result<Tokenizer> {
let tokenizer_bytes = model.tokenizer_bytes().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "Tokenizer file not found")
})?;
let tokenizer = Tokenizer::from_bytes(tokenizer_bytes).map_err(|e| anyhow::anyhow!(e))?;
Ok(tokenizer)
}
}
impl super::InferenceRunner for InferenceRunner {
fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
// Encode the input text into tokens.
let encoding = self
.tokenizer
.encode_fast(input.buffer_text.as_str(), true)
.map_err(|e| anyhow::anyhow!(e))?;
let input_ids = encoding.get_ids();
let attention_mask = encoding.get_attention_mask();
let input_ids = Value::from_array((
[1, input_ids.len()],
input_ids.iter().map(|&x| x as i64).collect_vec(),
))?;
let attention_mask = Value::from_array((
[1, attention_mask.len()],
attention_mask.iter().map(|&x| x as i64).collect_vec(),
))?;
let mut session = self.session.lock();
let outputs = session.run(ort::inputs![
"input_ids" => input_ids,
"attention_mask" => attention_mask,
])?;
let logits = &outputs[0];
let logits = logits.try_extract_array::<f32>()?;
let probabilities = logits.softmax(ndarray::Axis(1));
let probabilities = probabilities.view();
let probabilities = probabilities
.as_slice()
.ok_or_else(|| anyhow::anyhow!("failed to get probabilities"))?;
ensure!(probabilities.len() == 2, "expected 2 probabilities");
Ok(ClassificationResult {
p_ai: probabilities[0],
p_shell: probabilities[1],
})
}
}
crates/input_classifier/src/parser.rs
+154
View File
@@ -0,0 +1,154 @@
use std::{iter::Peekable, mem, str::Chars};
use itertools::Itertools;
/// Parses a query into tokens, taking into account delimiters.
pub fn parse_query_into_tokens(query: &str) -> Vec<String> {
let parser = SentenceParser {
chars: query.chars().peekable(),
active_delimiter: None,
active_token: String::new(),
};
parser.into_iter().collect_vec()
}
#[derive(PartialEq, Eq)]
enum WordDelimiter {
Separator,
DoubleQuote,
SingleQuote,
Backtick,
Whitespace,
}
fn convert_char_to_delimiter(c: char) -> Option<WordDelimiter> {
match c {
'\'' => Some(WordDelimiter::SingleQuote),
'"' => Some(WordDelimiter::DoubleQuote),
'`' => Some(WordDelimiter::Backtick),
',' | '.' | '!' | '?' => Some(WordDelimiter::Separator),
c if c.is_whitespace() => Some(WordDelimiter::Whitespace),
_ => None,
}
}
/// Parse a sentence into tokens for natural language classificiation.
/// The main difference from the default "split_whitespace" parser is
/// for matching double quotes and single quotes, we would keep their
/// closed content as one token.
struct SentenceParser<'a> {
chars: Peekable<Chars<'a>>,
active_delimiter: Option<WordDelimiter>,
active_token: String,
}
impl Iterator for SentenceParser<'_> {
type Item = String;
fn next(&mut self) -> Option<Self::Item> {
while let Some(c) = self.chars.next() {
let delimiter = convert_char_to_delimiter(c);
let next_delimiter = self.chars.peek().map(|c| convert_char_to_delimiter(*c));
match delimiter {
Some(WordDelimiter::Whitespace) if self.active_delimiter.is_none() => {
if self.active_token.is_empty() {
continue;
}
return Some(mem::take(&mut self.active_token));
}
Some(WordDelimiter::Separator) if self.active_delimiter.is_none() => {
if self.active_token.is_empty() {
continue;
}
// If next_delimiter is not whitespace or None, this means the delimiter
// is in the middle of a word. In this case, push the plain text character
// to the active token.
if next_delimiter
.map(|c| c == Some(WordDelimiter::Whitespace))
.unwrap_or(true)
{
return Some(mem::take(&mut self.active_token));
} else {
self.active_token.push(c);
}
}
Some(WordDelimiter::DoubleQuote) => {
let complete_quote =
if self.active_delimiter == Some(WordDelimiter::DoubleQuote) {
self.active_delimiter = None;
true
} else if !self.active_token.is_empty() || self.active_delimiter.is_some() {
false
} else {
self.active_delimiter = Some(WordDelimiter::DoubleQuote);
false
};
self.active_token.push(c);
if complete_quote {
let token = mem::take(&mut self.active_token);
// Skip empty quotes since this could be an in progress edit.
if token == "\"\"" {
continue;
}
return Some(token);
}
}
Some(WordDelimiter::Backtick) => {
let complete_quote = if self.active_delimiter == Some(WordDelimiter::Backtick) {
self.active_delimiter = None;
true
} else if !self.active_token.is_empty() || self.active_delimiter.is_some() {
false
} else {
self.active_delimiter = Some(WordDelimiter::Backtick);
false
};
self.active_token.push(c);
if complete_quote {
let token = mem::take(&mut self.active_token);
return Some(token);
}
}
Some(WordDelimiter::SingleQuote) => {
let complete_quote =
if self.active_delimiter == Some(WordDelimiter::SingleQuote) {
self.active_delimiter = None;
true
} else if !self.active_token.is_empty() || self.active_delimiter.is_some() {
false
} else {
self.active_delimiter = Some(WordDelimiter::SingleQuote);
false
};
self.active_token.push(c);
if complete_quote {
let token = mem::take(&mut self.active_token);
// Skip empty quotes since this could be an in progress edit.
if token == "''" {
continue;
}
return Some(token);
}
}
_ => self.active_token.push(c),
}
}
if self.active_token.is_empty() {
return None;
}
Some(mem::take(&mut self.active_token))
}
}
#[cfg(test)]
#[path = "parser_tests.rs"]
pub mod tests;
crates/input_classifier/src/parser_tests.rs
+64
View File
@@ -0,0 +1,64 @@
use super::*;
#[test]
fn test_sentence_parser() {
assert_eq!(
parse_query_into_tokens("This is a question?"),
vec![
"This".to_string(),
"is".to_string(),
"a".to_string(),
"question".to_string()
]
);
assert_eq!(
parse_query_into_tokens("No I can't!"),
vec!["No".to_string(), "I".to_string(), "can't".to_string()]
);
assert_eq!(
parse_query_into_tokens("A quote \"Inside quote\""),
vec![
"A".to_string(),
"quote".to_string(),
"\"Inside quote\"".to_string()
]
);
assert_eq!(
parse_query_into_tokens("A quote \"Inside ' quote\""),
vec![
"A".to_string(),
"quote".to_string(),
"\"Inside ' quote\"".to_string()
]
);
assert_eq!(
parse_query_into_tokens("A quote \"Inside 'something' quote\""),
vec![
"A".to_string(),
"quote".to_string(),
"\"Inside 'something' quote\"".to_string()
]
);
assert_eq!(
parse_query_into_tokens("Empty quote \"\"!?!"),
vec!["Empty".to_string(), "quote".to_string()]
);
assert_eq!(
parse_query_into_tokens("www.google.com"),
vec!["www.google.com".to_string(),]
);
assert_eq!(
parse_query_into_tokens("Command `mockery --name example_interface`"),
vec![
"Command".to_string(),
"`mockery --name example_interface`".to_string()
]
);
}
crates/input_classifier/src/test_utils.rs
+47
View File
@@ -0,0 +1,47 @@
use std::{collections::HashSet, sync::Arc};
use smol_str::SmolStr;
use warp_completer::{
completer::{GeneratorContext, PathCompletionContext},
signatures::CommandRegistry,
};
/// An implementation of `CompletionContext` for testing purposes.
pub struct CompletionContext {
command_registry: Arc<CommandRegistry>,
}
impl CompletionContext {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
Self {
command_registry: CommandRegistry::global_instance(),
}
}
}
impl warp_completer::completer::CompletionContext for CompletionContext {
fn top_level_commands(&self) -> Box<dyn Iterator<Item = &str> + '_> {
Box::new(self.command_registry.registered_commands())
}
fn command_registry(&self) -> &CommandRegistry {
&self.command_registry
}
fn environment_variable_names(&self) -> Option<&HashSet<SmolStr>> {
None
}
fn shell_supports_autocd(&self) -> Option<bool> {
None
}
fn path_completion_context(&self) -> Option<&dyn PathCompletionContext> {
None
}
fn generator_context(&self) -> Option<&dyn GeneratorContext> {
None
}
}
crates/input_classifier/src/util.rs
+120
View File
@@ -0,0 +1,120 @@
use std::collections::HashSet;
use lazy_static::lazy_static;
use natural_language_detection::check_if_token_has_shell_syntax;
use warp_completer::ParsedTokensSnapshot;
/// The percentage of input tokens that can be described by our completion engine before
/// we consider the input as a shell command. This could be tuned.
const DETECT_AS_COMMAND_THRESHOLD: f32 = 0.5;
/// Threshold for the case when we have a low number of input tokens and require a higher
/// confidence level. This could be tuned.
const DETECT_AS_COMMAND_LOW_TOKEN_THRESHOLD: f32 = 0.7;
lazy_static! {
/// One-off commands / keywords that should trigger a shell command classification.
///
/// `claude`, `codex`, and `gemini` are not actually _really_ one-off shell command keywords,
/// but false-positive NL classifications for these inputs (where the user was trying to use
/// claude code, codex CLI, or gemini CLI) suck, because the user often thinks we're
/// intentionally trying to push them away from those CLIs into Agent Mode, so we mitigate the
/// risk by always treating as shell.
static ref ONE_OFF_SHELL_COMMAND_KEYWORDS: HashSet<&'static str> = HashSet::from(["#", "echo", "man", "sudo", "claude", "codex", "gemini"]);
static ref ONE_OFF_NATURAL_LANGUAGE_WORDS: HashSet<&'static str> = HashSet::from(["hello", "hi", "hey", "hola", "thanks", "explain", "yes", "no", "what", "nice", "1. "]);
/// A set of words that should trigger an AI classification if they are the entire input
/// and the input is a follow-up to an agent response.
static ref AGENT_FOLLOW_UP_INPUTS: HashSet<&'static str> = HashSet::from(["yes", "continue", "do it"]);
}
pub fn is_agent_follow_up_input(input: &str) -> bool {
AGENT_FOLLOW_UP_INPUTS.contains(input)
}
pub fn is_one_off_shell_command_keyword(word: &str) -> bool {
ONE_OFF_SHELL_COMMAND_KEYWORDS.contains(word)
}
/// Returns true if the word is a one-off natural language word or a prefix of a one-off natural language word.
pub fn is_one_off_natural_language_word_or_prefix(word: &str) -> bool {
is_one_off_natural_language_word(word) || is_prefix_of_natural_language_word(word)
}
// Returns true if the word is a one-off natural language word.
pub fn is_one_off_natural_language_word(word: &str) -> bool {
ONE_OFF_NATURAL_LANGUAGE_WORDS.contains(word)
}
/// Checks if the input string is a prefix of any word in the ONE_OFF_NATURAL_LANGUAGE_WORDS set.
/// This helps with progressive typing detection to avoid mode flipping.
pub fn is_prefix_of_natural_language_word(input: &str) -> bool {
// input is already lowercase from caller
ONE_OFF_NATURAL_LANGUAGE_WORDS
.iter()
.any(|word| word.starts_with(input))
}
pub async fn is_likely_shell_command(
input: &ParsedTokensSnapshot,
word_tokens_count: usize,
) -> bool {
const YIELD_BATCH_SIZE: usize = 5;
let mut likely_command_token_count = 0;
let total_token_count = input.parsed_tokens.len();
let mut is_first_token_command = false;
for (idx, token) in input.parsed_tokens.iter().enumerate() {
// Periodically, yield to the executor so this task can be aborted if
// requested.
if idx % YIELD_BATCH_SIZE == 0 {
futures_lite::future::yield_now().await;
}
// Early return if we encounter a one-off command / keyword at the beginning of the line.
if token.token_index == 0 && ONE_OFF_SHELL_COMMAND_KEYWORDS.contains(&token.token.as_str())
{
return true;
}
if token.token_description.is_some()
|| check_if_token_has_shell_syntax(token.token.as_str())
{
likely_command_token_count += 1;
}
if token.token_index == 0 {
is_first_token_command = token.token_description.is_some();
}
}
// When token count is lower than 2, we should make sure all tokens
// are matching the target classification category.
let command_threshold = if total_token_count <= 2 {
1.0
} else if total_token_count <= 4 {
DETECT_AS_COMMAND_LOW_TOKEN_THRESHOLD
} else {
DETECT_AS_COMMAND_THRESHOLD
};
// Classify as shell if:
// 1) We hit significant threshold of likely shell command tokens.
// 2) When there are fewer than 3 tokens, the first token is a valid top-level command.
if likely_command_token_count >= (total_token_count as f32 * command_threshold) as usize
|| (word_tokens_count < 3 && is_first_token_command)
{
return true;
}
false
}
/// Returns true if the first token is a command that is installed on the system.
pub fn is_installed_binary(input: &ParsedTokensSnapshot) -> bool {
input
.parsed_tokens
.first()
.map(|token| token.token_description.is_some())
.unwrap_or(false)
}