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)]
30enum TermKind {
31 Direct,
33 Synonym,
35 Bigram,
37}
38
39struct Term {
41 text: String,
42 kind: TermKind,
43}
44
45fn normalize_text(s: &str) -> String {
48 s.chars()
49 .map(|c| {
50 let u = c as u32;
51 if (0xFF01..=0xFF5E).contains(&u) {
52 char::from_u32(u - 0xFEE0).unwrap_or(c)
53 } else {
54 c
55 }
56 })
57 .collect()
58}
59
60fn is_cjk(c: char) -> bool {
62 matches!(
63 c,
64 '\u{3400}'..='\u{4DBF}' | '\u{4E00}'..='\u{9FFF}' | '\u{F900}'..='\u{FAFF}'
65 )
66}
67
68fn cjk_bigrams(s: &str) -> Vec<String> {
71 let chars: Vec<char> = s.chars().filter(|c| is_cjk(*c)).collect();
72 if chars.len() < 2 {
73 return Vec::new();
74 }
75 chars.windows(2).map(|w| w.iter().collect()).collect()
76}
77
78pub struct KeywordMatcher {
93 pub name_weight: usize,
95 pub type_weight: usize,
97 pub desc_weight: usize,
99 pub synonyms: HashMap<String, Vec<String>>,
102 pub use_tfidf: bool,
104 pub synonym_weight: f64,
106 pub cjk_bigram_weight: f64,
108}
109
110impl Default for KeywordMatcher {
111 fn default() -> Self {
112 Self {
113 name_weight: 3,
114 type_weight: 2,
115 desc_weight: 1,
116 synonyms: HashMap::new(),
117 use_tfidf: true,
118 synonym_weight: 0.7,
119 cjk_bigram_weight: 0.5,
120 }
121 }
122}
123
124impl KeywordMatcher {
125 pub fn new() -> Self {
127 Self::default()
128 }
129
130 pub fn with_synonyms(mut self, synonyms: HashMap<String, Vec<String>>) -> Self {
132 self.synonyms = synonyms;
133 self
134 }
135
136 pub fn with_tfidf(mut self, enabled: bool) -> Self {
138 self.use_tfidf = enabled;
139 self
140 }
141
142 fn build_terms(&self, query: &str) -> Vec<Term> {
144 let normalized = normalize_text(query).to_lowercase();
145 let tokens: Vec<String> = normalized
146 .split_whitespace()
147 .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()).to_string())
148 .filter(|w| !w.is_empty())
149 .collect();
150
151 let mut terms: Vec<Term> = Vec::new();
152 let mut seen: HashSet<String> = HashSet::new();
153 for tok in tokens {
154 Self::push_term(&mut terms, &mut seen, tok.clone(), TermKind::Direct);
156
157 let syns = self
159 .synonyms
160 .iter()
161 .find(|(k, _)| normalize_text(k).to_lowercase() == tok)
162 .map(|(_, v)| v);
163 if let Some(syns) = syns {
164 for syn in syns {
165 Self::push_term(&mut terms, &mut seen, syn.clone(), TermKind::Synonym);
166 }
167 }
168
169 if tok.chars().any(is_cjk) {
171 for bg in cjk_bigrams(&tok) {
172 Self::push_term(&mut terms, &mut seen, bg, TermKind::Bigram);
173 }
174 }
175 }
176 terms
177 }
178
179 fn push_term(terms: &mut Vec<Term>, seen: &mut HashSet<String>, text: String, kind: TermKind) {
180 if seen.insert(text.clone()) {
181 terms.push(Term { text, kind });
182 }
183 }
184
185 fn compute_idf(&self, terms: &[Term], store: &GraphStore) -> HashMap<String, f64> {
190 let n = store.all_entities().len() as f64;
191 let mut df: HashMap<String, usize> = HashMap::new();
192 for term in terms {
193 df.entry(term.text.clone()).or_insert(0);
194 }
195 for entity in store.all_entities().values() {
196 let name = normalize_text(&entity.name).to_lowercase();
197 let desc = normalize_text(&entity.description).to_lowercase();
198 let typ = normalize_text(&entity.entity_type).to_lowercase();
199 for term in terms {
200 if name.contains(&term.text)
201 || desc.contains(&term.text)
202 || typ.contains(&term.text)
203 {
204 if let Some(c) = df.get_mut(&term.text) {
205 *c += 1;
206 }
207 }
208 }
209 }
210 let mut idf = HashMap::with_capacity(terms.len());
211 for (text, count) in &df {
212 let w = ((n + 1.0) / (*count as f64 + 1.0)).ln() + 1.0;
213 idf.insert(text.clone(), w);
214 }
215 idf
216 }
217
218 fn match_score(
220 &self,
221 term: &Term,
222 name: &str,
223 type_name: &str,
224 desc: &str,
225 idf_weight: f64,
226 ) -> f64 {
227 let mut score = 0.0f64;
228 if name.contains(&term.text) {
229 score += self.name_weight as f64 * idf_weight;
230 }
231 if type_name.contains(&term.text) {
232 score += self.type_weight as f64 * idf_weight;
233 }
234 if desc.contains(&term.text) {
235 score += self.desc_weight as f64 * idf_weight;
236 }
237 match term.kind {
238 TermKind::Direct => score,
239 TermKind::Synonym => score * self.synonym_weight,
240 TermKind::Bigram => score * self.cjk_bigram_weight,
241 }
242 }
243}
244
245impl EntityMatcher for KeywordMatcher {
246 fn find_relevant(&self, query: &str, store: &GraphStore, top_k: usize) -> Vec<String> {
247 let terms = self.build_terms(query);
248 if terms.is_empty() {
249 return Vec::new();
250 }
251
252 let idf = if self.use_tfidf {
254 Some(self.compute_idf(&terms, store))
255 } else {
256 None
257 };
258
259 let mut scored: Vec<(String, f64)> = Vec::new();
260
261 for (id, entity) in store.all_entities() {
262 let name_lower = normalize_text(&entity.name).to_lowercase();
263 let desc_lower = normalize_text(&entity.description).to_lowercase();
264 let type_lower = normalize_text(&entity.entity_type).to_lowercase();
265
266 let mut score = 0.0f64;
267 for term in &terms {
268 let idf_weight = idf
269 .as_ref()
270 .and_then(|m| m.get(&term.text))
271 .copied()
272 .unwrap_or(1.0);
273 score += self.match_score(term, &name_lower, &type_lower, &desc_lower, idf_weight);
274 }
275
276 if score > 0.0 {
277 scored.push((id.clone(), score));
278 }
279 }
280
281 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
282 scored.into_iter().take(top_k).map(|(id, _)| id).collect()
283 }
284}
285
286pub struct EmbeddingMatcher<E: Embeddings> {
296 embeddings: E,
297 cache: std::sync::Mutex<HashMap<String, Vec<f32>>>,
299 min_score: f64,
301}
302
303impl<E: Embeddings> EmbeddingMatcher<E> {
304 pub fn new(embeddings: E) -> Self {
306 Self {
307 embeddings,
308 cache: std::sync::Mutex::new(HashMap::new()),
309 min_score: 0.0,
310 }
311 }
312
313 pub fn with_min_score(mut self, min_score: f64) -> Self {
315 self.min_score = min_score;
316 self
317 }
318
319 async fn get_entity_embedding(&self, entity_id: &str, entity_text: &str) -> Option<Vec<f32>> {
321 {
323 let cache = self.cache.lock().unwrap_or_else(|e| e.into_inner());
324 if let Some(vec) = cache.get(entity_id) {
325 return Some(vec.clone());
326 }
327 }
328
329 match self.embeddings.embed_query(entity_text).await {
331 Ok(vec) => {
332 self.cache
333 .lock()
334 .unwrap_or_else(|e| e.into_inner())
335 .insert(entity_id.to_string(), vec.clone());
336 Some(vec)
337 }
338 Err(e) => {
339 log::warn!(
341 "实体 `{}` 嵌入失败,已从图匹配中排除: {}",
342 entity_id,
343 e
344 );
345 None
346 }
347 }
348 }
349}
350
351impl<E: Embeddings + 'static> EntityMatcher for EmbeddingMatcher<E> {
352 fn find_relevant(&self, query: &str, _store: &GraphStore, _top_k: usize) -> Vec<String> {
353 log::warn!(
361 "EmbeddingMatcher::find_relevant (sync) cannot run embedding matching and no longer \
362 silently falls back to keyword matching; returning empty results for query '{}'. \
363 Use find_relevant_async instead, or configure KeywordMatcher for sync matching.",
364 query
365 );
366 Vec::new()
367 }
368}
369
370impl<E: Embeddings + 'static> EmbeddingMatcher<E> {
371 pub async fn find_relevant_async(
380 &self,
381 query: &str,
382 store: &GraphStore,
383 top_k: usize,
384 ) -> Result<Vec<String>, GraphRAGError> {
385 let query_vec = self.embeddings.embed_query(query).await.map_err(|e| {
386 GraphRAGError::QueryError(format!("EmbeddingMatcher: query embedding failed: {}", e))
387 })?;
388
389 let mut scored: Vec<(String, f64)> = Vec::new();
390
391 for (id, entity) in store.all_entities() {
392 let entity_text = format!(
393 "{} {} {}",
394 entity.name, entity.entity_type, entity.description
395 );
396
397 if let Some(entity_vec) = self.get_entity_embedding(id, &entity_text).await {
398 match cosine_similarity(&query_vec, &entity_vec) {
399 Ok(score) => {
400 scored.push((id.clone(), score as f64));
401 }
402 Err(lc_core::math::MathError::LengthMismatch(a, b)) => {
405 return Err(GraphRAGError::QueryError(format!(
406 "EmbeddingMatcher: vector dimension mismatch {} vs {} (embedding model changed?)",
407 a, b
408 )));
409 }
410 }
411 }
412 }
413
414 let mut scored = filter_by_score(scored, self.min_score);
415 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
416 Ok(scored.into_iter().take(top_k).map(|(id, _)| id).collect())
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use super::*;
423 use crate::graph_rag::graph_store::{Entity, Relation};
424 use lc_embeddings::MockEmbeddings;
425
426 fn make_test_store() -> GraphStore {
427 let mut store = GraphStore::new();
428 store.add_entity(Entity {
429 id: "e1".into(),
430 name: "Rust".into(),
431 entity_type: "Technology".into(),
432 description: "A systems programming language".into(),
433 });
434 store.add_entity(Entity {
435 id: "e2".into(),
436 name: "Python".into(),
437 entity_type: "Technology".into(),
438 description: "A scripting language".into(),
439 });
440 store.add_entity(Entity {
441 id: "e3".into(),
442 name: "Alice".into(),
443 entity_type: "Person".into(),
444 description: "A developer who uses Rust".into(),
445 });
446 store.add_entity(Entity {
447 id: "e4".into(),
448 name: "Tokio".into(),
449 entity_type: "Library".into(),
450 description: "An async runtime for Rust".into(),
451 });
452 store.add_relation(Relation {
453 source: "e3".into(),
454 target: "e1".into(),
455 relation_type: "uses".into(),
456 description: "Alice uses Rust".into(),
457 doc_id: None,
458 });
459 store
460 }
461
462 #[test]
463 fn test_keyword_matcher_basic() {
464 let store = make_test_store();
465 let matcher = KeywordMatcher::new();
466 let results = matcher.find_relevant("Rust programming", &store, 10);
467 assert!(!results.is_empty());
468 assert_eq!(results[0], "e1");
470 }
471
472 #[test]
473 fn test_keyword_matcher_top_k() {
474 let store = make_test_store();
475 let matcher = KeywordMatcher::new();
476 let results = matcher.find_relevant("Technology", &store, 1);
477 assert_eq!(results.len(), 1);
478 }
479
480 #[test]
481 fn test_keyword_matcher_no_match() {
482 let store = make_test_store();
483 let matcher = KeywordMatcher::new();
484 let results = matcher.find_relevant("cooking recipe", &store, 10);
485 assert!(results.is_empty());
486 }
487
488 #[test]
492 fn test_embedding_matcher_sync_returns_empty_not_keyword_fallback() {
493 let store = make_test_store();
494 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
495 let results = matcher.find_relevant("Rust", &store, 10);
496 assert!(
497 results.is_empty(),
498 "sync find_relevant must NOT silently fall back to keyword matching"
499 );
500 }
501
502 #[tokio::test]
507 async fn test_embedding_matcher_async_still_works() {
508 let mut store = GraphStore::new();
509 store.add_entity(Entity {
510 id: "e1".into(),
511 name: "Rust".into(),
512 entity_type: "Technology".into(),
513 description: "A systems programming language".into(),
514 });
515 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
516 let query = "Rust Technology A systems programming language";
517 let results = matcher
518 .find_relevant_async(query, &store, 10)
519 .await
520 .unwrap();
521 assert_eq!(results, vec!["e1".to_string()]);
522 }
523
524 #[test]
525 fn test_keyword_matcher_custom_weights() {
526 let store = make_test_store();
527 let matcher = KeywordMatcher {
528 name_weight: 10,
529 type_weight: 1,
530 desc_weight: 0,
531 ..Default::default()
532 };
533 let results = matcher.find_relevant("Rust", &store, 10);
534 assert!(!results.is_empty());
535 assert_eq!(results[0], "e1");
536 }
537
538 #[test]
540 fn test_keyword_matcher_synonym_expansion() {
541 let mut store = GraphStore::new();
542 store.add_entity(Entity {
543 id: "e1".into(),
544 name: "PostgreSQL".into(),
545 entity_type: "Database".into(),
546 description: "relational database".into(),
547 });
548 store.add_entity(Entity {
549 id: "e2".into(),
550 name: "机器学习".into(),
551 entity_type: "Technology".into(),
552 description: "AI 领域".into(),
553 });
554
555 let synonyms: HashMap<String, Vec<String>> =
556 HashMap::from([("数据库".to_string(), vec!["database".to_string()])]);
557 let matcher = KeywordMatcher::new().with_synonyms(synonyms);
558
559 let results = matcher.find_relevant("数据库", &store, 10);
561 assert!(
562 results.contains(&"e1".to_string()),
563 "同义词 'database' 应能召回 PostgreSQL(e1)"
564 );
565 }
566
567 #[test]
569 fn test_keyword_matcher_fullwidth_normalization() {
570 let mut store = GraphStore::new();
571 store.add_entity(Entity {
572 id: "e1".into(),
573 name: "Rust".into(),
574 entity_type: "Technology".into(),
575 description: "systems language".into(),
576 });
577
578 let matcher = KeywordMatcher::new();
579 let results = matcher.find_relevant("Rust", &store, 10);
581 assert_eq!(results, vec!["e1".to_string()]);
582 }
583
584 #[test]
586 fn test_keyword_matcher_cjk_bigram_recall() {
587 let mut store = GraphStore::new();
588 store.add_entity(Entity {
589 id: "e1".into(),
590 name: "机器学习".into(),
591 entity_type: "Technology".into(),
592 description: "AI 领域".into(),
593 });
594
595 let matcher = KeywordMatcher::new();
596 let results = matcher.find_relevant("机器学习算法", &store, 10);
598 assert!(
599 results.contains(&"e1".to_string()),
600 "CJK 二元组应能召回 '机器学习' 实体"
601 );
602 }
603
604 #[test]
606 fn test_keyword_matcher_tfidf_common_lower_than_rare() {
607 let mut store = GraphStore::new();
608 for name in ["Rust", "Python", "Ruby"] {
609 store.add_entity(Entity {
610 id: name.to_lowercase(),
611 name: name.to_string(),
612 entity_type: "Technology".to_string(),
613 description: String::new(),
614 });
615 }
616
617 let matcher = KeywordMatcher::new();
618 let terms = matcher.build_terms("rust technology");
619 let idf = matcher.compute_idf(&terms, &store);
620
621 let tech_idf = idf["technology"];
623 let rust_idf = idf["rust"];
624 assert!(
625 rust_idf > tech_idf,
626 "稀有词 'rust' 的 IDF ({:.3}) 应高于常见词 'technology' ({:.3})",
627 rust_idf,
628 tech_idf
629 );
630 }
631
632 #[test]
634 fn test_keyword_matcher_tfidf_breaks_fixed_weight_tie() {
635 let mut store = GraphStore::new();
636 store.add_entity(Entity {
637 id: "e1".into(),
638 name: "data".into(),
639 entity_type: "Technology".into(),
640 description: "machine learning".into(),
641 });
642 store.add_entity(Entity {
643 id: "e2".into(),
644 name: "learning".into(),
645 entity_type: "Technology".into(),
646 description: "data".into(),
647 });
648 store.add_entity(Entity {
649 id: "e3".into(),
650 name: "extra".into(),
651 entity_type: "Technology".into(),
652 description: "data warehouse".into(),
653 });
654
655 let matcher = KeywordMatcher::new();
656 let results = matcher.find_relevant("data learning", &store, 10);
659 assert_eq!(results[0], "e2");
660 }
661
662 #[test]
664 fn test_cosine_similarity_identical() {
665 let v = vec![1.0, 0.0, 0.0];
666 let sim = cosine_similarity(&v, &v).unwrap();
667 assert!((sim - 1.0).abs() < 0.001);
668 }
669
670 #[test]
671 fn test_cosine_similarity_orthogonal() {
672 let a = vec![1.0, 0.0];
673 let b = vec![0.0, 1.0];
674 let sim = cosine_similarity(&a, &b).unwrap();
675 assert!((sim - 0.0).abs() < 0.001);
676 }
677
678 #[test]
679 fn test_cosine_similarity_opposite() {
680 let a = vec![1.0, 0.0];
681 let b = vec![-1.0, 0.0];
682 let sim = cosine_similarity(&a, &b).unwrap();
683 assert!((sim - (-1.0)).abs() < 0.001);
684 }
685
686 #[test]
687 fn test_cosine_similarity_zero_norm() {
688 let sim = cosine_similarity(&[], &[]).unwrap();
690 assert_eq!(sim, 0.0);
691 }
692
693 #[test]
695 fn test_cosine_similarity_different_lengths_errors() {
696 let a = vec![1.0];
697 let b = vec![1.0, 2.0];
698 assert!(cosine_similarity(&a, &b).is_err());
699 }
700}