use std::collections::{hash_map::Entry, HashMap};
use khive_fusion::FusionStrategy;
use khive_score::DeterministicScore;
use uuid::Uuid;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RankScoreKind {
Rrf,
Vector,
Keyword,
Weighted,
Union,
}
impl RankScoreKind {
pub fn of(strategy: &FusionStrategy) -> Result<Self, (&str, &serde_json::Value)> {
match strategy {
FusionStrategy::Rrf { .. } | FusionStrategy::WeightedRrf { .. } => Ok(Self::Rrf),
FusionStrategy::VectorOnly => Ok(Self::Vector),
FusionStrategy::KeywordOnly => Ok(Self::Keyword),
FusionStrategy::Weighted { .. } => Ok(Self::Weighted),
FusionStrategy::Union => Ok(Self::Union),
FusionStrategy::Custom { name, params } => Err((name, params)),
}
}
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Rrf => "rrf",
Self::Vector => "vector",
Self::Keyword => "keyword",
Self::Weighted => "weighted",
Self::Union => "union",
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct SearchSignals {
pub vector_similarity: Option<DeterministicScore>,
pub keyword_score: Option<DeterministicScore>,
}
#[derive(Clone, Debug)]
pub struct SearchHit {
pub entity_id: Uuid,
pub score: DeterministicScore,
pub rank_score_kind: RankScoreKind,
pub signals: SearchSignals,
pub source: SearchSource,
pub title: Option<String>,
pub snippet: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SignalMerge {
FirstPresent,
Maximum,
}
pub fn merge_hit_metadata(
metadata: &mut HashMap<Uuid, SearchHit>,
hit: SearchHit,
signals: SignalMerge,
) {
match metadata.entry(hit.entity_id) {
Entry::Occupied(mut entry) => {
let existing = entry.get_mut();
existing.source = existing.source.union(hit.source);
existing.signals.vector_similarity = match signals {
SignalMerge::FirstPresent => existing
.signals
.vector_similarity
.or(hit.signals.vector_similarity),
SignalMerge::Maximum => existing
.signals
.vector_similarity
.max(hit.signals.vector_similarity),
};
existing.signals.keyword_score = match signals {
SignalMerge::FirstPresent => {
existing.signals.keyword_score.or(hit.signals.keyword_score)
}
SignalMerge::Maximum => existing
.signals
.keyword_score
.max(hit.signals.keyword_score),
};
if existing.title.is_none() {
existing.title = hit.title;
}
if existing.snippet.is_none() {
existing.snippet = hit.snippet;
}
}
Entry::Vacant(entry) => {
entry.insert(hit);
}
}
}
#[derive(Clone, Debug)]
pub struct HybridSearchOutcome {
pub hits: Vec<SearchHit>,
pub vector_error: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SearchSource {
Vector,
Text,
Both,
}
impl SearchSource {
#[must_use]
pub const fn union(self, other: Self) -> Self {
match (self, other) {
(Self::Text, Self::Text) => Self::Text,
(Self::Vector, Self::Vector) => Self::Vector,
_ => Self::Both,
}
}
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Vector => "vector",
Self::Text => "text",
Self::Both => "both",
}
}
}