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

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/input_classifier/src/onnx/candle.rs
+3 -3
View File
@@ -7,9 +7,8 @@ use galaxy_completer::ParsedTokensSnapshot;
use prost::Message as _;
use tokenizers::Tokenizer;
use super::ClassificationResult;
use super::Model;
use super::{ClassificationResult, Model};
use crate::InputClassifierDecisionSource;
pub struct InferenceRunner {
model: ModelProto,
@@ -96,6 +95,7 @@ impl super::InferenceRunner for InferenceRunner {
Ok(ClassificationResult {
p_ai: probabilities[0],
p_shell: probabilities[1],
source: InputClassifierDecisionSource::InputClassifier,
})
}
}
crates/input_classifier/src/onnx/mod.rs
+52 -18
View File
@@ -10,21 +10,28 @@ use async_trait::async_trait;
use galaxy_completer::ParsedTokensSnapshot;
use rust_embed::RustEmbed;
use crate::parser::parse_query_into_tokens;
use crate::util::{
is_likely_shell_command, is_one_off_natural_language_word, is_one_off_shell_command_keyword,
};
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,
},
ClassificationResult, Context, InputClassificationResult, InputClassifier,
InputClassifierDecisionSource, InputType,
};
#[derive(Clone, Copy, RustEmbed)]
#[folder = "models/onnx"]
#[include = "bert_tiny_tokenizer.json"]
#[cfg_attr(feature = "nld_classifier_v1", include = "bert_tiny_v1.onnx")]
#[cfg_attr(feature = "nld_classifier_v2", include = "bert_tiny_v2.onnx")]
#[cfg_attr(feature = "nld_classifier_v3", include = "bert_tiny_v3.onnx")]
struct Models;
#[derive(Copy, Clone)]
#[derive(Copy, Clone, Debug)]
pub enum Model {
BertTiny,
BertTinyV1,
BertTinyV2,
BertTinyV3,
}
impl Model {
@@ -38,13 +45,15 @@ impl Model {
fn model_path(&self) -> &'static str {
match self {
Model::BertTiny => "bert_tiny.onnx",
Model::BertTinyV1 => "bert_tiny_v1.onnx",
Model::BertTinyV2 => "bert_tiny_v2.onnx",
Model::BertTinyV3 => "bert_tiny_v3.onnx",
}
}
fn tokenizer_path(&self) -> &'static str {
match self {
Model::BertTiny => "bert_tiny_tokenizer.json",
Model::BertTinyV1 | Model::BertTinyV2 | Model::BertTinyV3 => "bert_tiny_tokenizer.json",
}
}
}
@@ -85,7 +94,11 @@ impl OnnxClassifier {
#[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 {
async fn detect_input_type(
&self,
input: ParsedTokensSnapshot,
context: &Context,
) -> InputClassificationResult {
let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
let total_word_token_count = word_tokens.len();
@@ -96,36 +109,53 @@ impl InputClassifier for OnnxClassifier {
// 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;
return InputClassificationResult::new(
InputType::AI,
InputClassifierDecisionSource::NaturalLanguageOneOffAllowlist,
);
}
// 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;
return InputClassificationResult::new(
InputType::Shell,
InputClassifierDecisionSource::ShellHeuristic,
);
}
}
if is_likely_shell_command(&input, total_word_token_count).await {
return InputType::Shell;
return InputClassificationResult::new(
InputType::Shell,
InputClassifierDecisionSource::ShellHeuristic,
);
}
// 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)
.map(|classification| {
InputClassificationResult::new(
classification.to_input_type(),
classification.source,
)
})
.unwrap_or(InputClassificationResult::new(
context.current_input_type,
InputClassifierDecisionSource::InputClassifierFallbackCurrentInput,
))
}
async fn classify_input(
&self,
input: galaxy_completer::ParsedTokensSnapshot,
_context: &Context,
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)
.classify_input(input, context)
.await;
}
@@ -163,7 +193,7 @@ impl InputClassifier for OnnxClassifier {
);
self.has_panicked.on_panic();
crate::heuristic_classifier::HeuristicClassifier
.classify_input(input, _context)
.classify_input(input, context)
.await
}
}
@@ -196,3 +226,7 @@ impl HasPanicked {
self.inner.is_completed()
}
}
#[cfg(test)]
#[path = "mod_tests.rs"]
mod tests;
crates/input_classifier/src/onnx/mod_tests.rs
+64
View File
@@ -0,0 +1,64 @@
use anyhow::Result;
use futures::executor::block_on;
use warp_completer::meta::SpannedItem;
use warp_completer::{ParsedTokenData, ParsedTokensSnapshot};
use super::*;
struct FailingInferenceRunner;
impl InferenceRunner for FailingInferenceRunner {
fn run_inference(&self, _input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
Err(anyhow::anyhow!("inference failed"))
}
}
fn parsed_input_without_descriptions(buffer_text: &str) -> ParsedTokensSnapshot {
let mut next_search_start = 0;
let parsed_tokens = buffer_text
.split_whitespace()
.enumerate()
.map(|(token_index, token)| {
let token_start =
buffer_text[next_search_start..].find(token).unwrap() + next_search_start;
let token_end = token_start + token.len();
next_search_start = token_end;
ParsedTokenData {
token: token.to_string().spanned((token_start, token_end)),
token_index,
token_description: None,
}
})
.collect();
ParsedTokensSnapshot {
buffer_text: buffer_text.to_owned(),
parsed_tokens,
}
}
#[test]
fn test_inference_error_reports_current_input_fallback_source() {
block_on(async move {
let classifier = OnnxClassifier {
inference_runner: Box::new(FailingInferenceRunner),
has_panicked: HasPanicked::new(),
};
let context = Context {
current_input_type: InputType::AI,
is_agent_follow_up: false,
};
let input = parsed_input_without_descriptions("help migrate database");
let decision = classifier.detect_input_type(input, &context).await;
assert_eq!(
decision,
InputClassificationResult::new(
InputType::AI,
InputClassifierDecisionSource::InputClassifierFallbackCurrentInput,
)
);
});
}
crates/input_classifier/src/onnx/ort.rs
+7 -7
View File
@@ -1,16 +1,15 @@
use anyhow::{Result, ensure};
use galaxy_completer::ParsedTokensSnapshot;
use itertools::Itertools as _;
use ort::{
execution_providers::CPUExecutionProvider, session::Session, tensor::ArrayExtensions as _,
value::Value,
};
use ort::execution_providers::CPUExecutionProvider;
use ort::session::Session;
use ort::tensor::ArrayExtensions as _;
use ort::value::Value;
use parking_lot::Mutex;
use tokenizers::Tokenizer;
use super::ClassificationResult;
use super::Model;
use super::{ClassificationResult, Model};
use crate::InputClassifierDecisionSource;
pub struct InferenceRunner {
session: Mutex<Session>,
@@ -86,6 +85,7 @@ impl super::InferenceRunner for InferenceRunner {
Ok(ClassificationResult {
p_ai: probabilities[0],
p_shell: probabilities[1],
source: InputClassifierDecisionSource::InputClassifier,
})
}
}