Skip to main content

codemem_engine/
scoring.rs

1//! Hybrid scoring for memory recall.
2
3use crate::bm25;
4use chrono::{DateTime, Utc};
5use codemem_core::{MemoryNode, ScoreBreakdown};
6use codemem_storage::graph::GraphEngine;
7
8/// Compute graph strength for a memory node by combining raw graph metrics.
9///
10/// Blends two signal sources:
11/// - **Code neighbors** (`sym:`, `file:`, etc.): PageRank, betweenness, connectivity, edge weight
12/// - **Memory neighbors** (other memories linked via SHARES_THEME, PRECEDED_BY, etc.):
13///   connectivity count and average edge weight
14///
15/// A function with many linked memories ranks higher. A conversational memory
16/// connected to many session peers ranks higher than an isolated one.
17pub fn graph_strength_for_memory(graph: &GraphEngine, memory_id: &str) -> f64 {
18    let metrics = match graph.raw_graph_metrics_for_memory(memory_id) {
19        Some(m) => m,
20        None => return 0.0,
21    };
22
23    // Code-graph component (0.0-1.0): centrality of linked code nodes
24    let code_score = if metrics.code_neighbor_count > 0 {
25        let connectivity = (metrics.code_neighbor_count as f64 / 5.0).min(1.0);
26        let avg_edge_w = (metrics.total_edge_weight / metrics.code_neighbor_count as f64).min(1.0);
27        0.4 * metrics.max_pagerank
28            + 0.3 * metrics.max_betweenness
29            + 0.2 * connectivity
30            + 0.1 * avg_edge_w
31    } else {
32        0.0
33    };
34
35    // Memory-graph component (0.0-1.0): connectivity to other memories
36    let memory_score = if metrics.memory_neighbor_count > 0 {
37        let connectivity = (metrics.memory_neighbor_count as f64 / 10.0).min(1.0);
38        let avg_edge_w =
39            (metrics.memory_edge_weight / metrics.memory_neighbor_count as f64).min(1.0);
40        0.6 * connectivity + 0.4 * avg_edge_w
41    } else {
42        0.0
43    };
44
45    // Blend: when both exist, code neighbors dominate (70/30).
46    // When only one exists, use it fully.
47    let score = match (
48        metrics.code_neighbor_count > 0,
49        metrics.memory_neighbor_count > 0,
50    ) {
51        (true, true) => 0.7 * code_score + 0.3 * memory_score,
52        (true, false) => code_score,
53        (false, true) => memory_score,
54        (false, false) => 0.0,
55    };
56
57    score.min(1.0)
58}
59
60/// Truncate a string to `max` bytes, appending "..." if truncated.
61/// Handles multi-byte UTF-8 safely by finding the nearest char boundary.
62pub fn truncate_content(s: &str, max: usize) -> String {
63    if s.len() <= max {
64        s.to_string()
65    } else {
66        let mut end = max;
67        while end > 0 && !s.is_char_boundary(end) {
68            end -= 1;
69        }
70        format!("{}...", &s[..end])
71    }
72}
73
74/// Compute 9-component hybrid score for a memory against a query.
75/// The `graph` parameter is used to look up edge counts for graph strength scoring.
76/// The `bm25` parameter provides BM25-based token overlap scoring; if the memory
77/// is in the index it uses the indexed score, otherwise falls back to `score_text`.
78/// The `now` parameter makes scoring deterministic and testable by avoiding internal clock reads.
79pub fn compute_score(
80    memory: &MemoryNode,
81    query_tokens: &[&str],
82    vector_similarity: f64,
83    graph: &GraphEngine,
84    bm25: &bm25::Bm25Index,
85    now: DateTime<Utc>,
86) -> ScoreBreakdown {
87    // BM25 token overlap (replaces naive split+intersect)
88    // Use pre-tokenized query tokens to avoid re-tokenizing per document.
89    let token_overlap = if query_tokens.is_empty() {
90        0.0
91    } else {
92        // Try indexed score first (memory already in the BM25 index),
93        // fall back to scoring against raw text for unindexed documents.
94        let indexed_score = bm25.score_with_tokens_str(query_tokens, &memory.id);
95        if indexed_score > 0.0 {
96            indexed_score
97        } else {
98            bm25.score_text_with_tokens_str(query_tokens, &memory.content)
99        }
100    };
101
102    // Temporal: how recently updated (exponential decay over 30 days)
103    let age_hours = (now - memory.updated_at).num_hours().max(0) as f64;
104    let temporal = (-age_hours / (30.0 * 24.0)).exp();
105
106    // Tag matching: fraction of query tokens found in tags.
107    // Per-memory `tags.join().to_lowercase()` is O(tags) which is typically <10 strings,
108    // so allocation is negligible.
109    let tag_matching = if !query_tokens.is_empty() {
110        let tag_str: String = memory.tags.join(" ").to_lowercase();
111        let matches = query_tokens
112            .iter()
113            .filter(|qt| tag_str.contains(**qt))
114            .count();
115        matches as f64 / query_tokens.len() as f64
116    } else {
117        0.0
118    };
119
120    // Recency: based on last access time (decay over 7 days)
121    let access_hours = (now - memory.last_accessed_at).num_hours().max(0) as f64;
122    let recency = (-access_hours / (7.0 * 24.0)).exp();
123
124    // Enhanced graph scoring: bridge memory UUIDs to code-graph centrality.
125    // Memory nodes live in a separate ID space from code nodes (sym:, file:),
126    // so we collect raw metrics from code-graph neighbors and apply the
127    // scoring formula here in the engine.
128    let graph_strength = graph_strength_for_memory(graph, &memory.id);
129
130    ScoreBreakdown {
131        vector_similarity,
132        graph_strength,
133        token_overlap,
134        temporal,
135        tag_matching,
136        importance: memory.importance,
137        confidence: memory.confidence,
138        recency,
139    }
140}
141
142#[cfg(test)]
143#[path = "tests/scoring_tests.rs"]
144mod tests;