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();
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()
335 .insert(entity_id.to_string(), vec.clone());
336 Some(vec)
337 }
338 Err(_) => None,
339 }
340 }
341}
342
343impl<E: Embeddings + 'static> EntityMatcher for EmbeddingMatcher<E> {
344 fn find_relevant(&self, query: &str, _store: &GraphStore, _top_k: usize) -> Vec<String> {
345 log::warn!(
353 "EmbeddingMatcher::find_relevant (sync) cannot run embedding matching and no longer \
354 silently falls back to keyword matching; returning empty results for query '{}'. \
355 Use find_relevant_async instead, or configure KeywordMatcher for sync matching.",
356 query
357 );
358 Vec::new()
359 }
360}
361
362impl<E: Embeddings + 'static> EmbeddingMatcher<E> {
363 pub async fn find_relevant_async(
372 &self,
373 query: &str,
374 store: &GraphStore,
375 top_k: usize,
376 ) -> Result<Vec<String>, GraphRAGError> {
377 let query_vec = self.embeddings.embed_query(query).await.map_err(|e| {
378 GraphRAGError::QueryError(format!("EmbeddingMatcher: query embedding failed: {}", e))
379 })?;
380
381 let mut scored: Vec<(String, f64)> = Vec::new();
382
383 for (id, entity) in store.all_entities() {
384 let entity_text = format!(
385 "{} {} {}",
386 entity.name, entity.entity_type, entity.description
387 );
388
389 if let Some(entity_vec) = self.get_entity_embedding(id, &entity_text).await {
390 match cosine_similarity(&query_vec, &entity_vec) {
391 Ok(score) => {
392 scored.push((id.clone(), score as f64));
393 }
394 Err(lc_core::math::MathError::LengthMismatch(a, b)) => {
397 return Err(GraphRAGError::QueryError(format!(
398 "EmbeddingMatcher: vector dimension mismatch {} vs {} (embedding model changed?)",
399 a, b
400 )));
401 }
402 }
403 }
404 }
405
406 let mut scored = filter_by_score(scored, self.min_score);
407 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
408 Ok(scored.into_iter().take(top_k).map(|(id, _)| id).collect())
409 }
410}
411
412#[cfg(test)]
413mod tests {
414 use super::*;
415 use crate::graph_rag::graph_store::{Entity, Relation};
416 use lc_embeddings::MockEmbeddings;
417
418 fn make_test_store() -> GraphStore {
419 let mut store = GraphStore::new();
420 store.add_entity(Entity {
421 id: "e1".into(),
422 name: "Rust".into(),
423 entity_type: "Technology".into(),
424 description: "A systems programming language".into(),
425 });
426 store.add_entity(Entity {
427 id: "e2".into(),
428 name: "Python".into(),
429 entity_type: "Technology".into(),
430 description: "A scripting language".into(),
431 });
432 store.add_entity(Entity {
433 id: "e3".into(),
434 name: "Alice".into(),
435 entity_type: "Person".into(),
436 description: "A developer who uses Rust".into(),
437 });
438 store.add_entity(Entity {
439 id: "e4".into(),
440 name: "Tokio".into(),
441 entity_type: "Library".into(),
442 description: "An async runtime for Rust".into(),
443 });
444 store.add_relation(Relation {
445 source: "e3".into(),
446 target: "e1".into(),
447 relation_type: "uses".into(),
448 description: "Alice uses Rust".into(),
449 doc_id: None,
450 });
451 store
452 }
453
454 #[test]
455 fn test_keyword_matcher_basic() {
456 let store = make_test_store();
457 let matcher = KeywordMatcher::new();
458 let results = matcher.find_relevant("Rust programming", &store, 10);
459 assert!(!results.is_empty());
460 assert_eq!(results[0], "e1");
462 }
463
464 #[test]
465 fn test_keyword_matcher_top_k() {
466 let store = make_test_store();
467 let matcher = KeywordMatcher::new();
468 let results = matcher.find_relevant("Technology", &store, 1);
469 assert_eq!(results.len(), 1);
470 }
471
472 #[test]
473 fn test_keyword_matcher_no_match() {
474 let store = make_test_store();
475 let matcher = KeywordMatcher::new();
476 let results = matcher.find_relevant("cooking recipe", &store, 10);
477 assert!(results.is_empty());
478 }
479
480 #[test]
484 fn test_embedding_matcher_sync_returns_empty_not_keyword_fallback() {
485 let store = make_test_store();
486 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
487 let results = matcher.find_relevant("Rust", &store, 10);
488 assert!(
489 results.is_empty(),
490 "sync find_relevant must NOT silently fall back to keyword matching"
491 );
492 }
493
494 #[tokio::test]
499 async fn test_embedding_matcher_async_still_works() {
500 let mut store = GraphStore::new();
501 store.add_entity(Entity {
502 id: "e1".into(),
503 name: "Rust".into(),
504 entity_type: "Technology".into(),
505 description: "A systems programming language".into(),
506 });
507 let matcher = EmbeddingMatcher::new(MockEmbeddings::new(8));
508 let query = "Rust Technology A systems programming language";
509 let results = matcher
510 .find_relevant_async(query, &store, 10)
511 .await
512 .unwrap();
513 assert_eq!(results, vec!["e1".to_string()]);
514 }
515
516 #[test]
517 fn test_keyword_matcher_custom_weights() {
518 let store = make_test_store();
519 let matcher = KeywordMatcher {
520 name_weight: 10,
521 type_weight: 1,
522 desc_weight: 0,
523 ..Default::default()
524 };
525 let results = matcher.find_relevant("Rust", &store, 10);
526 assert!(!results.is_empty());
527 assert_eq!(results[0], "e1");
528 }
529
530 #[test]
532 fn test_keyword_matcher_synonym_expansion() {
533 let mut store = GraphStore::new();
534 store.add_entity(Entity {
535 id: "e1".into(),
536 name: "PostgreSQL".into(),
537 entity_type: "Database".into(),
538 description: "relational database".into(),
539 });
540 store.add_entity(Entity {
541 id: "e2".into(),
542 name: "机器学习".into(),
543 entity_type: "Technology".into(),
544 description: "AI 领域".into(),
545 });
546
547 let synonyms: HashMap<String, Vec<String>> =
548 HashMap::from([("数据库".to_string(), vec!["database".to_string()])]);
549 let matcher = KeywordMatcher::new().with_synonyms(synonyms);
550
551 let results = matcher.find_relevant("数据库", &store, 10);
553 assert!(
554 results.contains(&"e1".to_string()),
555 "同义词 'database' 应能召回 PostgreSQL(e1)"
556 );
557 }
558
559 #[test]
561 fn test_keyword_matcher_fullwidth_normalization() {
562 let mut store = GraphStore::new();
563 store.add_entity(Entity {
564 id: "e1".into(),
565 name: "Rust".into(),
566 entity_type: "Technology".into(),
567 description: "systems language".into(),
568 });
569
570 let matcher = KeywordMatcher::new();
571 let results = matcher.find_relevant("Rust", &store, 10);
573 assert_eq!(results, vec!["e1".to_string()]);
574 }
575
576 #[test]
578 fn test_keyword_matcher_cjk_bigram_recall() {
579 let mut store = GraphStore::new();
580 store.add_entity(Entity {
581 id: "e1".into(),
582 name: "机器学习".into(),
583 entity_type: "Technology".into(),
584 description: "AI 领域".into(),
585 });
586
587 let matcher = KeywordMatcher::new();
588 let results = matcher.find_relevant("机器学习算法", &store, 10);
590 assert!(
591 results.contains(&"e1".to_string()),
592 "CJK 二元组应能召回 '机器学习' 实体"
593 );
594 }
595
596 #[test]
598 fn test_keyword_matcher_tfidf_common_lower_than_rare() {
599 let mut store = GraphStore::new();
600 for name in ["Rust", "Python", "Ruby"] {
601 store.add_entity(Entity {
602 id: name.to_lowercase(),
603 name: name.to_string(),
604 entity_type: "Technology".to_string(),
605 description: String::new(),
606 });
607 }
608
609 let matcher = KeywordMatcher::new();
610 let terms = matcher.build_terms("rust technology");
611 let idf = matcher.compute_idf(&terms, &store);
612
613 let tech_idf = idf["technology"];
615 let rust_idf = idf["rust"];
616 assert!(
617 rust_idf > tech_idf,
618 "稀有词 'rust' 的 IDF ({:.3}) 应高于常见词 'technology' ({:.3})",
619 rust_idf,
620 tech_idf
621 );
622 }
623
624 #[test]
626 fn test_keyword_matcher_tfidf_breaks_fixed_weight_tie() {
627 let mut store = GraphStore::new();
628 store.add_entity(Entity {
629 id: "e1".into(),
630 name: "data".into(),
631 entity_type: "Technology".into(),
632 description: "machine learning".into(),
633 });
634 store.add_entity(Entity {
635 id: "e2".into(),
636 name: "learning".into(),
637 entity_type: "Technology".into(),
638 description: "data".into(),
639 });
640 store.add_entity(Entity {
641 id: "e3".into(),
642 name: "extra".into(),
643 entity_type: "Technology".into(),
644 description: "data warehouse".into(),
645 });
646
647 let matcher = KeywordMatcher::new();
648 let results = matcher.find_relevant("data learning", &store, 10);
651 assert_eq!(results[0], "e2");
652 }
653
654 #[test]
656 fn test_cosine_similarity_identical() {
657 let v = vec![1.0, 0.0, 0.0];
658 let sim = cosine_similarity(&v, &v).unwrap();
659 assert!((sim - 1.0).abs() < 0.001);
660 }
661
662 #[test]
663 fn test_cosine_similarity_orthogonal() {
664 let a = vec![1.0, 0.0];
665 let b = vec![0.0, 1.0];
666 let sim = cosine_similarity(&a, &b).unwrap();
667 assert!((sim - 0.0).abs() < 0.001);
668 }
669
670 #[test]
671 fn test_cosine_similarity_opposite() {
672 let a = vec![1.0, 0.0];
673 let b = vec![-1.0, 0.0];
674 let sim = cosine_similarity(&a, &b).unwrap();
675 assert!((sim - (-1.0)).abs() < 0.001);
676 }
677
678 #[test]
679 fn test_cosine_similarity_zero_norm() {
680 let sim = cosine_similarity(&[], &[]).unwrap();
682 assert_eq!(sim, 0.0);
683 }
684
685 #[test]
687 fn test_cosine_similarity_different_lengths_errors() {
688 let a = vec![1.0];
689 let b = vec![1.0, 2.0];
690 assert!(cosine_similarity(&a, &b).is_err());
691 }
692}