weavatrix_memory/extraction/
linker.rs

1use super::{
2    ExtractedEntity, ExtractionInput, LinkCandidate, LinkDecision, LinkMethod, LinkPolicy,
3    key::normalized,
4};
5use crate::{Confidence, EntityId, MemoryError, Result};
6use std::collections::{BTreeMap, HashMap};
7
8mod catalog;
9mod scoring;
10use catalog::IndexBucket;
11use scoring::{decision, exact_key, generated_id, scope_score};
12
13#[derive(Debug, Clone)]
14pub struct EntityLinker {
15    nodes: Vec<CatalogEntity>,
16    by_id: HashMap<EntityId, usize>,
17    by_label: HashMap<(String, String), IndexBucket>,
18    by_alias: HashMap<(String, String), IndexBucket>,
19    by_external_id: HashMap<(String, String, String), IndexBucket>,
20}
21
22#[derive(Debug, Clone)]
23struct CatalogEntity {
24    id: EntityId,
25    kind: String,
26    repository: Option<String>,
27    branch: Option<String>,
28}
29
30impl EntityLinker {
31    /// Resolves one extracted mention without silently choosing ambiguity.
32    ///
33    /// # Errors
34    ///
35    /// Rejects stable identifiers whose existing entity has a different kind,
36    /// generated identifier collisions, or invalid policy scores.
37    pub fn link(
38        &self,
39        mention: &ExtractedEntity,
40        input: &ExtractionInput,
41        policy: LinkPolicy,
42    ) -> Result<LinkDecision> {
43        if policy.minimum_score > 10_000 || policy.minimum_margin > 10_000 {
44            return Err(MemoryError::InvalidValue {
45                field: "link_policy",
46                reason: "scores must be between 0 and 10,000 basis points",
47            });
48        }
49        if let Some(stable_id) = &mention.stable_id {
50            return self.link_stable_id(mention, stable_id);
51        }
52        let candidates = self.collect_candidates(mention, input)?;
53        self.resolve_candidates(mention, input, policy, candidates)
54    }
55
56    fn link_stable_id(
57        &self,
58        mention: &ExtractedEntity,
59        stable_id: &EntityId,
60    ) -> Result<LinkDecision> {
61        if let Some(index) = self.by_id.get(stable_id) {
62            if self.nodes[*index].kind != normalized(&mention.kind) {
63                return Err(MemoryError::InvalidValue {
64                    field: "extracted_entity.stable_id",
65                    reason: "existing entity kind does not match extracted kind",
66                });
67            }
68            let score = mention.confidence;
69            return Ok(decision(
70                mention,
71                Some(stable_id.clone()),
72                score,
73                LinkMethod::StableId,
74                vec![LinkCandidate {
75                    entity_id: stable_id.clone(),
76                    score,
77                    method: LinkMethod::StableId,
78                }],
79            ));
80        }
81        Ok(decision(
82            mention,
83            Some(stable_id.clone()),
84            mention.confidence,
85            LinkMethod::Created,
86            Vec::new(),
87        ))
88    }
89
90    fn collect_candidates(
91        &self,
92        mention: &ExtractedEntity,
93        input: &ExtractionInput,
94    ) -> Result<Vec<LinkCandidate>> {
95        let kind = normalized(&mention.kind);
96        let mut candidates = BTreeMap::<EntityId, LinkCandidate>::new();
97        self.add_index_matches(
98            self.by_label
99                .get(&(kind.clone(), normalized(&mention.label))),
100            mention,
101            input,
102            9_000,
103            LinkMethod::Label,
104            &mut candidates,
105        )?;
106        self.add_index_matches(
107            self.by_alias
108                .get(&(kind.clone(), normalized(&mention.label))),
109            mention,
110            input,
111            8_500,
112            LinkMethod::Alias,
113            &mut candidates,
114        )?;
115        for alias in &mention.aliases {
116            self.add_index_matches(
117                self.by_alias.get(&(kind.clone(), normalized(alias))),
118                mention,
119                input,
120                8_500,
121                LinkMethod::Alias,
122                &mut candidates,
123            )?;
124        }
125        for (key, value) in &mention.attributes {
126            if key.starts_with("external_id.") {
127                self.add_index_matches(
128                    self.by_external_id
129                        .get(&(kind.clone(), key.clone(), exact_key(value))),
130                    mention,
131                    input,
132                    9_700,
133                    LinkMethod::ExternalId,
134                    &mut candidates,
135                )?;
136            }
137        }
138        for hint in &mention.hints {
139            if let Some(index) = self.by_id.get(&hint.entity_id)
140                && self.nodes[*index].kind == kind
141            {
142                self.add_candidate(
143                    *index,
144                    mention,
145                    input,
146                    hint.confidence.basis_points(),
147                    LinkMethod::ProviderHint,
148                    &mut candidates,
149                )?;
150            }
151        }
152        let mut candidates = candidates.into_values().collect::<Vec<_>>();
153        candidates.sort_by(|left, right| {
154            right
155                .score
156                .cmp(&left.score)
157                .then_with(|| left.entity_id.cmp(&right.entity_id))
158        });
159        Ok(candidates)
160    }
161
162    #[allow(clippy::too_many_arguments)]
163    fn add_index_matches(
164        &self,
165        indexes: Option<&IndexBucket>,
166        mention: &ExtractedEntity,
167        input: &ExtractionInput,
168        base_score: u16,
169        method: LinkMethod,
170        candidates: &mut BTreeMap<EntityId, LinkCandidate>,
171    ) -> Result<()> {
172        if let Some(indexes) = indexes {
173            match indexes {
174                IndexBucket::One(index) => {
175                    self.add_candidate(*index, mention, input, base_score, method, candidates)?;
176                }
177                IndexBucket::Many(indexes) => {
178                    for index in indexes {
179                        self.add_candidate(*index, mention, input, base_score, method, candidates)?;
180                    }
181                }
182            }
183        }
184        Ok(())
185    }
186
187    #[allow(clippy::too_many_arguments)]
188    fn add_candidate(
189        &self,
190        index: usize,
191        mention: &ExtractedEntity,
192        input: &ExtractionInput,
193        base_score: u16,
194        method: LinkMethod,
195        candidates: &mut BTreeMap<EntityId, LinkCandidate>,
196    ) -> Result<()> {
197        let node = &self.nodes[index];
198        let (score, scoped) = scope_score(base_score, node, input);
199        let score = score.min(mention.confidence.basis_points());
200        let method = if method == LinkMethod::Label && scoped {
201            LinkMethod::ScopedLabel
202        } else {
203            method
204        };
205        let candidate = LinkCandidate {
206            entity_id: node.id.clone(),
207            score: Confidence::from_basis_points(score)?,
208            method,
209        };
210        match candidates.get(&node.id) {
211            Some(existing) if existing.score >= candidate.score => {}
212            _ => {
213                candidates.insert(node.id.clone(), candidate);
214            }
215        }
216        Ok(())
217    }
218
219    fn resolve_candidates(
220        &self,
221        mention: &ExtractedEntity,
222        input: &ExtractionInput,
223        policy: LinkPolicy,
224        candidates: Vec<LinkCandidate>,
225    ) -> Result<LinkDecision> {
226        let Some(best) = candidates.first() else {
227            return self.unmatched(mention, input, policy, candidates);
228        };
229        if best.score.basis_points() < policy.minimum_score {
230            return self.unmatched(mention, input, policy, candidates);
231        }
232        if candidates.get(1).is_some_and(|second| {
233            best.score
234                .basis_points()
235                .saturating_sub(second.score.basis_points())
236                < policy.minimum_margin
237        }) {
238            return Ok(decision(
239                mention,
240                None,
241                best.score,
242                LinkMethod::Ambiguous,
243                candidates,
244            ));
245        }
246        Ok(decision(
247            mention,
248            Some(best.entity_id.clone()),
249            best.score,
250            best.method,
251            candidates,
252        ))
253    }
254
255    fn unmatched(
256        &self,
257        mention: &ExtractedEntity,
258        input: &ExtractionInput,
259        policy: LinkPolicy,
260        candidates: Vec<LinkCandidate>,
261    ) -> Result<LinkDecision> {
262        if !policy.create_unmatched {
263            return Ok(decision(
264                mention,
265                None,
266                Confidence::from_basis_points(0)?,
267                LinkMethod::Unresolved,
268                candidates,
269            ));
270        }
271        let id = generated_id(mention, input)?;
272        if self.by_id.contains_key(&id) {
273            return Err(MemoryError::Extraction {
274                provider: "entity-linker".to_owned(),
275                message: format!("generated entity identifier collides with {id}"),
276            });
277        }
278        Ok(decision(
279            mention,
280            Some(id),
281            mention.confidence,
282            LinkMethod::Created,
283            candidates,
284        ))
285    }
286}
weavatrix_memory/extraction/linker.rs

weavatrix_memory/extraction/
linker.rs
