122 lines
3.6 KiB
Rust
122 lines
3.6 KiB
Rust
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;
|
|
pub use heuristic_classifier::HeuristicClassifier;
|
|
pub use input_type::InputType;
|
|
#[cfg(feature = "onnx")]
|
|
pub use onnx::{Model as OnnxModel, OnnxClassifier};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Sources produced by the input classifier pipeline.
|
|
#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum InputClassifierDecisionSource {
|
|
// Classification result coming from Onnx classifier
|
|
InputClassifier,
|
|
// Classification result coming from fall back heuristic classifier when onnx classifier panicked
|
|
InputClassifierFallbackHeuristic,
|
|
// Classification result coming from current input type when onnx classifier failed
|
|
InputClassifierFallbackCurrentInput,
|
|
// Input match with ONE_OFF_NATURAL_LANGUAGE_WORDS
|
|
NaturalLanguageOneOffAllowlist,
|
|
// Input match with ONE_OFF_SHELL_COMMAND_KEYWORDS
|
|
ShellCommandAllowList,
|
|
// Classification result coming from is_likely_shell_command
|
|
ShellHeuristic,
|
|
}
|
|
|
|
/// The detected input type along with the decision source that produced it.
|
|
#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub struct InputClassificationResult {
|
|
/// The detected input type.
|
|
pub input_type: InputType,
|
|
/// The classifier source that produced this classification.
|
|
pub source: InputClassifierDecisionSource,
|
|
}
|
|
|
|
impl InputClassificationResult {
|
|
pub fn new(input_type: InputType, source: InputClassifierDecisionSource) -> Self {
|
|
Self { input_type, source }
|
|
}
|
|
}
|
|
|
|
/// 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: galaxy_completer::ParsedTokensSnapshot,
|
|
context: &Context,
|
|
) -> InputClassificationResult;
|
|
|
|
async fn classify_input(
|
|
&self,
|
|
input: galaxy_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,
|
|
/// The classifier source that produced this classification.
|
|
pub source: InputClassifierDecisionSource,
|
|
}
|
|
|
|
impl ClassificationResult {
|
|
fn pure_ai(source: InputClassifierDecisionSource) -> Self {
|
|
Self {
|
|
p_shell: 0.0,
|
|
p_ai: 1.0,
|
|
source,
|
|
}
|
|
}
|
|
|
|
fn pure_shell(source: InputClassifierDecisionSource) -> Self {
|
|
Self {
|
|
p_shell: 1.0,
|
|
p_ai: 0.0,
|
|
source,
|
|
}
|
|
}
|
|
|
|
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,
|
|
}
|