use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum QueryMode {
SemanticLookup,
EntityLookup {
mention: String,
},
TemporalLookup {
temporal_expr: String,
},
Mixed {
components: Vec<QueryMode>,
},
}
impl QueryMode {
pub fn kind(&self) -> &'static str {
match self {
Self::SemanticLookup => "semantic",
Self::EntityLookup { .. } => "entity",
Self::TemporalLookup { .. } => "temporal",
Self::Mixed { .. } => "mixed",
}
}
pub fn from_kind(kind: &str) -> Option<Self> {
match kind {
"semantic" => Some(Self::SemanticLookup),
"entity" => None, "temporal" => None, "mixed" => None, _ => None,
}
}
}
#[cfg(test)]
impl std::fmt::Display for QueryMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?}", self)
}
}
#[cfg(test)]
impl std::error::Error for QueryMode {}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifyResult {
pub mode: QueryMode,
pub confidence: f32,
pub reason: Option<String>,
}
pub fn classify(query: &str) -> ClassifyResult {
let mut components = Vec::new();
for mention in extract_entity_mentions(query) {
components.push(QueryMode::EntityLookup { mention });
}
if let Some(temporal_expr) = extract_temporal_expr(query) {
components.push(QueryMode::TemporalLookup { temporal_expr });
}
match components.len() {
0 => ClassifyResult {
mode: QueryMode::SemanticLookup,
confidence: 0.8,
reason: Some("no entity or temporal signals detected".into()),
},
1 => {
if let Some(mode) = components.pop() {
ClassifyResult {
confidence: 0.7,
reason: Some(format!("single signal: {}", mode.kind())),
mode,
}
} else {
ClassifyResult {
mode: QueryMode::SemanticLookup,
confidence: 0.8,
reason: Some("no entity or temporal signals detected".into()),
}
}
}
_ => ClassifyResult {
mode: QueryMode::Mixed { components },
confidence: 0.6,
reason: Some("multiple signal types detected".into()),
},
}
}
fn extract_entity_mentions(query: &str) -> Vec<String> {
let mut mentions = Vec::new();
for word in query.split_whitespace() {
if let Some(name) = word.strip_prefix('@') {
if !name.is_empty() {
mentions.push(
name.trim_matches(|c: char| !c.is_alphanumeric() && c != '_' && c != '-')
.to_string(),
);
}
}
}
let mut in_quote = false;
let mut escaped = false;
let mut current = String::new();
for ch in query.chars() {
if escaped {
if in_quote {
current.push(ch);
}
escaped = false;
continue;
}
if ch == '\\' {
escaped = true;
continue;
}
if ch == '"' {
if in_quote && !current.is_empty() {
mentions.push(std::mem::take(&mut current));
} else {
current.clear();
}
in_quote = !in_quote;
continue;
}
if in_quote {
current.push(ch);
}
}
mentions.retain(|mention| !mention.is_empty());
mentions.sort();
mentions.dedup();
mentions
}
fn extract_temporal_expr(query: &str) -> Option<String> {
let lower = query.to_lowercase();
let temporal_markers = [
"yesterday",
"last week",
"last month",
"last year",
"today",
"this week",
"this month",
"before",
"after",
"since",
"between",
"recent",
"recently",
"latest",
"oldest",
"earlier",
];
let mut markers = temporal_markers.to_vec();
markers.sort_by_key(|marker| std::cmp::Reverse(marker.len()));
for marker in markers {
if lower.contains(marker) {
return Some(marker.to_string());
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pure_semantic_query() {
let result = classify("how does the search pipeline work");
assert_eq!(result.mode.kind(), "semantic");
}
#[test]
fn at_mention_triggers_entity() {
let result = classify("what does @reembed_all do");
match &result.mode {
QueryMode::EntityLookup { mention } => assert_eq!(mention, "reembed_all"),
other => unreachable!("QueryMode::EntityLookup expected, got {other:?}"),
}
}
#[test]
fn quoted_string_triggers_entity() {
let result = classify("find \"MemoryStore\" usage");
match &result.mode {
QueryMode::EntityLookup { mention } => assert_eq!(mention, "MemoryStore"),
other => unreachable!("QueryMode::EntityLookup expected, got {other:?}"),
}
}
#[test]
fn temporal_keyword_triggers_temporal() {
let result = classify("what changed last week");
match &result.mode {
QueryMode::TemporalLookup { temporal_expr } => {
assert_eq!(temporal_expr, "last week");
}
other => unreachable!("QueryMode::TemporalLookup expected, got {other:?}"),
}
}
#[test]
fn mixed_entity_and_temporal() {
let result = classify("what did @alice do last week");
assert_eq!(result.mode.kind(), "mixed");
}
#[test]
fn multiple_entity_mentions_are_preserved() {
let result = classify(r#"compare @alice and @bob with "MemoryStore""#);
match &result.mode {
QueryMode::Mixed { components } => {
let mentions: Vec<_> = components
.iter()
.filter_map(|component| match component {
QueryMode::EntityLookup { mention } => Some(mention.as_str()),
_ => None,
})
.collect();
assert_eq!(mentions, vec!["MemoryStore", "alice", "bob"]);
}
other => unreachable!("QueryMode::Mixed expected, got {other:?}"),
}
}
#[test]
fn temporal_marker_prefers_more_specific_phrase() {
let result = classify("summarize last week and recent changes");
match &result.mode {
QueryMode::TemporalLookup { temporal_expr } => {
assert_eq!(temporal_expr, "last week");
}
other => unreachable!("QueryMode::TemporalLookup expected, got {other:?}"),
}
}
}