use std::str::FromStr;
use uuid::Uuid;
use khive_storage::types::PageRequest;
use khive_storage::EntityFilter;
use crate::error::{RuntimeError, RuntimeResult};
use crate::operations::Resolved;
use crate::reference_ring::ReferenceRing;
use crate::runtime::{KhiveRuntime, NamespaceToken};
#[derive(Clone, Debug, PartialEq)]
pub struct ReferenceCandidate {
pub id: Uuid,
pub name: Option<String>,
pub score: f64,
}
#[derive(Clone, Debug, PartialEq)]
pub enum ReferenceResolution {
Resolved { id: Uuid, confidence: f64 },
Ambiguous { candidates: Vec<ReferenceCandidate> },
NotFound,
}
const RING_EXACT_CONFIDENCE: f64 = 0.95;
const RING_SUBSTRING_CONFIDENCE: f64 = 0.7;
const RING_AUTO_RESOLVE_CONFIDENCE: f64 = 0.7;
const EXACT_NAME_CONFIDENCE: f64 = 0.98;
const SEARCH_MARGIN_RATIO: f64 = 2.0;
const SEARCH_VECTOR_SIMILARITY_FLOOR: f64 = 0.3;
const SEARCH_RESOLVED_CONFIDENCE: f64 = 0.6;
const STAGE4_MIN_SEARCH_LIMIT: u32 = 20;
pub async fn resolve_reference(
runtime: &KhiveRuntime,
ring: &ReferenceRing,
token: &NamespaceToken,
nl_ref: &str,
limit: u32,
entity_kind: Option<&str>,
) -> RuntimeResult<ReferenceResolution> {
resolve_reference_with_entity_type(runtime, ring, token, nl_ref, limit, entity_kind, None).await
}
#[allow(clippy::too_many_arguments)]
pub async fn resolve_reference_with_entity_type(
runtime: &KhiveRuntime,
ring: &ReferenceRing,
token: &NamespaceToken,
nl_ref: &str,
limit: u32,
entity_kind: Option<&str>,
entity_type: Option<&str>,
) -> RuntimeResult<ReferenceResolution> {
let trimmed = nl_ref.trim();
if trimmed.is_empty() {
return Ok(ReferenceResolution::NotFound);
}
if let Ok(uuid) = Uuid::from_str(trimmed) {
return match runtime.resolve_by_id(token, uuid).await? {
Some(Resolved::Entity(_)) => Ok(ReferenceResolution::Resolved {
id: uuid,
confidence: 1.0,
}),
Some(_) | None => Ok(ReferenceResolution::NotFound),
};
}
if is_hex_prefix(trimmed) {
return match runtime.resolve_prefix_unfiltered(trimmed).await {
Ok(Some(uuid)) => match runtime.resolve_by_id(token, uuid).await? {
Some(Resolved::Entity(_)) => Ok(ReferenceResolution::Resolved {
id: uuid,
confidence: 1.0,
}),
Some(_) | None => Ok(ReferenceResolution::NotFound),
},
Ok(None) => Ok(ReferenceResolution::NotFound),
Err(RuntimeError::AmbiguousPrefix { matches, .. }) => {
let mut entity_matches = Vec::with_capacity(matches.len());
for id in matches {
if matches!(
runtime.resolve_by_id(token, id).await?,
Some(Resolved::Entity(_))
) {
entity_matches.push(id);
}
}
match entity_matches.len() {
0 => Ok(ReferenceResolution::NotFound),
1 => Ok(ReferenceResolution::Resolved {
id: entity_matches[0],
confidence: 1.0,
}),
_ => Ok(ReferenceResolution::Ambiguous {
candidates: entity_matches
.into_iter()
.map(|id| ReferenceCandidate {
id,
name: None,
score: 1.0,
})
.collect(),
}),
}
}
Err(e) => Err(e),
};
}
let actor = token.actor();
let actor_key = format!("{}:{}", actor.kind, actor.id);
let ring_entries = ring.snapshot(token.namespace().as_str(), &actor_key);
let needle = trimmed.to_ascii_lowercase();
let exact: Vec<ReferenceCandidate> = ring_entries
.iter()
.filter(|e| {
e.name
.as_deref()
.is_some_and(|n| n.to_ascii_lowercase() == needle)
})
.map(|e| ReferenceCandidate {
id: e.id,
name: e.name.clone(),
score: RING_EXACT_CONFIDENCE,
})
.collect();
if let Some(resolution) = resolve_from_candidates(exact) {
return Ok(resolution);
}
let substring: Vec<ReferenceCandidate> = ring_entries
.iter()
.filter(|e| {
e.name.as_deref().is_some_and(|n| {
let n_lower = n.to_ascii_lowercase();
n_lower.contains(&needle) || needle.contains(&n_lower)
})
})
.map(|e| ReferenceCandidate {
id: e.id,
name: e.name.clone(),
score: RING_SUBSTRING_CONFIDENCE,
})
.collect();
if let Some(resolution) = resolve_from_candidates(substring) {
return Ok(resolution);
}
if let Some(resolution) =
exact_name_match(runtime, token, trimmed, entity_kind, entity_type).await?
{
return Ok(resolution);
}
let candidate_limit = limit.max(1);
let search_limit = candidate_limit.max(STAGE4_MIN_SEARCH_LIMIT);
let hits = runtime
.hybrid_search_with_vector_similarity_floor(
token,
trimmed,
None,
search_limit,
entity_kind,
entity_type,
&[],
None,
SEARCH_VECTOR_SIMILARITY_FLOOR,
)
.await?;
let candidates: Vec<ReferenceCandidate> = hits
.into_iter()
.map(|h| ReferenceCandidate {
id: h.entity_id,
name: h.title,
score: h.score.to_f64(),
})
.collect();
match candidates.len() {
0 => Ok(ReferenceResolution::NotFound),
1 => Ok(ReferenceResolution::Resolved {
id: candidates[0].id,
confidence: SEARCH_RESOLVED_CONFIDENCE,
}),
_ => {
let top_score = candidates[0].score;
let second_score = candidates[1].score;
let decisive =
second_score <= f64::EPSILON || top_score / second_score >= SEARCH_MARGIN_RATIO;
if decisive {
Ok(ReferenceResolution::Resolved {
id: candidates[0].id,
confidence: SEARCH_RESOLVED_CONFIDENCE,
})
} else {
let mut candidates = candidates;
candidates.truncate(candidate_limit as usize);
Ok(ReferenceResolution::Ambiguous { candidates })
}
}
}
}
fn resolve_from_candidates(candidates: Vec<ReferenceCandidate>) -> Option<ReferenceResolution> {
match candidates.len() {
0 => None,
1 => {
let top = &candidates[0];
Some(if top.score >= RING_AUTO_RESOLVE_CONFIDENCE {
ReferenceResolution::Resolved {
id: top.id,
confidence: top.score,
}
} else {
ReferenceResolution::Ambiguous { candidates }
})
}
_ => Some(ReferenceResolution::Ambiguous { candidates }),
}
}
async fn exact_name_match(
runtime: &KhiveRuntime,
token: &NamespaceToken,
name: &str,
entity_kind: Option<&str>,
entity_type: Option<&str>,
) -> RuntimeResult<Option<ReferenceResolution>> {
let mut filter = EntityFilter {
name_exact: Some(name.to_string()),
kinds: entity_kind.map(|k| vec![k.to_string()]).unwrap_or_default(),
..EntityFilter::default()
};
if let (Some(kind), Some(entity_type)) = (entity_kind, entity_type) {
filter
.entity_types_by_kind
.insert(kind.to_string(), vec![entity_type.to_string()]);
}
let page = runtime
.entities(token)?
.query_entities(
token.namespace().as_str(),
filter,
PageRequest {
offset: 0,
limit: 10,
},
)
.await
.map_err(RuntimeError::Storage)?;
let total = page.total.unwrap_or(page.items.len() as u64);
let exact: Vec<ReferenceCandidate> = page
.items
.into_iter()
.map(|e| ReferenceCandidate {
id: e.id,
name: Some(e.name),
score: EXACT_NAME_CONFIDENCE,
})
.collect();
Ok(match total {
0 => None,
1 => exact
.into_iter()
.next()
.map(|top| ReferenceResolution::Resolved {
id: top.id,
confidence: EXACT_NAME_CONFIDENCE,
}),
_ => Some(ReferenceResolution::Ambiguous { candidates: exact }),
})
}
fn is_hex_prefix(s: &str) -> bool {
s.len() >= 8 && s.chars().all(|c| c.is_ascii_hexdigit())
}
#[cfg(test)]
#[path = "reference_resolution_tests.rs"]
mod tests;