Initial public release of Warp.
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
David Stern
@@ -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 }
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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
|
||||
);
|
||||
});
|
||||
}
|
||||
@@ -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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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],
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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],
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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()
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user