cognee-search 0.1.3

Context retrieval (search) over the cognee knowledge graph and vector store.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

/// Blend a similarity score with feedback weight.
///
/// Rust uses cosine similarity (higher is better), unlike Python which uses cosine distance.
/// Blending formula adapted for similarity space [-1, 1]:
///   normalized_sim = (score + 1.0) / 2.0     -- map [-1,1] to [0,1]
///   blended = (1 - fi) * normalized_sim + fi * feedback_weight
///   result  = blended * 2.0 - 1.0             -- map back to [-1,1]
fn effective_score(score: f32, feedback_weight: f32, feedback_influence: f32) -> f32 {
    if feedback_influence <= 0.0 {
        return score;
    }
    let fw = feedback_weight.clamp(0.0, 1.0);
    let normalized_sim = (score + 1.0) / 2.0;
    let blended = (1.0 - feedback_influence) * normalized_sim + feedback_influence * fw;
    blended * 2.0 - 1.0
}

/// Computes the total distance for a triplet (source_node, edge, target_node).
///
/// Each component is a cosine distance (lower = better). The total distance is
/// the sum of all three, matching Python's `_calculate_query_top_triplet_importances`.
///
/// When `feedback_influence` is non-zero, each node's similarity score is blended
/// with its `feedback_weight` before summing.
pub fn rank_edge_score(
    source_distance: f32,
    target_distance: f32,
    edge_distance: f32,
    feedback_influence: f32,
    source_feedback_weight: f32,
    target_feedback_weight: f32,
) -> f32 {
    let s = effective_score(source_distance, source_feedback_weight, feedback_influence);
    let t = effective_score(target_distance, target_feedback_weight, feedback_influence);
    s + t + edge_distance
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sums_three_distance_components() {
        let score = rank_edge_score(0.1, 0.2, 0.3, 0.0, 0.5, 0.5);
        assert!((score - 0.6).abs() < 1e-6, "expected 0.6, got {score}");
    }

    #[test]
    fn unmatched_components_use_penalty() {
        // When all three components equal the default penalty (6.5), total is 19.5
        let score = rank_edge_score(6.5, 6.5, 6.5, 0.0, 0.5, 0.5);
        assert!((score - 19.5).abs() < 1e-6, "expected 19.5, got {score}");
    }

    #[test]
    fn effective_score_zero_influence_returns_input() {
        // with feedback_influence = 0.0, output == input regardless of feedback_weight
        assert_eq!(effective_score(0.5, 0.8, 0.0), 0.5);
        assert_eq!(effective_score(-0.3, 0.2, 0.0), -0.3);
    }

    #[test]
    fn effective_score_full_influence_returns_feedback_based() {
        // with feedback_influence = 1.0, output is purely based on feedback_weight
        // normalized_sim irrelevant, blended = fw, result = fw * 2 - 1
        let result = effective_score(0.0, 1.0, 1.0); // fw=1.0 -> blended=1.0 -> result=1.0
        assert!((result - 1.0).abs() < 1e-6);
        let result2 = effective_score(0.0, 0.0, 1.0); // fw=0.0 -> blended=0.0 -> result=-1.0
        assert!((result2 - (-1.0)).abs() < 1e-6);
    }
}