use std::collections::{HashMap, HashSet};
use std::time::Instant;
use crate::recall_feedback::{on_recall_hit, on_recall_miss};
use serde_json::{json, Value};
use uuid::Uuid;
use khive_fusion::FusionStrategy;
use khive_runtime::{micros_to_iso, NamespaceToken, RuntimeError, SearchSource, VerbRegistry};
use khive_storage::types::{EdgeFilter, PageRequest};
use khive_storage::EdgeRelation;
use crate::config::ScoreBreakdown;
use crate::rerank::{weighted_rerank, RerankFeatures};
use crate::scoring::{
calculate_score, contains_cjk, needs_multilingual, normalize_min_score,
normalize_rank_fusion_scores, normalize_rrf_scores, ScoreInput,
};
use crate::MemoryPack;
use super::common::{
compute_score, deser, fuse_candidates, make_pipeline, note_matches_tags, plog, plog_n,
recall_candidate_count, to_json, validate_memory_type, RecallCandidateParams, RecallParams,
TextSnippetPolicy, DEFAULT_DECAY_EPISODIC, DEFAULT_DECAY_SEMANTIC, DEFAULT_SALIENCE_EPISODIC,
DEFAULT_SALIENCE_SEMANTIC, PROF_CID, RECALL_CALL_ID,
};
impl MemoryPack {
pub(crate) async fn handle_recall(
&self,
token: &NamespaceToken,
params: Value,
_registry: &VerbRegistry,
) -> Result<Value, RuntimeError> {
use std::sync::atomic::Ordering;
let recall_start = Instant::now();
let p: RecallParams = deser(params)?;
let prof = super::common::recall_profile_enabled();
let call_id = if prof {
let id = RECALL_CALL_ID.fetch_add(1, Ordering::Relaxed);
PROF_CID.with(|c| c.set(id));
id
} else {
0
};
let t_total = if prof { Some(Instant::now()) } else { None };
let mut t_stage = if prof { Some(Instant::now()) } else { None };
let query_trimmed = p.query.trim();
if query_trimmed.is_empty() {
return Err(RuntimeError::InvalidInput("query must not be empty".into()));
}
if !crate::scoring::is_meaningful_query(query_trimmed) {
return Err(RuntimeError::InvalidInput(format!(
"query {query_trimmed:?} does not contain enough meaningful content \
(must have at least 2 alphabetic or CJK characters and not consist \
of repeated characters)"
)));
}
if let Some(mt) = &p.memory_type {
validate_memory_type(mt)?;
}
if let Some(ref fs) = p.fusion_strategy {
super::common::parse_fusion_strategy_str(fs)?;
}
let mut cfg = p.effective_config(self.active_config());
if let Some(ref fs) = p.fusion_strategy {
let mut new_strategy = super::common::parse_fusion_strategy_str(fs)?;
if let (
FusionStrategy::Weighted {
weights: ref mut new_w,
},
FusionStrategy::Weighted {
weights: ref existing_w,
},
) = (&mut new_strategy, &cfg.fuse_strategy)
{
*new_w = existing_w.clone();
}
cfg.fuse_strategy = new_strategy;
}
cfg.validate()?;
let effective_min_score: f32 = {
let raw = if let Some(floor) = p.score_floor {
floor as f64
} else {
cfg.min_score
};
normalize_min_score(raw).map_err(RuntimeError::from)?
};
let limit = if let Some(k) = p.top_k {
k.min(crate::scoring::MAX_RECALL_LIMIT)
} else {
p.limit
.map(|v| v as usize)
.unwrap_or(10)
.clamp(1, crate::scoring::MAX_RECALL_LIMIT)
};
let limit_u32 = u32::try_from(limit).unwrap_or(u32::MAX);
let mut scoring_cfg = cfg.scoring.clone().unwrap_or_default();
scoring_cfg.apply_dos_caps();
let cjk_fts_bypass = scoring_cfg.enable_multilingual_routing && contains_cjk(query_trimmed);
let use_multilingual =
scoring_cfg.enable_multilingual_routing && needs_multilingual(query_trimmed);
let candidate_limit =
recall_candidate_count(&cfg, limit_u32).min(scoring_cfg.max_recall_candidates as u32);
if prof {
if let Some(ref t) = t_stage {
plog(call_id, "setup", t.elapsed().as_micros());
}
t_stage = Some(Instant::now());
}
let effective_fts_gather = crate::config::RecallFtsGatherConfig::from_env()
.map_err(|e| RuntimeError::InvalidInput(format!("fts_gather env parse error: {e}")))?
.unwrap_or_else(|| cfg.fts_gather.clone());
let ann_overfetch_max_rounds = cfg
.ann_overfetch_max_rounds
.unwrap_or_else(super::common::ann_overfetch_max_rounds);
let candidates = self
.collect_recall_candidates(
query_trimmed,
token,
RecallCandidateParams {
candidate_limit,
embedding_model: p.embedding_model.as_deref(),
cjk_fts_bypass,
use_multilingual,
scoring_cfg: &scoring_cfg,
snippet_policy: TextSnippetPolicy::Omit,
fts_gather: &effective_fts_gather,
ann_overfetch_max_rounds,
},
)
.await?;
if prof {
if let Some(ref t) = t_stage {
plog_n(
call_id,
"candidates",
t.elapsed().as_micros(),
candidates.text_hits.len()
+ candidates
.vector_hits_per_model
.iter()
.map(|(_, h)| h.len())
.sum::<usize>(),
);
}
t_stage = Some(Instant::now());
}
let actual_multilingual_routed = candidates.multilingual_routed;
let (memory_ids, mut notes_by_id) =
self.load_memory_candidate_notes(token, &candidates).await?;
if prof {
if let Some(ref t) = t_stage {
plog_n(
call_id,
"hydration",
t.elapsed().as_micros(),
notes_by_id.len(),
);
}
t_stage = Some(Instant::now());
}
let raw_vec_scores: HashMap<Uuid, f32> = {
let mut map = HashMap::new();
for (_, hits) in &candidates.vector_hits_per_model {
for h in hits {
let score = h.score.to_f64() as f32;
map.entry(h.subject_id)
.and_modify(|s| {
if score > *s {
*s = score;
}
})
.or_insert(score);
}
}
map
};
let fused = fuse_candidates(&candidates, &memory_ids, &cfg, candidate_limit as usize);
if prof {
if let Some(ref t) = t_stage {
plog_n(call_id, "fusion", t.elapsed().as_micros(), fused.len());
}
t_stage = Some(Instant::now());
}
if fused.is_empty() {
if let Ok(mut state) = self.recall_state.lock() {
on_recall_miss(&mut state);
}
return to_json(&Vec::<Value>::new());
}
let fused_pairs: Vec<(Uuid, f32)> = fused
.iter()
.map(|h| (h.entity_id, h.score.to_f64() as f32))
.collect();
let is_rrf = matches!(&cfg.fuse_strategy, FusionStrategy::Rrf { .. });
let normalized_relevance: HashMap<Uuid, f32> = if is_rrf {
normalize_rrf_scores(fused_pairs, &scoring_cfg)
} else {
normalize_rank_fusion_scores(fused_pairs, &scoring_cfg)
};
let source_by_id: HashMap<Uuid, SearchSource> =
fused.iter().map(|h| (h.entity_id, h.source)).collect();
let now_micros = chrono::Utc::now().timestamp_micros();
let now_millis = now_micros / 1_000;
let entity_names: Vec<String> = p
.entity_names
.as_deref()
.unwrap_or(&[])
.iter()
.map(|s| s.to_lowercase())
.collect();
struct ScoredNote {
id: Uuid,
rank_score: f32,
score: f32,
raw_score: Option<f32>,
breakdown: ScoreBreakdown,
note: khive_storage::note::Note,
resolved_memory_type: String,
effective_salience: f64,
effective_decay_factor: f64,
}
let recall_pipeline = make_pipeline(&cfg);
let mut ranked: Vec<ScoredNote> = Vec::new();
for hit in &fused {
let id = hit.entity_id;
let norm_relevance = match normalized_relevance.get(&id) {
Some(&v) => v,
None => continue,
};
if let Some(&raw) = raw_vec_scores.get(&id) {
if raw < scoring_cfg.min_raw_relevance {
continue;
}
}
let note = match notes_by_id.remove(&id) {
Some(note) => note,
None => continue,
};
let note_memory_type: String = note
.properties
.as_ref()
.and_then(|pr| pr.get("memory_type"))
.and_then(|v| v.as_str())
.unwrap_or("episodic")
.to_owned();
if let Some(mt) = &p.memory_type {
if note_memory_type != mt.as_str() {
continue;
}
}
if let Some(filter_tags) = p.tags.as_ref().filter(|tags| !tags.is_empty()) {
if !note_matches_tags(note.properties.as_ref(), filter_tags, p.tag_mode) {
continue;
}
}
let salience = note.salience.unwrap_or(if note_memory_type == "semantic" {
DEFAULT_SALIENCE_SEMANTIC
} else {
DEFAULT_SALIENCE_EPISODIC
});
let decay_factor = note
.decay_factor
.unwrap_or(if note_memory_type == "semantic" {
DEFAULT_DECAY_SEMANTIC
} else {
DEFAULT_DECAY_EPISODIC
});
if salience < cfg.min_salience {
continue;
}
let rank_score = calculate_score(
&ScoreInput {
salience: salience as f32,
memory_type_str: ¬e_memory_type,
content: ¬e.content,
created_at_millis: note.created_at / 1_000,
decay_factor: decay_factor as f32,
now_millis,
relevance_score: norm_relevance,
entity_names: &entity_names,
},
&scoring_cfg,
);
let age_days_f64 =
((now_micros - note.created_at).max(0) as f64) / (1_000_000.0 * 86_400.0);
let (_, breakdown) = compute_score(
&cfg,
&recall_pipeline,
norm_relevance as f64,
salience,
decay_factor,
age_days_f64,
);
let source = source_by_id.get(&id).copied().unwrap_or(SearchSource::Text);
let final_score = if !cfg.reranker_weights.is_empty() {
let features = RerankFeatures {
relevance: norm_relevance as f64,
salience: breakdown.salience_decayed,
temporal: breakdown.temporal,
text_match: matches!(source, SearchSource::Text | SearchSource::Both),
vector_match: matches!(source, SearchSource::Vector | SearchSource::Both),
};
weighted_rerank(&features, &cfg.reranker_weights) as f32
} else {
rank_score
};
let raw_score_opt = raw_vec_scores.get(&id).copied();
let absolute_relevance = raw_score_opt.unwrap_or(final_score).clamp(0.0, 1.0);
debug_assert!(
absolute_relevance <= 1.0,
"score violates [0,1] contract: {absolute_relevance}"
);
if final_score < effective_min_score {
continue;
}
ranked.push(ScoredNote {
id,
rank_score: final_score,
score: absolute_relevance,
raw_score: raw_score_opt,
breakdown,
note,
resolved_memory_type: note_memory_type,
effective_salience: salience,
effective_decay_factor: decay_factor,
});
}
if prof {
if let Some(ref t) = t_stage {
plog_n(call_id, "scoring", t.elapsed().as_micros(), ranked.len());
}
t_stage = Some(Instant::now());
}
if scoring_cfg.mmr_penalty > 0.0 && scoring_cfg.mmr_prefix_len > 0 {
let prefix_len = scoring_cfg.mmr_prefix_len;
let prefixes: Vec<String> = ranked
.iter()
.map(|sn| sn.note.content.chars().take(prefix_len).collect::<String>())
.collect();
for i in 1..ranked.len() {
for j in 0..i {
if prefixes[i] == prefixes[j] {
ranked[i].rank_score =
(ranked[i].rank_score - scoring_cfg.mmr_penalty).max(0.0);
break;
}
}
}
}
if prof {
if let Some(ref t) = t_stage {
plog_n(call_id, "mmr", t.elapsed().as_micros(), ranked.len());
}
t_stage = Some(Instant::now());
}
if scoring_cfg.enable_supersedes_suppression {
let mut superseded_by_prop: HashSet<Uuid> = HashSet::new();
for sn in &ranked {
if let Some(target_str) = sn
.note
.properties
.as_ref()
.and_then(|pr| pr.get("supersedes"))
.and_then(|v| v.as_str())
{
if let Ok(uid) = target_str.parse::<Uuid>() {
superseded_by_prop.insert(uid);
} else {
let prefix = target_str.to_lowercase();
for sn2 in &ranked {
if sn2.id.as_hyphenated().to_string().starts_with(&prefix) {
superseded_by_prop.insert(sn2.id);
break;
}
}
}
}
}
let graph = self.runtime.graph(token)?;
let candidate_ids: Vec<Uuid> = ranked.iter().map(|sn| sn.id).collect();
let mut superseded_by_edge: HashSet<Uuid> = HashSet::new();
{
let limit = candidate_ids.len().max(1) as u32;
let edges = graph
.query_edges(
EdgeFilter {
target_ids: candidate_ids.clone(),
relations: vec![EdgeRelation::Supersedes],
..EdgeFilter::default()
},
vec![],
PageRequest { limit, offset: 0 },
)
.await?;
for edge in &edges.items {
superseded_by_edge.insert(edge.target_id);
}
}
let superseded_ids: HashSet<Uuid> = superseded_by_prop
.union(&superseded_by_edge)
.copied()
.collect();
if !superseded_ids.is_empty() {
ranked.retain(|sn| !superseded_ids.contains(&sn.id));
}
}
if prof {
if let Some(ref t) = t_stage {
plog_n(call_id, "supersedes", t.elapsed().as_micros(), ranked.len());
}
t_stage = Some(Instant::now());
}
ranked.sort_by(|a, b| {
b.rank_score
.partial_cmp(&a.rank_score)
.unwrap_or(std::cmp::Ordering::Equal)
.then(a.id.cmp(&b.id))
});
ranked.truncate(limit);
let token_budget_chars = scoring_cfg.default_token_budget * scoring_cfg.chars_per_token;
let pre_budget_count = ranked.len();
let mut total_chars = 0usize;
let mut budget_cutoff: Option<usize> = None;
for (i, sn) in ranked.iter().enumerate() {
let entry_chars = sn.note.content.len();
if total_chars + entry_chars > token_budget_chars {
budget_cutoff = Some(i);
break;
}
total_chars += entry_chars;
}
if let Some(cut) = budget_cutoff {
ranked.truncate(cut);
}
let budget_capped = ranked.len() < pre_budget_count;
let is_verbose = cfg.include_breakdown || p.include_breakdown.unwrap_or(false);
let full_content = p.full_content.unwrap_or(true);
const PREVIEW_CHARS: usize = 200;
let results: Vec<Value> = ranked
.into_iter()
.map(|sn| {
let content_out =
if !full_content && sn.note.content.chars().count() > PREVIEW_CHARS {
let preview: String = sn.note.content.chars().take(PREVIEW_CHARS).collect();
format!("{preview}…")
} else {
sn.note.content.clone()
};
let mut result = json!({
"id": sn.id.to_string(),
"score": sn.score,
"rank_score": sn.rank_score,
"raw_score": sn.raw_score,
"content": content_out,
"salience": sn.effective_salience,
"decay_factor": sn.effective_decay_factor,
"memory_type": sn.resolved_memory_type,
"created_at": micros_to_iso(sn.note.created_at),
});
if is_verbose {
result["breakdown"] = json!(sn.breakdown);
}
if actual_multilingual_routed {
result["multilingual_routed"] = json!(true);
}
result
})
.collect();
{
let latency_us = recall_start.elapsed().as_micros() as i64;
let top_id = results.first().and_then(|r| {
r.get("id")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<Uuid>().ok())
});
if let Ok(mut state) = self.recall_state.lock() {
if let Some(tid) = top_id {
on_recall_hit(&mut state, tid, latency_us);
} else {
on_recall_miss(&mut state);
}
}
}
if is_verbose && candidates.vector_hits_per_model.len() > 1 {
let per_model: Vec<Value> = candidates
.vector_hits_per_model
.iter()
.map(|(model, hits)| {
let hits_json: Vec<Value> = hits
.iter()
.map(|h| {
json!({
"id": h.subject_id.to_string(),
"score": h.score.to_f64(),
"rank": h.rank,
})
})
.collect();
json!({ "model": model, "hits": hits_json })
})
.collect();
let truncated_for_budget = if budget_capped {
pre_budget_count - results.len()
} else {
0
};
return to_json(&json!({
"results": results,
"candidates": {
"vector_candidates_per_model": per_model,
},
"budget_capped": budget_capped,
"truncated_for_budget": truncated_for_budget,
}));
}
if prof {
if let Some(ref t) = t_stage {
plog_n(call_id, "serialize", t.elapsed().as_micros(), results.len());
}
if let Some(ref t) = t_total {
plog(call_id, "total", t.elapsed().as_micros());
}
}
to_json(&results)
}
}