1use super::graph_store::GraphStore;
9use crate::graph_rag::GraphRAGError;
10use crate::hybrid::filter_by_score;
11use lc_core::math::cosine_similarity;
12use lc_embeddings::Embeddings;
13use std::collections::{HashMap, HashSet};
14
15pub trait EntityMatcher: Send + Sync {
20 fn find_relevant(&self, query: &str, store: &GraphStore, top_k: usize) -> Vec<String>;
22}
23
24#[derive(Debug, Clone, Copy)]
31enum TermKind {
32 Direct,
34 Synonym,
36 Bigram,
38}
39
40struct Term {
42 text: String,
43 kind: TermKind,
44}
45
46fn normalize_text(s: &str) -> String {
50 s.chars()
51 .map(|c| {
52 let u = c as u32;
53 if (0xFF01..=0xFF5E).contains(&u) {
54 char::from_u32(u - 0xFEE0).unwrap_or(c)
55 } else {
56 c
57 }
58 })
59 .collect()
60}
61
62fn is_cjk(c: char) -> bool {
64 matches!(
65 c,
66 '\u{3400}'..='\u{4DBF}' | '\u{4E00}'..='\u{9FFF}' | '\u{F900}'..='\u{FAFF}'
67 )
68}
69
70fn cjk_bigrams(s: &str) -> Vec<String> {
74 let chars: Vec<char> = s.chars().filter(|c| is_cjk(*c)).collect();
75 if chars.len() < 2 {
76 return Vec::new();
77 }
78 chars.windows(2).map(|w| w.iter().collect()).collect()
79}
80
81pub struct KeywordMatcher {
99 pub name_weight: usize,
101 pub type_weight: usize,
103 pub desc_weight: usize,
105 pub synonyms: HashMap<String, Vec<String>>,
109 pub use_tfidf: bool,
112 pub synonym_weight: f64,
114 pub cjk_bigram_weight: f64,
116}
117
118impl Default for KeywordMatcher {
119 fn default() -> Self {
120 Self {
121 name_weight: 3,
122 type_weight: 2,
123 desc_weight: 1,
124 synonyms: HashMap::new(),
125 use_tfidf: true,
126 synonym_weight: 0.7,
127 cjk_bigram_weight: 0.5,
128 }
129 }
130}
131
132impl KeywordMatcher {
133 pub fn new() -> Self {
135 Self::default()
136 }
137
138 pub fn with_synonyms(mut self, synonyms: HashMap<String, Vec<String>>) -> Self {
140 self.synonyms = synonyms;
141 self
142 }
143
144 pub fn with_tfidf(mut self, enabled: bool) -> Self {
146 self.use_tfidf = enabled;
147 self
148 }
149
150 fn build_terms(&self, query: &str) -> Vec<Term> {
153 let normalized = normalize_text(query).to_lowercase();
154 let tokens: Vec<String> = normalized
155 .split_whitespace()
156 .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()).to_string())
157 .filter(|w| !w.is_empty())
158 .collect();
159
160 let mut terms: Vec<Term> = Vec::new();
161 let mut seen: HashSet<String> = HashSet::new();
162 for tok in tokens {
163 Self::push_term(&mut terms, &mut seen, tok.clone(), TermKind::Direct);
166
167 let syns = self
169 .synonyms
170 .iter()
171 .find(|(k, _)| normalize_text(k).to_lowercase() == tok)
172 .map(|(_, v)| v);
173 if let Some(syns) = syns {
174 for syn in syns {
175 Self::push_term(&mut terms, &mut seen, syn.clone(), TermKind::Synonym);
176 }
177 }
178
179 if tok.chars().any(is_cjk) {
181 for bg in cjk_bigrams(&tok) {
182 Self::push_term(&mut terms, &mut seen, bg, TermKind::Bigram);
183 }
184 }
185 }
186 terms
187 }
188
189 fn push_term(terms: &mut Vec<Term>, seen: &mut HashSet<String>, text: String, kind: TermKind) {
190 if seen.insert(text.clone()) {
191 terms.push(Term { text, kind });
192 }
193 }
194
195 fn compute_idf(&self, terms: &[Term], store: &GraphStore) -> HashMap<String, f64> {
201 let n = store.all_entities().len() as f64;
202 let mut df: HashMap<String, usize> = HashMap::new();
203 for term in terms {
204 df.entry(term.text.clone()).or_insert(0);
205 }
206 for entity in store.all_entities().values() {
207 let name = normalize_text(&entity.name).to_lowercase();
208 let desc = normalize_text(&entity.description).to_lowercase();
209 let typ = normalize_text(&entity.entity_type).to_lowercase();
210 for term in terms {
211 if name.contains(&term.text)
212 || desc.contains(&term.text)
213 || typ.contains(&term.text)
214 {
215 if let Some(c) = df.get_mut(&term.text) {
216 *c += 1;
217 }
218 }
219 }
220 }
221 let mut idf = HashMap::with_capacity(terms.len());
222 for (text, count) in &df {
223 let w = ((n + 1.0) / (*count as f64 + 1.0)).ln() + 1.0;
224 idf.insert(text.clone(), w);
225 }
226 idf
227 }
228
229 fn match_score(
232 &self,
233 term: &Term,
234 name: &str,
235 type_name: &str,
236 desc: &str,
237 idf_weight: f64,
238 ) -> f64 {
239 let mut score = 0.0f64;
240 if name.contains(&term.text) {
241 score += self.name_weight as f64 * idf_weight;
242 }
243 if type_name.contains(&term.text) {
244 score += self.type_weight as f64 * idf_weight;
245 }
246 if desc.contains(&term.text) {
247 score += self.desc_weight as f64 * idf_weight;
248 }
249 match term.kind {
250 TermKind::Direct => score,
251 TermKind::Synonym => score * self.synonym_weight,
252 TermKind::Bigram => score * self.cjk_bigram_weight,
253 }
254 }
255}
256
257impl EntityMatcher for KeywordMatcher {
258 fn find_relevant(&self, query: &str, store: &GraphStore, top_k: usize) -> Vec<String> {
259 let terms = self.build_terms(query);
260 if terms.is_empty() {
261 return Vec::new();
262 }
263
264 let idf = if self.use_tfidf {
267 Some(self.compute_idf(&terms, store))
268 } else {
269 None
270 };
271
272 let mut scored: Vec<(String, f64)> = Vec::new();
273
274 for (id, entity) in store.all_entities() {
275 let name_lower = normalize_text(&entity.name).to_lowercase();
276 let desc_lower = normalize_text(&entity.description).to_lowercase();
277 let type_lower = normalize_text(&entity.entity_type).to_lowercase();
278
279 let mut score = 0.0f64;
280 for term in &terms {
281 let idf_weight = idf
282 .as_ref()
283 .and_then(|m| m.get(&term.text))
284 .copied()
285 .unwrap_or(1.0);
286 score += self.match_score(term, &name_lower, &type_lower, &desc_lower, idf_weight);
287 }
288
289 if score > 0.0 {
290 scored.push((id.clone(), score));
291 }
292 }
293
294 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
295 scored.into_iter().take(top_k).map(|(id, _)| id).collect()
296 }
297}
298
299pub struct EmbeddingMatcher<E: Embeddings> {
309 embeddings: E,
310 cache: std::sync::Mutex<HashMap<String, Vec<f32>>>,
312 min_score: f64,
314}
315
316impl<E: Embeddings> EmbeddingMatcher<E> {
317 pub fn new(embeddings: E) -> Self {
319 Self {
320 embeddings,
321 cache: std::sync::Mutex::new(HashMap::new()),
322 min_score: 0.0,
323 }
324 }
325
326 pub fn with_min_score(mut self, min_score: f64) -> Self {
329 self.min_score = min_score;
330 self
331 }
332
333 async fn get_entity_embedding(&self, entity_id: &str, entity_text: &str) -> Option<Vec<f32>> {
335 {
337 let cache = self.cache.lock().unwrap_or_else(|e| e.into_inner());
338 if let Some(vec) = cache.get(entity_id) {
339 return Some(vec.clone());
340 }
341 }
342
343 match self.embeddings.embed_query(entity_text).await {
345 Ok(vec) => {
346 self.cache
347 .lock()
348 .unwrap_or_else(|e| e.into_inner())
349 .insert(entity_id.to_string(), vec.clone());
350 Some(vec)
351 }
352 Err(e) => {
353 log::warn!(
356 "entity `{}` embedding failed; excluded from graph matching: {}",
357 entity_id,
358 e
359 );
360 None
361 }
362 }
363 }
364}
365
366impl<E: Embeddings + 'static> EntityMatcher for EmbeddingMatcher<E> {
367 fn find_relevant(&self, query: &str, _store: &GraphStore, _top_k: usize) -> Vec<String> {
368 log::warn!(
378 "EmbeddingMatcher::find_relevant (sync) cannot run embedding matching and no longer \
379 silently falls back to keyword matching; returning empty results for query '{}'. \
380 Use find_relevant_async instead, or configure KeywordMatcher for sync matching.",
381 query
382 );
383 Vec::new()
384 }
385}
386
387impl<E: Embeddings + 'static> EmbeddingMatcher<E> {
388 pub async fn find_relevant_async(
398 &self,
399 query: &str,
400 store: &GraphStore,
401 top_k: usize,
402 ) -> Result<Vec<String>, GraphRAGError> {
403 let query_vec = self.embeddings.embed_query(query).await.map_err(|e| {
404 GraphRAGError::QueryError(format!("EmbeddingMatcher: query embedding failed: {}", e))
405 })?;
406
407 let mut scored: Vec<(String, f64)> = Vec::new();
408
409 for (id, entity) in store.all_entities() {
410 let entity_text = format!(
411 "{} {} {}",
412 entity.name, entity.entity_type, entity.description
413 );
414
415 if let Some(entity_vec) = self.get_entity_embedding(id, &entity_text).await {
416 match cosine_similarity(&query_vec, &entity_vec) {
417 Ok(score) => {
418 scored.push((id.clone(), score as f64));
419 }
420 Err(lc_core::math::MathError::LengthMismatch(a, b)) => {
423 return Err(GraphRAGError::QueryError(format!(
424 "EmbeddingMatcher: vector dimension mismatch {} vs {} (embedding model changed?)",
425 a, b
426 )));
427 }
428 Err(other) => {
431 return Err(GraphRAGError::QueryError(format!(
432 "EmbeddingMatcher: similarity computation failed: {}",
433 other
434 )));
435 }
436 }
437 }
438 }
439
440 let mut scored = filter_by_score(scored, self.min_score);
441 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
442 Ok(scored.into_iter().take(top_k).map(|(id, _)| id).collect())
443 }
444}
445
446#[cfg(test)]
447mod tests {
448 use super::*;
449 use crate::graph_rag::graph_store::{Entity, Relation};
450 use lc_embeddings::MockEmbeddings;
451
452 fn make_test_store() -> GraphStore {
453 let mut store = GraphStore::new();
454 store.add_entity(Entity {
455 id: "e1".into(),
456 name: "Rust".into(),
457 entity_type: "Technology".into(),
458 description: "A systems programming language".into(),
459 });
460 store.add_entity(Entity {
461 id: "e2".into(),
462 name: "Python".into(),
463 entity_type: "Technology".into(),
464 description: "A scripting language".into(),
465 });
466 store.add_entity(Entity {
467 id: "e3".into(),
468 name: "Alice".into(),
469 entity_type: "Person".into(),
470 description: "A developer who uses Rust".into(),
471 });
472 store.add_entity(Entity {
473 id: "e4".into(),
474 name: "Tokio".into(),
475 entity_type: "Library".into(),
476 description: "An async runtime for Rust".into(),
477 });
478 store.add_relation(Relation {
479 source: "e3".into(),
480 target: "e1".into(),
481 relation_type: "uses".into(),
482 description: "Alice uses Rust".into(),
483 doc_id: None,
484 });
485 store
486 }
487
488 #[test]
489 fn test_keyword_matcher_basic() {
490 let store = make_test_store();
491 let matcher = KeywordMatcher::new();
492 let results = matcher.find_relevant("Rust programming", &store, 10);
493 assert!(!results.is_empty());
494 assert_eq!(results[0], "e1");
496 }
497
498 #[test]
499 fn test_keyword_matcher_top_k() {
500 let store = make_test_store();
501 let matcher = KeywordMatcher::new();
502 let results = matcher.find_relevant("Technology", &store, 1);
503 assert_eq!(results.len(), 1);
504 }
505
506 #[test]
507 fn test_keyword_matcher_no_match() {
508 let store = make_test_store();
509 let matcher = KeywordMatcher::new();
510 let results = matcher.find_relevant("cooking recipe", &store, 10);
511 assert!(results.is_empty());
512 }
513
514 #[test]
519 fn test_embedding_matcher_sync_returns_empty_not_keyword_fallback() {
520 let store = make_test_store();
521 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
522 let results = matcher.find_relevant("Rust", &store, 10);
523 assert!(
524 results.is_empty(),
525 "sync find_relevant must NOT silently fall back to keyword matching"
526 );
527 }
528
529 #[tokio::test]
536 async fn test_embedding_matcher_async_still_works() {
537 let mut store = GraphStore::new();
538 store.add_entity(Entity {
539 id: "e1".into(),
540 name: "Rust".into(),
541 entity_type: "Technology".into(),
542 description: "A systems programming language".into(),
543 });
544 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
545 let query = "Rust Technology A systems programming language";
546 let results = matcher
547 .find_relevant_async(query, &store, 10)
548 .await
549 .unwrap();
550 assert_eq!(results, vec!["e1".to_string()]);
551 }
552
553 #[test]
554 fn test_keyword_matcher_custom_weights() {
555 let store = make_test_store();
556 let matcher = KeywordMatcher {
557 name_weight: 10,
558 type_weight: 1,
559 desc_weight: 0,
560 ..Default::default()
561 };
562 let results = matcher.find_relevant("Rust", &store, 10);
563 assert!(!results.is_empty());
564 assert_eq!(results[0], "e1");
565 }
566
567 #[test]
570 fn test_keyword_matcher_synonym_expansion() {
571 let mut store = GraphStore::new();
572 store.add_entity(Entity {
573 id: "e1".into(),
574 name: "PostgreSQL".into(),
575 entity_type: "Database".into(),
576 description: "relational database".into(),
577 });
578 store.add_entity(Entity {
579 id: "e2".into(),
580 name: "机器学习".into(),
581 entity_type: "Technology".into(),
582 description: "AI 领域".into(),
583 });
584
585 let synonyms: HashMap<String, Vec<String>> =
586 HashMap::from([("数据库".to_string(), vec!["database".to_string()])]);
587 let matcher = KeywordMatcher::new().with_synonyms(synonyms);
588
589 let results = matcher.find_relevant("数据库", &store, 10);
592 assert!(
593 results.contains(&"e1".to_string()),
594 "同义词 'database' 应能召回 PostgreSQL(e1)"
595 );
596 }
597
598 #[test]
601 fn test_keyword_matcher_fullwidth_normalization() {
602 let mut store = GraphStore::new();
603 store.add_entity(Entity {
604 id: "e1".into(),
605 name: "Rust".into(),
606 entity_type: "Technology".into(),
607 description: "systems language".into(),
608 });
609
610 let matcher = KeywordMatcher::new();
611 let results = matcher.find_relevant("Rust", &store, 10);
614 assert_eq!(results, vec!["e1".to_string()]);
615 }
616
617 #[test]
620 fn test_keyword_matcher_cjk_bigram_recall() {
621 let mut store = GraphStore::new();
622 store.add_entity(Entity {
623 id: "e1".into(),
624 name: "机器学习".into(),
625 entity_type: "Technology".into(),
626 description: "AI 领域".into(),
627 });
628
629 let matcher = KeywordMatcher::new();
630 let results = matcher.find_relevant("机器学习算法", &store, 10);
632 assert!(
633 results.contains(&"e1".to_string()),
634 "CJK 二元组应能召回 '机器学习' 实体"
635 );
636 }
637
638 #[test]
640 fn test_keyword_matcher_tfidf_common_lower_than_rare() {
641 let mut store = GraphStore::new();
642 for name in ["Rust", "Python", "Ruby"] {
643 store.add_entity(Entity {
644 id: name.to_lowercase(),
645 name: name.to_string(),
646 entity_type: "Technology".to_string(),
647 description: String::new(),
648 });
649 }
650
651 let matcher = KeywordMatcher::new();
652 let terms = matcher.build_terms("rust technology");
653 let idf = matcher.compute_idf(&terms, &store);
654
655 let tech_idf = idf["technology"];
657 let rust_idf = idf["rust"];
658 assert!(
659 rust_idf > tech_idf,
660 "稀有词 'rust' 的 IDF ({:.3}) 应高于常见词 'technology' ({:.3})",
661 rust_idf,
662 tech_idf
663 );
664 }
665
666 #[test]
668 fn test_keyword_matcher_tfidf_breaks_fixed_weight_tie() {
669 let mut store = GraphStore::new();
670 store.add_entity(Entity {
671 id: "e1".into(),
672 name: "data".into(),
673 entity_type: "Technology".into(),
674 description: "machine learning".into(),
675 });
676 store.add_entity(Entity {
677 id: "e2".into(),
678 name: "learning".into(),
679 entity_type: "Technology".into(),
680 description: "data".into(),
681 });
682 store.add_entity(Entity {
683 id: "e3".into(),
684 name: "extra".into(),
685 entity_type: "Technology".into(),
686 description: "data warehouse".into(),
687 });
688
689 let matcher = KeywordMatcher::new();
690 let results = matcher.find_relevant("data learning", &store, 10);
693 assert_eq!(results[0], "e2");
694 }
695
696 #[test]
699 fn test_cosine_similarity_identical() {
700 let v = vec![1.0, 0.0, 0.0];
701 let sim = cosine_similarity(&v, &v).unwrap();
702 assert!((sim - 1.0).abs() < 0.001);
703 }
704
705 #[test]
706 fn test_cosine_similarity_orthogonal() {
707 let a = vec![1.0, 0.0];
708 let b = vec![0.0, 1.0];
709 let sim = cosine_similarity(&a, &b).unwrap();
710 assert!((sim - 0.0).abs() < 0.001);
711 }
712
713 #[test]
714 fn test_cosine_similarity_opposite() {
715 let a = vec![1.0, 0.0];
716 let b = vec![-1.0, 0.0];
717 let sim = cosine_similarity(&a, &b).unwrap();
718 assert!((sim - (-1.0)).abs() < 0.001);
719 }
720
721 #[test]
722 fn test_cosine_similarity_zero_norm() {
723 let sim = cosine_similarity(&[], &[]).unwrap();
725 assert_eq!(sim, 0.0);
726 }
727
728 #[test]
731 fn test_cosine_similarity_different_lengths_errors() {
732 let a = vec![1.0];
733 let b = vec![1.0, 2.0];
734 assert!(cosine_similarity(&a, &b).is_err());
735 }
736}