relay-knowledge 1.1.9

Graph-database-based knowledge graph project.
Documentation
use std::collections::BTreeMap;

use crate::domain::{RECIPROCAL_RANK_FUSION_K, RankingSignal, RetrievalHit, RetrieverSource};

use super::ScoredHit;

pub(super) fn merge_ranked(
    candidates: &mut BTreeMap<String, Candidate>,
    hits: Vec<ScoredHit>,
    fallback_source: RetrieverSource,
    explanation: &'static str,
) {
    for (index, scored) in hits.into_iter().enumerate() {
        let rank = index + 1;
        let rrf_score = 1.0 / (RECIPROCAL_RANK_FUSION_K + rank as f64);
        let source = match scored.source {
            RetrieverSource::CodeGraph => RetrieverSource::CodeGraph,
            _ => fallback_source,
        };
        let explanation_text = scored
            .explanation
            .unwrap_or_else(|| format!("{explanation}; modality={}", scored.modality));
        let candidate = candidates
            .entry(scored.key)
            .and_modify(|candidate| candidate.merge_hit(&scored.hit))
            .or_insert_with(|| Candidate::new(scored.hit));
        if !candidate.hit.retriever_sources.contains(&source) {
            candidate.hit.retriever_sources.push(source);
        }
        candidate.hit.ranking.push(RankingSignal {
            source,
            rank,
            score: scored.source_score,
            explanation: explanation_text,
        });
        candidate.rrf_score += rrf_score;
    }
}

pub(super) struct Candidate {
    hit: RetrievalHit,
    rrf_score: f64,
}

impl Candidate {
    fn new(hit: RetrievalHit) -> Self {
        Self {
            hit,
            rrf_score: 0.0,
        }
    }

    pub(super) fn into_hit(mut self) -> RetrievalHit {
        self.hit.score = self.rrf_score;
        self.hit
    }

    fn merge_hit(&mut self, hit: &RetrievalHit) {
        if !hit.content.is_empty() && !self.hit.content.contains(&hit.content) {
            if !self.hit.content.is_empty() {
                self.hit.content.push_str("\n\n");
            }
            self.hit.content.push_str(&hit.content);
        }
        if self.hit.source_path.is_none() {
            self.hit.source_path = hit.source_path.clone();
        }
        if self.hit.source_span.is_none() {
            self.hit.source_span = hit.source_span;
        }
        if self.hit.code_artifact.is_none() {
            self.hit.code_artifact = hit.code_artifact.clone();
        }
        for label in &hit.entity_labels {
            if !self.hit.entity_labels.contains(label) {
                self.hit.entity_labels.push(label.clone());
            }
        }
        for entity in &hit.entities {
            if !self
                .hit
                .entities
                .iter()
                .any(|existing| existing.id == entity.id)
            {
                self.hit.entities.push(entity.clone());
            }
        }
        for fact in &hit.graph_facts {
            if !self
                .hit
                .graph_facts
                .iter()
                .any(|existing| existing.fact_id == fact.fact_id && existing.kind == fact.kind)
            {
                self.hit.graph_facts.push(fact.clone());
            }
        }
    }
}